package native_test

import (
	"bufio"
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	omni "git.yuntuo.hk/k2safe/go-omni"
	"git.yuntuo.hk/k2safe/go-omni/transport/native"
	netrouter "git.yuntuo.hk/k2safe/go-omni/transport/nethttp"
)

func TestServerKeepAliveAndApplicationBridge(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/hello", func(ctx *omni.Context) error {
		return ctx.OK(map[string]string{"message": "hello"})
	})
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	reader := bufio.NewReader(conn)

	for index, closeAfter := range []bool{false, true} {
		connectionHeader := "keep-alive"
		if closeAfter {
			connectionHeader = "close"
		}
		fmt.Fprintf(conn, "GET /hello HTTP/1.1\r\nHost: test\r\nConnection: %s\r\n\r\n", connectionHeader)
		request := &http.Request{Method: http.MethodGet}
		response, err := http.ReadResponse(reader, request)
		if err != nil {
			t.Fatalf("response %d: %v", index, err)
		}
		body, err := io.ReadAll(response.Body)
		response.Body.Close()
		if err != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), "\"message\":\"hello\"") {
			t.Fatalf("response %d: status=%d body=%s error=%v", index, response.StatusCode, body, err)
		}
	}
}

func TestServerRejectsMalformedBodyAndLimits(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.POST("/body", func(ctx *omni.Context) error { return ctx.OK("ok") })

	server := native.New(application)
	server.MaxBodyBytes = 4
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	defer stopServer(t, server, listener, done)

	tests := []struct {
		request string
		status  int
	}{
		{
			request: "POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\nZ\r\nbad\r\n0\r\n\r\n",
			status:  http.StatusBadRequest,
		},
		{
			request: "POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 5\r\n\r\ntest!",
			status:  http.StatusRequestEntityTooLarge,
		},
	}
	for _, test := range tests {
		conn, err := net.Dial("tcp", listener.Addr().String())
		if err != nil {
			t.Fatal(err)
		}
		_, _ = io.WriteString(conn, test.request)
		response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodPost})
		if err != nil {
			conn.Close()
			t.Fatal(err)
		}
		response.Body.Close()
		if response.StatusCode != test.status {
			t.Fatalf("expected status %d, got %d", test.status, response.StatusCode)
		}
		conn.Close()
	}
}

func TestServerStreamsContentLengthBodyToHandler(t *testing.T) {
	entered := make(chan struct{})
	application := omni.NewWithRouter(netrouter.New())
	application.POST("/body", func(ctx *omni.Context) error {
		close(entered)
		body, err := io.ReadAll(ctx.Request().Body)
		if err != nil {
			return err
		}
		return ctx.OK(map[string]string{"body": string(body)})
	})
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	if _, err := io.WriteString(conn, "POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 5\r\n\r\nhel"); err != nil {
		t.Fatal(err)
	}
	select {
	case <-entered:
	case <-time.After(time.Second):
		t.Fatal("handler did not start before the complete body arrived")
	}
	if _, err := io.WriteString(conn, "lo"); err != nil {
		t.Fatal(err)
	}
	response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodPost})
	if err != nil {
		t.Fatal(err)
	}
	body, err := io.ReadAll(response.Body)
	response.Body.Close()
	if err != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), `"body":"hello"`) {
		t.Fatalf("unexpected streamed body response: status=%d body=%s error=%v", response.StatusCode, body, err)
	}
}

func TestServerChunkedBodyAndTrailers(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.POST("/body", func(ctx *omni.Context) error {
		body, err := io.ReadAll(ctx.Request().Body)
		if err != nil {
			return err
		}
		return ctx.OK(map[string]string{
			"body":     string(body),
			"trailer":  ctx.Request().Trailer.Get("X-Trace"),
			"transfer": strings.Join(ctx.Request().TransferEncoding, ","),
		})
	})
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	request := "POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n" +
		"4;source=wire\r\nWiki\r\n5\r\npedia\r\n0\r\nX-Trace: yes\r\n\r\n"
	response := rawResponseMethod(t, listener.Addr().String(), http.MethodPost, request)
	defer response.Body.Close()
	body, err := io.ReadAll(response.Body)
	if err != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), `"body":"Wikipedia"`) || !strings.Contains(string(body), `"trailer":"yes"`) || !strings.Contains(string(body), `"transfer":"chunked"`) {
		t.Fatalf("unexpected chunked response: status=%d body=%s error=%v", response.StatusCode, body, err)
	}
}

