go-omni/transport/native/server.go

// Package native provides GoOmni's first project-owned HTTP/1.1 server.
//
// It deliberately keeps the public handler boundary as net/http so existing
// GoOmni business code remains unchanged while connection management and
// protocol I/O move out of net/http.Server. The parser supports Content-Length
// and chunked request bodies with configurable size limits.
package native

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	defaultReadHeaderTimeout = 5 * time.Second
	defaultReadTimeout       = 30 * time.Second
	defaultWriteTimeout      = 30 * time.Second
	defaultIdleTimeout       = 60 * time.Second
	defaultMaxHeaderBytes    = 1 << 20
	defaultMaxBodyBytes      = 8 << 20
	defaultReadBufferSize    = 4 << 10
	maxPooledResponseBytes   = 64 << 10
	nativeALPNProtocol       = "http/1.1"
)

// Server is a small HTTP/1.1 server with project-owned connection lifecycle.
// It supports keep-alive, graceful shutdown, bounded concurrency and standard
// http.Handler values. Zero MaxConnections and MaxConcurrentRequests mean
// unlimited; pooled protocol objects are reused automatically.
type Server struct {
	ReadHeaderTimeout     time.Duration
	ReadTimeout           time.Duration
	WriteTimeout          time.Duration
	IdleTimeout           time.Duration
	ReadBufferSize        int
	MaxHeaderBytes        int
	MaxBodyBytes          int64
	MaxConnections        int
	MaxConcurrentRequests int
	// TLSConfig enables TLS for Run and App.RunWith. ServeTLS can be used
	// when the listener and TLS configuration are assembled by the caller.
	TLSConfig *tls.Config

	mu               sync.Mutex
	handler          http.Handler
	listener         net.Listener
	conns            map[*connection]struct{}
	closing          bool
	cancel           context.CancelFunc
	wg               sync.WaitGroup
	connSlots        chan struct{}
	requestSlots     chan struct{}
	accepted         atomic.Uint64
	rejectedConns    atomic.Uint64
	rejectedRequests atomic.Uint64
	activeConns      atomic.Int64
	activeRequests   atomic.Int64
	readerPool       sync.Pool
	writerPool       sync.Pool
	connectionPool   sync.Pool
}

// Stats is a point-in-time view of the native server lifecycle. Counters are
// monotonic unless their name starts with Active.
type Stats struct {
	AcceptedConnections uint64 `json:"acceptedConnections"`
	// Rejections count admission decisions, even if writing a 503 fails.
	RejectedConnections uint64 `json:"rejectedConnections"`
	RejectedRequests    uint64 `json:"rejectedRequests"`
	ActiveConnections   int64  `json:"activeConnections"`
	ActiveRequests      int64  `json:"activeRequests"`
}

// Stats returns native connection and request counters without taking the
// server lifecycle mutex, so it is safe to call from a metrics endpoint.
func (s *Server) Stats() Stats {
	if s == nil {
		return Stats{}
	}
	return Stats{
		AcceptedConnections: s.accepted.Load(),
		RejectedConnections: s.rejectedConns.Load(),
		RejectedRequests:    s.rejectedRequests.Load(),
		ActiveConnections:   s.activeConns.Load(),
		ActiveRequests:      s.activeRequests.Load(),
	}
}

type connection struct {
	server        *Server
	conn          net.Conn
	active        atomic.Bool
	closeWhenIdle atomic.Bool
	closed        atomic.Bool
	counted       bool
	slotHeld      bool
	closeDone     chan struct{}
}

// New creates a native server. The handler may also be supplied later through
// App.RunWith, which is useful when the App owns the lifecycle.
func New(handler http.Handler) *Server {
	server := &Server{
		handler:           handler,
		ReadHeaderTimeout: defaultReadHeaderTimeout,
		ReadTimeout:       defaultReadTimeout,
		WriteTimeout:      defaultWriteTimeout,
		IdleTimeout:       defaultIdleTimeout,
		ReadBufferSize:    defaultReadBufferSize,
		MaxHeaderBytes:    defaultMaxHeaderBytes,
		MaxBodyBytes:      defaultMaxBodyBytes,
		conns:             make(map[*connection]struct{}),
	}
	server.readerPool.New = func() any { return bufio.NewReaderSize(nil, server.readBufferSize()) }
	server.writerPool.New = func() any { return newResponseWriter("") }
	server.connectionPool.New = func() any { return &connection{server: server} }
	return server
}

func (s *Server) readBufferSize() int {
	if s.ReadBufferSize > 0 {
		return s.ReadBufferSize
	}
	return defaultReadBufferSize
}

