feat: Add SetMaxResponseSize to limit response body size (#406)

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
+}

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 {

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

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" {

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

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.