func TestServerChunkedBodyLimit(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.POST("/body", func(ctx *omni.Context) error { return ctx.OK("unexpected") })
	server := native.New(application)
	server.MaxBodyBytes = 4
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	defer stopServer(t, server, listener, done)

	request := "POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\n5\r\nhello\r\n0\r\n\r\n"
	response := rawResponseMethod(t, listener.Addr().String(), http.MethodPost, request)
	defer response.Body.Close()
	if response.StatusCode != http.StatusRequestEntityTooLarge {
		t.Fatalf("expected chunked body limit status 413, got %d", response.StatusCode)
	}
}

func TestServerSlowBodyTimesOutWithoutSuccessfulHandler(t *testing.T) {
	var handlerCalls atomic.Int32
	application := omni.NewWithRouter(netrouter.New())
	application.POST("/body", func(ctx *omni.Context) error {
		body, err := io.ReadAll(ctx.Request().Body)
		if err != nil {
			return err
		}
		handlerCalls.Add(1)
		if len(body) != 4 {
			return fmt.Errorf("unexpected body length %d", len(body))
		}
		return ctx.OK("ok")
	})
	server := native.New(application)
	server.ReadTimeout = 30 * time.Millisecond
	server.IdleTimeout = time.Second
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	_, _ = io.WriteString(conn, "POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 4\r\n\r\na")
	requirePeerClose(t, conn)
	if handlerCalls.Load() != 0 {
		t.Fatal("handler reported success before the request body was complete")
	}
}

func TestServerPipelinedRequestsAndLargeBody(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.POST("/body", func(ctx *omni.Context) error {
		body, err := io.ReadAll(ctx.Request().Body)
		if err != nil {
			return err
		}
		return ctx.OK(map[string]int{"length": len(body)})
	})
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	body := strings.Repeat("x", 64<<10)
	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	reader := bufio.NewReader(conn)
	request := fmt.Sprintf("POST /body HTTP/1.1\r\nHost: test\r\nConnection: keep-alive\r\nContent-Length: %d\r\n\r\n%sGET /health HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n", len(body), body)
	if _, err := io.WriteString(conn, request); err != nil {
		t.Fatal(err)
	}
	first, err := http.ReadResponse(reader, &http.Request{Method: http.MethodPost})
	if err != nil {
		t.Fatal(err)
	}
	firstBody, err := io.ReadAll(first.Body)
	first.Body.Close()
	if err != nil || first.StatusCode != http.StatusOK || !strings.Contains(string(firstBody), `"length":65536`) {
		t.Fatalf("unexpected large body response: status=%d body=%s error=%v", first.StatusCode, firstBody, err)
	}
	second, err := http.ReadResponse(reader, &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	secondBody, err := io.ReadAll(second.Body)
	second.Body.Close()
	if err != nil || second.StatusCode != http.StatusOK || !strings.Contains(string(secondBody), `"ok"`) {
		t.Fatalf("unexpected pipelined response: status=%d body=%s error=%v", second.StatusCode, secondBody, err)
	}
}

func TestServerFlushesChunkedResponses(t *testing.T) {
	flushed := make(chan struct{})
	release := make(chan struct{})
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/stream", func(ctx *omni.Context) error {
		flusher, ok := ctx.Raw().(http.Flusher)
		if !ok {
			return fmt.Errorf("native response writer does not implement http.Flusher")
		}
		ctx.Raw().Header().Set("Content-Type", "text/plain")
		_, _ = io.WriteString(ctx.Raw(), "one\n")
		flusher.Flush()
		close(flushed)
		<-release
		_, _ = io.WriteString(ctx.Raw(), "two\n")
		return nil
	})
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	reader := bufio.NewReader(conn)
	_, _ = io.WriteString(conn, "GET /stream HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	response, err := http.ReadResponse(reader, &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	if response.StatusCode != http.StatusOK || response.TransferEncoding == nil || response.TransferEncoding[0] != "chunked" {
		response.Body.Close()
		t.Fatalf("expected chunked response, status=%d transfer=%v", response.StatusCode, response.TransferEncoding)
	}
	select {
	case <-flushed:
	case <-time.After(time.Second):
		response.Body.Close()
		t.Fatal("handler did not flush response")
	}
	first := make([]byte, len("one\n"))
	if _, err := io.ReadFull(response.Body, first); err != nil || string(first) != "one\n" {
		response.Body.Close()
		t.Fatalf("unexpected first stream chunk: %q error=%v", first, err)
	}
	close(release)
	rest, err := io.ReadAll(response.Body)
	response.Body.Close()
	if err != nil || string(rest) != "two\n" {
		t.Fatalf("unexpected remaining stream body: %q error=%v", rest, err)
	}
}

