MakyuuIchaival/httputils/server.go

package httputils

import (
	"crypto/tls"
	"git.gammaspectra.live/S.O.N.G/MakyuuIchaival/tlsutils"
	"github.com/dgrr/http2"
	"github.com/lucas-clemente/quic-go/http3"
	"github.com/valyala/fasthttp"
	"io"
	"log"
	"mime"
	"net"
	"net/http"
	"sync"
	"time"
)

type Server struct {
	ListenAddress         string
	TLSConfig             *tlsutils.Configuration
	EnableHTTP2           bool
	FastHTTPRequestServer FastHTTPRequestServer
	NetHTTPRequestServer  NetHTTPRequestServer
	EnableHTTP3           bool
	Debug                 bool
	Handler               RequestHandler

	extraHeadersMutex sync.RWMutex
	extraHeaders      map[string]string
}

const ServeStreamChunked = -1

type RequestContext interface {
	GetPath() string
	GetConnectionTime() time.Time
	GetRequestTime() time.Time
	GetTLSServerName() string
	GetBody() io.Reader
	GetRequestHeader(name string) string
	GetResponseHeader(name string) string
	AddResponseHeader(name string, value string)
	SetResponseHeader(name string, value string)

	ServeStream(stream Stream)
	ServeBytes(content []byte)
	SetResponseCode(code int)
	DoRedirect(location string, code int)

	IsGet() bool
	IsPost() bool
	IsOptions() bool
	IsHead() bool

	GetExtraHeaders() map[string]string

	AddTiming(name string, desc string, d time.Duration)
	AddTimingInformational(name string, desc string)
}

type RequestHandler func(ctx RequestContext)

type FastHTTPRequestServer func(ctx *fasthttp.RequestCtx) RequestContext
type NetHTTPRequestServer func(w http.ResponseWriter, r *http.Request) RequestContext

func (server *Server) Serve() {
	var wg sync.WaitGroup

	mime.AddExtensionType(".bin", "application/octet-stream")

	if server.EnableHTTP2 {
		server.TLSConfig.Config.NextProtos = []string{
			"http/1.1",
			http2.H2TLSProto,
		}
	} else {
		server.AddExtraHeader("Connection", "close")
	}

	wg.Add(1)
	go func() {
		defer wg.Done()

		handler := func(ctx *fasthttp.RequestCtx) RequestContext {
			return NewRequestContextFromFastHttp(server, ctx)
		}

		if server.FastHTTPRequestServer != nil {
			handler = server.FastHTTPRequestServer
		}

		s := &fasthttp.Server{
			ReadTimeout: 5 * time.Second,
			IdleTimeout: 15 * time.Second,
			Handler: func(ctx *fasthttp.RequestCtx) {
				server.Handler(handler(ctx))
			},
			NoDefaultServerHeader: true,
			NoDefaultDate:         true,
			DisableKeepalive:      !server.EnableHTTP2,
			TCPKeepalive:          server.EnableHTTP2,
			TLSConfig:             server.TLSConfig.Config,
		}

		if server.EnableHTTP2 {
			http2.ConfigureServer(s, http2.ServerConfig{
				Debug: server.Debug,
			})
		}

		log.Printf("[Server] Serving TCP on %s", server.ListenAddress)
		ln, err := net.Listen("tcp", server.ListenAddress)
		if err != nil {
			log.Panic(err)
		}
		defer ln.Close()
		err = s.Serve(tls.NewListener(ln, s.TLSConfig.Clone()))
		if err != nil {
			log.Panic(err)
		}
	}()

	if server.EnableHTTP3 {
		wg.Add(1)
		go func() {
			defer wg.Done()

			handler := func(w http.ResponseWriter, r *http.Request) RequestContext {
				return NewRequestContextFromHttp(server, w, r)
			}

			if server.NetHTTPRequestServer != nil {
				handler = server.NetHTTPRequestServer
			}

			s := &http3.Server{
				Server: &http.Server{
					Addr:      server.ListenAddress,
					TLSConfig: server.TLSConfig.QUICConfig,
					Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
						if server.Debug {
							log.Print("Received HTTP/3 request")
						}
						server.Handler(handler(w, r))
					}),
				},
			}

			h := http.Header{}
			s.SetQuicHeaders(h)
			server.AddExtraHeader("Alt-Svc", h.Get("Alt-Svc"))

			log.Printf("[Server] Serving UDP on %s", server.ListenAddress)
			err := s.ListenAndServe()
			if err != nil {
				log.Panic(err)
			}
		}()
	}

	wg.Wait()
}

func (server *Server) GetExtraHeaders() map[string]string {
	server.extraHeadersMutex.RLock()

	r := make(map[string]string)
	if server.extraHeaders != nil {
		for key, value := range server.extraHeaders {
			r[key] = value
		}
	}
	server.extraHeadersMutex.RUnlock()
	return r
}

func (server *Server) AddExtraHeader(key, value string) {
	server.extraHeadersMutex.Lock()
	if server.extraHeaders == nil {
		server.extraHeaders = make(map[string]string)
	}
	server.extraHeaders[key] = value
	server.extraHeadersMutex.Unlock()
}