Fix redirect request body mismatch with origin request body caused by…

… buffer reused by mistake (#568) * Fix request body buffer reused by mistake when redirect * Add testcase for POST redirect with original method and body * Use buffer from the pool --------- Co-authored-by: Jeevanandam M <jeeva@myjeeva.com>
go-resty · Mar 8, 2023 · 5ebbbb1 · 5ebbbb1
1 parent 15c3640
commit 5ebbbb1
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 7 deletions.
diff --git a/client_test.go b/client_test.go
@@ -17,6 +17,7 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 )
@@ -858,3 +859,27 @@ func TestHostURLForGH318AndGH407(t *testing.T) {
 	assertNil(t, err)
 	assertNotNil(t, resp)
 }
+
+func TestPostRedirectWithBody(t *testing.T) {
+	ts := createPostServer(t)
+	defer ts.Close()
+
+	targetURL, _ := url.Parse(ts.URL)
+	t.Log("ts.URL:", ts.URL)
+	t.Log("targetURL.Host:", targetURL.Host)
+
+	c := dc()
+	wg := sync.WaitGroup{}
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			resp, err := c.R().
+				SetBody([]byte(strconv.Itoa(newRnd().Int()))).
+				Post(targetURL.String() + "/redirect-with-body")
+			assertError(t, err)
+			assertNotNil(t, resp)
+		}()
+	}
+	wg.Wait()
+}
diff --git a/middleware.go b/middleware.go
@@ -514,7 +514,12 @@ func saveResponseIntoFile(c *Client, res *Response) error {
 func getBodyCopy(r *Request) (*bytes.Buffer, error) {
 	// If r.bodyBuf present, return the copy
 	if r.bodyBuf != nil {
-		return bytes.NewBuffer(r.bodyBuf.Bytes()), nil
+		bodyCopy := acquireBuffer()
+		if _, err := io.Copy(bodyCopy, bytes.NewReader(r.bodyBuf.Bytes())); err != nil {
+			// cannot use io.Copy(bodyCopy, r.bodyBuf) because io.Copy reset r.bodyBuf
+			return nil, err
+		}
+		return bodyCopy, nil
 	}
 
 	// Maybe body is `io.Reader`.

diff --git a/resty_test.go b/resty_test.go
@@ -14,6 +14,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"os"
 	"path/filepath"
 	"reflect"
@@ -240,15 +241,13 @@ func createPostServer(t *testing.T) *httptest.Server {
 			handleLoginEndpoint(t, w, r)
 
 			handleUsersEndpoint(t, w, r)
-
-			if r.URL.Path == "/login-json-html" {
+			switch r.URL.Path {
+			case "/login-json-html":
 				w.Header().Set(hdrContentTypeKey, "text/html")
 				w.WriteHeader(http.StatusOK)
 				_, _ = w.Write([]byte(`<htm><body>Test JSON request with HTML response</body></html>`))
 				return
-			}
-
-			if r.URL.Path == "/usersmap" {
+			case "/usersmap":
 				// JSON
 				if IsJSONType(r.Header.Get(hdrContentTypeKey)) {
 					if r.URL.Query().Get("status") == "500" {
@@ -287,9 +286,19 @@ func createPostServer(t *testing.T) *httptest.Server {
 
 					return
 				}
-			} else if r.URL.Path == "/redirect" {
+			case "/redirect":
 				w.Header().Set(hdrLocationKey, "/login")
 				w.WriteHeader(http.StatusTemporaryRedirect)
+			case "/redirect-with-body":
+				body, _ := ioutil.ReadAll(r.Body)
+				query := url.Values{}
+				query.Add("body", string(body))
+				w.Header().Set(hdrLocationKey, "/redirected-with-body?"+query.Encode())
+				w.WriteHeader(http.StatusTemporaryRedirect)
+			case "/redirected-with-body":
+				body, _ := ioutil.ReadAll(r.Body)
+				assertEqual(t, r.URL.Query().Get("body"), string(body))
+				w.WriteHeader(http.StatusOK)
 			}
 		}
 	})