func TestServerProtocolBoundaries(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	tests := []struct {
		name    string
		request string
		status  int
	}{
		{
			name:    "missing host",
			request: "GET /health HTTP/1.1\r\n\r\n",
			status:  http.StatusBadRequest,
		},
		{
			name:    "duplicate host",
			request: "GET /health HTTP/1.1\r\nHost: one\r\nHost: two\r\n\r\n",
			status:  http.StatusBadRequest,
		},
		{
			name:    "unsupported version",
			request: "GET /health HTTP/2.0\r\nHost: test\r\n\r\n",
			status:  http.StatusHTTPVersionNotSupported,
		},
		{
			name:    "malformed header",
			request: "GET /health HTTP/1.1\r\nHost: test\r\nBroken\r\n\r\n",
			status:  http.StatusBadRequest,
		},
		{
			name:    "folded header",
			request: "GET /health HTTP/1.1\r\nHost: test\r\n X-Bad: value\r\n\r\n",
			status:  http.StatusBadRequest,
		},
		{
			name:    "control header value",
			request: "GET /health HTTP/1.1\r\nHost: test\r\nX-Bad: value\x01\r\n\r\n",
			status:  http.StatusBadRequest,
		},
		{
			name:    "malformed chunked",
			request: "POST /health HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\nZ\r\n\r\n",
			status:  http.StatusBadRequest,
		},
		{
			name:    "expectation",
			request: "GET /health HTTP/1.1\r\nHost: test\r\nExpect: 100-continue\r\n\r\n",
			status:  http.StatusExpectationFailed,
		},
		{
			name:    "conflicting content length",
			request: "POST /health HTTP/1.1\r\nHost: test\r\nContent-Length: 1\r\nContent-Length: 2\r\n\r\na",
			status:  http.StatusBadRequest,
		},
		{
			name:    "invalid content length",
			request: "POST /health HTTP/1.1\r\nHost: test\r\nContent-Length: nope\r\n\r\n",
			status:  http.StatusBadRequest,
		},
		{
			name:    "malformed target",
			request: "GET health HTTP/1.1\r\nHost: test\r\n\r\n",
			status:  http.StatusBadRequest,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			response := rawResponse(t, listener.Addr().String(), test.request)
			defer response.Body.Close()
			if response.StatusCode != test.status {
				t.Fatalf("expected status %d, got %d", test.status, response.StatusCode)
			}
			if !response.Close {
				t.Fatal("protocol error should close the connection")
			}
		})
	}
}

