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