go-omni/transport/http2/server.go

// Package http2 provides an opt-in HTTP/2 transport experiment.
//
// It owns the listener, TLS/ALPN setup, connection admission and global
// request backpressure while using x/net/http2 for frame parsing and stream
// flow control. It is deliberately separate from transport/native, whose
// production preview remains HTTP/1.1.
package http2

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	xhttp2 "golang.org/x/net/http2"
)

const (
	defaultShutdownTimeout = 10 * time.Second
	defaultTLSHandshake    = 10 * time.Second
	defaultMaxBodyBytes    = 8 << 20
	http2ALPNProtocol      = "h2"
)

// Server is an opt-in HTTP/2 transport. Run uses cleartext h2c for local
// experiments; RunTLS or ServeTLS should be used for deployed TLS traffic.
// HTTP/2 stream limits and global request admission are independent: a client
// may open many streams, but only admitted handlers run concurrently.
type Server struct {
	MaxConnections        int
	MaxConcurrentRequests int
	MaxConcurrentStreams  uint32
	MaxHeaderBytes        int
	MaxBodyBytes          int64
	IdleTimeout           time.Duration
	TLSHandshakeTimeout   time.Duration
	ShutdownTimeout       time.Duration
	TLSConfig             *tls.Config

	mu             sync.Mutex
	handler        http.Handler
	listener       net.Listener
	shutdownConfig *http.Server
	shutdownDone   chan struct{}
	shutdownOnce   sync.Once
	closing        bool
	cancel         context.CancelFunc
	conns          map[net.Conn]struct{}
	protocolConns  map[net.Conn]struct{}
	connSlots      chan struct{}
	requestSlots   chan struct{}
	accepted       atomic.Uint64
	rejectedConns  atomic.Uint64
	rejectedReqs   atomic.Uint64
	protocolErrors atomic.Uint64
	goAwaySent     atomic.Uint64
	rstStreams     atomic.Uint64
	goAwayFrames   atomic.Uint64
	activeConns    atomic.Int64
	activeRequests atomic.Int64
	wg             sync.WaitGroup
}

// Stats is a point-in-time view of HTTP/2 transport admission and lifecycle.
// Accepted and rejected counters are monotonic; active counters are gauges.
type Stats struct {
	AcceptedConnections uint64 `json:"acceptedConnections"`
	RejectedConnections uint64 `json:"rejectedConnections"`
	RejectedRequests    uint64 `json:"rejectedRequests"`
	ProtocolErrors      uint64 `json:"protocolErrors"`
	GoAwaySent          uint64 `json:"goAwaySent"`
	RSTStreams          uint64 `json:"rstStreams"`
	GoAwayFrames        uint64 `json:"goAwayFrames"`
	ActiveConnections   int64  `json:"activeConnections"`
	ActiveRequests      int64  `json:"activeRequests"`
}

// New creates an HTTP/2 transport with the supplied handler. The handler may
// also be supplied by App.RunWith through the Runner interface.
func New(handler http.Handler) *Server {
	return &Server{
		handler:      handler,
		conns:        make(map[net.Conn]struct{}),
		MaxBodyBytes: defaultMaxBodyBytes,
	}
}

// Stats returns a lock-free snapshot suitable for a metrics handler.
func (s *Server) Stats() Stats {
	if s == nil {
		return Stats{}
	}
	return Stats{
		AcceptedConnections: s.accepted.Load(),
		RejectedConnections: s.rejectedConns.Load(),
		RejectedRequests:    s.rejectedReqs.Load(),
		ProtocolErrors:      s.protocolErrors.Load(),
		GoAwaySent:          s.goAwaySent.Load(),
		RSTStreams:          s.rstStreams.Load(),
		GoAwayFrames:        s.goAwayFrames.Load(),
		ActiveConnections:   s.activeConns.Load(),
		ActiveRequests:      s.activeRequests.Load(),
	}
}

