go-omni/transport/http2/proxy_test.go

package http2

import (
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"testing"
	"time"

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

// This opt-in suite runs against a real Nginx executable. The Linux wrapper
// builds the test binary separately, then runs it in an isolated Nginx image.
func TestNginxHTTP2Proxy(t *testing.T) {
	if os.Getenv("GO_OMNI_HTTP2_PROXY_TEST") != "1" {
		t.Skip("use scripts/verifyHTTP2Proxy.py for real Nginx HTTP/2 upstream validation")
	}
	if output, err := exec.Command("nginx", "-v").CombinedOutput(); err != nil {
		t.Fatalf("nginx unavailable: %v: %s", err, output)
	} else {
		t.Logf("proxyVersion=%s", strings.TrimSpace(string(output)))
	}
	t.Run("TLSAndConnectionReuse", testProxyTLSAndReuse)
	t.Run("TimeoutAndCancellation", testProxyTimeoutAndCancellation)
	t.Run("GoAwayAndRollingDrain", testProxyGoAwayAndDrain)
}

type proxyReply struct {
	Instance   string `json:"instance"`
	RequestID  string `json:"requestId"`
	Protocol   string `json:"protocol"`
	ALPN       string `json:"alpn"`
	Method     string `json:"method"`
	BodyBytes  int    `json:"bodyBytes"`
	BodySHA256 string `json:"bodySha256"`
}

type proxyOrigin struct {
	server     *Server
	address    string
	done       chan struct{}
	serveErr   error
	requests   atomic.Uint64
	duplicates atomic.Uint64
	seen       sync.Map
}

func newProxyOrigin(t *testing.T, certificate tls.Certificate, instance string, beforeReply func(*http.Request)) *proxyOrigin {
	t.Helper()
	o := &proxyOrigin{done: make(chan struct{})}
	o.server = New(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if request.ProtoMajor != 2 || request.TLS == nil || request.TLS.NegotiatedProtocol != "h2" || request.TLS.ServerName != "localhost" {
			t.Error("proxy upstream did not negotiate verified-name TLS/h2")
			http.Error(writer, "wrong upstream protocol", http.StatusInternalServerError)
			return
		}
		o.requests.Add(1)
		id := request.Header.Get("X-Request-ID")
		if _, loaded := o.seen.LoadOrStore(id, struct{}{}); loaded {
			o.duplicates.Add(1)
		}
		body, err := io.ReadAll(request.Body)
		if err != nil {
			http.Error(writer, "body read failed", http.StatusBadRequest)
			return
		}
		if beforeReply != nil {
			beforeReply(request)
		}
		if request.Context().Err() != nil {
			return
		}
		writer.Header().Set("Content-Type", "application/json")
		writer.Header().Set("X-Request-ID", id)
		_ = json.NewEncoder(writer).Encode(proxyReply{
			Instance: instance, RequestID: id, Protocol: request.Proto, ALPN: request.TLS.NegotiatedProtocol,
			Method: request.Method, BodyBytes: len(body), BodySHA256: fmt.Sprintf("%x", sha256.Sum256(body)),
		})
	}))
	o.server.MaxConnections = 16
	o.server.MaxConcurrentRequests = 32
	o.server.MaxConcurrentStreams = 32
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	o.address = listener.Addr().String()
	go func() {
		o.serveErr = o.server.ServeTLS(context.Background(), listener, &tls.Config{
			Certificates: []tls.Certificate{certificate}, MinVersion: tls.VersionTLS12,
		})
		close(o.done)
	}()
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		defer cancel()
		if err := o.shutdown(ctx); err != nil {
			t.Errorf("origin cleanup: %v", err)
		}
		stats := o.server.Stats()
		if stats.ActiveConnections != 0 || stats.ActiveRequests != 0 || stats.ProtocolErrors != 0 || o.duplicates.Load() != 0 {
			t.Errorf("origin lifecycle/duplicate failure: stats=%+v duplicates=%d", stats, o.duplicates.Load())
		}
	})
	return o
}

func (o *proxyOrigin) shutdown(ctx context.Context) error {
	if err := o.server.Shutdown(ctx); err != nil {
		return err
	}
	select {
	case <-o.done:
		return o.serveErr
	case <-ctx.Done():
		return ctx.Err()
	}
}

