From 2ecb93d8801e49b5cf59cf5d6fb8463984456919 Mon Sep 17 00:00:00 2001
From: roc <roc@imroc.cc>
Date: Tue, 8 Aug 2023 15:30:22 +0800
Subject: [PATCH] Improve cookie jar.

* Add SetCookeJarFactory.
* Use memoryCookieJarFactory to create cookie jar by default when create Client.
* Add some comments.
---
 client.go      | 46 ++++++++++++++++++++++++++++++++++++++--------
 client_test.go | 16 ++++++++++++++++
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/client.go b/client.go
index 4bf54009..c140d1b5 100644
--- a/client.go
+++ b/client.go
@@ -50,6 +50,7 @@ type Client struct {
 	AllowGetMethodPayload bool
 	*Transport
 
+	cookiejarFactory        func() *cookiejar.Jar
 	trace                   bool
 	disableAutoReadResponse bool
 	commonErrorType         reflect.Type
@@ -1022,9 +1023,13 @@ func (c *Client) EnableTraceAll() *Client {
 	return c
 }
 
-// SetCookieJar set the `CookeJar` to the underlying `http.Client`, set to nil if you
-// want to disable cookie.
+// SetCookieJar set the cookie jar to the underlying `http.Client`, set to nil if you
+// want to disable cookies.
+// Note: If you use Client.Clone to clone a new Client, the new client will share the same
+// cookie jar as the old Client after cloning. Use SetCookieJarFactory instead if you want
+// to create a new CookieJar automatically when cloning a client.
 func (c *Client) SetCookieJar(jar http.CookieJar) *Client {
+	c.cookiejarFactory = nil
 	c.httpClient.Jar = jar
 	return c
 }
@@ -1042,11 +1047,10 @@ func (c *Client) GetCookies(url string) ([]*http.Cookie, error) {
 }
 
 // ClearCookies clears all cookies if cookie is enabled.
+// Note: It has no effect if you called SetCookieJar instead of
+// SetCookieJarFactory.
 func (c *Client) ClearCookies() *Client {
-	if c.httpClient.Jar != nil {
-		jar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
-		c.httpClient.Jar = jar
-	}
+	c.initCookieJar()
 	return c
 }
 
@@ -1446,6 +1450,7 @@ func (c *Client) Clone() *Client {
 	client := *c.httpClient
 	client.Transport = cc.Transport
 	cc.httpClient = &client
+	cc.initCookieJar()
 
 	// clone client middleware
 	if len(cc.roundTripWrappers) > 0 {
@@ -1467,14 +1472,17 @@ func (c *Client) Clone() *Client {
 	return &cc
 }
 
+func memoryCookieJarFactory() *cookiejar.Jar {
+	jar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
+	return jar
+}
+
 // C create a new client.
 func C() *Client {
 	t := T()
 
-	jar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
 	httpClient := &http.Client{
 		Transport: t,
-		Jar:       jar,
 		Timeout:   2 * time.Minute,
 	}
 	beforeRequest := []RequestMiddleware{
@@ -1496,13 +1504,35 @@ func C() *Client {
 		jsonUnmarshal:         json.Unmarshal,
 		xmlMarshal:            xml.Marshal,
 		xmlUnmarshal:          xml.Unmarshal,
+		cookiejarFactory:      memoryCookieJarFactory,
 	}
 	httpClient.CheckRedirect = c.defaultCheckRedirect
+	c.initCookieJar()
 
 	c.initTransport()
 	return c
 }
 
+// SetCookieJarFactory set the functional factory of cookie jar, which creates
+// cookie jar that store cookies for underlying `http.Client`. After client clone,
+// the cookie jar of the new client will also be regenerated using this factory
+// function.
+func (c *Client) SetCookieJarFactory(factory func() *cookiejar.Jar) *Client {
+	c.cookiejarFactory = factory
+	c.initCookieJar()
+	return c
+}
+
+func (c *Client) initCookieJar() {
+	if c.cookiejarFactory == nil {
+		return
+	}
+	jar := c.cookiejarFactory()
+	if jar != nil {
+		c.httpClient.Jar = jar
+	}
+}
+
 func (c *Client) initTransport() {
 	c.Debugf = func(format string, v ...interface{}) {
 		if c.DebugLog {
diff --git a/client_test.go b/client_test.go
index 2204f8c2..98df9f55 100644
--- a/client_test.go
+++ b/client_test.go
@@ -7,9 +7,11 @@ import (
 	"errors"
 	"github.com/imroc/req/v3/internal/header"
 	"github.com/imroc/req/v3/internal/tests"
+	"golang.org/x/net/publicsuffix"
 	"io"
 	"net"
 	"net/http"
+	"net/http/cookiejar"
 	"net/url"
 	"os"
 	"strings"
@@ -632,3 +634,17 @@ func TestSetResultStateCheckFunc(t *testing.T) {
 	tests.AssertNoError(t, err)
 	tests.AssertEqual(t, ErrorState, resp.ResultState())
 }
+func TestCloneCookieJar(t *testing.T) {
+	c1 := C()
+	c2 := c1.Clone()
+	tests.AssertEqual(t, true, c1.httpClient.Jar != c2.httpClient.Jar)
+
+	jar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
+	c1.SetCookieJar(jar)
+	c2 = c1.Clone()
+	tests.AssertEqual(t, true, c1.httpClient.Jar == c2.httpClient.Jar)
+
+	c2.SetCookieJar(nil)
+	tests.AssertEqual(t, true, c2.cookiejarFactory == nil)
+	tests.AssertEqual(t, true, c2.httpClient.Jar == nil)
+}