go-omni/httpc/request_test.go

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)
	}
}