func testProxyTLSAndReuse(t *testing.T) {
	certificate, roots := testCertificate(t)
	origin := newProxyOrigin(t, certificate, "primary", nil)
	endpoint := startTestNginx(t, certificate, origin.address, "")
	client := proxyClient(t, roots)
	for index := 0; index < 32; index++ {
		payload := []byte(fmt.Sprintf("payload-%d", index))
		result, err := proxyCall(context.Background(), client, endpoint+"/echo", http.MethodPost, fmt.Sprintf("reuse-%d", index), payload, 200)
		if err != nil || result.Instance != "primary" {
			t.Fatalf("reuse request %d: %+v %v", index, result, err)
		}
	}
	stats := origin.server.Stats()
	if stats.AcceptedConnections != 1 || stats.ActiveConnections != 1 || origin.requests.Load() != 32 {
		t.Fatalf("upstream TLS connection was not reused: %+v requests=%d", stats, origin.requests.Load())
	}
	for _, path := range []string{"/badTrust", "/badName"} {
		if _, err := proxyCall(context.Background(), client, endpoint+path, http.MethodGet, path, nil, 502); err != nil {
			t.Fatal(err)
		}
	}
	if origin.requests.Load() != 32 {
		t.Fatal("an upstream with invalid certificate trust or hostname reached the handler")
	}
	if _, err := proxyCall(context.Background(), client, endpoint+"/echo", http.MethodGet, "afterTLSFailures", nil, 200); err != nil {
		t.Fatal(err)
	}
	t.Logf("proxyResult=tlsAndReuse reusedRequests=32 upstreamConnections=1 invalidTrustStatus=502 invalidNameStatus=502 recovered=true")
}

func testProxyTimeoutAndCancellation(t *testing.T) {
	certificate, roots := testCertificate(t)
	started, canceled := make(chan string, 2), make(chan string, 2)
	origin := newProxyOrigin(t, certificate, "primary", func(request *http.Request) {
		if request.URL.Path == "/timeout" || request.URL.Path == "/cancel" {
			started <- request.URL.Path
			<-request.Context().Done()
			canceled <- request.URL.Path
		}
	})
	endpoint := startTestNginx(t, certificate, origin.address, "")
	client := proxyClient(t, roots)
	begin := time.Now()
	if _, err := proxyCall(context.Background(), client, endpoint+"/timeout", http.MethodGet, "timeout", nil, 504); err != nil {
		t.Fatal(err)
	}
	if time.Since(begin) > 2*time.Second {
		t.Fatal("upstream timeout was not bounded")
	}
	awaitProxyPath(t, started, "/timeout")
	awaitProxyPath(t, canceled, "/timeout")
	waitHTTP2Stats(t, origin.server, func(stats Stats) bool { return stats.ActiveRequests == 0 })

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	result := make(chan error, 1)
	go func() {
		_, err := proxyCall(ctx, client, endpoint+"/cancel", http.MethodGet, "cancel", nil, 200)
		result <- err
	}()
	awaitProxyPath(t, started, "/cancel")
	cancel()
	if err := serveResult(t, result); err == nil {
		t.Fatal("canceled proxy request unexpectedly succeeded")
	}
	awaitProxyPath(t, canceled, "/cancel")
	waitHTTP2Stats(t, origin.server, func(stats Stats) bool { return stats.ActiveRequests == 0 })
	if _, err := proxyCall(context.Background(), client, endpoint+"/echo", http.MethodGet, "afterCancellation", nil, 200); err != nil {
		t.Fatal(err)
	}
	t.Logf("proxyResult=timeoutAndCancellation timeoutStatus=504 canceledHandlers=2 recovered=true stats=%+v", origin.server.Stats())
}

