Repository: go-resty/resty
Pull request #833 (2024-08-28 16:36:05)

Feature/update bazel config
I updated the bazel configuration
- Use bazel version 6.5.0
note:
When there are changes to the go file, it is recommended to run the following command
``` shell
bazelisk run gazelle
```
If add some deps, run
``` shell
bazelisk run gazelle-update-repos
```

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/go-resty/resty/pull/833?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-resty) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 96.68%. 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 [(`7a0d79e`)](https://app.codecov.io/gh/go-resty/resty/commit/7a0d79e6a9e59ed0b19bf1ed0a11d47801094154?dropdown=coverage&el=desc&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     #833   +/-   ##
=======================================
  Coverage   96.68%   96.68%           
=======================================
  Files          14       14           
  Lines        1748     1748           
=======================================
  Hits         1690     1690           
  Misses         37       37           
  Partials       21       21           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/go-resty/resty/pull/833?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).


Patch:
diff --git a/BUILD.bazel b/BUILD.bazel
index e41515b..47da6c2 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -1,5 +1,5 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 load("@bazel_gazelle//:def.bzl", "gazelle")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 # gazelle:prefix github.com/go-resty/resty/v2
 # gazelle:go_naming_convention import_alias
@@ -17,17 +17,19 @@ go_library(
         "resty.go",
         "retry.go",
         "trace.go",
-        "transport_js.go",
-        "transport_other.go",
         "transport.go",
         "transport112.go",
+        "transport_js.go",
+        "transport_other.go",
         "util.go",
         "util_curl.go",
-        "shellescape/shellescape.go"
     ],
     importpath = "github.com/go-resty/resty/v2",
     visibility = ["//visibility:public"],
-    deps = ["@org_golang_x_net//publicsuffix:go_default_library"],
+    deps = [
+        "//shellescape",
+        "@org_golang_x_net//publicsuffix",
+    ],
 )
 
 go_test(
@@ -36,6 +38,7 @@ go_test(
         "client_test.go",
         "context_test.go",
         "example_test.go",
+        "middleware_test.go",
         "request_test.go",
         "resty_test.go",
         "retry_test.go",
@@ -43,7 +46,10 @@ go_test(
     ],
     data = glob([".testdata/*"]),
     embed = [":resty"],
-    deps = ["@org_golang_x_net//proxy:go_default_library"],
+    deps = [
+        "@org_golang_x_net//proxy",
+        "@org_golang_x_time//rate",
+    ],
 )
 
 alias(
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/examples/BUILD.bazel b/examples/BUILD.bazel
new file mode 100644
index 0000000..849ea4e
--- /dev/null
+++ b/examples/BUILD.bazel
@@ -0,0 +1,10 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_test")
+
+go_test(
+    name = "examples_test",
+    srcs = [
+        "debug_curl_test.go",
+        "server_test.go",
+    ],
+    deps = ["//:resty"],
+)
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.
 //
diff --git a/shellescape/BUILD.bazel b/shellescape/BUILD.bazel
new file mode 100644
index 0000000..fe829e3
--- /dev/null
+++ b/shellescape/BUILD.bazel
@@ -0,0 +1,14 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "shellescape",
+    srcs = ["shellescape.go"],
+    importpath = "github.com/go-resty/resty/v2/shellescape",
+    visibility = ["//visibility:public"],
+)
+
+alias(
+    name = "go_default_library",
+    actual = ":shellescape",
+    visibility = ["//visibility:public"],
+)


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