go-omni/transport/native/tls_test.go

package native_test

import (
	"bufio"
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync/atomic"
	"testing"
	"time"

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

func startNativeTLS(t *testing.T, server *native.Server) (string, *tls.Config) {
	t.Helper()
	certificate := testCertificate(t)
	leaf, err := x509.ParseCertificate(certificate.Certificate[0])
	if err != nil {
		t.Fatal(err)
	}
	roots := x509.NewCertPool()
	roots.AddCert(leaf)
	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.ServeTLS(context.Background(), listener, &tls.Config{
			Certificates: []tls.Certificate{certificate}, MinVersion: tls.VersionTLS12,
			NextProtos: []string{"http/1.1"},
		})
	}()
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := server.Shutdown(ctx); err != nil {
			t.Errorf("TLS shutdown: %v", err)
		}
		select {
		case err := <-done:
			if err != nil {
				t.Errorf("TLS serve: %v", err)
			}
		case <-ctx.Done():
			t.Error("TLS serve did not exit")
		}
	})
	return listener.Addr().String(), &tls.Config{
		RootCAs: roots, ServerName: "localhost", MinVersion: tls.VersionTLS12,
		NextProtos: []string{"http/1.1"},
	}
}

func requirePeerClose(t *testing.T, conn net.Conn) {
	t.Helper()
	conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	_, err := conn.Read(make([]byte, 1))
	if err == nil {
		t.Fatal("expected server to close the connection")
	}
	var netErr net.Error
	if errors.As(err, &netErr) && netErr.Timeout() {
		t.Fatalf("client timed out; server did not close the connection: %v", err)
	}
}

func TestNativeTLSSlowHandshakeTimesOut(t *testing.T) {
	for _, prefix := range []string{"", "\x16\x03\x01\x00\x64\x01"} {
		t.Run(fmt.Sprintf("partialBytes=%d", len(prefix)), func(t *testing.T) {
			var calls atomic.Int32
			server := native.New(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { calls.Add(1) }))
			server.ReadHeaderTimeout = 50 * time.Millisecond
			address, _ := startNativeTLS(t, server)
			conn, err := net.DialTimeout("tcp", address, time.Second)
			if err != nil {
				t.Fatal(err)
			}
			defer conn.Close()
			if _, err := io.WriteString(conn, prefix); err != nil {
				t.Fatal(err)
			}
			requirePeerClose(t, conn)
			if calls.Load() != 0 {
				t.Fatal("unfinished TLS handshake reached the handler")
			}
		})
	}
}

func TestNativeTLSConnectionLimitAndKeepAlive(t *testing.T) {
	var calls atomic.Int32
	server := native.New(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		calls.Add(1)
		_, _ = io.WriteString(w, "healthy")
	}))
	server.MaxConnections = 1
	address, config := startNativeTLS(t, server)
	dial := func() (*tls.Conn, error) {
		return tls.DialWithDialer(&net.Dialer{Timeout: time.Second}, "tcp", address, config)
	}
	first, err := dial()
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()
	first.SetDeadline(time.Now().Add(3 * time.Second))
	reader := bufio.NewReader(first)
	for index := 0; index < 3; index++ {
		if _, err := io.WriteString(first, "GET /health HTTP/1.1\r\nHost: localhost\r\n\r\n"); 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 || string(body) != "healthy" || response.Close {
			t.Fatalf("Keep-Alive response body=%q err=%v close=%v", body, err, response.Close)
		}
	}
	// A completed TLS handshake proves the first slot is held before testing
	// rejection. TCP accept can still succeed; the TLS handshake must fail.
	rejected, err := dial()
	if err == nil {
		rejected.Close()
		t.Fatal("TLS connection exceeded MaxConnections")
	}
	var netErr net.Error
	if errors.As(err, &netErr) && netErr.Timeout() {
		t.Fatalf("over-limit handshake hung: %v", err)
	}
	first.Close()
	deadline := time.Now().Add(2 * time.Second)
	var next *tls.Conn
	for time.Now().Before(deadline) {
		next, err = dial()
		if err == nil {
			break
		}
		time.Sleep(time.Millisecond)
	}
	if err != nil {
		t.Fatalf("TLS connection capacity not released: %v", err)
	}
	defer next.Close()
	next.SetDeadline(time.Now().Add(time.Second))
	if _, err := io.WriteString(next, "GET /health HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"); err != nil {
		t.Fatal(err)
	}
	response, err := http.ReadResponse(bufio.NewReader(next), &http.Request{Method: http.MethodGet})
	if err != nil {
		t.Fatal(err)
	}
	body, err := io.ReadAll(response.Body)
	response.Body.Close()
	if err != nil || string(body) != "healthy" || calls.Load() != 4 {
		t.Fatalf("recovered response=%q calls=%d err=%v", body, calls.Load(), err)
	}
}

func TestNativeTLSRequestMetadata(t *testing.T) {
	server := native.New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.TLS == nil || !r.TLS.HandshakeComplete || r.TLS.NegotiatedProtocol != "http/1.1" {
			http.Error(w, "missing TLS request metadata", http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusNoContent)
	}))
	address, config := startNativeTLS(t, server)
	client := &http.Client{Transport: &http.Transport{TLSClientConfig: config}, Timeout: 2 * time.Second}
	defer client.CloseIdleConnections()
	response, err := client.Get("https://" + address + "/")
	if err != nil {
		t.Fatal(err)
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusNoContent {
		t.Fatalf("TLS metadata unavailable to handler: status=%d", response.StatusCode)
	}
}