Use Hooks to avoid data races

go-resty · Mar 9, 2023 · 2a97c22 · 2a97c22
1 parent d250cba
commit 2a97c22
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 44 deletions.
diff --git a/client.go b/client.go
@@ -128,30 +128,29 @@ type Client struct {
 	// value when `SetAuthToken` option is used.
 	HeaderAuthorizationKey string
 
-	digestCredentials  *digestCredentials
-	jsonEscapeHTML     bool
-	setContentLength   bool
-	closeConnection    bool
-	notParseResponse   bool
-	trace              bool
-	debugBodySizeLimit int64
-	outputDirectory    string
-	scheme             string
-	log                Logger
-	httpClient         *http.Client
-	proxyURL           *url.URL
-	beforeRequest      []RequestMiddleware
-	udBeforeRequest    []RequestMiddleware
-  udBeforeRequestLock sync.RWMutex
-	preReqHook         PreRequestHook
-	successHooks       []SuccessHook
-	afterResponse      []ResponseMiddleware
-  afterResponseLock   sync.RWMutex
-	requestLog         RequestLogCallback
-	responseLog        ResponseLogCallback
-	errorHooks         []ErrorHook
-	invalidHooks       []ErrorHook
-	panicHooks         []ErrorHook
+	jsonEscapeHTML      bool
+	setContentLength    bool
+	closeConnection     bool
+	notParseResponse    bool
+	trace               bool
+	debugBodySizeLimit  int64
+	outputDirectory     string
+	scheme              string
+	log                 Logger
+	httpClient          *http.Client
+	proxyURL            *url.URL
+	beforeRequest       []RequestMiddleware
+	udBeforeRequest     []RequestMiddleware
+	udBeforeRequestLock sync.RWMutex
+	preReqHook          PreRequestHook
+	successHooks        []SuccessHook
+	afterResponse       []ResponseMiddleware
+	afterResponseLock   sync.RWMutex
+	requestLog          RequestLogCallback
+	responseLog         ResponseLogCallback
+	errorHooks          []ErrorHook
+	invalidHooks        []ErrorHook
+	panicHooks          []ErrorHook
 }
 
 // User type is to hold an username and password information
@@ -406,14 +405,27 @@ func (c *Client) SetAuthScheme(scheme string) *Client {
 // a Digest challenge in the WWW-Authenticate Header, requests will be resent with the appropriate Authorization Header.
 //
 // For Example: To set the Digest scheme with user "Mufasa" and password "Circle Of Life"
-// 		client.SetDigestAuth("Mufasa", "Circle Of Life")
+//
+//	client.SetDigestAuth("Mufasa", "Circle Of Life")
 //
 // Information about Digest Access Authentication can be found in RFC7616:
-//     https://datatracker.ietf.org/doc/html/rfc7616
+//
+//	https://datatracker.ietf.org/doc/html/rfc7616
 //
 // See `Request.SetDigestAuth`.
 func (c *Client) SetDigestAuth(username, password string) *Client {
-	c.digestCredentials = &digestCredentials{username: username, password: password}
+	oldTransport := c.httpClient.Transport
+	c.OnBeforeRequest(func(c *Client, _ *Request) error {
+		c.httpClient.Transport = &digestTransport{
+			digestCredentials: digestCredentials{username, password},
+			transport:         oldTransport,
+		}
+		return nil
+	})
+	c.OnAfterResponse(func(c *Client, _ *Response) error {
+		c.httpClient.Transport = oldTransport
+		return nil
+	})
 	return c
 }
 
@@ -1049,16 +1061,6 @@ func (c *Client) execute(req *Request) (*Response, error) {
 
 	req.RawRequest.Body = newRequestBodyReleaser(req.RawRequest.Body, req.bodyBuf)
 
-	transport := c.httpClient.Transport
-	if req.digestCredentials != nil {
-		c.httpClient.Transport = &digestTransport{digestCredentials: *req.digestCredentials, transport: transport}
-	} else if c.digestCredentials != nil {
-		c.httpClient.Transport = &digestTransport{digestCredentials: *c.digestCredentials, transport: transport}
-	}
-	defer func() {
-		c.httpClient.Transport = transport
-	}()
-
 	req.Time = time.Now()
 	resp, err := c.httpClient.Do(req.RawRequest)
 

diff --git a/request.go b/request.go
@@ -50,7 +50,6 @@ type Request struct {
 	// Since v2.4.0
 	Attempt int
 
-	digestCredentials   *digestCredentials
 	isMultiPart         bool
 	isFormData          bool
 	setContentLength    bool
@@ -507,17 +506,27 @@ func (r *Request) SetAuthScheme(scheme string) *Request {
 // a Digest challenge in the WWW-Authenticate Header, the request will be resent with the appropriate Authorization Header.
 //
 // For Example: To set the Digest scheme with username "Mufasa" and password "Circle Of Life"
-// 		client.R().SetDigestAuth("Mufasa", "Circle Of Life")
+//
+//	client.R().SetDigestAuth("Mufasa", "Circle Of Life")
 //
 // Information about Digest Access Authentication can be found in RFC7616:
-//     https://datatracker.ietf.org/doc/html/rfc7616
+//
+//	https://datatracker.ietf.org/doc/html/rfc7616
 //
 // This method overrides the username and password set by method `Client.SetDigestAuth`.
 func (r *Request) SetDigestAuth(username, password string) *Request {
-	r.digestCredentials = &digestCredentials{
-		username: username,
-		password: password,
-	}
+	oldTransport := r.client.httpClient.Transport
+	r.client.OnBeforeRequest(func(c *Client, _ *Request) error {
+		c.httpClient.Transport = &digestTransport{
+			digestCredentials: digestCredentials{username, password},
+			transport:         oldTransport,
+		}
+		return nil
+	})
+	r.client.OnAfterResponse(func(c *Client, _ *Response) error {
+		c.httpClient.Transport = oldTransport
+		return nil
+	})
 
 	return r
 }