func TestServerHTTP10KeepAliveAndRequestMetadata(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/metadata", func(ctx *omni.Context) error {
		return ctx.OK(map[string]string{
			"method": ctx.Request().Method,
			"host":   ctx.Request().Host,
		})
	})
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	reader := bufio.NewReader(conn)
	_, _ = io.WriteString(conn, "GET /metadata HTTP/1.0\r\nHost: legacy\r\nConnection: keep-alive\r\n\r\n")
	response, err := http.ReadResponse(reader, &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	body, err := io.ReadAll(response.Body)
	response.Body.Close()
	if err != nil || response.StatusCode != http.StatusOK || response.Close || !strings.Contains(string(body), `"host":"legacy"`) {
		t.Fatalf("unexpected HTTP/1.0 response: close=%v status=%d body=%s error=%v", response.Close, response.StatusCode, body, err)
	}
	_, _ = io.WriteString(conn, "GET /metadata HTTP/1.0\r\nHost: legacy\r\n\r\n")
	response, err = http.ReadResponse(reader, &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if !response.Close {
		t.Fatal("HTTP/1.0 response without keep-alive should close")
	}
}

func TestServerSuppressesBodiesForHeadAndNoContent(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.Handle(http.MethodHead, "/head", func(ctx *omni.Context) error { return ctx.OK("hidden") })
	application.GET("/empty", func(ctx *omni.Context) error {
		ctx.Raw().WriteHeader(http.StatusNoContent)
		return nil
	})
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	headResponse := rawResponseMethod(t, listener.Addr().String(), http.MethodHead, "HEAD /head HTTP/1.1\r\nHost: test\r\n\r\n")
	headBody, err := io.ReadAll(headResponse.Body)
	headResponse.Body.Close()
	if err != nil || len(headBody) != 0 || headResponse.StatusCode != http.StatusOK || headResponse.ContentLength <= 0 {
		t.Fatalf("unexpected HEAD response: status=%d contentLength=%d body=%q error=%v", headResponse.StatusCode, headResponse.ContentLength, headBody, err)
	}
	emptyResponse := rawResponse(t, listener.Addr().String(), "GET /empty HTTP/1.1\r\nHost: test\r\n\r\n")
	emptyBody, err := io.ReadAll(emptyResponse.Body)
	emptyResponse.Body.Close()
	if err != nil || len(emptyBody) != 0 || emptyResponse.StatusCode != http.StatusNoContent || emptyResponse.ContentLength != 0 {
		t.Fatalf("unexpected 204 response: status=%d contentLength=%d body=%q error=%v", emptyResponse.StatusCode, emptyResponse.ContentLength, emptyBody, err)
	}
}

func rawResponse(t *testing.T, address, request string) *http.Response {
	return rawResponseMethod(t, address, http.MethodGet, request)
}

func rawResponseMethod(t *testing.T, address, method, request string) *http.Response {
	t.Helper()
	conn, err := net.Dial("tcp", address)
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.WriteString(conn, request)
	response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: method})
	if err != nil {
		conn.Close()
		t.Fatal(err)
	}
	response.Body = &closeBody{ReadCloser: response.Body, conn: conn}
	return response
}

type closeBody struct {
	io.ReadCloser
	conn net.Conn
}

func (b *closeBody) Close() error {
	err := b.ReadCloser.Close()
	return errors.Join(err, b.conn.Close())
}

func TestServerShutdownCancelsActiveHandler(t *testing.T) {
	started := make(chan struct{})
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/wait", func(ctx *omni.Context) error {
		close(started)
		<-ctx.Context().Done()
		return ctx.Context().Err()
	})
	server, listener, done := startServer(t, application)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.WriteString(conn, "GET /wait HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("handler did not start")
	}
	shutdownContext, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := server.Shutdown(shutdownContext); err != nil {
		t.Fatal(err)
	}
	response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusInternalServerError {
		t.Fatalf("expected canceled handler to become 500, got %d", response.StatusCode)
	}
	conn.Close()
	select {
	case err := <-done:
		if err != nil && !errors.Is(err, net.ErrClosed) {
			t.Fatal(err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("server did not stop")
	}
}

func TestServerCancelsHandlerWhenClientDisconnects(t *testing.T) {
	for _, stream := range []bool{false, true} {
		t.Run(fmt.Sprintf("stream=%t", stream), func(t *testing.T) {
			started := make(chan struct{})
			canceled := make(chan struct{})
			application := omni.NewWithRouter(netrouter.New())
			application.GET("/wait", func(ctx *omni.Context) error {
				if stream {
					_, _ = io.WriteString(ctx.Raw(), "ready\n")
					ctx.Raw().(http.Flusher).Flush()
				}
				close(started)
				<-ctx.Context().Done()
				close(canceled)
				return nil
			})
			server, listener, done := startServer(t, application)
			defer stopServer(t, server, listener, done)

			conn, err := net.Dial("tcp", listener.Addr().String())
			if err != nil {
				t.Fatal(err)
			}
			defer conn.Close()
			conn.SetDeadline(time.Now().Add(2 * time.Second))
			_, _ = io.WriteString(conn, "GET /wait HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
			select {
			case <-started:
			case <-time.After(time.Second):
				t.Fatal("handler did not start")
			}
			if stream {
				response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodGet})
				if err != nil {
					t.Fatal(err)
				}
				first := make([]byte, len("ready\n"))
				if _, err := io.ReadFull(response.Body, first); err != nil || string(first) != "ready\n" {
					t.Fatalf("failed reading stream: %q err=%v", first, err)
				}
				// Close the socket before closing the unfinished response body.
				conn.Close()
				response.Body.Close()
			} else {
				conn.Close()
			}
			select {
			case <-canceled:
			case <-time.After(time.Second):
				t.Fatal("client disconnect did not cancel request context")
			}
		})
	}
}

