Skip to content

Commit 05fa854

Browse files
committed
Stream oversized ETag responses
Avoid fully buffering responses that exceed the ETag cache entry budget when Content-Length is unknown. Preserve the consumed prefix with the remaining response stream so callers still receive the complete body. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: c4eabbcd-abb9-4c5d-a1f5-c98618623c5b
1 parent 3d894b8 commit 05fa854

2 files changed

Lines changed: 70 additions & 11 deletions

File tree

pkg/http/transport/etag.go

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -166,27 +166,27 @@ func (t *ETagTransport) RoundTrip(req *http.Request) (*http.Response, error) {
166166
return resp, nil
167167
}
168168

169-
// Skip caching bodies that exceed the per-entry byte budget. When the
170-
// length is known up front, avoid buffering the body at all.
171169
maxEntry := t.maxEntryBytes()
172-
if resp.ContentLength > int64(maxEntry) {
173-
t.remove(key)
174-
return resp, nil
175-
}
176-
177-
body, readErr := io.ReadAll(resp.Body)
178-
resp.Body.Close()
170+
limited := io.LimitReader(resp.Body, int64(maxEntry)+1)
171+
body, readErr := io.ReadAll(limited)
179172
if readErr != nil {
173+
resp.Body.Close()
180174
return nil, readErr
181175
}
182-
resp.Body = io.NopCloser(bytes.NewReader(body))
183-
resp.ContentLength = int64(len(body))
184176

185177
if len(body) > maxEntry {
186178
t.remove(key)
179+
resp.Body = struct {
180+
io.Reader
181+
io.Closer
182+
}{io.MultiReader(bytes.NewReader(body), resp.Body), resp.Body}
187183
return resp, nil
188184
}
189185

186+
resp.Body.Close()
187+
resp.Body = io.NopCloser(bytes.NewReader(body))
188+
resp.ContentLength = int64(len(body))
189+
190190
t.add(key, etagEntry{
191191
etag: etag,
192192
status: resp.StatusCode,

pkg/http/transport/etag_test.go

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package transport
22

33
import (
4+
"bytes"
45
"context"
56
"io"
67
"net/http"
@@ -14,6 +15,27 @@ import (
1415
"github.com/stretchr/testify/require"
1516
)
1617

18+
type roundTripFunc func(*http.Request) (*http.Response, error)
19+
20+
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
21+
return f(req)
22+
}
23+
24+
type countingReadCloser struct {
25+
reader io.Reader
26+
bytesRead int
27+
}
28+
29+
func (c *countingReadCloser) Read(p []byte) (int, error) {
30+
n, err := c.reader.Read(p)
31+
c.bytesRead += n
32+
return n, err
33+
}
34+
35+
func (c *countingReadCloser) Close() error {
36+
return nil
37+
}
38+
1739
// TestETagTransport_ServesCachedBodyOn304 verifies the core conditional-request
1840
// flow: the first GET carries no If-None-Match and is cached with its ETag; the
1941
// second GET sends the cached ETag and, on a 304 Not Modified, is served the
@@ -236,6 +258,43 @@ func TestETagTransport_SkipsBodiesOverEntryByteBudget(t *testing.T) {
236258
assert.Empty(t, lastIfNoneMatch, "an oversized response must not be cached or revalidated")
237259
}
238260

261+
func TestETagTransport_StreamsOversizedUnknownLengthBody(t *testing.T) {
262+
t.Parallel()
263+
264+
const maxEntry = 16
265+
body := []byte("0123456789abcdef-streamed-remainder")
266+
stream := &countingReadCloser{reader: bytes.NewReader(body)}
267+
transport := roundTripFunc(func(req *http.Request) (*http.Response, error) {
268+
assert.Empty(t, req.Header.Get(headers.IfNoneMatchHeader), "an oversized response must not be cached or revalidated")
269+
return &http.Response{
270+
StatusCode: http.StatusOK,
271+
Status: "200 OK",
272+
Header: http.Header{headers.ETagHeader: []string{`"streamed"`}},
273+
Body: stream,
274+
ContentLength: -1,
275+
Request: req,
276+
}, nil
277+
})
278+
rt := &ETagTransport{Transport: transport, MaxEntryBytes: maxEntry}
279+
280+
req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.test/stream", nil)
281+
require.NoError(t, err)
282+
resp, err := rt.RoundTrip(req)
283+
require.NoError(t, err)
284+
require.LessOrEqual(t, stream.bytesRead, maxEntry+1, "RoundTrip must not read the entire oversized body before returning")
285+
286+
data, err := io.ReadAll(resp.Body)
287+
require.NoError(t, err)
288+
require.NoError(t, resp.Body.Close())
289+
assert.Equal(t, body, data, "caller receives the buffered prefix plus the streamed remainder")
290+
291+
stream = &countingReadCloser{reader: bytes.NewReader(body)}
292+
resp, err = rt.RoundTrip(req)
293+
require.NoError(t, err)
294+
_, _ = io.Copy(io.Discard, resp.Body)
295+
require.NoError(t, resp.Body.Close())
296+
}
297+
239298
// TestETagTransport_EvictsByTotalByteBudget verifies that inserting a second
240299
// entry that pushes the cache over its total-byte budget evicts the
241300
// least-recently-used entry, which is then re-fetched in full.

0 commit comments

Comments
 (0)