// Run implements app.Runner. Without TLSConfig it starts an h2c listener for
// local testing. Use RunTLS or TLSConfig for real HTTPS traffic.
func (s *Server) Run(ctx context.Context, handler http.Handler, address string) error {
	if handler != nil {
		s.handler = handler
	}
	if s.handler == nil {
		return errors.New("go-omni/http2: nil handler")
	}
	if address == "" {
		address = ":8080"
	}
	listener, err := net.Listen("tcp", address)
	if err != nil {
		return err
	}
	if s.TLSConfig != nil {
		return s.ServeTLS(ctx, listener, s.TLSConfig)
	}
	return s.Serve(ctx, listener)
}

// RunTLS loads a certificate pair and serves HTTP/2 over TLS with ALPN h2.
func (s *Server) RunTLS(ctx context.Context, handler http.Handler, address, certFile, keyFile string) error {
	if handler != nil {
		s.handler = handler
	}
	if s.handler == nil {
		return errors.New("go-omni/http2: nil handler")
	}
	if address == "" {
		address = ":8443"
	}
	certificate, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}
	listener, err := net.Listen("tcp", address)
	if err != nil {
		return err
	}
	return s.ServeTLS(ctx, listener, &tls.Config{Certificates: []tls.Certificate{certificate}, MinVersion: tls.VersionTLS12})
}

// ServeTLS wraps listener with TLS and requires ALPN h2. A caller-provided
// config is cloned, so the original value is never mutated.
func (s *Server) ServeTLS(ctx context.Context, listener net.Listener, config *tls.Config) error {
	if listener == nil {
		return errors.New("go-omni/http2: nil listener")
	}
	if config == nil {
		return errors.New("go-omni/http2: nil TLS config")
	}
	prepared, err := prepareTLSConfig(config)
	if err != nil {
		_ = listener.Close()
		return err
	}
	return s.Serve(ctx, tls.NewListener(listener, prepared))
}

func prepareTLSConfig(config *tls.Config) (*tls.Config, error) {
	clone := config.Clone()
	if len(clone.NextProtos) > 0 {
		for _, protocol := range clone.NextProtos {
			if protocol != http2ALPNProtocol {
				return nil, fmt.Errorf("go-omni/http2: unsupported TLS ALPN protocol %q; only %q is supported", protocol, http2ALPNProtocol)
			}
		}
	}
	clone.NextProtos = []string{http2ALPNProtocol}
	return clone, nil
}