func TestServerLongHandlerPreservesKeepAliveAndPipeline(t *testing.T) {
	for _, pipeline := range []bool{false, true} {
		t.Run(fmt.Sprintf("pipeline=%t", pipeline), func(t *testing.T) {
			application := omni.New()
			application.GET("/slow", func(ctx *omni.Context) error {
				timer := time.NewTimer(25 * time.Millisecond)
				defer timer.Stop()
				select {
				case <-timer.C:
					return ctx.OK("slow done")
				case <-ctx.Context().Done():
					return ctx.Context().Err()
				}
			})
			application.POST("/next", func(ctx *omni.Context) error {
				body, err := io.ReadAll(ctx.Request().Body)
				if err != nil {
					return err
				}
				return ctx.OK(string(body))
			})
			server, listener, done := startServer(t, application)
			defer stopServer(t, server, listener, done)
			conn, err := net.Dial("tcp", listener.Addr().String())
			if err != nil {
				t.Fatal(err)
			}
			defer conn.Close()
			conn.SetDeadline(time.Now().Add(3 * time.Second))
			reader := bufio.NewReader(conn)
			for attempt := 0; attempt < 3; attempt++ {
				slow := "GET /slow HTTP/1.1\r\nHost: test\r\n\r\n"
				next := "POST /next HTTP/1.1\r\nHost: test\r\nContent-Length: 7\r\n\r\npayload"
				if pipeline {
					slow += next
				}
				if _, err := io.WriteString(conn, slow); err != nil {
					t.Fatal(err)
				}
				for index, want := range []string{"slow done", "payload"} {
					if index == 1 && !pipeline {
						if _, err := io.WriteString(conn, next); err != nil {
							t.Fatal(err)
						}
					}
					response, err := http.ReadResponse(reader, &http.Request{Method: http.MethodGet})
					if err != nil {
						t.Fatal(err)
					}
					body, err := io.ReadAll(response.Body)
					response.Body.Close()
					if err != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), want) {
						t.Fatalf("attempt=%d response=%d status=%d body=%s err=%v", attempt, index, response.StatusCode, body, err)
					}
				}
			}
		})
	}
}

func TestServerConcurrentRequestLimit(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/work", func(ctx *omni.Context) error {
		close(started)
		<-release
		return ctx.OK("done")
	})
	server := native.New(application)
	server.MaxConcurrentRequests = 1
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	defer func() {
		stopServer(t, server, listener, done)
	}()

	first, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()
	_, _ = io.WriteString(first, "GET /work HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("first handler did not start")
	}

	second, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer second.Close()
	_, _ = io.WriteString(second, "GET /work HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	response, err := http.ReadResponse(bufio.NewReader(second), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("expected 503 at concurrency limit, got %d", response.StatusCode)
	}
	if stats := server.Stats(); stats.RejectedRequests != 1 {
		t.Fatalf("expected one rejected request, got %+v", stats)
	}
	close(release)
	response, err = http.ReadResponse(bufio.NewReader(first), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("expected first request to complete, got %d", response.StatusCode)
	}
}

func TestServerConcurrentRequestLimitRecovers(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	var first sync.Once
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/work", func(ctx *omni.Context) error {
		first.Do(func() {
			close(started)
			<-release
		})
		return ctx.OK("done")
	})
	server := native.New(application)
	server.MaxConcurrentRequests = 1
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	defer stopServer(t, server, listener, done)

	firstConn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer firstConn.Close()
	_, _ = io.WriteString(firstConn, "GET /work HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("first handler did not start")
	}

	secondConn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	_, _ = io.WriteString(secondConn, "GET /work HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	response, err := http.ReadResponse(bufio.NewReader(secondConn), &http.Request{Method: http.MethodGet})
	if err != nil {
		secondConn.Close()
		t.Fatal(err)
	}
	response.Body.Close()
	secondConn.Close()
	if response.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("expected 503 at concurrency limit, got %d", response.StatusCode)
	}

	close(release)
	response, err = http.ReadResponse(bufio.NewReader(firstConn), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("expected held request to complete, got %d", response.StatusCode)
	}
	deadline := time.Now().Add(time.Second)
	for server.Stats().ActiveRequests != 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if stats := server.Stats(); stats.ActiveRequests != 0 {
		t.Fatalf("request slot was not released after response: %+v", stats)
	}

	thirdConn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer thirdConn.Close()
	_, _ = io.WriteString(thirdConn, "GET /work HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	response, err = http.ReadResponse(bufio.NewReader(thirdConn), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("request slot was not reusable after overload, got %d", response.StatusCode)
	}
}

