Optimize digest auth

* Use transport middleware instead of client middleware to compatible
  with response auto unmarshal (SetSuccessResult / SetErrorResult).
* Support cache, no need to re-authenticate the same site.
* Deprecate request level digest auth.

Signed-off-by: roc <roc@imroc.cc>

Author: imroc

client.go

@@ -49,7 +49,7 @@ type Client struct {
 	DebugLog              bool
 	AllowGetMethodPayload bool
 	*Transport
-
+	digestAuth              *digestAuth
 	cookiejarFactory        func() *cookiejar.Jar
 	trace                   bool
 	disableAutoReadResponse bool
@@ -842,10 +842,14 @@ func (c *Client) SetCommonBasicAuth(username, password string) *Client {
 // Information about Digest Access Authentication can be found in RFC7616:
 //
 //	https://datatracker.ietf.org/doc/html/rfc7616
-//
-// See `Request.SetDigestAuth`
 func (c *Client) SetCommonDigestAuth(username, password string) *Client {
-	c.OnAfterResponse(handleDigestAuthFunc(username, password))
+	c.digestAuth = &digestAuth{
+		Username:   username,
+		Password:   password,
+		HttpClient: c.httpClient,
+		cache:      make(map[string]*cchal),
+	}
+	c.Transport.WrapRoundTripFunc(c.digestAuth.HttpRoundTripWrapperFunc)
 	return c
 }
 

digest.go

@@ -1,12 +1,140 @@
 package req
 
 import (
+	"bytes"
+	"io"
 	"net/http"
+	"sync"
 
 	"github.com/icholy/digest"
 	"github.com/imroc/req/v3/internal/header"
 )
 
+// cchal is a cached challenge and the number of times it's been used.
+type cchal struct {
+	c *digest.Challenge
+	n int
+}
+
+type digestAuth struct {
+	Username   string
+	Password   string
+	HttpClient *http.Client
+	cache      map[string]*cchal
+	cacheMu    sync.Mutex
+}
+
+func (da *digestAuth) digest(req *http.Request, chal *digest.Challenge, count int) (*digest.Credentials, error) {
+	opt := digest.Options{
+		Method:   req.Method,
+		URI:      req.URL.RequestURI(),
+		GetBody:  req.GetBody,
+		Count:    count,
+		Username: da.Username,
+		Password: da.Password,
+	}
+	return digest.Digest(chal, opt)
+}
+
+// challenge returns a cached challenge and count for the provided request
+func (da *digestAuth) challenge(req *http.Request) (*digest.Challenge, int, bool) {
+	da.cacheMu.Lock()
+	defer da.cacheMu.Unlock()
+	host := req.URL.Hostname()
+	cc, ok := da.cache[host]
+	if !ok {
+		return nil, 0, false
+	}
+	cc.n++
+	return cc.c, cc.n, true
+}
+
+// prepare attempts to find a cached challenge that matches the
+// requested domain, and use it to set the Authorization header
+func (da *digestAuth) prepare(req *http.Request) error {
+	// add cookies
+	if da.HttpClient.Jar != nil {
+		for _, cookie := range da.HttpClient.Jar.Cookies(req.URL) {
+			req.AddCookie(cookie)
+		}
+	}
+	// add auth
+	chal, count, ok := da.challenge(req)
+	if !ok {
+		return nil
+	}
+	cred, err := da.digest(req, chal, count)
+	if err != nil {
+		return err
+	}
+	if cred != nil {
+		req.Header.Set("Authorization", cred.String())
+	}
+	return nil
+}
+
+func (da *digestAuth) HttpRoundTripWrapperFunc(rt http.RoundTripper) HttpRoundTripFunc {
+	return func(req *http.Request) (resp *http.Response, err error) {
+		clone, err := cloner(req)
+		if err != nil {
+			return nil, err
+		}
+
+		// make a copy of the request
+		first, err := clone()
+		if err != nil {
+			return nil, err
+		}
+
+		// prepare the first request using a cached challenge
+		if err := da.prepare(first); err != nil {
+			return nil, err
+		}
+
+		// the first request will either succeed or return a 401
+		res, err := rt.RoundTrip(first)
+		if err != nil || res.StatusCode != http.StatusUnauthorized {
+			return res, err
+		}
+
+		// drain and close the first message body
+		_, _ = io.Copy(io.Discard, res.Body)
+		_ = res.Body.Close()
+
+		// find and cache the challenge
+		host := req.URL.Hostname()
+		chal, err := digest.FindChallenge(res.Header)
+		if err != nil {
+			// existing cached challenge didn't work, so remove it
+			da.cacheMu.Lock()
+			delete(da.cache, host)
+			da.cacheMu.Unlock()
+			if err == digest.ErrNoChallenge {
+				return res, nil
+			}
+			return nil, err
+		} else {
+			// found new challenge, so cache it
+			da.cacheMu.Lock()
+			da.cache[host] = &cchal{c: chal}
+			da.cacheMu.Unlock()
+		}
+
+		// make a second copy of the request
+		second, err := clone()
+		if err != nil {
+			return nil, err
+		}
+
+		// prepare the second request based on the new challenge
+		if err := da.prepare(second); err != nil {
+			return nil, err
+		}
+
+		return rt.RoundTrip(second)
+	}
+}
+
 // create response middleware for http digest authentication.
 func handleDigestAuthFunc(username, password string) ResponseMiddleware {
 	return func(client *Client, resp *Response) error {
@@ -60,3 +188,38 @@ func createDigestAuth(req *http.Request, resp *http.Response, username, password
 	}
 	return cred.String(), nil
 }
+
+// cloner returns a function which makes clones of the provided request
+func cloner(req *http.Request) (func() (*http.Request, error), error) {
+	getbody := req.GetBody
+	// if there's no GetBody function set we have to copy the body
+	// into memory to use for future clones
+	if getbody == nil {
+		if req.Body == nil || req.Body == http.NoBody {
+			getbody = func() (io.ReadCloser, error) {
+				return http.NoBody, nil
+			}
+		} else {
+			body, err := io.ReadAll(req.Body)
+			if err != nil {
+				return nil, err
+			}
+			if err := req.Body.Close(); err != nil {
+				return nil, err
+			}
+			getbody = func() (io.ReadCloser, error) {
+				return io.NopCloser(bytes.NewReader(body)), nil
+			}
+		}
+	}
+	return func() (*http.Request, error) {
+		clone := req.Clone(req.Context())
+		body, err := getbody()
+		if err != nil {
+			return nil, err
+		}
+		clone.Body = body
+		clone.GetBody = getbody
+		return clone, nil
+	}, nil
+}

request.go

@@ -431,7 +431,7 @@ func (r *Request) SetBasicAuth(username, password string) *Request {
 //
 //	https://datatracker.ietf.org/doc/html/rfc7616
 //
-// This method overrides the username and password set by method `Client.SetCommonDigestAuth`.
+// Deprecated: Use Client.SetCommonDigestAuth instead. Request level digest auth is not recommended,
 func (r *Request) SetDigestAuth(username, password string) *Request {
 	r.OnAfterResponse(handleDigestAuthFunc(username, password))
 	return r