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 currently supports
// Content-Length request bodies and rejects chunked requests.
package native

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"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
)

// 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

	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{}
	readerPool     sync.Pool
	writerPool     sync.Pool
	connectionPool sync.Pool
}

type connection struct {
	server        *Server
	conn          net.Conn
	active        atomic.Bool
	closeWhenIdle atomic.Bool
	closed        atomic.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
	}
	return s.Serve(ctx, listener)
}

// 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
		}
		if slots := s.connSlots; slots != nil {
			select {
			case slots <- struct{}{}:
			default:
				_ = 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.slotHeld = s.connSlots != nil
		state.closeDone = make(chan struct{})
		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)
		writer := s.writerPool.Get().(*responseWriter)
		writer.reset(request.Method)
		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)
				return
			default:
				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)
				if closeAfter || closing {
					return
				}
				continue
			}
		}
		handler := s.handler
		handler.ServeHTTP(writer, request)
		cancel()

		closing := s.isClosing() || state.closeWhenIdle.Load()
		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)
			return
		}
		if requestSlot {
			<-s.requestSlots
		}
		writer.reset("")
		s.writerPool.Put(writer)
		state.active.Store(false)
		if closeAfter || closing {
			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()
	}
}

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) || strings.ContainsAny(value, "\r\n") {
			return nil, true, &protocolError{status: http.StatusBadRequest, message: "malformed header"}
		}
		header.Add(name, value)
	}

	if value := header.Get("Expect"); value != "" {
		return nil, true, &protocolError{status: http.StatusExpectationFailed, message: "unsupported expectation"}
	}
	if value := header.Get("Transfer-Encoding"); value != "" && !strings.EqualFold(value, "identity") {
		return nil, true, &protocolError{status: http.StatusNotImplemented, message: "chunked requests are not supported"}
	}
	contentLength := int64(0)
	if value := header.Get("Content-Length"); value != "" {
		contentLength, err = strconv.ParseInt(strings.TrimSpace(value), 10, 64)
		if err != nil || contentLength < 0 {
			return nil, true, &protocolError{status: http.StatusBadRequest, message: "invalid content length"}
		}
	}
	if contentLength > maxBodyBytes {
		return nil, true, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
	}
	body := make([]byte, contentLength)
	if contentLength > 0 {
		if setBodyDeadline != nil {
			setBodyDeadline()
		}
		if _, err := io.ReadFull(reader, body); err != nil {
			return nil, true, err
		}
	}
	request := &http.Request{
		Method:        parts[0],
		URL:           requestURL,
		Proto:         parts[2],
		ProtoMajor:    protoMajor,
		ProtoMinor:    protoMinor,
		Header:        header,
		Body:          io.NopCloser(bytes.NewReader(body)),
		ContentLength: contentLength,
		Host:          header.Get("Host"),
		RequestURI:    target,
		RemoteAddr:    conn.RemoteAddr().String(),
	}
	closeAfter := strings.EqualFold(header.Get("Connection"), "close") ||
		(protoMajor == 1 && protoMinor == 0 && !strings.Contains(strings.ToLower(header.Get("Connection")), "keep-alive"))
	request.Close = closeAfter
	return request, closeAfter, nil
}

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 _, char := range value {
		if char <= ' ' || strings.ContainsRune("()<>@,;:\\\"/[]?={} \t", char) {
			return false
		}
	}
	return true
}

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

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
}

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
}

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.wroteHeader {
		w.WriteHeader(http.StatusOK)
	}
	return w.body.Write(data)
}

func (w *responseWriter) writeTo(conn net.Conn, closeAfter bool, timeout time.Duration) error {
	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.StatusNotModified || (w.status >= 100 && w.status < 200)
	if noBody {
		body = nil
	}
	if w.header.Get("Content-Length") == "" && w.status != http.StatusNoContent && w.status != http.StatusNotModified {
		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 {
		for _, value := range w.header.Values(key) {
			if strings.ContainsAny(value, "\r\n") {
				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 := conn.Write(output.Bytes())
	_ = conn.SetWriteDeadline(time.Time{})
	return err
}