support update request in retry hook (#98)

client.go

@@ -929,43 +929,43 @@ func (c *Client) do(r *Request) (resp *Response, err error) {
 		Request: r,
 	}
 
-	for _, f := range r.client.udBeforeRequest {
-		if err = f(r.client, r); err != nil {
-			return
+	for {
+		for _, f := range r.client.udBeforeRequest {
+			if err = f(r.client, r); err != nil {
+				return
+			}
 		}
-	}
-	for _, f := range r.client.beforeRequest {
-		if err = f(r.client, r); err != nil {
-			return
+		for _, f := range r.client.beforeRequest {
+			if err = f(r.client, r); err != nil {
+				return
+			}
 		}
-	}
 
-	// setup trace
-	if r.trace == nil && r.client.trace {
-		r.trace = &clientTrace{}
-	}
-	if r.trace != nil {
-		r.ctx = r.trace.createContext(r.Context())
-	}
+		// setup trace
+		if r.trace == nil && r.client.trace {
+			r.trace = &clientTrace{}
+		}
+		if r.trace != nil {
+			r.ctx = r.trace.createContext(r.Context())
+		}
 
-	// setup url and host
-	var host string
-	if h := r.getHeader("Host"); h != "" {
-		host = h // Host header override
-	} else {
-		host = r.URL.Host
-	}
+		// setup url and host
+		var host string
+		if h := r.getHeader("Host"); h != "" {
+			host = h // Host header override
+		} else {
+			host = r.URL.Host
+		}
 
-	// setup header
-	var header http.Header
-	if r.Headers == nil {
-		header = make(http.Header)
-	} else {
-		header = r.Headers.Clone()
-	}
-	contentLength := int64(len(r.body))
+		// setup header
+		var header http.Header
+		if r.Headers == nil {
+			header = make(http.Header)
+		} else {
+			header = r.Headers.Clone()
+		}
+		contentLength := int64(len(r.body))
 
-	for {
 		var reqBody io.ReadCloser
 		if r.getBody != nil {
 			reqBody, err = r.getBody()
@@ -1036,10 +1036,12 @@ func (c *Client) do(r *Request) (resp *Response, err error) {
 		}
 		time.Sleep(r.retryOption.GetRetryInterval(resp, r.RetryAttempt))
 
-		// clean buffers
+		// clean up before retry
 		if r.dumpBuffer != nil {
 			r.dumpBuffer.Reset()
 		}
+		r.trace = nil
+		r.ctx = nil
 		resp.body = nil
 	}
 

req_test.go

@@ -297,9 +297,6 @@ func handleGet(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusTooManyRequests)
 		w.Header().Set(hdrContentTypeKey, jsonContentType)
 		w.Write([]byte(`{"errMsg":"too many requests"}`))
-	case "/retry":
-		r.ParseForm()
-		r.Form.Get("attempt")
 	case "/chunked":
 		w.Header().Add("Trailer", "Expires")
 		w.Write([]byte(`This is a chunked body`))
@@ -356,6 +353,14 @@ func handleGet(w http.ResponseWriter, r *http.Request) {
 		w.Write([]byte(r.URL.RawQuery))
 	case "/search":
 		handleSearch(w, r)
+	case "/protected":
+		auth := r.Header.Get("Authorization")
+		if auth == "Bearer goodtoken" {
+			w.Write([]byte("good"))
+		} else {
+			w.WriteHeader(http.StatusUnauthorized)
+			w.Write([]byte(`bad`))
+		}
 	default:
 		if strings.HasPrefix(r.URL.Path, "/user") {
 			handleGetUserProfile(w, r)

retry_test.go

@@ -139,3 +139,22 @@ func TestRetryWithSetResult(t *testing.T) {
 	assertSuccess(t, resp, err)
 	assertEqual(t, "test=test", headers.Get("Cookie"))
 }
+
+func TestRetryWithModify(t *testing.T) {
+	tokens := []string{"badtoken1", "badtoken2", "goodtoken"}
+	tokenIndex := 0
+	c := tc().EnableDumpAll().
+		SetCommonRetryCount(2).
+		SetCommonRetryHook(func(resp *Response, err error) {
+			tokenIndex++
+			resp.Request.SetBearerAuthToken(tokens[tokenIndex])
+		}).SetCommonRetryCondition(func(resp *Response, err error) bool {
+		return err != nil || resp.StatusCode == http.StatusUnauthorized
+	})
+
+	resp, err := c.R().
+		SetBearerAuthToken(tokens[tokenIndex]).
+		Get("/protected")
+	assertSuccess(t, resp, err)
+	assertEqual(t, 2, resp.Request.RetryAttempt)
+}