// Serve accepts h2c connections from listener until ctx is canceled.
func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
	if listener == nil {
		return errors.New("go-omni/http2: nil listener")
	}
	if s.handler == nil {
		return errors.New("go-omni/http2: nil handler")
	}
	if ctx == nil {
		ctx = context.Background()
	}

	s.mu.Lock()
	if s.cancel != nil {
		s.mu.Unlock()
		return errors.New("go-omni/http2: server is already running")
	}
	// Keep the connection context independent from the caller's context. A
	// canceled parent starts graceful shutdown; it must not cancel active
	// streams before the GOAWAY/drain phase has completed.
	baseCtx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel
	s.listener = listener
	s.shutdownDone = make(chan struct{})
	s.shutdownOnce = sync.Once{}
	s.closing = false
	if s.conns == nil {
		s.conns = make(map[net.Conn]struct{})
	}
	if s.protocolConns == nil {
		s.protocolConns = make(map[net.Conn]struct{})
	}
	if s.MaxConnections > 0 {
		s.connSlots = make(chan struct{}, s.MaxConnections)
	} else {
		s.connSlots = nil
	}
	if s.MaxConcurrentRequests > 0 {
		s.requestSlots = make(chan struct{}, s.MaxConcurrentRequests)
	} else {
		s.requestSlots = nil
	}
	s.mu.Unlock()

	maxBodyBytes := s.MaxBodyBytes
	if maxBodyBytes <= 0 {
		maxBodyBytes = defaultMaxBodyBytes
	}
	requestHandler := http.Handler(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.ContentLength > maxBodyBytes {
			_ = request.Body.Close()
			http.Error(writer, "request body too large\n", http.StatusRequestEntityTooLarge)
			return
		}
		request.Body = http.MaxBytesReader(writer, request.Body, maxBodyBytes)
		s.handler.ServeHTTP(writer, request)
	}))
	gate := &requestGate{server: s, next: requestHandler, slots: s.requestSlots}
	protocolServer := &xhttp2.Server{
		MaxConcurrentStreams: s.MaxConcurrentStreams,
		IdleTimeout:          s.IdleTimeout,
		CountError: func(string) {
			s.protocolErrors.Add(1)
		},
	}
	shutdownConfig := &http.Server{IdleTimeout: s.IdleTimeout, MaxHeaderBytes: s.MaxHeaderBytes}
	if err := xhttp2.ConfigureServer(shutdownConfig, protocolServer); err != nil {
		cancel()
		s.mu.Lock()
		s.listener = nil
		s.cancel = nil
		s.mu.Unlock()
		return err
	}
	s.mu.Lock()
	s.shutdownConfig = shutdownConfig
	s.mu.Unlock()
	contextWakeDone := make(chan struct{})
	go func() {
		select {
		case <-ctx.Done():
			// Wake Accept. Serve performs the graceful drain synchronously so
			// its return value reports a shutdown timeout to the caller.
			_ = listener.Close()
		case <-contextWakeDone:
		}
	}()

	defer func() {
		close(contextWakeDone)
		s.mu.Lock()
		closing := s.closing
		shutdownDone := s.shutdownDone
		s.listener = nil
		s.shutdownConfig = nil
		s.cancel = nil
		s.mu.Unlock()
		if !closing {
			cancel()
		} else if shutdownDone != nil {
			<-shutdownDone
		}
	}()

	for {
		conn, err := listener.Accept()
		if err != nil {
			if ctx.Err() != nil || s.isClosing() {
				break
			}
			return err
		}
		s.accepted.Add(1)
		if slots := s.connSlots; slots != nil {
			select {
			case slots <- struct{}{}:
			default:
				s.rejectedConns.Add(1)
				_ = conn.Close()
				continue
			}
		}
		s.mu.Lock()
		if s.closing {
			s.mu.Unlock()
			if s.connSlots != nil {
				<-s.connSlots
			}
			_ = conn.Close()
			continue
		}
		s.conns[conn] = struct{}{}
		s.activeConns.Add(1)
		s.wg.Add(1)
		s.mu.Unlock()
		go s.serveConn(baseCtx, protocolServer, shutdownConfig, gate, conn)
	}
	if ctx.Err() != nil && !s.isClosing() {
		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), s.shutdownTimeout())
		err := s.Shutdown(shutdownCtx)
		shutdownCancel()
		return err
	}
	return nil
}

func (s *Server) shutdownTimeout() time.Duration {
	if s.ShutdownTimeout > 0 {
		return s.ShutdownTimeout
	}
	return defaultShutdownTimeout
}

