diff --git a/internal/httpfetcher/fetch_internal_test.go b/internal/httpfetcher/fetch_internal_test.go index 83c78c0..c443987 100644 --- a/internal/httpfetcher/fetch_internal_test.go +++ b/internal/httpfetcher/fetch_internal_test.go @@ -9,6 +9,7 @@ import ( "net/http" "net/http/httptest" "slices" + "strconv" "strings" "sync" "testing" @@ -297,6 +298,139 @@ func TestFetchMaxResponseSizeEnforced(t *testing.T) { } } +// fetchWithMaxResponseSize serves handler as the fake upstream and fetches +// its /image route with MaxResponseSize set to maxResponseSize. +func fetchWithMaxResponseSize( + t *testing.T, + maxResponseSize int64, + handler http.HandlerFunc, +) *FetchResult { + t.Helper() + + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + + cfg := DefaultConfig() + cfg.MaxResponseSize = maxResponseSize + + f, _ := newServerFetcher(t, srv, cfg) + + res := fetchImage(t, f, "/image") + t.Cleanup(func() { _ = res.Content.Close() }) + + return res +} + +// TestFetchContentLengthBodyAtMaxResponseSize checks that a body sent with a +// Content-Length is read whole at exactly MaxResponseSize bytes and refused +// at one byte more. +func TestFetchContentLengthBodyAtMaxResponseSize(t *testing.T) { + t.Parallel() + + const maxResponseSize = 8 + + for _, tc := range []struct { + size int + wantErr error + }{ + {maxResponseSize, nil}, + {maxResponseSize + 1, ErrResponseTooLarge}, + } { + t.Run(strconv.Itoa(tc.size)+" bytes", func(t *testing.T) { + t.Parallel() + + body := strings.Repeat("x", tc.size) + + res := fetchWithMaxResponseSize(t, maxResponseSize, + func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", contentTypeJPEG) + w.Header().Set("Content-Length", strconv.Itoa(tc.size)) + _, _ = io.WriteString(w, body) + }) + + if res.ContentLength != int64(tc.size) { + t.Fatalf("ContentLength = %d, want %d", res.ContentLength, tc.size) + } + + data, err := io.ReadAll(res.Content) + t.Logf("body of %d bytes: read %d, error %v", tc.size, len(data), err) + + if !errors.Is(err, tc.wantErr) { + t.Fatalf("read error = %v, want %v", err, tc.wantErr) + } + + if err == nil && string(data) != body { + t.Errorf("read %q, want %q", data, body) + } + }) + } +} + +// TestFetchChunkedBodyAtMaxResponseSize checks that a chunked body is read +// whole at exactly MaxResponseSize bytes and refused at one byte more. The +// upstream sends the end of the body only once the test has read +// MaxResponseSize bytes, so the read after them comes before the end of the +// body has arrived, as it can from a real upstream. +func TestFetchChunkedBodyAtMaxResponseSize(t *testing.T) { + t.Parallel() + + const maxResponseSize = 8 + + for _, tc := range []struct { + size int + wantErr error + }{ + {maxResponseSize, nil}, + {maxResponseSize + 1, ErrResponseTooLarge}, + } { + t.Run(strconv.Itoa(tc.size)+" bytes", func(t *testing.T) { + t.Parallel() + + body := strings.Repeat("x", tc.size) + endBody := make(chan struct{}) + + res := fetchWithMaxResponseSize(t, maxResponseSize, + func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", contentTypeJPEG) + _, _ = io.WriteString(w, body) + // With no Content-Length, flushing sends the body + // chunked; its end is sent when the handler returns. + _ = http.NewResponseController(w).Flush() + + select { + case <-endBody: + case <-r.Context().Done(): + } + }) + + if res.ContentLength != -1 { + t.Fatalf("ContentLength = %d, want -1 (chunked)", res.ContentLength) + } + + head := make([]byte, maxResponseSize) + + _, err := io.ReadFull(res.Content, head) + if err != nil { + t.Fatalf("read the first %d bytes: %v", maxResponseSize, err) + } + + close(endBody) + + rest, err := io.ReadAll(res.Content) + data := string(head) + string(rest) + t.Logf("body of %d bytes: read %d, error %v", tc.size, len(data), err) + + if !errors.Is(err, tc.wantErr) { + t.Fatalf("read error = %v, want %v", err, tc.wantErr) + } + + if err == nil && data != body { + t.Errorf("read %q, want %q", data, body) + } + }) + } +} + func TestFetchSemaphoreReleasedOnError(t *testing.T) { t.Parallel()