// Run implements app.Runner without importing the app package.
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/native: 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, listens on address and serves native HTTP
// over TLS. The server still owns the listener and follows the same shutdown
// and connection-limit rules as Run.
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/native: 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 a TLS listener and serves native HTTP. The
// caller owns the TLS configuration value; it must contain a certificate or
// a certificate callback accepted by crypto/tls.
func (s *Server) ServeTLS(ctx context.Context, listener net.Listener, config *tls.Config) error {
	if listener == nil {
		return errors.New("go-omni/native: nil listener")
	}
	if config == nil {
		return errors.New("go-omni/native: nil TLS config")
	}
	config, err := prepareTLSConfig(config)
	if err != nil {
		_ = listener.Close()
		return err
	}
	return s.Serve(ctx, tls.NewListener(listener, config))
}

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

// Serve accepts connections from listener until the context is canceled or an
// unrecoverable accept error occurs. The listener is owned by the Server.
func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
	if listener == nil {
		return errors.New("go-omni/native: nil listener")
	}
	if s.handler == nil {
		return errors.New("go-omni/native: nil handler")
	}
	if ctx == nil {
		ctx = context.Background()
	}

	s.mu.Lock()
	if s.listener != nil {
		s.mu.Unlock()
		return errors.New("go-omni/native: server is already running")
	}
	s.listener = listener
	s.closing = false
	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
	}
	baseCtx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel
	s.mu.Unlock()
	defer cancel()

	shutdownDone := make(chan struct{})
	go func() {
		select {
		case <-ctx.Done():
			timeout := s.shutdownTimeout()
			shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
			_ = s.Shutdown(shutdownCtx)
			cancel()
		case <-shutdownDone:
		}
	}()

	var serveErr error
	for {
		conn, err := listener.Accept()
		if err != nil {
			if ctx.Err() != nil || s.isClosing() {
				break
			}
			serveErr = err
			break
		}
		s.accepted.Add(1)
		if slots := s.connSlots; slots != nil {
			select {
			case slots <- struct{}{}:
			default:
				s.rejectedConns.Add(1)
				_ = conn.Close()
				continue
			}
		}
		state := s.connectionPool.Get().(*connection)
		state.server = s
		state.conn = conn
		state.active.Store(false)
		state.closeWhenIdle.Store(false)
		state.closed.Store(false)
		state.counted = true
		state.slotHeld = s.connSlots != nil
		state.closeDone = make(chan struct{})
		s.activeConns.Add(1)
		s.mu.Lock()
		if s.closing {
			s.mu.Unlock()
			state.close()
			s.recycleConnection(state)
			continue
		}
		s.conns[state] = struct{}{}
		s.wg.Add(1)
		s.mu.Unlock()
		go s.serveConnection(baseCtx, state)
	}

	close(shutdownDone)
	s.mu.Lock()
	if s.listener == listener {
		s.listener = nil
	}
	s.mu.Unlock()
	if ctx.Err() != nil || s.isClosing() {
		timeout := s.shutdownTimeout()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
		shutdownErr := s.Shutdown(shutdownCtx)
		cancel()
		return errors.Join(shutdownErr, serveErr)
	}
	_ = listener.Close()
	timeout := s.shutdownTimeout()
	shutdownCtx, cancel := context.WithTimeout(context.Background(), timeout)
	shutdownErr := s.Shutdown(shutdownCtx)
	cancel()
	return errors.Join(shutdownErr, serveErr)
}