func (s *Server) serveConn(ctx context.Context, protocolServer *xhttp2.Server, baseConfig *http.Server, handler http.Handler, conn net.Conn) {
	defer s.wg.Done()
	defer func() {
		s.mu.Lock()
		delete(s.conns, conn)
		delete(s.protocolConns, conn)
		s.activeConns.Add(-1)
		s.mu.Unlock()
		if s.connSlots != nil {
			<-s.connSlots
		}
		_ = conn.Close()
	}()
	if tlsConn, ok := conn.(*tls.Conn); ok {
		_ = conn.SetDeadline(time.Now().Add(s.tlsHandshakeTimeout()))
		if err := tlsConn.HandshakeContext(ctx); err != nil {
			return
		}
		_ = conn.SetDeadline(time.Time{})
	}
	s.mu.Lock()
	if s.closing {
		s.mu.Unlock()
		return
	}
	s.protocolConns[conn] = struct{}{}
	s.mu.Unlock()
	observed := &observedConn{
		Conn:       conn,
		readParser: frameParser{prefaceRemain: 24},
		onFrame: func(direction frameDirection, frameType byte) {
			if direction == frameRead && frameType == frameTypeRSTStream {
				s.rstStreams.Add(1)
			}
			if direction == frameWrite && frameType == frameTypeGoAway {
				s.goAwayFrames.Add(1)
			}
		},
	}
	var protocolConn net.Conn = observed
	if tlsConn, ok := conn.(*tls.Conn); ok {
		protocolConn = &observedTLSConn{observedConn: observed, state: tlsConn.ConnectionState()}
	}
	protocolServer.ServeConn(protocolConn, &xhttp2.ServeConnOpts{
		Context: ctx, Handler: handler, BaseConfig: baseConfig,
	})
}

func (s *Server) tlsHandshakeTimeout() time.Duration {
	if s.TLSHandshakeTimeout > 0 {
		return s.TLSHandshakeTimeout
	}
	return defaultTLSHandshake
}

// Shutdown stops accepting connections and waits for active HTTP/2 sessions.
// If the context expires, remaining connections are closed so handlers can
// observe cancellation through their request contexts.
func (s *Server) Shutdown(ctx context.Context) error {
	if s == nil {
		return nil
	}
	if ctx == nil {
		timeout := s.ShutdownTimeout
		if timeout <= 0 {
			timeout = defaultShutdownTimeout
		}
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(context.Background(), timeout)
		defer cancel()
	}
	s.mu.Lock()
	if s.cancel == nil {
		s.mu.Unlock()
		return nil
	}
	alreadyClosing := s.closing
	s.closing = true
	listener := s.listener
	shutdownConfig := s.shutdownConfig
	cancel := s.cancel
	connections := make([]net.Conn, 0, len(s.conns))
	for conn := range s.conns {
		connections = append(connections, conn)
	}
	protocolConnectionCount := len(s.protocolConns)
	s.mu.Unlock()
	if listener != nil {
		_ = listener.Close()
	}
	// x/net/http2's configured shutdown callback sends GOAWAY and lets active
	// streams finish. The base request context remains live until the timeout
	// path below, where connections are forcibly closed.
	if shutdownConfig != nil && !alreadyClosing {
		s.goAwaySent.Add(uint64(protocolConnectionCount))
		_ = shutdownConfig.Shutdown(context.Background())
	}

	done := make(chan struct{})
	go func() {
		s.wg.Wait()
		close(done)
	}()
	select {
	case <-done:
		cancel()
		s.finishShutdown()
		return nil
	case <-ctx.Done():
		for _, conn := range connections {
			_ = conn.Close()
		}
		cancel()
		s.finishShutdown()
		return ctx.Err()
	}
}

func (s *Server) finishShutdown() {
	s.mu.Lock()
	done := s.shutdownDone
	s.mu.Unlock()
	if done != nil {
		s.shutdownOnce.Do(func() { close(done) })
	}
}

func (s *Server) isClosing() bool {
	s.mu.Lock()
	closing := s.closing
	s.mu.Unlock()
	return closing
}

type requestGate struct {
	server *Server
	next   http.Handler
	slots  chan struct{}
}

func (g *requestGate) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	if g.slots != nil {
		select {
		case g.slots <- struct{}{}:
			g.server.activeRequests.Add(1)
			defer func() {
				<-g.slots
				g.server.activeRequests.Add(-1)
			}()
		default:
			g.server.rejectedReqs.Add(1)
			http.Error(writer, "server busy\n", http.StatusServiceUnavailable)
			return
		}
	} else {
		g.server.activeRequests.Add(1)
		defer g.server.activeRequests.Add(-1)
	}
	g.next.ServeHTTP(writer, request)
}