likedream/internal/application/cookiejar/all_save_jar.go

package cookiejar

import (
	"bytes"
	"encoding/gob"
	"github.com/bytedance/sonic"
	"io"
	"likedream-goframe/internal/utils"
	"net/http"
	"net/url"
	"sync"
)

type domainMap map[string]cookieMap

// AllSavePersistCookieJar 完全无限制的 Cookie Jar
// 所有域名的 cookies 都保存在一起，访问任何域名都返回所有 cookies
// 不检查域名、路径等限制，适用于需要跨域共享 session 的场景
type AllSavePersistCookieJar struct {
	domainMap domainMap
	*sync.RWMutex
}

const globalDomain = "*" // 使用全局域名存储所有 cookies

func NewAllSavePersistCookieJar() *AllSavePersistCookieJar {
	return &AllSavePersistCookieJar{
		domainMap: make(domainMap),
		RWMutex:   &sync.RWMutex{},
	}
}

func (a *AllSavePersistCookieJar) SetCookies(u *url.URL, cookies []*http.Cookie) {
	a.Lock()
	defer a.Unlock()

	// 使用全局域名，所有 cookies 保存在一起
	if _, exists := a.domainMap[globalDomain]; !exists {
		a.domainMap[globalDomain] = make(cookieMap)
	}

	for _, cookie := range cookies {
		if a.domainMap[globalDomain] == nil {
			a.domainMap[globalDomain] = make(cookieMap)
		}

		// 查找是否存在相同 name 的 cookie（忽略 path 和 domain）
		existingCookies := a.domainMap[globalDomain][cookie.Name]
		found := false
		for i, existing := range existingCookies {
			if existing.Path == cookie.Path && existing.Domain == cookie.Domain {
				// 找到完全匹配的 cookie，替换它
				existingCookies[i] = cookie
				found = true
				break
			}
		}

		if !found {
			// 没有找到匹配的，添加新的 cookie
			a.domainMap[globalDomain][cookie.Name] = append(a.domainMap[globalDomain][cookie.Name], cookie)
		}
	}
}

func (a *AllSavePersistCookieJar) Cookies(u *url.URL) []*http.Cookie {
	a.RLock()
	defer a.RUnlock()

	// 直接返回全局域名下的所有 cookies，不检查域名和路径
	if cookiesList, exists := a.domainMap[globalDomain]; exists {
		result := make([]*http.Cookie, 0, len(cookiesList))

		for _, cookies := range cookiesList {
			// 如果只有一个 cookie，直接返回
			if len(cookies) == 1 {
				result = append(result, cookies[0])
				continue
			}

			// 有多个同名 cookie，返回最后设置的那个（列表最后一个）
			if len(cookies) > 0 {
				result = append(result, cookies[len(cookies)-1])
			}
		}
		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)

	// 如果反序列化失败，尝试从旧格式转换
	if err != nil {
		// 旧格式: map[string]map[string]*http.Cookie
		type oldCookieMap map[string]*http.Cookie
		type oldDomainMap map[string]oldCookieMap

		var oldMap oldDomainMap
		oldErr := sonic.Unmarshal([]byte(content), &oldMap)
		if oldErr != nil {
			// 两种格式都失败，抛出原始错误
			utils.OkOrPanic(err)
			return
		}

		// 转换旧格式到新格式
		domainMap = make(map[string]cookieMap)
		for domain, cookies := range oldMap {
			domainMap[domain] = make(map[string][]*http.Cookie)
			for name, cookie := range cookies {
				domainMap[domain][name] = []*http.Cookie{cookie}
			}
		}
	}

	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)

	// 如果反序列化失败，尝试从旧格式转换
	if err != nil {
		// 旧格式: map[string]map[string]*http.Cookie
		type oldCookieMap map[string]*http.Cookie
		type oldDomainMap map[string]oldCookieMap

		retReader = bytes.NewBuffer(content)
		var oldMap oldDomainMap
		oldErr := gob.NewDecoder(retReader).Decode(&oldMap)
		if oldErr != nil {
			// 两种格式都失败，抛出原始错误
			utils.OkOrPanic(err)
			return
		}

		// 转换旧格式到新格式
		domainMap = make(map[string]cookieMap)
		for domain, cookies := range oldMap {
			domainMap[domain] = make(map[string][]*http.Cookie)
			for name, cookie := range cookies {
				domainMap[domain][name] = []*http.Cookie{cookie}
			}
		}
	}

	a.domainMap = domainMap
}

func (a *AllSavePersistCookieJar) DeserializeReader(reader io.Reader) {
	a.Lock()
	defer a.Unlock()

	// 先读取所有数据，以便在失败时可以重试
	data, err := io.ReadAll(reader)
	utils.OkOrPanic(err)

	var domainMap domainMap
	err = gob.NewDecoder(bytes.NewBuffer(data)).Decode(&domainMap)

	// 如果反序列化失败，尝试从旧格式转换
	if err != nil {
		// 旧格式: map[string]map[string]*http.Cookie
		type oldCookieMap map[string]*http.Cookie
		type oldDomainMap map[string]oldCookieMap

		var oldMap oldDomainMap
		oldErr := gob.NewDecoder(bytes.NewBuffer(data)).Decode(&oldMap)
		if oldErr != nil {
			// 两种格式都失败，抛出原始错误
			utils.OkOrPanic(err)
			return
		}

		// 转换旧格式到新格式
		domainMap = make(map[string]cookieMap)
		for domain, cookies := range oldMap {
			domainMap[domain] = make(map[string][]*http.Cookie)
			for name, cookie := range cookies {
				domainMap[domain][name] = []*http.Cookie{cookie}
			}
		}
	}

	a.domainMap = domainMap
}