// Shutdown stops accepting requests, cancels active request contexts and waits
// for connections to finish. On timeout, active sockets are closed.
func (s *Server) Shutdown(ctx context.Context) error {
	if ctx == nil {
		ctx = context.Background()
	}
	s.mu.Lock()
	s.closing = true
	if s.cancel != nil {
		s.cancel()
	}
	listener := s.listener
	states := make([]*connection, 0, len(s.conns))
	for state := range s.conns {
		state.closeWhenIdle.Store(true)
		states = append(states, state)
	}
	s.listener = nil
	s.mu.Unlock()
	if listener != nil {
		_ = listener.Close()
	}
	for _, state := range states {
		if !state.active.Load() {
			state.close()
		}
	}

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

func (s *Server) shutdownTimeout() time.Duration {
	if value := s.WriteTimeout; value > 0 {
		return value
	}
	return defaultWriteTimeout
}

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

func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
	defer s.wg.Done()
	defer func() {
		state.close()
		s.recycleConnection(state)
	}()
	reader := s.readerPool.Get().(*bufio.Reader)
	reader.Reset(state.conn)
	defer func() {
		reader.Reset(nil)
		s.readerPool.Put(reader)
	}()
	firstRequest := true
	for {
		readTimeout := s.ReadHeaderTimeout
		if readTimeout <= 0 {
			readTimeout = defaultReadHeaderTimeout
		}
		if !firstRequest && s.IdleTimeout > 0 {
			readTimeout = s.IdleTimeout
		} else if s.IdleTimeout > 0 && s.IdleTimeout < readTimeout {
			readTimeout = s.IdleTimeout
		}
		_ = state.conn.SetReadDeadline(time.Now().Add(readTimeout))
		request, closeAfter, err := readRequest(reader, state.conn, s.MaxHeaderBytes, s.MaxBodyBytes, func() {
			if s.ReadTimeout > 0 {
				_ = state.conn.SetReadDeadline(time.Now().Add(s.ReadTimeout))
			} else {
				_ = state.conn.SetReadDeadline(time.Time{})
			}
		})
		if err != nil {
			if !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) && !isTimeout(err) {
				status := http.StatusBadRequest
				message := "bad request\n"
				var protocolErr *protocolError
				if errors.As(err, &protocolErr) {
					status = protocolErr.status
					message = protocolErr.message + "\n"
				}
				_ = writeSimple(state.conn, status, message, true, s.WriteTimeout)
			}
			return
		}
		_ = state.conn.SetReadDeadline(time.Time{})
		firstRequest = false
		requestContext, cancel := context.WithCancel(baseCtx)
		request = request.WithContext(requestContext)
		state.active.Store(true)
		s.activeRequests.Add(1)
		writer := s.writerPool.Get().(*responseWriter)
		writer.reset(request.Method)
		writer.configure(request, state.conn, closeAfter, s.WriteTimeout, cancel)
		requestSlot := false
		if slots := s.requestSlots; slots != nil {
			select {
			case slots <- struct{}{}:
				requestSlot = true
			case <-baseCtx.Done():
				cancel()
				writer.reset("")
				s.writerPool.Put(writer)
				state.active.Store(false)
				s.activeRequests.Add(-1)
				return
			default:
				s.rejectedRequests.Add(1)
				writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
				writer.WriteHeader(http.StatusServiceUnavailable)
				writer.body.WriteString("server busy\n")
				closing := s.isClosing() || state.closeWhenIdle.Load()
				_ = writer.writeTo(state.conn, closeAfter || closing, s.WriteTimeout)
				cancel()
				writer.reset("")
				s.writerPool.Put(writer)
				state.active.Store(false)
				s.activeRequests.Add(-1)
				if closeAfter || closing {
					return
				}
				continue
			}
		}
		handler := s.handler
		var disconnectWatch *clientDisconnectWatch
		if body, ok := request.Body.(*requestBody); ok {
			// A body reader and the disconnect watcher must never read the
			// bufio.Reader concurrently. The body starts the watcher once it
			// reaches EOF; body read errors cancel the request directly.
			body.cancel = cancel
			body.onComplete = func() {
				disconnectWatch = newClientDisconnectWatch(reader, state.conn, cancel)
			}
		} else {
			disconnectWatch = newClientDisconnectWatch(reader, state.conn, cancel)
		}
		handler.ServeHTTP(writer, request)
		bodyErr := request.Body.Close()
		if disconnectWatch != nil {
			disconnectWatch.stop()
		}
		cancel()

		closing := s.isClosing() || state.closeWhenIdle.Load() || bodyErr != nil
		if bodyErr != nil {
			var protocolErr *protocolError
			if !errors.As(bodyErr, &protocolErr) {
				if requestSlot {
					<-s.requestSlots
				}
				writer.reset("")
				s.writerPool.Put(writer)
				state.active.Store(false)
				s.activeRequests.Add(-1)
				return
			}
			// Request bodies are streamed into the handler, so a malformed
			// chunk or an over-limit body can be discovered while draining
			// after the handler returns. Replace the buffered response with
			// the same protocol error that eager parsing used to return.
			if !writer.streaming {
				writer.reset(request.Method)
				writer.configure(request, state.conn, true, s.WriteTimeout, cancel)
				writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
				writer.WriteHeader(protocolErr.status)
				writer.body.WriteString(protocolErr.message + "\n")
			}
		}
		if err := writer.writeTo(state.conn, closeAfter || closing, s.WriteTimeout); err != nil {
			if requestSlot {
				<-s.requestSlots
			}
			writer.reset("")
			s.writerPool.Put(writer)
			state.active.Store(false)
			s.activeRequests.Add(-1)
			return
		}
		if requestSlot {
			<-s.requestSlots
		}
		shouldClose := closeAfter || writer.closeAfter || closing
		writer.reset("")
		s.writerPool.Put(writer)
		state.active.Store(false)
		s.activeRequests.Add(-1)
		if shouldClose {
			return
		}
	}
}