func TestServerConnectionLimitReleasesAfterRequest(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/hold", func(ctx *omni.Context) error {
		close(started)
		<-release
		return ctx.OK("done")
	})
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	server := native.New(application)
	server.MaxConnections = 1
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	defer stopServer(t, server, listener, done)

	first, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()
	_, _ = io.WriteString(first, "GET /hold HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("first request did not start")
	}

	second, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	second.SetReadDeadline(time.Now().Add(time.Second))
	_, _ = io.WriteString(second, "GET /health HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	if _, err := second.Read(make([]byte, 1)); err == nil {
		second.Close()
		t.Fatal("connection over MaxConnections was not rejected")
	}
	second.Close()
	if stats := server.Stats(); stats.RejectedConnections < 1 {
		t.Fatalf("expected rejected connection counter, got %+v", stats)
	}

	close(release)
	firstResponse, err := http.ReadResponse(bufio.NewReader(first), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	firstResponse.Body.Close()
	deadline := time.Now().Add(time.Second)
	for server.Stats().ActiveConnections != 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if stats := server.Stats(); stats.ActiveConnections != 0 {
		t.Fatalf("connection slot was not released after response: %+v", stats)
	}

	third, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer third.Close()
	_, _ = io.WriteString(third, "GET /health HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
	thirdResponse, err := http.ReadResponse(bufio.NewReader(third), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	thirdResponse.Body.Close()
	if thirdResponse.StatusCode != http.StatusOK {
		t.Fatalf("connection slot was not released, got status %d", thirdResponse.StatusCode)
	}
}

func TestServerStatsTrackLifecycle(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/hold", func(ctx *omni.Context) error {
		close(started)
		<-release
		return ctx.OK("done")
	})
	server, listener, done := startServer(t, application)
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	waitForStats := func(check func(native.Stats) bool) {
		t.Helper()
		deadline := time.Now().Add(time.Second)
		for time.Now().Before(deadline) {
			if check(server.Stats()) {
				return
			}
			time.Sleep(time.Millisecond)
		}
		t.Fatalf("stats condition not reached: %+v", server.Stats())
	}
	waitForStats(func(stats native.Stats) bool { return stats.ActiveConnections == 1 && stats.AcceptedConnections >= 1 })
	if _, err := io.WriteString(conn, "GET /hold HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n"); err != nil {
		t.Fatal(err)
	}
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("hold handler did not start")
	}
	waitForStats(func(stats native.Stats) bool { return stats.ActiveRequests == 1 })
	close(release)
	response, err := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("unexpected hold response: %d", response.StatusCode)
	}
	waitForStats(func(stats native.Stats) bool { return stats.ActiveRequests == 0 })
	waitForStats(func(stats native.Stats) bool { return stats.ActiveConnections == 0 })
}

func TestServerConnectionChurnReusesStateSafely(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	server := native.New(application)
	server.MaxConnections = 32
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	defer stopServer(t, server, listener, done)

	// Repeated short-lived connections exercise accept/close, connection-pool
	// reuse and the MaxConnections gate under a bursty workload.
	const workers = 16
	const requestsPerWorker = 200
	errors := make(chan error, workers)
	var group sync.WaitGroup
	for worker := 0; worker < workers; worker++ {
		group.Add(1)
		go func() {
			defer group.Done()
			for index := 0; index < requestsPerWorker; index++ {
				conn, err := net.Dial("tcp", listener.Addr().String())
				if err != nil {
					errors <- err
					return
				}
				_, writeErr := io.WriteString(conn, "GET /health HTTP/1.1\r\nHost: test\r\nConnection: close\r\n\r\n")
				if writeErr == nil {
					response, readErr := http.ReadResponse(bufio.NewReader(conn), &http.Request{Method: http.MethodGet})
					if readErr != nil {
						writeErr = readErr
					} else {
						response.Body.Close()
						if response.StatusCode != http.StatusOK {
							writeErr = fmt.Errorf("unexpected status %d", response.StatusCode)
						}
					}
				}
				conn.Close()
				if writeErr != nil {
					errors <- writeErr
					return
				}
			}
		}()
	}
	group.Wait()
	close(errors)
	for err := range errors {
		t.Fatal(err)
	}
}

