go-omni/transport/native/server_test.go

package native_test

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"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 TestServerRejectsUnsupportedBodyAndLimits(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) }()
	waitForServer(t, listener.Addr().String())
	defer stopServer(t, server, listener, done)

	for _, requestText := range []string{
		"POST /body HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\n4\r\ntest\r\n0\r\n\r\n",
		"POST /body HTTP/1.1\r\nHost: test\r\nContent-Length: 5\r\n\r\ntest!",
	} {
		conn, err := net.Dial("tcp", listener.Addr().String())
		if err != nil {
			t.Fatal(err)
		}
		_, _ = io.WriteString(conn, requestText)
		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 != http.StatusNotImplemented && response.StatusCode != http.StatusRequestEntityTooLarge {
			t.Fatalf("unexpected protocol status %d", response.StatusCode)
		}
		conn.Close()
	}
}

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:    "chunked",
			request: "POST /health HTTP/1.1\r\nHost: test\r\nTransfer-Encoding: chunked\r\n\r\n",
			status:  http.StatusNotImplemented,
		},
		{
			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 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)
	}
	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 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 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)
	}
}