feat: add max reponse body limit (#830)

client.go

@@ -128,6 +128,7 @@ type Client struct {
 	// 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 @@ func (c *Client) R() *Request {
 		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 @@ func (c *Client) SetJSONEscapeHTML(b bool) *Client {
 	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 @@ func (c *Client) execute(req *Request) (*Response, error) {
 			}
 		}
 
-		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,39 @@ func (c *Client) execute(req *Request) (*Response, error) {
 	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)
+	}
+
+	var buf [512]byte // make buf stack allocated
+	result := make([]byte, 0, 512)
+	total := 0
+	for {
+		n, err := r.Read(buf[:])
+		total += n
+		if total > maxSize {
+			return nil, ErrResponseBodyTooLarge
+		}
+
+		if err != nil {
+			if err == io.EOF {
+				result = append(result, buf[:n]...)
+				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()

client_test.go

@@ -6,6 +6,8 @@ package resty
 
 import (
 	"bytes"
+	"compress/gzip"
+	"crypto/rand"
 	"crypto/tls"
 	"errors"
 	"fmt"
@@ -1097,3 +1099,56 @@ func TestClone(t *testing.T) {
 	assertEqual(t, "clone", parent.UserInfo.Username)
 	assertEqual(t, "clone", clone.UserInfo.Username)
 }
+
+func TestResponseBodyLimit(t *testing.T) {
+	ts := createTestServer(func(w http.ResponseWriter, r *http.Request) {
+		io.CopyN(w, rand.Reader, 100*800)
+	})
+	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("body less than limit", func(t *testing.T) {
+		c := dc()
+
+		res, err := c.R().SetResponseBodyLimit(800*100 + 10).Get(ts.URL + "/")
+		assertNil(t, err)
+		assertEqual(t, 800*100, len(res.body))
+	})
+
+	t.Run("no body limit", func(t *testing.T) {
+		c := dc()
+
+		res, err := c.R().Get(ts.URL + "/")
+		assertNil(t, err)
+		assertEqual(t, 800*100, len(res.body))
+	})
+
+	t.Run("read error", func(t *testing.T) {
+		tse := createTestServer(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set(hdrContentEncodingKey, "gzip")
+			var buf [1024]byte
+			w.Write(buf[:])
+		})
+		defer tse.Close()
+
+		c := dc()
+
+		_, err := c.R().SetResponseBodyLimit(10240).Get(tse.URL + "/")
+		assertErrorIs(t, err, gzip.ErrHeader)
+	})
+}

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.
 //

resty_test.go

@@ -809,6 +809,7 @@ func dclr() *Request {
 }
 
 func assertNil(t *testing.T, v interface{}) {
+	t.Helper()
 	if !isNil(v) {
 		t.Errorf("[%v] was expected to be nil", v)
 	}
@@ -841,6 +842,7 @@ func assertErrorIs(t *testing.T, e, g error) (r bool) {
 }
 
 func assertEqual(t *testing.T, e, g interface{}) (r bool) {
+	t.Helper()
 	if !equal(e, g) {
 		t.Errorf("Expected [%v], got [%v]", e, g)
 	}