go-omni/transport/http2/server_test.go

package http2

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	xhttp2 "golang.org/x/net/http2"
)

func TestServerHandlesH2CRequest(t *testing.T) {
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.ProtoMajor != 2 {
			http.Error(writer, "not HTTP/2", http.StatusInternalServerError)
			return
		}
		_, _ = writer.Write([]byte("ok"))
	}))
	address, stop := startServer(t, server)
	defer stop()

	transport := &xhttp2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, address string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, address)
		},
	}
	client := &http.Client{Transport: transport, Timeout: 2 * time.Second}
	defer transport.CloseIdleConnections()
	response, err := client.Get("http://" + address + "/health")
	if err != nil {
		t.Fatal(err)
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("unexpected status: %d", response.StatusCode)
	}
	if stats := server.Stats(); stats.AcceptedConnections != 1 {
		t.Fatalf("unexpected HTTP/2 stats: %+v", stats)
	}
}

func TestServerEnforcesGlobalRequestLimitAndRecovers(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	var once sync.Once
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.URL.Path == "/hold" {
			once.Do(func() { close(started) })
			<-release
		}
		_, _ = writer.Write([]byte("ok"))
	}))
	server.MaxConcurrentRequests = 1
	address, stop := startServer(t, server)
	defer stop()
	transport := &xhttp2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, address string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, address)
		},
	}
	client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
	defer transport.CloseIdleConnections()

	firstDone := make(chan *http.Response, 1)
	firstErr := make(chan error, 1)
	go func() {
		response, err := client.Get("http://" + address + "/hold")
		if err != nil {
			firstErr <- err
			return
		}
		firstDone <- response
	}()
	select {
	case <-started:
	case err := <-firstErr:
		t.Fatal(err)
	case <-time.After(time.Second):
		t.Fatal("first HTTP/2 stream did not start")
	}

	response, err := client.Get("http://" + address + "/busy")
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusServiceUnavailable {
		t.Fatalf("expected 503 at stream limit, got %d", response.StatusCode)
	}
	if stats := server.Stats(); stats.RejectedRequests != 1 {
		t.Fatalf("expected one rejected stream, got %+v", stats)
	}

	close(release)
	select {
	case response = <-firstDone:
		response.Body.Close()
	case err := <-firstErr:
		t.Fatal(err)
	case <-time.After(time.Second):
		t.Fatal("held HTTP/2 stream did not finish")
	}
	response, err = client.Get("http://" + address + "/recovered")
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("request limit did not recover, got %d", response.StatusCode)
	}
}

func TestServerRejectsBodyAboveLimitAndRecovers(t *testing.T) {
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		body, err := io.ReadAll(request.Body)
		if err != nil {
			http.Error(writer, err.Error(), http.StatusRequestEntityTooLarge)
			return
		}
		_, _ = writer.Write(body)
	}))
	server.MaxBodyBytes = 4
	address, stop := startServer(t, server)
	defer stop()
	transport := &xhttp2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, address string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, address)
		},
	}
	client := &http.Client{Transport: transport, Timeout: 3 * time.Second}
	defer transport.CloseIdleConnections()

	request, err := http.NewRequest(http.MethodPost, "http://"+address+"/too-large", bytes.NewReader([]byte("12345")))
	if err != nil {
		t.Fatal(err)
	}
	response, err := client.Do(request)
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusRequestEntityTooLarge {
		t.Fatalf("expected 413 for oversized body, got %d", response.StatusCode)
	}

	request, err = http.NewRequest(http.MethodPost, "http://"+address+"/ok", bytes.NewReader([]byte("1234")))
	if err != nil {
		t.Fatal(err)
	}
	response, err = client.Do(request)
	if err != nil {
		t.Fatal(err)
	}
	body, readErr := io.ReadAll(response.Body)
	response.Body.Close()
	if readErr != nil || response.StatusCode != http.StatusOK || string(body) != "1234" {
		t.Fatalf("body limit did not recover: status=%d body=%q err=%v", response.StatusCode, body, readErr)
	}
}

