修改 cookie 获取逻辑 ⚡ (#168)

bot.go

@@ -6,7 +6,6 @@ import (
 	"errors"
 	"io"
 	"log"
-	"net/url"
 	"os/exec"
 	"runtime"
 	"sync"
@@ -111,14 +110,7 @@ func (b *Bot) hotLogin(storage HotReloadStorage) error {
 
 // 热登陆初始化
 func (b *Bot) hotLoginInit(item *HotReloadStorageItem) error {
-	cookies := item.Cookies
-	for u, ck := range cookies {
-		path, err := url.Parse(u)
-		if err != nil {
-			return err
-		}
-		b.Caller.Client.Jar.SetCookies(path, ck)
-	}
+	b.Caller.Client.Jar = item.Jar.AsCookieJar()
 	b.Storage.LoginInfo = item.LoginInfo
 	b.Storage.Request = item.BaseRequest
 	b.Caller.Client.Domain = item.WechatDomain
@@ -365,10 +357,10 @@ func (b *Bot) DumpHotReloadStorage() error {
 // DumpTo 将热登录需要的数据写入到指定的 io.Writer 中
 // 注: 写之前最好先清空之前的数据
 func (b *Bot) DumpTo(writer io.Writer) error {
-	cookies := b.Caller.Client.GetCookieMap()
+	cookies := b.Caller.Client.GetCookieJar()
 	item := HotReloadStorageItem{
 		BaseRequest:  b.Storage.Request,
-		Cookies:      cookies,
+		Jar:          cookies,
 		LoginInfo:    b.Storage.LoginInfo,
 		WechatDomain: b.Caller.Client.Domain,
 		UUID:         b.uuid,

client.go

@@ -9,7 +9,6 @@ import (
 	"io"
 	"mime/multipart"
 	"net/http"
-	"net/http/cookiejar"
 	"net/url"
 	"os"
 	"strconv"
@@ -48,14 +47,13 @@ type Client struct {
 }
 
 func NewClient() *Client {
-	jar, _ := cookiejar.New(nil)
 	timeout := 30 * time.Second
 	return &Client{
 		Client: &http.Client{
 			CheckRedirect: func(req *http.Request, via []*http.Request) error {
 				return http.ErrUseLastResponse
 			},
-			Jar:     jar,
+			Jar:     newCookieJar(),
 			Timeout: timeout,
 		}}
 }
@@ -96,30 +94,15 @@ func (c *Client) do(req *http.Request) (*http.Response, error) {
 	return resp, err
 }
 
-func (c *Client) setCookie(resp *http.Response) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	cookies := resp.Cookies()
-	if c.cookies == nil {
-		c.cookies = make(map[string][]*http.Cookie)
-	}
-	path := fmt.Sprintf("%s://%s%s", resp.Request.URL.Scheme, resp.Request.URL.Host, resp.Request.URL.Path)
-	c.cookies[path] = cookies
-}
-
 // Do 抽象Do方法,将所有的有效的cookie存入Client.cookies
 // 方便热登陆时获取
 func (c *Client) Do(req *http.Request) (*http.Response, error) {
-	resp, err := c.do(req)
-	if err == nil {
-		c.setCookie(resp)
-	}
-	return resp, err
+	return c.do(req)
 }
 
-// GetCookieMap 获取当前client的所有的有效的client
-func (c *Client) GetCookieMap() map[string][]*http.Cookie {
-	return c.cookies
+// GetCookieJar 获取当前client的所有的有效的client
+func (c *Client) GetCookieJar() *Jar {
+	return fromCookieJar(c.Client.Jar)
 }
 
 // GetLoginUUID 获取登录的uuid

jar.go

@@ -0,0 +1,59 @@
+package openwechat
+
+import (
+	"net/http"
+	"net/http/cookiejar"
+	"sync"
+	"time"
+	"unsafe"
+)
+
+// Jar is a struct which as same as cookiejar.Jar
+type Jar struct {
+	PsList cookiejar.PublicSuffixList
+
+	// mu locks the remaining fields.
+	mu sync.Mutex
+
+	// Entries is a set of entries, keyed by their eTLD+1 and subkeyed by
+	// their name/domain/path.
+	Entries map[string]map[string]entry
+
+	// nextSeqNum is the next sequence number assigned to a new cookie
+	// created SetCookies.
+	NextSeqNum uint64
+}
+
+// AsCookieJar unsafe convert to http.CookieJar
+func (j *Jar) AsCookieJar() http.CookieJar {
+	return (*cookiejar.Jar)(unsafe.Pointer(j))
+}
+
+func newCookieJar() http.CookieJar {
+	jar, _ := cookiejar.New(nil)
+	return jar
+}
+
+func fromCookieJar(jar http.CookieJar) *Jar {
+	return (*Jar)(unsafe.Pointer(jar.(*cookiejar.Jar)))
+}
+
+type entry struct {
+	Name       string
+	Value      string
+	Domain     string
+	Path       string
+	SameSite   string
+	Secure     bool
+	HttpOnly   bool
+	Persistent bool
+	HostOnly   bool
+	Expires    time.Time
+	Creation   time.Time
+	LastAccess time.Time
+
+	// seqNum is a sequence number so that Jar returns cookies in a
+	// deterministic order, even for cookies that have equal Path length and
+	// equal Creation time. This simplifies testing.
+	seqNum uint64
+}

stroage.go

@@ -2,7 +2,6 @@ package openwechat
 
 import (
 	"io"
-	"net/http"
 	"os"
 )
 
@@ -14,7 +13,7 @@ type Storage struct {
 }
 
 type HotReloadStorageItem struct {
-	Cookies      map[string][]*http.Cookie
+	Jar          *Jar
 	BaseRequest  *BaseRequest
 	LoginInfo    *LoginInfo
 	WechatDomain WechatDomain