func (c *connection) close() {
	if c.closed.CompareAndSwap(false, true) {
		defer close(c.closeDone)
		_ = c.conn.Close()
		c.server.mu.Lock()
		delete(c.server.conns, c)
		c.server.mu.Unlock()
		if c.counted {
			c.server.activeConns.Add(-1)
			c.counted = false
		}
	}
}

func (s *Server) recycleConnection(c *connection) {
	<-c.closeDone
	if c.slotHeld {
		if slots := s.connSlots; slots != nil {
			<-slots
		}
		c.slotHeld = false
	}
	s.connectionPool.Put(c)
}

type protocolError struct {
	status  int
	message string
}

func (e *protocolError) Error() string { return e.message }

func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBodyBytes int64, setBodyDeadline func()) (*http.Request, bool, error) {
	if maxHeaderBytes <= 0 {
		maxHeaderBytes = defaultMaxHeaderBytes
	}
	if maxBodyBytes <= 0 {
		maxBodyBytes = defaultMaxBodyBytes
	}
	used := 0
	readLine := func() (string, error) {
		line, err := reader.ReadString('\n')
		used += len(line)
		if used > maxHeaderBytes {
			return "", &protocolError{status: http.StatusRequestHeaderFieldsTooLarge, message: "request headers too large"}
		}
		if err != nil {
			return "", err
		}
		if !strings.HasSuffix(line, "\n") {
			return "", &protocolError{status: http.StatusBadRequest, message: "malformed request line"}
		}
		return strings.TrimSuffix(strings.TrimSuffix(line, "\n"), "\r"), nil
	}

	line, err := readLine()
	if err != nil {
		return nil, false, err
	}
	parts := strings.Split(line, " ")
	if len(parts) != 3 || !validToken(parts[0]) {
		return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed request line"}
	}
	protoMajor, protoMinor, ok := parseProtocol(parts[2])
	if !ok {
		return nil, true, &protocolError{status: http.StatusHTTPVersionNotSupported, message: "unsupported HTTP version"}
	}
	target := parts[1]
	var requestURL *url.URL
	if target == "*" {
		requestURL = &url.URL{Path: "*"}
	} else {
		requestURL, err = url.ParseRequestURI(target)
		if err != nil || requestURL.Path == "" || !strings.HasPrefix(requestURL.Path, "/") {
			return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed request target"}
		}
	}

	header := make(http.Header)
	for {
		line, err = readLine()
		if err != nil {
			return nil, true, err
		}
		if line == "" {
			break
		}
		if strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t") {
			return nil, true, &protocolError{status: http.StatusBadRequest, message: "folded headers are not supported"}
		}
		index := strings.IndexByte(line, ':')
		if index <= 0 {
			return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed header"}
		}
		name := strings.TrimSpace(line[:index])
		value := strings.TrimSpace(line[index+1:])
		if !validToken(name) || !validHeaderValue(value) {
			return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed header"}
		}
		header.Add(name, value)
	}
	if protoMajor == 1 && protoMinor == 1 && (len(header.Values("Host")) != 1 || header.Get("Host") == "") {
		return nil, true, &protocolError{status: http.StatusBadRequest, message: "missing host header"}
	}

	for _, value := range header.Values("Expect") {
		if strings.TrimSpace(value) != "" {
			return nil, true, &protocolError{status: http.StatusExpectationFailed, message: "unsupported expectation"}
		}
	}
	chunked := false
	for _, value := range header.Values("Transfer-Encoding") {
		for _, coding := range strings.Split(value, ",") {
			coding = strings.TrimSpace(coding)
			if coding == "" || strings.EqualFold(coding, "identity") {
				continue
			}
			if strings.EqualFold(coding, "chunked") && !chunked {
				chunked = true
				continue
			}
			return nil, true, &protocolError{status: http.StatusNotImplemented, message: "transfer encoding is not supported"}
		}
	}
	if chunked && len(header.Values("Content-Length")) > 0 {
		return nil, true, &protocolError{status: http.StatusBadRequest, message: "content length with chunked encoding"}
	}
	contentLength := int64(0)
	contentLengthSeen := false
	for _, value := range header.Values("Content-Length") {
		for _, part := range strings.Split(value, ",") {
			part = strings.TrimSpace(part)
			parsed, parseErr := strconv.ParseInt(part, 10, 64)
			if parseErr != nil || parsed < 0 || (contentLengthSeen && parsed != contentLength) {
				return nil, true, &protocolError{status: http.StatusBadRequest, message: "invalid content length"}
			}
			contentLength = parsed
			contentLengthSeen = true
		}
	}
	if !chunked && contentLength > maxBodyBytes {
		return nil, true, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
	}
	var body io.ReadCloser = http.NoBody
	var trailer http.Header
	if chunked {
		trailer = make(http.Header)
		body = &requestBody{
			reader: reader, maxBytes: maxBodyBytes, chunked: true,
			trailer: trailer, setDeadline: setBodyDeadline, readLine: readLine,
		}
		// net/http uses -1 to indicate that the decoded length was not known
		// from the request headers.
		contentLength = -1
	} else {
		if contentLength > 0 {
			if setBodyDeadline != nil {
				setBodyDeadline()
			}
			body = &requestBody{reader: reader, remaining: contentLength, setDeadline: setBodyDeadline}
		}
	}
	request := &http.Request{
		Method:        parts[0],
		URL:           requestURL,
		Proto:         parts[2],
		ProtoMajor:    protoMajor,
		ProtoMinor:    protoMinor,
		Header:        header,
		Body:          body,
		ContentLength: contentLength,
		Host:          header.Get("Host"),
		RequestURI:    target,
		RemoteAddr:    conn.RemoteAddr().String(),
		Trailer:       trailer,
	}
	if tlsConn, ok := conn.(*tls.Conn); ok {
		state := tlsConn.ConnectionState()
		request.TLS = &state
	}
	if chunked {
		request.TransferEncoding = []string{"chunked"}
	}
	closeAfter := headerHasToken(header.Values("Connection"), "close") ||
		(protoMajor == 1 && protoMinor == 0 && !headerHasToken(header.Values("Connection"), "keep-alive"))
	request.Close = closeAfter
	return request, closeAfter, nil
}

