Repository: go-resty/resty
Pull request #830 (2024-08-27 23:26:18)

feat: add max reponse body limit
Client not will be able to check resposne body size limit.

If server return a large response than expected, it will discard the response and return a err `ErrResponseBodyTooLarge`. 

```golang
c := Client().SetResponseBodyLimit(1024*1024*1024) // do not allow body larger than 1GiB
res, err := c.R().Get(...)
if err != nil{
  if errors.Is(err, resty.ErrResponseBodyTooLarge){
    ...
  }
}
```

close https://github.com/go-resty/resty/issues/819

jeevatkm: @trim21 Thanks for creating a PR.

I wanted to confirm before I start reviewing it. If you don't mind, your use case is to fail the request if the response body size exceeds the specific size. Can you update the PR description to one use case and its example?
trim21: > @trim21 Thanks for creating a PR.
> 
> I wanted to confirm before I start reviewing it. If you don't mind, your use case is to fail the request if the response body size exceeds the specific size. Can you update the PR description to one use case and its example?

is this ok?
trim21: I'm not sure size limit should be `int` or `int64`.

It's a length limit of `[]byte`, and `len([]byte)` is int, so it would make more sense to be int?
jeevatkm: > I'm not sure size limit should be `int` or `int64`.
> 
> It's a length limit of `[]byte`, and `len([]byte)` is int, so it would make more sense to be int?

@trim21 Yes, for size limit, keep type `int`
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/go-resty/resty/pull/830?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) Report
Attention: Patch coverage is `83.33333%` with `5 lines` in your changes missing coverage. Please review.
> Project coverage is 96.61%. Comparing base [(`f575bf6`)](https://app.codecov.io/gh/go-resty/resty/commit/f575bf6ff178ccf6b23dec8eee2a2484eb25f8f2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) to head [(`aa3fdd1`)](https://app.codecov.io/gh/go-resty/resty/commit/aa3fdd1726124f7a189fe39095b9379ad463efd0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty).

| [Files with missing lines](https://app.codecov.io/gh/go-resty/resty/pull/830?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) | Patch % | Lines |
|---|---|---|
| [client.go](https://app.codecov.io/gh/go-resty/resty/pull/830?src=pr&el=tree&filepath=client.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty#diff-Y2xpZW50Lmdv) | 81.48% | [4 Missing and 1 partial :warning: ](https://app.codecov.io/gh/go-resty/resty/pull/830?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##               v2     #830      +/-   ##
==========================================
- Coverage   96.68%   96.61%   -0.07%     
==========================================
  Files          14       14              
  Lines        1748     1772      +24     
==========================================
+ Hits         1690     1712      +22     
- Misses         37       39       +2     
  Partials       21       21              
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/go-resty/resty/pull/830?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty).

jeevatkm: @trim21 Can you please check this https://github.com/go-resty/resty/pull/830/checks?check_run_id=29386261232 and improve it?
trim21: > @trim21 Can you please check this https://github.com/go-resty/resty/pull/830/checks?check_run_id=29386261232 and improve it?

Wow can I make test client read a error from response?

Patch:
diff --git a/client.go b/client.go
index aeb5507..d33d629 100644
--- a/client.go
+++ b/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()
diff --git a/request.go b/request.go
index 91d432c..cfbe89b 100644
--- a/request.go
+++ b/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.
 //


Test patch:
diff --git a/client_test.go b/client_test.go
index c89d932..acd31d4 100644
--- a/client_test.go
+++ b/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)
+	})
+}
diff --git a/resty_test.go b/resty_test.go
index 22b483c..95ef0b5 100644
--- a/resty_test.go
+++ b/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)
 	}