Something went wrong. Try again.
An ATProto Powered Reputation Network
Something went wrong. Try again.
13 kB · 418 lines
Go
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419package upstream
import ( "bytes" "context" "errors" "fmt" "io" "log/slog" "net/http" "net/http/httptest" "strings" "sync/atomic" "testing" "time"
"github.com/coder/websocket" "provenance/internal/logging")
type transport func(*http.Request) (*http.Response, error)
func (f transport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
func response(status int) *http.Response { return &http.Response{StatusCode: status, Header: make(http.Header), ContentLength: unknownBodyLength, Body: io.NopCloser(strings.NewReader("body"))}}
func TestWorkerBound(t *testing.T) { q, err := New(context.Background(), 2) if err != nil { t.Fatal(err) } defer q.Close() started := make(chan struct{}, 8) client := &http.Client{Transport: q.Transport(transport(func(r *http.Request) (*http.Response, error) { if r.Header.Get("User-Agent") != UserAgent { t.Error("missing project User-Agent") } started <- struct{}{} return response(http.StatusOK), nil }))} responses := make(chan *http.Response, 8) for range 4 { go func() { resp, err := client.Get("https://example.com") if err != nil { t.Error(err) return } responses <- resp }() } first, second := <-responses, <-responses <-started <-started select { case <-started: t.Fatal("workers released before response bodies closed") case <-time.After(20 * time.Millisecond): } first.Body.Close() second.Body.Close() for range 2 { (<-responses).Body.Close() }}
func TestCompleteReadReleasesWorker(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() q, err := New(ctx, 1) if err != nil { t.Fatal(err) } defer q.Close() const payload = "range" client := &http.Client{Transport: q.Transport(transport(func(*http.Request) (*http.Response, error) { return &http.Response{StatusCode: http.StatusPartialContent, Header: make(http.Header), ContentLength: int64(len(payload)), Body: io.NopCloser(strings.NewReader(payload))}, nil }))} first, err := client.Get("https://archive.example/first") if err != nil { t.Fatal(err) } defer first.Body.Close() buf := make([]byte, len(payload)) if _, err := io.ReadFull(first.Body, buf); err != nil { t.Fatal(err) } second, err := client.Get("https://archive.example/second") if err != nil { t.Fatalf("complete response read retained the worker: %v", err) } defer second.Body.Close() if _, err := io.ReadFull(second.Body, buf); err != nil { t.Fatal(err) } if _, err := first.Body.Read(buf); !errors.Is(err, io.EOF) { t.Fatalf("completed body returned %v instead of EOF", err) }}
func TestRequestLogs(t *testing.T) { var output bytes.Buffer previous := slog.Default() slog.SetDefault(logging.New(&output, slog.LevelDebug, "appview", nil)) t.Cleanup(func() { slog.SetDefault(previous) }) const payload = "private-record-content" server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Write([]byte(payload)) })) defer server.Close() q, err := New(context.Background(), 1) if err != nil { t.Fatal(err) } defer q.Close() client := &http.Client{Transport: q.Transport(server.Client().Transport)} req, _ := http.NewRequest(http.MethodGet, server.URL+"/download?access_token=query-secret", nil) req.Header.Set("Authorization", "Bearer header-secret") resp, err := client.Do(req) if err != nil { t.Fatal(err) } if _, err := io.ReadAll(resp.Body); err != nil { t.Fatal(err) } resp.Body.Close() q.Close() for _, secret := range []string{payload, "query-secret", "header-secret"} { if strings.Contains(output.String(), secret) { t.Fatalf("request log exposed %s", secret) } } for _, detail := range []string{`"path":"/download"`, `"status":200`, `"read_bytes":22`, `"queue_wait":`, `"duration":`} { if !strings.Contains(output.String(), detail) { t.Errorf("missing request detail: %s", detail) } }}
func TestIncompleteReadHoldsWorker(t *testing.T) { const payload = "range" for _, length := range []int64{int64(len(payload)), unknownBodyLength} { t.Run(fmt.Sprint(length), func(t *testing.T) { q, err := New(context.Background(), 1) if err != nil { t.Fatal(err) } defer q.Close() client := &http.Client{Transport: q.Transport(transport(func(*http.Request) (*http.Response, error) { return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), ContentLength: length, Body: io.NopCloser(strings.NewReader(payload))}, nil }))} first, err := client.Get("https://archive.example/first") if err != nil { t.Fatal(err) } defer first.Body.Close() size := 1 if length < 0 { size = len(payload) } if _, err := io.ReadFull(first.Body, make([]byte, size)); err != nil { t.Fatal(err) } ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) defer cancel() req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "https://archive.example/second", nil) if resp, err := client.Do(req); !errors.Is(err, context.DeadlineExceeded) { if resp != nil { resp.Body.Close() } t.Fatalf("incomplete response released its worker: %v", err) } if _, err := io.ReadAll(first.Body); err != nil { t.Fatal(err) } ctx, cancel = context.WithTimeout(context.Background(), time.Second) defer cancel() req, _ = http.NewRequestWithContext(ctx, http.MethodGet, "https://archive.example/third", nil) resp, err := client.Do(req) if err != nil { t.Fatal(err) } resp.Body.Close() }) }}
func TestCallerRetryOwnership(t *testing.T) { q, err := New(context.Background(), 1) if err != nil { t.Fatal(err) } defer q.Close() var calls atomic.Int64 client := &http.Client{Transport: q.TransportWithPolicy(transport(func(*http.Request) (*http.Response, error) { calls.Add(1) resp := response(http.StatusTooManyRequests) resp.Header.Set("Retry-After", "60") return resp, nil }), CallerRetries)} resp, err := client.Get("https://archive.example.com") if err != nil { t.Fatal(err) } resp.Body.Close() if calls.Load() != 1 { t.Fatal("queue multiplied caller retries") } ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) defer cancel() req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "https://archive.example.com", nil) if _, err := client.Do(req); !errors.Is(err, context.DeadlineExceeded) || calls.Load() != 1 { t.Fatal("caller retries bypassed origin cooldown") }}
func TestRetries(t *testing.T) { for _, method := range []string{http.MethodGet, http.MethodPost, http.MethodPut, http.MethodDelete} { t.Run(method, func(t *testing.T) { q, _ := New(context.Background(), 1) defer q.Close() now := time.Now() q.now = func() time.Time { return now } q.wait = func(ctx context.Context, delay time.Duration) error { now = now.Add(delay); return ctx.Err() } calls := 0 client := &http.Client{Transport: q.Transport(transport(func(*http.Request) (*http.Response, error) { calls++ return response(http.StatusServiceUnavailable), nil }))} req, _ := http.NewRequest(method, "https://example.com", nil) resp, err := client.Do(req) if err != nil { t.Fatal(err) } resp.Body.Close() want := MaxAttempts if method != http.MethodGet { want = 1 } if calls != want { t.Fatalf("attempts: %d, want %d", calls, want) } }) }}
func TestSharedCooldown(t *testing.T) { q, _ := New(context.Background(), 1) defer q.Close() now := time.Now() q.now = func() time.Time { return now } var delays []time.Duration q.wait = func(ctx context.Context, delay time.Duration) error { delays = append(delays, delay) now = now.Add(delay) return ctx.Err() } first := &http.Client{Transport: q.Transport(transport(func(*http.Request) (*http.Response, error) { resp := response(http.StatusTooManyRequests) resp.Header.Set("Retry-After", "60") return resp, nil }))} // Writes do not retry, but their cooldown applies to a different client on the same origin. req, _ := http.NewRequest(http.MethodPost, "https://EXAMPLE.com/a", nil) resp, err := first.Do(req) if err != nil { t.Fatal(err) } resp.Body.Close() second := &http.Client{Transport: q.Transport(transport(func(*http.Request) (*http.Response, error) { return response(http.StatusOK), nil }))} resp, err = second.Get("https://other.example.com/b") if err != nil { t.Fatal(err) } resp.Body.Close() if len(delays) != 0 { t.Fatal("cooldown affected a different origin") } resp, err = second.Get("https://example.com:443/b") if err != nil { t.Fatal(err) } resp.Body.Close() if len(delays) != 1 || delays[0] != time.Minute { t.Fatalf("shared cooldown: %v", delays) } date := now.Add(2 * time.Minute).UTC().Truncate(time.Second) if got := retryAfter(date.Format(http.TimeFormat), now); got != date.Sub(now) { t.Fatalf("HTTP date: %v", got) }}
func TestCanceledQueue(t *testing.T) { q, _ := New(context.Background(), 1) defer q.Close() var calls atomic.Int32 client := &http.Client{Transport: q.Transport(transport(func(*http.Request) (*http.Response, error) { calls.Add(1); return response(http.StatusOK), nil }))} resp, err := client.Get("https://example.com") if err != nil { t.Fatal(err) } ctx, cancel := context.WithCancel(context.Background()) cancel() req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "https://example.com", nil) if _, err := client.Do(req); err == nil { t.Fatal("canceled request succeeded") } resp.Body.Close() q.Close() if calls.Load() != 1 { t.Fatal("canceled request reached upstream") }}
func TestWebSocketWorker(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/fetch" { w.Write([]byte("ok")) return } conn, err := websocket.Accept(w, r, nil) if err != nil { t.Error(err) return } defer conn.CloseNow() kind, data, err := conn.Read(ctx) if err != nil { return } if err := conn.Write(ctx, kind, data); err != nil { t.Error(err) } })) defer server.Close() q, _ := New(ctx, 1) defer q.Close() client := &http.Client{Transport: q.Transport(http.DefaultTransport)} conn, _, err := websocket.Dial(ctx, server.URL, &websocket.DialOptions{HTTPClient: client}) if err != nil { t.Fatal(err) } defer conn.CloseNow() if err := conn.Write(ctx, websocket.MessageText, []byte("hello")); err != nil { t.Fatal(err) } _, data, err := conn.Read(ctx) if err != nil || string(data) != "hello" { t.Fatalf("upgrade lost its read/write body: %s %v", data, err) } req, _ := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+"/fetch", nil) resp, err := client.Do(req) if err != nil { t.Fatalf("WebSocket occupied the only fetch worker: %v", err) } resp.Body.Close()}
func TestRetryAfter(t *testing.T) { now := time.Now() for _, raw := range []string{"", "invalid", "-1", "0", now.Add(-time.Hour).Format(http.TimeFormat)} { if got := retryAfter(raw, now); got != 0 { t.Fatalf("invalid or expired header %q: %v", raw, got) } } for range 100 { delay := Jitter(time.Second) if delay < time.Second/2 || delay > time.Second { t.Fatalf("jitter outside bounds: %v", delay) } }}
func TestRetryCancellation(t *testing.T) { q, _ := New(context.Background(), 1) defer q.Close() ctx, cancel := context.WithCancel(context.Background()) defer cancel() calls := 0 q.wait = func(ctx context.Context, delay time.Duration) error { cancel(); return ctx.Err() } client := &http.Client{Transport: q.Transport(transport(func(*http.Request) (*http.Response, error) { calls++ resp := response(http.StatusTooManyRequests) resp.Header.Set("Retry-After", "3600") return resp, nil }))} req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "https://example.com", nil) if _, err := client.Do(req); err == nil { t.Fatal("retry wait ignored cancellation") } q.Close() if calls != 1 { t.Fatalf("canceled retry made %d requests", calls) }}
func TestTransportErrors(t *testing.T) { for _, test := range []struct { err error attempts int }{{errors.New("invalid transport configuration"), 1}, {io.ErrUnexpectedEOF, MaxAttempts}} { q, _ := New(context.Background(), 1) now := time.Now() q.now = func() time.Time { return now } q.wait = func(ctx context.Context, delay time.Duration) error { now = now.Add(delay); return ctx.Err() } calls := 0 client := &http.Client{Transport: q.Transport(transport(func(*http.Request) (*http.Response, error) { calls++; return nil, test.err }))} if _, err := client.Get("https://example.com"); !errors.Is(err, test.err) { t.Fatalf("lost transport error: %v", err) } q.Close() if calls != test.attempts { t.Fatalf("error %v made %d attempts, want %d", test.err, calls, test.attempts) } }}