type requestBody struct {
	reader      *bufio.Reader
	remaining   int64
	maxBytes    int64
	readBytes   int64
	chunked     bool
	chunkRemain int64
	needSuffix  bool
	trailer     http.Header
	readLine    func() (string, error)
	setDeadline func()
	started     bool
	cancel      context.CancelFunc
	onComplete  func()
	done        bool
	closed      bool
	err         error
}

func (b *requestBody) Read(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	if b.err != nil {
		return 0, b.err
	}
	if b.done {
		return 0, io.EOF
	}
	if !b.chunked {
		return b.readContentLength(p)
	}
	return b.readChunked(p)
}

func (b *requestBody) readContentLength(p []byte) (int, error) {
	if b.remaining == 0 {
		b.complete()
		return 0, io.EOF
	}
	if !b.started {
		b.started = true
		if b.setDeadline != nil {
			b.setDeadline()
		}
	}
	if int64(len(p)) > b.remaining {
		p = p[:b.remaining]
	}
	n, err := b.reader.Read(p)
	b.remaining -= int64(n)
	b.readBytes += int64(n)
	if b.remaining == 0 {
		b.complete()
		if n > 0 {
			return n, nil
		}
		return 0, io.EOF
	}
	if err != nil {
		b.fail(err)
	}
	return n, err
}

func (b *requestBody) readChunked(p []byte) (int, error) {
	for {
		if b.chunkRemain == 0 {
			if b.needSuffix {
				var suffix [2]byte
				if _, err := io.ReadFull(b.reader, suffix[:]); err != nil {
					b.fail(err)
					return 0, err
				}
				if suffix[0] != '\r' || suffix[1] != '\n' {
					err := &protocolError{status: http.StatusBadRequest, message: "malformed chunk terminator"}
					b.fail(err)
					return 0, err
				}
				b.needSuffix = false
			}
			if err := b.beginChunk(); err != nil {
				b.fail(err)
				return 0, err
			}
			if b.done {
				return 0, io.EOF
			}
		}
		if int64(len(p)) > b.chunkRemain {
			p = p[:b.chunkRemain]
		}
		n, err := b.reader.Read(p)
		b.chunkRemain -= int64(n)
		b.readBytes += int64(n)
		if b.chunkRemain == 0 {
			b.needSuffix = true
		}
		if n > 0 {
			if err != nil && b.chunkRemain > 0 {
				b.fail(err)
				return n, err
			}
			return n, nil
		}
		if err != nil {
			b.fail(err)
			return 0, err
		}
	}
}

