diff --git a/client.go b/client.go index 2c8acd0..d802fe2 100644 --- a/client.go +++ b/client.go @@ -55,6 +55,7 @@ type Client struct { cookiejarFactory func() http.CookieJar trace bool disableAutoReadResponse bool + maxResponseSize int64 // 0 means no limit commonErrorType reflect.Type retryOption *retryOption jsonMarshal func(v any) ([]byte, error) @@ -810,6 +811,31 @@ func (c *Client) EnableAutoReadResponse() *Client { return c } +// SetMaxResponseSize sets the maximum allowed size of a response body in bytes. +// +// Enforcement: +// - When Response.ContentLength is known and greater than the limit, the body +// is closed without reading and a ResponseBodyTooLargeError is returned. +// Closing early may prevent connection reuse for that request. +// - Otherwise the body is wrapped so application reads stop at the limit. +// - HEAD requests never fail the Content-Length early check (there is no body). +// +// The limit is applied to bytes delivered to the application after the transport +// has handled Content-Encoding (e.g. gzip decompression). For auto-decompressed +// responses ContentLength is typically -1, so only the streaming limit applies. +// Charset auto-decode, if enabled, also runs underneath the limit. +// +// A value of 0 or less disables the limit (default). This is useful for bounding +// memory use and network bandwidth when talking to untrusted or unexpectedly +// large endpoints. +func (c *Client) SetMaxResponseSize(max int64) *Client { + if max < 0 { + max = 0 + } + c.maxResponseSize = max + return c +} + // SetAutoDecodeContentType set the content types that will be auto-detected and decode to utf-8 // (e.g. "json", "xml", "html", "text"). func (c *Client) SetAutoDecodeContentType(contentTypes ...string) *Client { @@ -1802,6 +1828,13 @@ func (c *Client) roundTrip(r *Request) (resp *Response, err error) { httpResponse, resp.Err = c.httpClient.Do(r.RawRequest) resp.Response = httpResponse + // Enforce response body size limit before any body consumption. + if resp.Err == nil { + if err := applyMaxResponseSize(r, resp); err != nil { + resp.Err = err + } + } + // auto-read response body if possible if resp.Err == nil && !c.disableAutoReadResponse && !r.isSaveResponse && !r.disableAutoReadResponse && resp.StatusCode > 199 { resp.ToBytes() @@ -1816,3 +1849,32 @@ func (c *Client) roundTrip(r *Request) (resp *Response, err error) { } return } + +// applyMaxResponseSize rejects oversized responses early when Content-Length is +// known, and otherwise wraps the body so reads stop at the configured limit. +func applyMaxResponseSize(r *Request, resp *Response) error { + max := r.getMaxResponseSize() + if max <= 0 || resp.Response == nil || resp.Body == nil { + return nil + } + + // HEAD keeps Content-Length from the resource header but has no body + // (see transfer.go). Do not treat that advertised length as a body limit + // violation — ParallelDownload relies on Head() + ContentLength for sizing. + if r.Method != http.MethodHead { + // Known Content-Length over the limit: reject without reading the body so + // bandwidth and memory are not wasted. Early close may prevent keep-alive reuse. + if cl := resp.ContentLength; cl > max { + _ = resp.Body.Close() + resp.Body = http.NoBody + return &ResponseBodyTooLargeError{Limit: max, ContentLength: cl} + } + } + + resp.Body = &maxResponseBodyReader{ + r: resp.Body, + n: max, + limit: max, + } + return nil +} diff --git a/client_wrapper.go b/client_wrapper.go index 74c5312..cd4a9ca 100644 --- a/client_wrapper.go +++ b/client_wrapper.go @@ -363,6 +363,12 @@ func EnableAutoReadResponse() *Client { return defaultClient.EnableAutoReadResponse() } +// SetMaxResponseSize is a global wrapper methods which delegated +// to the default client's Client.SetMaxResponseSize. +func SetMaxResponseSize(max int64) *Client { + return defaultClient.SetMaxResponseSize(max) +} + // SetAutoDecodeContentType is a global wrapper methods which delegated // to the default client's Client.SetAutoDecodeContentType. func SetAutoDecodeContentType(contentTypes ...string) *Client { diff --git a/middleware.go b/middleware.go index dde3f0e..8686c57 100644 --- a/middleware.go +++ b/middleware.go @@ -473,6 +473,13 @@ func handleDownload(c *Client, r *Response) (err error) { if r.Response == nil || !r.Request.isSaveResponse { return nil } + // Content-Length early reject replaces the body with http.NoBody. Skip + // download in that case so we do not create an empty output file. Other + // prior errors (e.g. result unmarshalling) still allow the download path + // when the body was already buffered. + if errors.Is(r.Err, ErrResponseBodyTooLarge) && r.body == nil { + return nil + } var body io.ReadCloser if r.body != nil { // already read diff --git a/req_test.go b/req_test.go index 9942bcb..6563b1d 100644 --- a/req_test.go +++ b/req_test.go @@ -1,6 +1,7 @@ package req import ( + "bytes" "encoding/json" "encoding/xml" "fmt" @@ -47,7 +48,7 @@ func createTestServer() *httptest.Server { func handleHTTP(w http.ResponseWriter, r *http.Request) { w.Header().Add("Method", r.Method) switch r.Method { - case http.MethodGet: + case http.MethodGet, http.MethodHead: handleGet(w, r) case http.MethodPost: handlePost(w, r) @@ -330,6 +331,56 @@ func handleGet(w http.ResponseWriter, r *http.Request) { } i += n } + case "/fixed-size": + // Returns a body of exactly the given size with Content-Length set. + // Query: size (default 1024). HEAD advertises Content-Length without a body. + r.ParseForm() + size := 1024 + if s := r.FormValue("size"); s != "" { + if n, err := strconv.Atoi(s); err == nil && n >= 0 { + size = n + } + } + w.Header().Set("Content-Length", strconv.Itoa(size)) + w.Header().Set(header.ContentType, "application/octet-stream") + if r.Method == http.MethodHead || size == 0 { + return + } + buf := bytes.Repeat([]byte{'x'}, min(size, 4096)) + for written := 0; written < size; { + n := min(len(buf), size-written) + wn, err := w.Write(buf[:n]) + if err != nil { + return + } + written += wn + } + case "/chunked-size": + // Returns a body of the given size without Content-Length (chunked). + r.ParseForm() + size := 1024 + if s := r.FormValue("size"); s != "" { + if n, err := strconv.Atoi(s); err == nil && n >= 0 { + size = n + } + } + w.Header().Set(header.ContentType, "application/octet-stream") + if r.Method == http.MethodHead { + return + } + // Ensure chunked encoding by not setting Content-Length and flushing. + flusher, _ := w.(http.Flusher) + buf := bytes.Repeat([]byte{'y'}, min(size, 4096)) + for written := 0; written < size; { + n := min(len(buf), size-written) + if _, err := w.Write(buf[:n]); err != nil { + return + } + if flusher != nil { + flusher.Flush() + } + written += n + } case "/protected": auth := r.Header.Get("Authorization") if auth == "Bearer goodtoken" { diff --git a/request.go b/request.go index 7435d1f..380bf84 100644 --- a/request.go +++ b/request.go @@ -55,20 +55,23 @@ type Request struct { uploadCallbackInterval time.Duration downloadCallback DownloadCallback downloadCallbackInterval time.Duration - unReplayableBody io.ReadCloser - retryOption *retryOption - bodyReadCloser io.ReadCloser - dumpOptions *DumpOptions - marshalBody any - ctx context.Context - uploadFiles []*FileUpload - uploadReader []io.ReadCloser - outputFile string - output io.Writer - trace *clientTrace - dumpBuffer *bytes.Buffer - responseReturnTime time.Time - afterResponse []ResponseMiddleware + // maxResponseSize, when non-nil, overrides Client.maxResponseSize for this + // request. A pointed-to value of 0 means no limit for this request. + maxResponseSize *int64 + unReplayableBody io.ReadCloser + retryOption *retryOption + bodyReadCloser io.ReadCloser + dumpOptions *DumpOptions + marshalBody any + ctx context.Context + uploadFiles []*FileUpload + uploadReader []io.ReadCloser + outputFile string + output io.Writer + trace *clientTrace + dumpBuffer *bytes.Buffer + responseReturnTime time.Time + afterResponse []ResponseMiddleware } type GetContentFunc func() (io.ReadCloser, error) @@ -1047,6 +1050,31 @@ func (r *Request) EnableAutoReadResponse() *Request { return r } +// SetMaxResponseSize sets the maximum allowed size of the response body in bytes +// for this request, overriding Client.SetMaxResponseSize. A value of 0 or less +// disables the limit for this request even if the client has a limit configured. +// +// See Client.SetMaxResponseSize for behavior details. +func (r *Request) SetMaxResponseSize(max int64) *Request { + if max < 0 { + max = 0 + } + r.maxResponseSize = &max + return r +} + +// getMaxResponseSize returns the effective max response body size for this +// request (request override, else client setting). 0 means no limit. +func (r *Request) getMaxResponseSize() int64 { + if r.maxResponseSize != nil { + return *r.maxResponseSize + } + if r.client != nil { + return r.client.maxResponseSize + } + return 0 +} + // DisableTrace disables trace. func (r *Request) DisableTrace() *Request { r.trace = nil diff --git a/response.go b/response.go index 2fb3ca1..556ddc2 100644 --- a/response.go +++ b/response.go @@ -2,6 +2,7 @@ package req import ( "errors" + "fmt" "io" "net/http" "strings" @@ -11,6 +12,73 @@ import ( "github.com/imroc/req/v3/internal/util" ) +// ErrResponseBodyTooLarge is returned when a response body exceeds the limit +// configured via Client.SetMaxResponseSize or Request.SetMaxResponseSize. +// Use errors.Is(err, ErrResponseBodyTooLarge) or errors.As with +// *ResponseBodyTooLargeError for inspection. +var ErrResponseBodyTooLarge = errors.New("req: response body too large") + +// ResponseBodyTooLargeError is returned when a response body exceeds the +// configured maximum size. Limit is the configured max in bytes. ContentLength +// is the response's Content-Length when the body was rejected based on headers +// without reading; it is -1 when the limit was exceeded while reading. +type ResponseBodyTooLargeError struct { + Limit int64 + ContentLength int64 +} + +func (e *ResponseBodyTooLargeError) Error() string { + if e.ContentLength >= 0 { + return fmt.Sprintf("req: response body too large: Content-Length %d exceeds limit %d", e.ContentLength, e.Limit) + } + return fmt.Sprintf("req: response body too large: exceeds limit of %d bytes", e.Limit) +} + +// Is reports whether target is ErrResponseBodyTooLarge. +func (e *ResponseBodyTooLargeError) Is(target error) bool { + return target == ErrResponseBodyTooLarge +} + +// maxResponseBodyReader limits how many bytes can be read from a response body. +// It is similar in spirit to http.MaxBytesReader but for client response bodies: +// when the limit is exceeded it returns *ResponseBodyTooLargeError and subsequent +// reads return the same sticky error. Close closes the underlying body. +type maxResponseBodyReader struct { + r io.ReadCloser + n int64 // bytes remaining + limit int64 // original limit, for the error + err error // sticky error +} + +func (l *maxResponseBodyReader) Read(p []byte) (n int, err error) { + if l.err != nil { + return 0, l.err + } + if len(p) == 0 { + return 0, nil + } + // Read at most remaining+1 so we can detect crossing the limit. + if int64(len(p))-1 > l.n { + p = p[:l.n+1] + } + n, err = l.r.Read(p) + + if int64(n) <= l.n { + l.n -= int64(n) + l.err = err + return n, err + } + + n = int(l.n) + l.n = 0 + l.err = &ResponseBodyTooLargeError{Limit: l.limit, ContentLength: -1} + return n, l.err +} + +func (l *maxResponseBodyReader) Close() error { + return l.r.Close() +} + // Response is the http response. type Response struct { // The underlying http.Response is embed into Response. diff --git a/response_maxsize_test.go b/response_maxsize_test.go new file mode 100644 index 0000000..815c5c0 --- /dev/null +++ b/response_maxsize_test.go @@ -0,0 +1,396 @@ +package req + +import ( + "bytes" + "errors" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/imroc/req/v3/internal/tests" +) + +func TestMaxResponseSizeWithinLimit(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + // "TestGet: text response" is 23 bytes + resp, err := c.SetMaxResponseSize(100).R().Get("/") + assertSuccess(t, resp, err) + tests.AssertEqual(t, "TestGet: text response", resp.String()) + }) +} + +func TestMaxResponseSizeExactLimit(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + const size = 64 + resp, err := c.SetMaxResponseSize(size).R().Get("/fixed-size?size=64") + assertSuccess(t, resp, err) + tests.AssertEqual(t, size, len(resp.Bytes())) + }) +} + +func TestMaxResponseSizeContentLengthExceeded(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + resp, err := c.SetMaxResponseSize(10).R().Get("/fixed-size?size=100") + if err == nil { + t.Fatal("expected error when Content-Length exceeds limit") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + var tooLarge *ResponseBodyTooLargeError + if !errors.As(err, &tooLarge) { + t.Fatalf("expected *ResponseBodyTooLargeError, got %T", err) + } + tests.AssertEqual(t, int64(10), tooLarge.Limit) + tests.AssertEqual(t, int64(100), tooLarge.ContentLength) + // Response headers/status should still be available. + tests.AssertEqual(t, http.StatusOK, resp.StatusCode) + // Body must not have been buffered. + tests.AssertEqual(t, 0, len(resp.Bytes())) + }) +} + +func TestMaxResponseSizeChunkedExceeded(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + // Chunked response has no Content-Length; limit is enforced while reading. + resp, err := c.SetMaxResponseSize(50).R().Get("/chunked-size?size=200") + if err == nil { + t.Fatal("expected error when chunked body exceeds limit") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + var tooLarge *ResponseBodyTooLargeError + if !errors.As(err, &tooLarge) { + t.Fatalf("expected *ResponseBodyTooLargeError, got %T", err) + } + tests.AssertEqual(t, int64(50), tooLarge.Limit) + tests.AssertEqual(t, int64(-1), tooLarge.ContentLength) + tests.AssertEqual(t, http.StatusOK, resp.StatusCode) + // Partial body may have been read up to the limit. + if len(resp.Bytes()) > 50 { + t.Fatalf("buffered body larger than limit: %d", len(resp.Bytes())) + } + }) +} + +func TestMaxResponseSizeChunkedWithinLimit(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + resp, err := c.SetMaxResponseSize(500).R().Get("/chunked-size?size=100") + assertSuccess(t, resp, err) + tests.AssertEqual(t, 100, len(resp.Bytes())) + }) +} + +func TestMaxResponseSizeRequestOverridesClient(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + c.SetMaxResponseSize(10) + // Request-level higher limit allows the larger body. + resp, err := c.R().SetMaxResponseSize(200).Get("/fixed-size?size=100") + assertSuccess(t, resp, err) + tests.AssertEqual(t, 100, len(resp.Bytes())) + + // Request-level lower limit rejects a body the client would accept. + c.SetMaxResponseSize(1000) + resp, err = c.R().SetMaxResponseSize(10).Get("/fixed-size?size=100") + if err == nil { + t.Fatal("expected request-level limit to reject response") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + _ = resp + }) +} + +func TestMaxResponseSizeRequestDisablesClientLimit(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + c.SetMaxResponseSize(10) + // 0 on the request disables the limit for this request. + resp, err := c.R().SetMaxResponseSize(0).Get("/fixed-size?size=100") + assertSuccess(t, resp, err) + tests.AssertEqual(t, 100, len(resp.Bytes())) + }) +} + +func TestMaxResponseSizeDisableAutoRead(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + c.SetMaxResponseSize(50).DisableAutoReadResponse() + resp, err := c.R().Get("/chunked-size?size=200") + // Headers succeed; body is not auto-read. + tests.AssertNoError(t, err) + tests.AssertEqual(t, http.StatusOK, resp.StatusCode) + + // Manual read hits the limit. + _, readErr := io.ReadAll(resp.Body) + if readErr == nil { + t.Fatal("expected error when reading oversized body manually") + } + if !errors.Is(readErr, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", readErr) + } + _ = resp.Body.Close() + }) +} + +func TestMaxResponseSizeToBytes(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + c.SetMaxResponseSize(50).DisableAutoReadResponse() + resp, err := c.R().Get("/chunked-size?size=200") + tests.AssertNoError(t, err) + + _, err = resp.ToBytes() + if err == nil { + t.Fatal("expected ToBytes to fail when body exceeds limit") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + if !errors.Is(resp.Err, ErrResponseBodyTooLarge) { + t.Fatalf("expected resp.Err to be ErrResponseBodyTooLarge, got %v", resp.Err) + } + }) +} + +func TestMaxResponseSizeDownload(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + dir := t.TempDir() + out := filepath.Join(dir, "out.bin") + + // Content-Length over limit: reject without writing a full file. + resp, err := c.SetMaxResponseSize(10).R(). + SetOutputFile(out). + Get("/fixed-size?size=100") + if err == nil { + t.Fatal("expected download to fail when Content-Length exceeds limit") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + _ = resp + // File should not have been created (handleDownload skipped on prior error). + if _, statErr := os.Stat(out); !os.IsNotExist(statErr) { + t.Fatalf("expected output file not to be created, stat err=%v", statErr) + } + + // Chunked body over limit: download fails mid-stream. + out2 := filepath.Join(dir, "out2.bin") + resp, err = c.SetMaxResponseSize(50).R(). + SetOutputFile(out2). + Get("/chunked-size?size=200") + if err == nil { + t.Fatal("expected download to fail when body exceeds limit during read") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + _ = resp + }) +} + +func TestMaxResponseSizeDownloadWriter(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + var buf bytes.Buffer + resp, err := c.SetMaxResponseSize(50).R(). + SetOutput(&buf). + Get("/chunked-size?size=200") + if err == nil { + t.Fatal("expected error when writing oversized body to output") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + if buf.Len() > 50 { + t.Fatalf("wrote more than limit: %d", buf.Len()) + } + _ = resp + }) +} + +func TestMaxResponseSizeLargeContentLengthNoFullRead(t *testing.T) { + // Regression for the original bandwidth concern: a huge Content-Length must + // be rejected without buffering the body. + testWithAllTransport(t, func(t *testing.T, c *Client) { + resp, err := c.SetMaxResponseSize(1024).R().Get("/download") + if err == nil { + t.Fatal("expected error for 100MiB Content-Length with 1KiB limit") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + var tooLarge *ResponseBodyTooLargeError + if !errors.As(err, &tooLarge) { + t.Fatalf("expected *ResponseBodyTooLargeError, got %T", err) + } + tests.AssertEqual(t, int64(1024), tooLarge.Limit) + if tooLarge.ContentLength <= 1024 { + t.Fatalf("expected ContentLength > limit, got %d", tooLarge.ContentLength) + } + tests.AssertEqual(t, 0, len(resp.Bytes())) + }) +} + +func TestMaxResponseSizeZeroMeansUnlimited(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + resp, err := c.SetMaxResponseSize(0).R().Get("/fixed-size?size=100") + assertSuccess(t, resp, err) + tests.AssertEqual(t, 100, len(resp.Bytes())) + }) +} + +func TestMaxResponseSizeNegativeTreatedAsUnlimited(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + resp, err := c.SetMaxResponseSize(-1).R().Get("/fixed-size?size=100") + assertSuccess(t, resp, err) + tests.AssertEqual(t, 100, len(resp.Bytes())) + }) +} + +func TestMaxResponseSizeClone(t *testing.T) { + c := tc().SetMaxResponseSize(42) + cc := c.Clone() + tests.AssertEqual(t, int64(42), cc.maxResponseSize) + + // Changing the clone must not affect the original. + cc.SetMaxResponseSize(99) + tests.AssertEqual(t, int64(42), c.maxResponseSize) + tests.AssertEqual(t, int64(99), cc.maxResponseSize) +} + +func TestMaxResponseSizeErrorMessage(t *testing.T) { + e1 := &ResponseBodyTooLargeError{Limit: 10, ContentLength: 100} + if !strings.Contains(e1.Error(), "Content-Length 100") { + t.Fatalf("unexpected error message: %s", e1.Error()) + } + if !strings.Contains(e1.Error(), "limit 10") { + t.Fatalf("unexpected error message: %s", e1.Error()) + } + + e2 := &ResponseBodyTooLargeError{Limit: 50, ContentLength: -1} + if !strings.Contains(e2.Error(), "exceeds limit of 50") { + t.Fatalf("unexpected error message: %s", e2.Error()) + } + if !errors.Is(e1, ErrResponseBodyTooLarge) || !errors.Is(e2, ErrResponseBodyTooLarge) { + t.Fatal("errors.Is should match ErrResponseBodyTooLarge") + } +} + +func TestMaxResponseBodyReaderStickyError(t *testing.T) { + r := &maxResponseBodyReader{ + r: io.NopCloser(bytes.NewReader(bytes.Repeat([]byte{'a'}, 100))), + n: 10, + limit: 10, + } + buf := make([]byte, 64) + n, err := r.Read(buf) + if err == nil { + t.Fatal("expected error on first oversized read") + } + if n > 10 { + t.Fatalf("read more than limit: %d", n) + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + // Sticky: subsequent reads return the same error. + _, err2 := r.Read(buf) + if !errors.Is(err2, ErrResponseBodyTooLarge) { + t.Fatalf("expected sticky ErrResponseBodyTooLarge, got %v", err2) + } +} + +func TestMaxResponseBodyReaderExactLimit(t *testing.T) { + data := bytes.Repeat([]byte{'b'}, 10) + r := &maxResponseBodyReader{ + r: io.NopCloser(bytes.NewReader(data)), + n: 10, + limit: 10, + } + got, err := io.ReadAll(r) + tests.AssertNoError(t, err) + tests.AssertEqual(t, data, got) +} + +func TestMaxResponseSizeWithSuccessResult(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + var result map[string]string + resp, err := c.SetMaxResponseSize(5).R(). + SetSuccessResult(&result). + Get("/json") + if err == nil { + t.Fatal("expected error when JSON body exceeds limit") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + _ = resp + }) +} + +func TestMaxResponseSizeHEADDoesNotEarlyReject(t *testing.T) { + // HEAD responses advertise Content-Length of the resource but carry no body. + // The limit must not reject them, or ParallelDownload size probes break. + testWithAllTransport(t, func(t *testing.T, c *Client) { + resp, err := c.SetMaxResponseSize(10).R().Head("/fixed-size?size=100") + assertSuccess(t, resp, err) + tests.AssertEqual(t, int64(100), resp.ContentLength) + tests.AssertEqual(t, 0, len(resp.Bytes())) + }) +} + +func TestMaxResponseSizeDownloadStillRunsAfterUnmarshalError(t *testing.T) { + // Combining SetOutput with SetSuccessResult: if unmarshalling fails, the + // download path should still write the buffered body (pre-existing contract). + testWithAllTransport(t, func(t *testing.T, c *Client) { + var buf bytes.Buffer + var result struct { + Missing string `json:"missing_required_shape"` + } + // Valid JSON that unmarshals into the empty struct without error... use + // invalid JSON endpoint would 404. Use /fixed-size which is not JSON so + // unmarshal fails, while body is auto-read and available for download. + resp, err := c.R(). + SetSuccessResult(&result). + SetOutput(&buf). + Get("/fixed-size?size=32") + // Unmarshal of non-JSON body should fail. + if err == nil { + // Some configs may not set Result when status is success but content is not JSON — + // ensure we at least got a response and the body was still written to output. + t.Logf("unmarshal did not error (err=nil); body written=%d status=%d", buf.Len(), resp.StatusCode) + } + if buf.Len() != 32 { + // If handleDownload was skipped incorrectly, buf would be empty. + // parseResponseBody runs first and may leave body in resp.body; + // handleDownload should still copy it. + if buf.Len() == 0 && resp != nil && len(resp.Bytes()) == 32 { + t.Fatal("download was skipped after non-size error; body was buffered but not written to output") + } + if buf.Len() != 32 { + t.Fatalf("expected 32 bytes written to output, got %d (err=%v)", buf.Len(), err) + } + } + }) +} + +func TestMaxResponseSizeContentLengthEarlyRejectSkipsDownload(t *testing.T) { + testWithAllTransport(t, func(t *testing.T, c *Client) { + var buf bytes.Buffer + resp, err := c.SetMaxResponseSize(10).R(). + SetOutput(&buf). + Get("/fixed-size?size=100") + if err == nil { + t.Fatal("expected size limit error") + } + if !errors.Is(err, ErrResponseBodyTooLarge) { + t.Fatalf("expected ErrResponseBodyTooLarge, got %v", err) + } + if buf.Len() != 0 { + t.Fatalf("expected no download output on Content-Length early reject, got %d bytes", buf.Len()) + } + _ = resp + }) +}