func testProxyGoAwayAndDrain(t *testing.T) {
	certificate, roots := testCertificate(t)
	started := make(chan context.Context, 1)
	release := make(chan struct{})
	var releaseOnce sync.Once
	releaseOld := func() { releaseOnce.Do(func() { close(release) }) }
	defer releaseOld()
	old := newProxyOrigin(t, certificate, "old", func(request *http.Request) {
		if request.URL.Path == "/hold" {
			started <- request.Context()
			select {
			case <-release:
			case <-request.Context().Done():
			}
		}
	})
	next := newProxyOrigin(t, certificate, "new", nil)
	endpoint := startTestNginx(t, certificate, old.address, next.address)
	client := proxyClient(t, roots)
	if reply, err := proxyCall(context.Background(), client, endpoint+"/echo", http.MethodGet, "beforeRollout", nil, 200); err != nil || reply.Instance != "old" {
		t.Fatalf("initial request: %+v %v", reply, err)
	}
	held := make(chan error, 1)
	go func() {
		reply, err := proxyCall(context.Background(), client, endpoint+"/hold", http.MethodPost, "heldOld", []byte("finish exactly once"), 200)
		if err == nil && reply.Instance != "old" {
			err = fmt.Errorf("held request switched instance: %+v", reply)
		}
		held <- err
	}()
	var heldContext context.Context
	select {
	case heldContext = <-started:
	case <-time.After(3 * time.Second):
		t.Fatal("held request never reached old instance")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
	defer cancel()
	stopped := make(chan error, 1)
	go func() { stopped <- old.shutdown(ctx) }()
	waitHTTP2Stats(t, old.server, func(stats Stats) bool { return stats.GoAwayFrames > 0 })
	if heldContext.Err() != nil {
		t.Fatal("GOAWAY canceled the in-flight request")
	}
	// The primary listener is closed, so Nginx must route new requests to the
	// backup while the old HTTP/2 stream is still draining. POSTs also prove
	// request bodies and IDs survive this transition without duplicate handlers.
	for index := 0; index < 20; index++ {
		reply, err := proxyCall(context.Background(), client, endpoint+"/echo", http.MethodPost, fmt.Sprintf("rollout-%d", index), []byte(fmt.Sprint(index)), 200)
		if err != nil || reply.Instance != "new" {
			t.Fatalf("request during drain %d: %+v %v", index, reply, err)
		}
	}
	select {
	case err := <-held:
		t.Fatalf("held request completed before release: %v", err)
	default:
	}
	if old.requests.Load() != 2 || next.requests.Load() != 20 {
		t.Fatalf("rollout request accounting: old=%d new=%d", old.requests.Load(), next.requests.Load())
	}
	releaseOld()
	if err := serveResult(t, held); err != nil {
		t.Fatal(err)
	}
	if err := serveResult(t, stopped); err != nil {
		t.Fatal(err)
	}
	stats := old.server.Stats()
	if stats.ActiveRequests != 0 || stats.ActiveConnections != 0 || stats.GoAwayFrames == 0 {
		t.Fatalf("old instance failed to drain: %+v", stats)
	}
	t.Logf("proxyResult=goAwayAndRollingDrain oldRequests=2 newRequests=20 heldPostCompleted=true duplicateHandlers=0 oldStats=%+v", stats)
}

func awaitProxyPath(t *testing.T, events <-chan string, expected string) {
	t.Helper()
	select {
	case path := <-events:
		if path != expected {
			t.Fatalf("handler event: got %s, want %s", path, expected)
		}
	case <-time.After(3 * time.Second):
		t.Fatalf("timed out waiting for handler event: %s", expected)
	}
}

func proxyClient(t *testing.T, roots *x509.CertPool) *http.Client {
	t.Helper()
	transport := &xhttp2.Transport{TLSClientConfig: &tls.Config{
		RootCAs: roots, ServerName: "localhost", MinVersion: tls.VersionTLS12, NextProtos: []string{"h2"},
	}, StrictMaxConcurrentStreams: true}
	t.Cleanup(transport.CloseIdleConnections)
	return &http.Client{Transport: transport, Timeout: 10 * time.Second}
}

func proxyCall(ctx context.Context, client *http.Client, endpoint, method, id string, body []byte, expectedStatus int) (proxyReply, error) {
	var reply proxyReply
	request, err := http.NewRequestWithContext(ctx, method, endpoint, bytes.NewReader(body))
	if err != nil {
		return reply, err
	}
	request.Header.Set("X-Request-ID", id)
	response, err := client.Do(request)
	if err != nil {
		return reply, err
	}
	defer response.Body.Close()
	data, err := io.ReadAll(io.LimitReader(response.Body, 8192))
	if err != nil {
		return reply, err
	}
	if response.StatusCode != expectedStatus || response.ProtoMajor != 2 || response.TLS == nil || response.TLS.NegotiatedProtocol != "h2" || len(response.TLS.VerifiedChains) == 0 {
		return reply, fmt.Errorf("unverified edge TLS/h2 or status: %s body=%q", response.Status, data)
	}
	if expectedStatus != 200 {
		return reply, nil
	}
	if err := json.Unmarshal(data, &reply); err != nil {
		return reply, err
	}
	if reply.RequestID != id || response.Header.Get("X-Request-ID") != id || reply.Protocol != "HTTP/2.0" || reply.ALPN != "h2" || reply.Method != method || reply.BodyBytes != len(body) || reply.BodySHA256 != fmt.Sprintf("%x", sha256.Sum256(body)) {
		return reply, fmt.Errorf("upstream protocol, request ID or body mismatch: %+v", reply)
	}
	return reply, nil
}

func startTestNginx(t *testing.T, certificate tls.Certificate, primary, backup string) string {
	t.Helper()
	directory := t.TempDir()
	certFile, keyFile := filepath.Join(directory, "cert.pem"), filepath.Join(directory, "key.pem")
	writeProxyFile(t, certFile, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certificate.Certificate[0]}))
	key, err := x509.MarshalPKCS8PrivateKey(certificate.PrivateKey)
	if err != nil {
		t.Fatal(err)
	}
	writeProxyFile(t, keyFile, pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: key}))
	wrong, _ := testCertificate(t)
	wrongFile := filepath.Join(directory, "wrong.pem")
	writeProxyFile(t, wrongFile, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: wrong.Certificate[0]}))
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	address := listener.Addr().String()
	_ = listener.Close()
	backupLine := ""
	if backup != "" {
		backupLine = "server " + backup + " backup;"
	}
	config := fmt.Sprintf(`daemon off;
master_process off;
pid %q;
error_log stderr info;
events { worker_connections 128; }
http {
    access_log off;
    client_body_temp_path %q;
    proxy_temp_path %q;
    fastcgi_temp_path %q;
    uwsgi_temp_path %q;
    scgi_temp_path %q;
    upstream origins { server %s max_fails=1 fail_timeout=30s; %s keepalive 16; }
    upstream badTrust { server %s; }
    upstream badName { server %s; }
    server {
        listen %s ssl;
        http2 on;
        ssl_certificate %q;
        ssl_certificate_key %q;
        proxy_http_version 2;
        proxy_set_header Connection "";
        proxy_set_header X-Request-ID $http_x_request_id;
        proxy_ssl_verify on;
        proxy_ssl_trusted_certificate %q;
        proxy_ssl_server_name on;
        proxy_ssl_name localhost;
        proxy_connect_timeout 500ms;
        proxy_send_timeout 2s;
        proxy_read_timeout 10s;
        proxy_buffering off;
        proxy_request_buffering off;
        proxy_next_upstream error timeout;
        proxy_next_upstream_tries 2;
        location / { proxy_pass https://origins; }
        location = /timeout { proxy_read_timeout 250ms; proxy_next_upstream off; proxy_pass https://origins; }
        location = /badTrust { proxy_ssl_trusted_certificate %q; proxy_next_upstream off; proxy_pass https://badTrust; }
        location = /badName { proxy_ssl_name wrong.invalid; proxy_next_upstream off; proxy_pass https://badName; }
    }
}
`, filepath.Join(directory, "nginx.pid"), filepath.Join(directory, "clientBody"), filepath.Join(directory, "proxyBody"),
		filepath.Join(directory, "fastcgi"), filepath.Join(directory, "uwsgi"), filepath.Join(directory, "scgi"),
		primary, backupLine, primary, primary, address, certFile, keyFile, certFile, wrongFile)
	configFile := filepath.Join(directory, "nginx.conf")
	writeProxyFile(t, configFile, []byte(config))
	logFile := filepath.Join(directory, "nginx.log")
	output, err := os.Create(logFile)
	if err != nil {
		t.Fatal(err)
	}
	command := exec.Command("nginx", "-p", directory+"/", "-c", configFile)
	command.Stdout, command.Stderr = output, output
	if err := command.Start(); err != nil {
		_ = output.Close()
		t.Fatal(err)
	}
	done := make(chan struct{})
	var processErr error
	go func() { processErr = command.Wait(); close(done) }()
	t.Cleanup(func() {
		_ = command.Process.Signal(syscall.SIGQUIT)
		select {
		case <-done:
		case <-time.After(3 * time.Second):
			_ = command.Process.Kill()
			<-done
			t.Error("nginx did not exit after SIGQUIT")
		}
		_ = output.Close()
		if processErr != nil || t.Failed() {
			log, _ := os.ReadFile(logFile)
			t.Logf("nginx exit=%v log:\n%s", processErr, log)
		}
		if processErr != nil {
			t.Errorf("nginx exit: %v", processErr)
		}
	})
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		select {
		case <-done:
			t.Fatalf("nginx exited during startup: %v", processErr)
		default:
		}
		conn, err := net.DialTimeout("tcp", address, 100*time.Millisecond)
		if err == nil {
			_ = conn.Close()
			return "https://" + address
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("nginx did not start")
	return ""
}

func writeProxyFile(t *testing.T, path string, data []byte) {
	t.Helper()
	if err := os.WriteFile(path, data, 0600); err != nil {
		t.Fatal(err)
	}
}