func (b *requestBody) beginChunk() error {
	if b.setDeadline != nil {
		b.setDeadline()
	}
	line, err := b.readLine()
	if err != nil {
		return err
	}
	chunkSizeText := line
	if extension := strings.IndexByte(chunkSizeText, ';'); extension >= 0 {
		chunkSizeText = chunkSizeText[:extension]
	}
	chunkSizeText = strings.TrimSpace(chunkSizeText)
	if chunkSizeText == "" || !validHexNumber(chunkSizeText) {
		return &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
	}
	chunkSize, parseErr := strconv.ParseUint(chunkSizeText, 16, 63)
	if parseErr != nil {
		return &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
	}
	if chunkSize == 0 {
		for {
			line, err = b.readLine()
			if err != nil {
				return err
			}
			if line == "" {
				b.complete()
				return nil
			}
			if strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t") {
				return &protocolError{status: http.StatusBadRequest, message: "folded trailers are not supported"}
			}
			index := strings.IndexByte(line, ':')
			if index <= 0 {
				return &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
			}
			name := strings.TrimSpace(line[:index])
			value := strings.TrimSpace(line[index+1:])
			if !validToken(name) || !validHeaderValue(value) {
				return &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
			}
			b.trailer.Add(name, value)
		}
	}
	if int64(chunkSize) > b.maxBytes-b.readBytes {
		return &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
	}
	b.chunkRemain = int64(chunkSize)
	return nil
}

func (b *requestBody) Close() error {
	if b.closed {
		return b.err
	}
	b.closed = true
	if !b.done && b.err == nil {
		_, _ = io.Copy(io.Discard, b)
	}
	return b.err
}

func (b *requestBody) complete() {
	if b.done {
		return
	}
	b.done = true
	if b.onComplete != nil {
		b.onComplete()
	}
}

func (b *requestBody) fail(err error) {
	if b.err == nil {
		b.err = err
	}
	b.done = true
	if b.cancel != nil {
		b.cancel()
	}
}

func validHexNumber(value string) bool {
	for index := 0; index < len(value); index++ {
		char := value[index]
		if !((char >= '0' && char <= '9') || (char >= 'a' && char <= 'f') || (char >= 'A' && char <= 'F')) {
			return false
		}
	}
	return true
}

func parseProtocol(value string) (int, int, bool) {
	if value != "HTTP/1.1" && value != "HTTP/1.0" {
		return 0, 0, false
	}
	return 1, int(value[len(value)-1] - '0'), true
}

func validToken(value string) bool {
	if value == "" {
		return false
	}
	for index := 0; index < len(value); index++ {
		char := value[index]
		if !isTokenByte(char) {
			return false
		}
	}
	return true
}

func isTokenByte(char byte) bool {
	if char >= '0' && char <= '9' || char >= 'A' && char <= 'Z' || char >= 'a' && char <= 'z' {
		return true
	}
	return strings.ContainsRune("!#$%&'*+-.^_`|~", rune(char))
}

func validHeaderValue(value string) bool {
	for index := 0; index < len(value); index++ {
		char := value[index]
		if (char < 0x20 && char != '\t') || char == 0x7f {
			return false
		}
	}
	return true
}

func headerHasToken(values []string, token string) bool {
	for _, value := range values {
		for _, item := range strings.Split(value, ",") {
			if strings.EqualFold(strings.TrimSpace(item), token) {
				return true
			}
		}
	}
	return false
}

func isTimeout(err error) bool {
	var netErr net.Error
	return errors.As(err, &netErr) && netErr.Timeout()
}

// Most handlers complete before this delay. A delayed, blocking read avoids a
// watcher goroutine for those requests and does not poll idle client sockets.
const clientDisconnectWatchDelay = 5 * time.Millisecond

type clientDisconnectWatch struct {
	conn    net.Conn
	done    chan struct{}
	stopped atomic.Bool
	timer   *time.Timer
}

func newClientDisconnectWatch(reader *bufio.Reader, conn net.Conn, cancel context.CancelFunc) *clientDisconnectWatch {
	watch := &clientDisconnectWatch{conn: conn, done: make(chan struct{})}
	watch.timer = time.AfterFunc(clientDisconnectWatchDelay, func() {
		defer close(watch.done)
		if watch.stopped.Load() {
			return
		}
		// Only this goroutine may read until stop joins it. Peek preserves
		// pipelined bytes for the next request and bounds read-ahead to the
		// existing reader buffer. Once data is available, stop watching:
		// draining it to detect EOF would require an unbounded pipeline queue.
		_, err := reader.Peek(1)
		if err != nil && !(watch.stopped.Load() && isTimeout(err)) {
			cancel()
		}
	})
	return watch
}

