feat: simplify HTTP client proxy pools and authenticated JSON calls

2026-09-26 15:28:19 +08:00
parent 068e85132f
commit 96e7d9629b
27 changed files with 1611 additions and 143 deletions
--- a/httpc/request_test.go
+++ b/httpc/request_test.go
@@ -0,0 +1,316 @@
+package httpc
+
+import (
+	"bytes"
+	"compress/gzip"
+	"context"
+	"errors"
+	"fmt"
+	"git.yuntuo.hk/k2safe/go-omni/logx"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestJSONHeadersConcurrentSnapshot(t *testing.T) {
+	var leaked atomic.Bool
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("Authorization") != "Bearer "+r.URL.Query().Get("caller") {
+			leaked.Store(true)
+		}
+		if r.Header.Get("X-Shared") != "original" || len(r.Header.Values("X-Tag")) != 2 {
+			leaked.Store(true)
+		}
+		if r.Header.Get("Content-Type") != "application/vnd.example+json" {
+			leaked.Store(true)
+		}
+		io.WriteString(w, `{"ok":true}`)
+	}))
+	defer s.Close()
+	c := testClient(t, Config{})
+	headers := http.Header{"X-Shared": []string{"original"}, "X-Tag": []string{"one", "two"}}
+	snapshot := WithHeaders(headers)
+	headers.Set("X-Shared", "changed")
+	headers["X-Tag"][0] = "changed"
+	var wg sync.WaitGroup
+	for i := 0; i < 24; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			id := fmt.Sprint(i)
+			var output struct {
+				OK bool `json:"ok"`
+			}
+			if err := c.JSON(context.Background(), "POST", s.URL+"?caller="+id, map[string]bool{"hello": true}, &output, snapshot, WithHeader("Authorization", "Bearer "+id), WithHeader("content-type", "application/vnd.example+json")); err != nil || !output.OK {
+				t.Errorf("JSON failed: %v", err)
+			}
+		}(i)
+	}
+	wg.Wait()
+	if leaked.Load() {
+		t.Fatal("headers leaked across requests or changed with source map")
+	}
+}
+func TestRequestOptionsValidationBeforeNetwork(t *testing.T) {
+	var calls atomic.Int32
+	c := testClient(t, Config{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
+		calls.Add(1)
+		return benchmarkTransport{`{}`, 200}.RoundTrip(r)
+	})})
+	invalid := []RequestOption{nil, WithHeader("bad name", "private-token"), WithHeader("Authorization", "private-token\r\nInjected: yes"), WithHeader("Host", "example.invalid"), WithHeader("Content-Length", "1"), WithHeader("Transfer-Encoding", "chunked"), WithHeader("Trailer", "X-Signature"), WithHeaders(http.Header{"X-Name": []string{"one"}, "x-name": []string{"two"}}), WithMaxResponseBytes(0), WithMaxResponseBytes(maxResponseBytes + 1)}
+	for _, option := range invalid {
+		err := c.JSON(context.Background(), "GET", "http://example.invalid", nil, nil, option)
+		if err == nil || strings.Contains(err.Error(), "private-token") {
+			t.Fatal("missing or unsafe error", err)
+		}
+	}
+	if calls.Load() != 0 {
+		t.Fatal("invalid options reached network")
+	}
+}
+func TestJSONResponseLimitsAndStreaming(t *testing.T) {
+	var calls atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		calls.Add(1)
+		switch r.URL.Path {
+		case "/empty":
+			w.WriteHeader(204)
+		case "/chunked":
+			w.(http.Flusher).Flush()
+			io.WriteString(w, strings.Repeat("x", 17))
+		case "/status-stream":
+			w.WriteHeader(401)
+			w.(http.Flusher).Flush()
+			<-r.Context().Done()
+		case "/multiple":
+			io.WriteString(w, `{} {}`)
+		default:
+			io.WriteString(w, `{"ok":true}`)
+		}
+	}))
+	defer server.Close()
+	c := testClient(t, Config{MaxResponseBytes: 8, Timeout: time.Second})
+	var out struct {
+		OK bool `json:"ok"`
+	}
+	if err := c.JSON(context.Background(), "GET", server.URL, nil, &out, WithMaxResponseBytes(32)); err != nil || !out.OK {
+		t.Fatal(err)
+	}
+	if err := c.JSON(context.Background(), "GET", server.URL, nil, nil); !errors.Is(err, ErrResponseTooLarge) {
+		t.Fatal("limit changed permanently", err)
+	}
+	for _, target := range []any{nil, &out} {
+		if err := c.JSON(context.Background(), "GET", server.URL+"/chunked", nil, target, WithMaxResponseBytes(16)); !errors.Is(err, ErrResponseTooLarge) {
+			t.Fatal("unbounded chunked body", err)
+		}
+	}
+	if err := c.JSON(context.Background(), "GET", server.URL+"/empty", nil, &out); err != nil || !out.OK {
+		t.Fatal("empty success changed output", err)
+	}
+	if err := c.JSON(context.Background(), "GET", server.URL+"/multiple", nil, &out); err == nil {
+		t.Fatal("multiple JSON documents accepted")
+	}
+	start := time.Now()
+	err := c.JSON(context.Background(), "GET", server.URL+"/status-stream", nil, nil)
+	var status *StatusError
+	if !errors.As(err, &status) || status.StatusCode != 401 || time.Since(start) > 500*time.Millisecond {
+		t.Fatal("waited for unknown status body", err)
+	}
+	if calls.Load() != 7 {
+		t.Fatal(calls.Load())
+	}
+}
+func TestErrorBodyPoolReuseAndDrainBound(t *testing.T) {
+	var connections atomic.Int32
+	s := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(401)
+		io.WriteString(w, "private-error-body")
+	}))
+	s.Config.ConnState = func(_ net.Conn, state http.ConnState) {
+		if state == http.StateNew {
+			connections.Add(1)
+		}
+	}
+	s.Start()
+	defer s.Close()
+	c := testClient(t, Config{})
+	for i := 0; i < 10; i++ {
+		err := c.JSON(context.Background(), "GET", s.URL, nil, nil)
+		var status *StatusError
+		if !errors.As(err, &status) || strings.Contains(err.Error(), "private-error-body") {
+			t.Fatal(err)
+		}
+	}
+	if connections.Load() != 1 {
+		t.Fatal("small status response prevented reuse", connections.Load())
+	}
+	body := &countedBody{Reader: strings.NewReader(strings.Repeat("x", 1<<20))}
+	custom := testClient(t, Config{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
+		return &http.Response{StatusCode: 401, ContentLength: 1, Body: body, Header: make(http.Header)}, nil
+	})})
+	err := custom.JSON(context.Background(), "GET", "http://example.invalid", nil, nil)
+	var status *StatusError
+	if !errors.As(err, &status) || !body.closed || body.read > (32<<10)+1 {
+		t.Fatal("unbounded/missing drain", body.read, body.closed, err)
+	}
+}
+
+type countedBody struct {
+	io.Reader
+	read   int
+	closed bool
+}
+
+func (b *countedBody) Read(p []byte) (int, error) {
+	n, err := b.Reader.Read(p)
+	b.read += n
+	return n, err
+}
+func (b *countedBody) Close() error { b.closed = true; return nil }
+
+type roundTripperFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
+
+type borrowedTransport struct{ closed atomic.Bool }
+
+func (t *borrowedTransport) RoundTrip(r *http.Request) (*http.Response, error) {
+	return benchmarkTransport{`{}`, 200}.RoundTrip(r)
+}
+func (t *borrowedTransport) CloseIdleConnections() { t.closed.Store(true) }
+func TestBorrowedTransportAndHeaderRetry(t *testing.T) {
+	owned := &borrowedTransport{}
+	c := testClient(t, Config{Transport: owned})
+	c.Close()
+	if owned.closed.Load() {
+		t.Fatal("closed caller's transport")
+	}
+	var attempts atomic.Int32
+	retry := testClient(t, Config{MaxRetries: 1, Backoff: time.Millisecond, Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
+		if r.Header.Get("Authorization") != "Bearer test" {
+			t.Error("retry lost header")
+		}
+		status := 200
+		if attempts.Add(1) == 1 {
+			status = 503
+		}
+		return benchmarkTransport{`{}`, status}.RoundTrip(r)
+	})})
+	if err := retry.JSON(context.Background(), "GET", "http://example.invalid", nil, nil, WithHeader("Authorization", "Bearer test")); err != nil || attempts.Load() != 2 {
+		t.Fatal(err, attempts.Load())
+	}
+}
+func FuzzRequestHeaderOptions(f *testing.F) {
+	f.Add("Authorization", "Bearer example")
+	f.Add("X-Name", "one\r\nX-Test: bad")
+	f.Fuzz(func(t *testing.T, name, value string) {
+		if len(name)+len(value) > 8192 {
+			t.Skip()
+		}
+		option := WithHeader(name, value)
+		var settings requestOptions
+		if err := option(&settings); err != nil {
+			return
+		}
+		req, err := http.NewRequest("GET", "http://example.invalid", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		req.Header = settings.headers
+		// The standard Request writer must be able to serialize accepted headers.
+		if err = req.Write(io.Discard); err != nil {
+			t.Fatal("accepted malformed header", err)
+		}
+	})
+}
+
+func TestHTTPLogDoesNotContainCredentials(t *testing.T) {
+	var buffer bytes.Buffer
+	logger, err := logx.New(logx.Config{Writer: &buffer})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer logger.Close()
+	c, err := New(Config{Transport: benchmarkTransport{"private-response", 401}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+	ctx := logx.WithLogger(context.Background(), logger)
+	err = c.JSON(ctx, "GET", "http://private-user:private-password@example.invalid/private-path?token=private-query", nil, nil, WithHeader("Authorization", "Bearer private-token"))
+	if err == nil {
+		t.Fatal("expected status error")
+	}
+	for _, secret := range []string{"private-user", "private-password", "private-path", "private-query", "private-token", "private-response"} {
+		if strings.Contains(buffer.String(), secret) || strings.Contains(err.Error(), secret) {
+			t.Fatal("sensitive field leaked")
+		}
+	}
+	if !strings.Contains(buffer.String(), "http.client") {
+		t.Fatal("missing access log")
+	}
+}
+func TestCloseAllowsInflightCallToFinish(t *testing.T) {
+	entered := make(chan struct{})
+	release := make(chan struct{})
+	var once sync.Once
+	releaseNow := func() { once.Do(func() { close(release) }) }
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { close(entered); <-release; io.WriteString(w, `{}`) }))
+	defer s.Close()
+	c := testClient(t, Config{Timeout: time.Second})
+	done := make(chan error, 1)
+	go func() { done <- c.JSON(context.Background(), "GET", s.URL, nil, nil) }()
+	defer releaseNow()
+	select {
+	case <-entered:
+	case err := <-done:
+		t.Fatal("request failed before handler", err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("handler not entered")
+	}
+	c.Close()
+	releaseNow()
+	if err := <-done; err != nil {
+		t.Fatal("Close interrupted in-flight call", err)
+	}
+	if err := c.JSON(context.Background(), "GET", s.URL, nil, nil); !errors.Is(err, ErrClosed) {
+		t.Fatal("accepted new call after Close", err)
+	}
+}
+
+func TestSlowKnownErrorBodyHasShortDrainBudget(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Length", "100")
+		w.WriteHeader(503)
+		w.(http.Flusher).Flush()
+		<-r.Context().Done()
+	}))
+	defer server.Close()
+	c := testClient(t, Config{Timeout: 2 * time.Second})
+	start := time.Now()
+	err := c.JSON(context.Background(), "GET", server.URL, nil, nil)
+	var status *StatusError
+	if !errors.As(err, &status) || status.StatusCode != 503 || time.Since(start) > 500*time.Millisecond {
+		t.Fatal("slow error body delayed status", err, time.Since(start))
+	}
+}
+
+func TestJSONLimitAppliesAfterGzipExpansion(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Encoding", "gzip")
+		writer := gzip.NewWriter(w)
+		_, _ = io.WriteString(writer, strings.Repeat("x", 4096))
+		_ = writer.Close()
+	}))
+	defer server.Close()
+	c := testClient(t, Config{MaxResponseBytes: 64})
+	if err := c.JSON(context.Background(), "GET", server.URL, nil, nil); !errors.Is(err, ErrResponseTooLarge) {
+		t.Fatal("gzip bypassed response limit", err)
+	}
+}