MakyuuIchaival/httputils/server.go

package httputils

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

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

	ReadTimeout  time.Duration
	WriteTimeout time.Duration
	IdleTimeout  time.Duration

	extraHeadersMutex sync.RWMutex
	extraHeaders      map[string]string
}

type RequestContext interface {
	GetPath() string
	GetConnectionTime() time.Time
	GetRequestTime() time.Time
	GetTLSServerName() string
	GetBody() io.Reader
	GetHost() string
	GetProtocol() string
	GetTLSVersion() uint16
	GetTLSCipher() uint16
	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{
			"h2",
			"http/1.1",
		}
	} else {
		server.AddExtraHeader("Connection", "close")
	}

	if server.UseFastHTTP {
		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:  server.ReadTimeout,
				WriteTimeout: server.WriteTimeout,
				IdleTimeout:  server.IdleTimeout,
				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)
			}
		}()
	} else {
		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 := &http.Server{
				ReadTimeout:  server.ReadTimeout,
				WriteTimeout: server.WriteTimeout,
				IdleTimeout:  server.IdleTimeout,
				Addr:         server.ListenAddress,
				TLSConfig:    server.TLSConfig.Config,
				Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					server.Handler(handler(w, r))
				}),
			}
			s.SetKeepAlivesEnabled(!server.EnableHTTP2)

			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{
				Addr: server.ListenAddress,
				QuicConfig: &quic.Config{
					HandshakeIdleTimeout: server.ReadTimeout,
					MaxIdleTimeout:       server.IdleTimeout,
					KeepAlivePeriod:      server.IdleTimeout,
				},
				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))
				}),
			}

			//TODO HACK: SetQuicHeaders does not work before server starts
			go func() {
				h := http.Header{}
				for {
					if err := s.SetQuicHeaders(h); err == nil {
						server.AddExtraHeader("Alt-Svc", h.Get("Alt-Svc"))
						break
					}
					runtime.Gosched()
				}
			}()

			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()
}