func TestServerReadHeaderTimeout(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	server := native.New(application)
	server.ReadHeaderTimeout = 25 * time.Millisecond
	server.IdleTimeout = 0
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	defer stopServer(t, server, listener, done)

	conn, err := net.Dial("tcp", listener.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	conn.SetReadDeadline(time.Now().Add(time.Second))
	_, _ = io.WriteString(conn, "GET /health HTTP/1.1\r\nHost: test")
	buffer := make([]byte, 1)
	if _, err := conn.Read(buffer); err == nil {
		t.Fatal("expected incomplete headers to time out and close")
	}
}

func TestAppRunWithNativeRunner(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	address := freeAddress(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- application.RunWith(ctx, native.New(nil), address) }()
	waitForServer(t, address)
	response, err := http.Get("http://" + address + "/health")
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("unexpected status %d", response.StatusCode)
	}
	cancel()
	select {
	case err := <-done:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("RunWith did not stop")
	}
}

func TestAppRunWithNativeTLS(t *testing.T) {
	application := omni.NewWithRouter(netrouter.New())
	application.GET("/health", func(ctx *omni.Context) error { return ctx.OK("ok") })
	certificate := testCertificate(t)
	server := native.New(nil)
	server.TLSConfig = &tls.Config{
		Certificates: []tls.Certificate{certificate},
		MinVersion:   tls.VersionTLS12,
	}
	address := freeAddress(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- application.RunWith(ctx, server, address) }()

	client := &http.Client{Transport: &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, // test certificate is self-signed
	}}
	defer client.CloseIdleConnections()
	var response *http.Response
	var err error
	deadline := time.Now().Add(2 * time.Second)
	for response == nil && time.Now().Before(deadline) {
		response, err = client.Get("https://" + address + "/health")
		if err != nil {
			time.Sleep(10 * time.Millisecond)
		}
	}
	if err != nil || response == nil {
		t.Fatalf("TLS request failed: response=%v error=%v", response, err)
	}
	body, readErr := io.ReadAll(response.Body)
	response.Body.Close()
	if readErr != nil || response.StatusCode != http.StatusOK || !strings.Contains(string(body), `"data":"ok"`) {
		t.Fatalf("unexpected TLS response: status=%d body=%q error=%v", response.StatusCode, body, readErr)
	}
	cancel()
	select {
	case err := <-done:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("TLS RunWith did not stop")
	}
}

func testCertificate(t *testing.T) tls.Certificate {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 120))
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now()
	template := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: "localhost"},
		NotBefore:    now.Add(-time.Minute),
		NotAfter:     now.Add(time.Hour),
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSNames:     []string{"localhost"},
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	certificate, err := tls.X509KeyPair(
		pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}),
		pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: mustMarshalECKey(t, key)}),
	)
	if err != nil {
		t.Fatal(err)
	}
	return certificate
}

func mustMarshalECKey(t *testing.T, key *ecdsa.PrivateKey) []byte {
	t.Helper()
	der, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		t.Fatal(err)
	}
	return der
}

func startServer(t *testing.T, handler http.Handler) (*native.Server, net.Listener, <-chan error) {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	server := native.New(handler)
	done := make(chan error, 1)
	go func() { done <- server.Serve(context.Background(), listener) }()
	waitForServer(t, listener.Addr().String())
	return server, listener, done
}

func stopServer(t *testing.T, server *native.Server, listener net.Listener, done <-chan error) {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	_ = server.Shutdown(ctx)
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("server did not stop")
	}
	_ = listener.Close()
}

func freeAddress(t *testing.T) string {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	address := listener.Addr().String()
	listener.Close()
	return address
}

func waitForServer(t *testing.T, address string) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for {
		conn, err := net.DialTimeout("tcp", address, 50*time.Millisecond)
		if err == nil {
			conn.Close()
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("server did not start: %v", err)
		}
		time.Sleep(10 * time.Millisecond)
	}
}