feat: add max reponse body limit

client.go

@@ -128,6 +128,7 @@
 	// HeaderAuthorizationKey is used to set/access Request Authorization header
 	// value when `SetAuthToken` option is used.
 	HeaderAuthorizationKey string
+	ResponseBodyLimit      int
 
 	jsonEscapeHTML      bool
 	setContentLength    bool
@@ -442,11 +443,12 @@
 		RawPathParams: map[string]string{},
 		Debug:         c.Debug,
 
-		client:          c,
-		multipartFiles:  []*File{},
-		multipartFields: []*MultipartField{},
-		jsonEscapeHTML:  c.jsonEscapeHTML,
-		log:             c.log,
+		client:            c,
+		multipartFiles:    []*File{},
+		multipartFields:   []*MultipartField{},
+		jsonEscapeHTML:    c.jsonEscapeHTML,
+		log:               c.log,
+		responseBodyLimit: c.ResponseBodyLimit,
 	}
 	return r
 }
@@ -1089,6 +1091,20 @@
 	return c
 }
 
+// SetResponseBodyLimit set a max body size limit on response, avoid reading too many data to memory.
+//
+// Client will return [resty.ErrResponseBodyTooLarge] if uncompressed response body size if larger than limit.
+// Body size limit will not be enforced in following case:
+//   - ResponseBodyLimit <= 0, which is the default behavior.
+//   - [Request.SetOutput] is called to save a response data to file.
+//   - "DoNotParseResponse" is set for client or request.
+//
+// this can be overridden at client level with [Request.SetResponseBodyLimit]
+func (c *Client) SetResponseBodyLimit(v int) *Client {
+	c.ResponseBodyLimit = v
+	return c
+}
+
 // EnableTrace method enables the Resty client trace for the requests fired from
 // the client using `httptrace.ClientTrace` and provides insights.
 //
@@ -1238,7 +1254,7 @@
 			}
 		}
 
-		if response.body, err = io.ReadAll(body); err != nil {
+		if response.body, err = readAllWithLimit(body, req.responseBodyLimit); err != nil {
 			response.setReceivedAt()
 			return response, err
 		}
@@ -1258,6 +1274,38 @@
 	return response, wrapNoRetryErr(err)
 }
 
+var ErrResponseBodyTooLarge = errors.New("resty: response body too large")
+
+// https://github.com/golang/go/issues/51115
+// [io.LimitedReader] can only return [io.EOF]
+func readAllWithLimit(r io.Reader, maxSize int) ([]byte, error) {
+	if maxSize <= 0 {
+		return io.ReadAll(r)
+	}
+
+	result := make([]byte, 0, 512)
+	total := 0
+	buf := make([]byte, 512)
+	for {
+		n, err := r.Read(buf)
+		total += n
+		if total > maxSize {
+			return nil, ErrResponseBodyTooLarge
+		}
+
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return nil, err
+		}
+
+		result = append(result, buf[:n]...)
+	}
+
+	return result, nil
+}
+
 // getting TLS client config if not exists then create one
 func (c *Client) tlsConfig() (*tls.Config, error) {
 	transport, err := c.Transport()

request.go

@@ -73,6 +73,7 @@ type Request struct {
 	multipartFiles      []*File
 	multipartFields     []*MultipartField
 	retryConditions     []RetryConditionFunc
+	responseBodyLimit   int
 }
 
 // Generate curl command for the request.
@@ -600,6 +601,20 @@ func (r *Request) SetDoNotParseResponse(parse bool) *Request {
 	return r
 }
 
+// SetResponseBodyLimit set a max body size limit on response, avoid reading too many data to memory.
+//
+// Request will return [resty.ErrResponseBodyTooLarge] if uncompressed response body size if larger than limit.
+// Body size limit will not be enforced in following case:
+//   - ResponseBodyLimit <= 0, which is the default behavior.
+//   - [Request.SetOutput] is called to save a response data to file.
+//   - "DoNotParseResponse" is set for client or request.
+//
+// This will override Client config.
+func (r *Request) SetResponseBodyLimit(v int) *Request {
+	r.responseBodyLimit = v
+	return r
+}
+
 // SetPathParam method sets single URL path key-value pair in the
 // Resty current request instance.
 //

retry_test.go

@@ -7,6 +7,7 @@ package resty
 import (
 	"bytes"
 	"context"
+	"crypto/rand"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -823,3 +824,33 @@ func TestResetMultipartReaders(t *testing.T) {
 	assertEqual(t, 500, resp.StatusCode())
 	assertNil(t, err)
 }
+
+func TestResponseBodyLimit(t *testing.T) {
+	ts := createTestServer(func(w http.ResponseWriter, r *http.Request) {
+		io.CopyN(w, rand.Reader, 100*1024)
+	})
+	defer ts.Close()
+
+	t.Run("Client body limit", func(t *testing.T) {
+		c := dc().SetResponseBodyLimit(1024)
+
+		_, err := c.R().Get(ts.URL + "/")
+		assertNotNil(t, err)
+		assertEqual(t, err, ErrResponseBodyTooLarge)
+	})
+
+	t.Run("request body limit", func(t *testing.T) {
+		c := dc()
+
+		_, err := c.R().SetResponseBodyLimit(1024).Get(ts.URL + "/")
+		assertNotNil(t, err)
+		assertEqual(t, err, ErrResponseBodyTooLarge)
+	})
+
+	t.Run("no body limit", func(t *testing.T) {
+		c := dc()
+
+		_, err := c.R().Get(ts.URL + "/")
+		assertNil(t, err)
+	})
+}