func (w *clientDisconnectWatch) stop() {
	if w.timer.Stop() {
		return
	}
	w.stopped.Store(true)
	// The watcher never changes deadlines: this interrupt cannot be
	// overwritten by a callback starting concurrently with stop.
	_ = w.conn.SetReadDeadline(time.Now())
	<-w.done
	_ = w.conn.SetReadDeadline(time.Time{})
}

func writeSimple(conn net.Conn, status int, body string, closeAfter bool, timeout time.Duration) error {
	writer := newResponseWriter("")
	writer.status = status
	writer.wroteHeader = true
	writer.body.WriteString(body)
	writer.header.Set("Content-Type", "text/plain; charset=utf-8")
	return writer.writeTo(conn, closeAfter, timeout)
}

type responseWriter struct {
	header         http.Header
	status         int
	wroteHeader    bool
	body           bytes.Buffer
	requestMethod  string
	protoMajor     int
	protoMinor     int
	conn           net.Conn
	closeAfter     bool
	writeTimeout   time.Duration
	streaming      bool
	chunked        bool
	streamingClose bool
	writeErr       error
	cancel         context.CancelFunc
}

func newResponseWriter(method string) *responseWriter {
	return &responseWriter{header: make(http.Header), requestMethod: method}
}

func (w *responseWriter) reset(method string) {
	for key := range w.header {
		delete(w.header, key)
	}
	w.status = 0
	w.wroteHeader = false
	if w.body.Cap() > maxPooledResponseBytes {
		w.body = bytes.Buffer{}
	} else {
		w.body.Reset()
	}
	w.requestMethod = method
	w.protoMajor = 0
	w.protoMinor = 0
	w.conn = nil
	w.closeAfter = false
	w.writeTimeout = 0
	w.streaming = false
	w.chunked = false
	w.streamingClose = false
	w.writeErr = nil
	w.cancel = nil
}

func (w *responseWriter) configure(request *http.Request, conn net.Conn, closeAfter bool, timeout time.Duration, cancel context.CancelFunc) {
	w.cancel = cancel
	w.protoMajor = request.ProtoMajor
	w.protoMinor = request.ProtoMinor
	w.conn = conn
	w.closeAfter = closeAfter
	w.writeTimeout = timeout
}

func (w *responseWriter) Header() http.Header { return w.header }

func (w *responseWriter) WriteHeader(status int) {
	if w.wroteHeader {
		return
	}
	w.status = status
	w.wroteHeader = true
}

func (w *responseWriter) Write(data []byte) (int, error) {
	if w.writeErr != nil {
		return 0, w.writeErr
	}
	if !w.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	if w.streaming {
		written, err := w.writeStream(data)
		if err != nil {
			w.writeErr = err
		}
		return written, err
	}
	return w.body.Write(data)
}

// Flush implements http.Flusher. Once a response is flushed, HTTP/1.1 uses
// chunked framing unless the handler supplied Content-Length. HTTP/1.0 falls
// back to close-delimited output because it has no chunked framing.
func (w *responseWriter) Flush() {
	if w.writeErr != nil || w.streaming || w.requestMethod == http.MethodHead || w.status == http.StatusNoContent ||
		w.status == http.StatusResetContent || w.status == http.StatusNotModified || (w.status >= 100 && w.status < 200) {
		return
	}
	_ = w.recordWriteError(w.beginStream())
}

func (w *responseWriter) beginStream() error {
	if w.conn == nil {
		return errors.New("go-omni/native: response stream is not connected")
	}
	if !w.wroteHeader {
		w.status = http.StatusOK
		w.wroteHeader = true
	}
	if w.requestMethod == http.MethodHead || w.status == http.StatusNoContent || w.status == http.StatusResetContent ||
		w.status == http.StatusNotModified || (w.status >= 100 && w.status < 200) {
		return nil
	}
	contentLength := w.header.Get("Content-Length")
	w.chunked = contentLength == "" && w.protoMajor == 1 && w.protoMinor == 1
	if w.chunked {
		w.header.Set("Transfer-Encoding", "chunked")
	} else if contentLength == "" {
		w.streamingClose = true
	}
	closeAfter := w.closeAfter || w.streamingClose
	if closeAfter {
		w.header.Set("Connection", "close")
	}
	w.closeAfter = closeAfter
	if err := w.writeHeaders(w.conn); err != nil {
		return err
	}
	w.streaming = true
	buffered := w.body.Bytes()
	w.body.Reset()
	if len(buffered) == 0 {
		return nil
	}
	_, err := w.writeStream(buffered)
	return err
}