func TestServerMultiplexes128Streams(t *testing.T) {
	started := make(chan struct{}, 128)
	release := make(chan struct{})
	var active, maximum atomic.Int32
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		current := active.Add(1)
		for {
			seen := maximum.Load()
			if current <= seen || maximum.CompareAndSwap(seen, current) {
				break
			}
		}
		started <- struct{}{}
		<-release
		active.Add(-1)
		_, _ = io.WriteString(writer, request.URL.Path)
	}))
	server.MaxConcurrentStreams = 128
	address, stop := startServer(t, server)
	defer stop()
	transport := &xhttp2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, address string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, address)
		},
	}
	client := &http.Client{Transport: transport, Timeout: 5 * time.Second}
	defer transport.CloseIdleConnections()

	type multiplexedResponse struct {
		response *http.Response
		expected string
	}
	responses := make(chan multiplexedResponse, 128)
	errorsCh := make(chan error, 128)
	for i := 0; i < 128; i++ {
		go func(i int) {
			response, err := client.Get(fmt.Sprintf("http://%s/stream-%d", address, i))
			if err != nil {
				errorsCh <- err
				return
			}
			responses <- multiplexedResponse{response: response, expected: fmt.Sprintf("/stream-%d", i)}
		}(i)
	}
	for i := 0; i < 128; i++ {
		select {
		case <-started:
		case err := <-errorsCh:
			close(release)
			t.Fatal(err)
		case <-time.After(5 * time.Second):
			close(release)
			t.Fatal("HTTP/2 streams did not reach the handler")
		}
	}
	if got := maximum.Load(); got < 128 {
		close(release)
		t.Fatalf("expected 128 concurrent streams, got %d", got)
	}
	if stats := server.Stats(); stats.ActiveConnections != 1 {
		close(release)
		t.Fatalf("expected one multiplexed connection, got %+v", stats)
	}
	close(release)
	for i := 0; i < 128; i++ {
		select {
		case result := <-responses:
			body, err := io.ReadAll(result.response.Body)
			result.response.Body.Close()
			if err != nil || result.response.StatusCode != http.StatusOK || string(body) != result.expected {
				t.Fatalf("unexpected multiplexed response: status=%d body=%q expected=%q err=%v", result.response.StatusCode, body, result.expected, err)
			}
		case err := <-errorsCh:
			t.Fatal(err)
		case <-time.After(5 * time.Second):
			t.Fatal("HTTP/2 streams did not finish")
		}
	}
}

func TestServerStreamsLargeBodyWithFlowControl(t *testing.T) {
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		var total int64
		buffer := make([]byte, 16<<10)
		for {
			n, err := request.Body.Read(buffer)
			total += int64(n)
			if n > 0 {
				time.Sleep(time.Millisecond)
			}
			if err == io.EOF {
				break
			}
			if err != nil {
				http.Error(writer, err.Error(), http.StatusBadRequest)
				return
			}
		}
		_, _ = fmt.Fprintf(writer, "%d", total)
	}))
	server.MaxBodyBytes = 4 << 20
	address, stop := startServer(t, server)
	defer stop()
	transport := &xhttp2.Transport{
		AllowHTTP: true,
		DialTLSContext: func(ctx context.Context, network, address string, _ *tls.Config) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, address)
		},
	}
	client := &http.Client{Transport: transport, Timeout: 10 * time.Second}
	defer transport.CloseIdleConnections()
	payload := bytes.Repeat([]byte("x"), 2<<20)
	request, err := http.NewRequest(http.MethodPost, "http://"+address+"/large", bytes.NewReader(payload))
	if err != nil {
		t.Fatal(err)
	}
	response, err := client.Do(request)
	if err != nil {
		t.Fatal(err)
	}
	body, readErr := io.ReadAll(response.Body)
	response.Body.Close()
	if readErr != nil || response.StatusCode != http.StatusOK || string(body) != "2097152" {
		t.Fatalf("large body flow-control result: status=%d body=%q err=%v", response.StatusCode, body, readErr)
	}
}

func TestServerTLSUsesH2ALPN(t *testing.T) {
	certificate, roots := testCertificate(t)
	server := New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.TLS == nil || request.TLS.NegotiatedProtocol != "h2" {
			http.Error(writer, "missing h2 ALPN", http.StatusInternalServerError)
			return
		}
		writer.WriteHeader(http.StatusNoContent)
	}))
	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,
		})
	}()
	stop := func() {
		t.Helper()
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := server.Shutdown(ctx); err != nil && !errors.Is(err, context.DeadlineExceeded) {
			t.Errorf("shutdown: %v", err)
		}
		if err := <-done; err != nil {
			t.Errorf("serve: %v", err)
		}
	}
	defer stop()

	transport := &xhttp2.Transport{TLSClientConfig: &tls.Config{
		RootCAs: roots, ServerName: "localhost", MinVersion: tls.VersionTLS12,
		NextProtos: []string{"h2"},
	}}
	client := &http.Client{Transport: transport, Timeout: 2 * time.Second}
	defer transport.CloseIdleConnections()
	response, err := client.Get("https://" + listener.Addr().String() + "/health")
	if err != nil {
		t.Fatal(err)
	}
	response.Body.Close()
	if response.StatusCode != http.StatusNoContent {
		t.Fatalf("unexpected TLS status: %d", response.StatusCode)
	}
}

func startServer(t *testing.T, server *Server) (string, func()) {
	t.Helper()
	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) }()
	stop := func() {
		t.Helper()
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := server.Shutdown(ctx); err != nil && !errors.Is(err, context.DeadlineExceeded) {
			t.Errorf("shutdown: %v", err)
		}
		if err := <-done; err != nil {
			t.Errorf("serve: %v", err)
		}
	}
	return listener.Addr().String(), stop
}

func testCertificate(t *testing.T) (tls.Certificate, *x509.CertPool) {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	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)
	}
	template := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: "localhost"},
		DNSNames:     []string{"localhost"},
		NotBefore:    time.Now().Add(-time.Minute),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, &key.PublicKey, key)
	if err != nil {
		t.Fatal(err)
	}
	leaf, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatal(err)
	}
	certificate := tls.Certificate{
		Certificate: [][]byte{der},
		PrivateKey:  key,
	}
	certificate.Leaf = leaf
	roots := x509.NewCertPool()
	roots.AddCert(leaf)
	return certificate, roots
}