sugarBecomer
161 lines
3.3 KiB
Go
161 lines
3.3 KiB
Go
package cookiejar
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/gob"
|
|
"github.com/bytedance/sonic"
|
|
"io"
|
|
"likedream-goframe/internal/utils"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
type domainMap map[string]cookieMap
|
|
|
|
// AllSavePersistCookieJar 支持泛域名的 Cookie Jar
|
|
// 不同的二级域名会共享相同主域名的 cookies
|
|
// 例如: passport2.chaoxing.com 和 mooc1-api.chaoxing.com 共享 .chaoxing.com 的 cookies
|
|
type AllSavePersistCookieJar struct {
|
|
domainMap domainMap
|
|
*sync.RWMutex
|
|
}
|
|
|
|
func NewAllSavePersistCookieJar() *AllSavePersistCookieJar {
|
|
return &AllSavePersistCookieJar{
|
|
domainMap: make(domainMap),
|
|
RWMutex: &sync.RWMutex{},
|
|
}
|
|
}
|
|
|
|
// extractMainDomain 提取主域名
|
|
// 例如: passport2.chaoxing.com -> chaoxing.com
|
|
//
|
|
// mooc1-api.chaoxing.com -> chaoxing.com
|
|
// www.example.co.uk -> example.co.uk
|
|
func extractMainDomain(host string) string {
|
|
// 移除端口号
|
|
if colonIndex := strings.Index(host, ":"); colonIndex != -1 {
|
|
host = host[:colonIndex]
|
|
}
|
|
|
|
parts := strings.Split(host, ".")
|
|
if len(parts) <= 2 {
|
|
// 已经是主域名或单级域名
|
|
return host
|
|
}
|
|
|
|
// 处理特殊的二级顶级域名，如 .co.uk, .com.cn 等
|
|
// 简化处理：取最后两部分
|
|
// 更复杂的情况可以使用 publicsuffix 包
|
|
return strings.Join(parts[len(parts)-2:], ".")
|
|
}
|
|
|
|
func (a *AllSavePersistCookieJar) SetCookies(u *url.URL, cookies []*http.Cookie) {
|
|
a.Lock()
|
|
defer a.Unlock()
|
|
|
|
// 提取主域名
|
|
mainDomain := extractMainDomain(u.Host)
|
|
|
|
if _, exists := a.domainMap[mainDomain]; !exists {
|
|
a.domainMap[mainDomain] = make(cookieMap)
|
|
}
|
|
|
|
for _, cookie := range cookies {
|
|
if a.domainMap[mainDomain] == nil {
|
|
a.domainMap[mainDomain] = make(cookieMap)
|
|
}
|
|
a.domainMap[mainDomain][cookie.Name] = cookie
|
|
}
|
|
}
|
|
|
|
func (a *AllSavePersistCookieJar) Cookies(u *url.URL) []*http.Cookie {
|
|
a.RLock()
|
|
defer a.RUnlock()
|
|
|
|
// 提取主域名
|
|
mainDomain := extractMainDomain(u.Host)
|
|
|
|
if cookies, exists := a.domainMap[mainDomain]; exists {
|
|
result := make([]*http.Cookie, 0, len(cookies))
|
|
for _, cookie := range cookies {
|
|
result = append(result, cookie)
|
|
}
|
|
return result
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (a *AllSavePersistCookieJar) Serialize() string {
|
|
a.RLock()
|
|
defer a.RUnlock()
|
|
|
|
json, err := sonic.Marshal(a.domainMap)
|
|
utils.OkOrPanic(err)
|
|
|
|
return string(json)
|
|
}
|
|
|
|
func (a *AllSavePersistCookieJar) Deserialize(content string) {
|
|
a.Lock()
|
|
defer a.Unlock()
|
|
|
|
if content == "" {
|
|
a.domainMap = make(domainMap)
|
|
return
|
|
}
|
|
|
|
var domainMap domainMap
|
|
err := sonic.Unmarshal([]byte(content), &domainMap)
|
|
utils.OkOrPanic(err)
|
|
|
|
a.domainMap = domainMap
|
|
}
|
|
|
|
func (a *AllSavePersistCookieJar) SerializeBytes() []byte {
|
|
a.RLock()
|
|
defer a.RUnlock()
|
|
|
|
ret := make([]byte, 0)
|
|
retWriter := bytes.NewBuffer(ret)
|
|
|
|
err := gob.NewEncoder(retWriter).
|
|
Encode(a.domainMap)
|
|
utils.OkOrPanic(err)
|
|
|
|
return retWriter.Bytes()
|
|
}
|
|
|
|
func (a *AllSavePersistCookieJar) DeserializeBytes(content []byte) {
|
|
a.Lock()
|
|
defer a.Unlock()
|
|
|
|
if len(content) == 0 {
|
|
a.domainMap = make(domainMap)
|
|
return
|
|
}
|
|
|
|
retReader := bytes.NewBuffer(content)
|
|
|
|
var domainMap domainMap
|
|
err := gob.NewDecoder(retReader).
|
|
Decode(&domainMap)
|
|
utils.OkOrPanic(err)
|
|
|
|
a.domainMap = domainMap
|
|
}
|
|
|
|
func (a *AllSavePersistCookieJar) DeserializeReader(reader io.Reader) {
|
|
a.Lock()
|
|
defer a.Unlock()
|
|
|
|
var domainMap domainMap
|
|
err := gob.NewDecoder(reader).
|
|
Decode(&domainMap)
|
|
utils.OkOrPanic(err)
|
|
|
|
a.domainMap = domainMap
|
|
}
|