func (w *responseWriter) writeStream(data []byte) (int, error) {
	if len(data) == 0 {
		return 0, nil
	}
	if w.conn == nil {
		return 0, errors.New("go-omni/native: response stream is not connected")
	}
	if w.chunked {
		var prefix [32]byte
		framed := strconv.AppendInt(prefix[:0], int64(len(data)), 16)
		framed = append(framed, '\r', '\n')
		if err := w.writeBytes(framed); err != nil {
			return 0, err
		}
		if err := w.writeBytes(data); err != nil {
			return 0, err
		}
		if err := w.writeBytes([]byte("\r\n")); err != nil {
			return 0, err
		}
		return len(data), nil
	}
	if err := w.writeBytes(data); err != nil {
		return 0, err
	}
	return len(data), nil
}

func (w *responseWriter) writeBytes(data []byte) error {
	if w.writeTimeout > 0 {
		_ = w.conn.SetWriteDeadline(time.Now().Add(w.writeTimeout))
	}
	return w.recordWriteError(writeAll(w.conn, data))
}

func (w *responseWriter) recordWriteError(err error) error {
	if err != nil {
		w.writeErr = err
		if w.cancel != nil {
			w.cancel()
		}
	}
	return err
}

func (w *responseWriter) writeTo(conn net.Conn, closeAfter bool, timeout time.Duration) error {
	if w.writeErr != nil {
		_ = conn.SetWriteDeadline(time.Time{})
		return w.writeErr
	}
	if w.streaming {
		if w.chunked {
			if err := w.writeBytes([]byte("0\r\n\r\n")); err != nil {
				w.writeErr = err
			}
		}
		_ = conn.SetWriteDeadline(time.Time{})
		return w.writeErr
	}
	if timeout > 0 {
		_ = conn.SetWriteDeadline(time.Now().Add(timeout))
	}
	if !w.wroteHeader {
		w.status = http.StatusOK
		w.wroteHeader = true
	}
	body := w.body.Bytes()
	noBody := w.requestMethod == http.MethodHead || w.status == http.StatusNoContent || w.status == http.StatusResetContent ||
		w.status == http.StatusNotModified || (w.status >= 100 && w.status < 200)
	if noBody {
		body = nil
	}
	if w.header.Get("Content-Length") == "" && w.header.Get("Transfer-Encoding") == "" && w.status >= 200 && w.status != http.StatusNoContent && w.status != http.StatusNotModified && w.status != http.StatusResetContent {
		w.header.Set("Content-Length", strconv.Itoa(len(w.body.Bytes())))
	}
	if closeAfter {
		w.header.Set("Connection", "close")
	}
	keys := make([]string, 0, len(w.header))
	for key := range w.header {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	var output bytes.Buffer
	output.WriteString("HTTP/1.1 ")
	output.WriteString(strconv.Itoa(w.status))
	output.WriteByte(' ')
	output.WriteString(http.StatusText(w.status))
	output.WriteString("\r\n")
	for _, key := range keys {
		if !validToken(key) {
			return errors.New("invalid response header name")
		}
		for _, value := range w.header.Values(key) {
			if !validHeaderValue(value) {
				return errors.New("invalid response header")
			}
			output.WriteString(key)
			output.WriteString(": ")
			output.WriteString(value)
			output.WriteString("\r\n")
		}
	}
	output.WriteString("\r\n")
	output.Write(body)
	err := writeAll(conn, output.Bytes())
	_ = conn.SetWriteDeadline(time.Time{})
	return err
}

func (w *responseWriter) writeHeaders(conn net.Conn) error {
	if w.writeTimeout > 0 {
		_ = conn.SetWriteDeadline(time.Now().Add(w.writeTimeout))
	}
	keys := make([]string, 0, len(w.header))
	for key := range w.header {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	var output bytes.Buffer
	output.WriteString("HTTP/1.1 ")
	output.WriteString(strconv.Itoa(w.status))
	output.WriteByte(' ')
	output.WriteString(http.StatusText(w.status))
	output.WriteString("\r\n")
	for _, key := range keys {
		if !validToken(key) {
			return errors.New("invalid response header name")
		}
		for _, value := range w.header.Values(key) {
			if !validHeaderValue(value) {
				return errors.New("invalid response header")
			}
			output.WriteString(key)
			output.WriteString(": ")
			output.WriteString(value)
			output.WriteString("\r\n")
		}
	}
	output.WriteString("\r\n")
	return w.recordWriteError(writeAll(conn, output.Bytes()))
}

func writeAll(conn net.Conn, data []byte) error {
	for len(data) > 0 {
		written, err := conn.Write(data)
		if written > 0 {
			data = data[written:]
		}
		if err != nil {
			return err
		}
		if written == 0 {
			return io.ErrShortWrite
		}
	}
	return nil
}