likedream/internal/application/cookiejar/jar.go

package cookiejar

import (
	"bytes"
	"encoding/gob"
	"io"
	"likedream-goframe/internal/utils"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/bytedance/sonic"
)

type hostMap map[string]cookieMap

type cookieMap map[string][]*http.Cookie
type PersistCookieJar struct {
	hostMap hostMap
	*sync.RWMutex
}

func NewPersistCookieJar() *PersistCookieJar {
	return &PersistCookieJar{
		hostMap: make(hostMap),
		RWMutex: &sync.RWMutex{},
	}
}

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

	host := u.Host

	if _, exists := p.hostMap[host]; !exists {
		p.hostMap[host] = make(cookieMap)
	}

	for _, cookie := range cookies {
		if p.hostMap[host] == nil {
			p.hostMap[host] = make(cookieMap)
		}

		// 检查并删除全局 cookie 中同名的 cookie
		if globalCookies, exists := p.hostMap[""]; exists {
			delete(globalCookies, cookie.Name)
		}

		// 查找是否存在相同 name, path 和 domain 的 cookie
		existingCookies := p.hostMap[host][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
			p.hostMap[host][cookie.Name] = append(p.hostMap[host][cookie.Name], cookie)
		}
	}

}

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

	host := u.Host
	requestPath := u.Path
	if requestPath == "" {
		requestPath = "/"
	}

	result := make([]*http.Cookie, 0)
	addedNames := make(map[string]bool)

	// 先添加匹配 host 的 cookie
	if cookiesList, exists := p.hostMap[host]; exists {
		for cookieName, cookies := range cookiesList {
			// 如果只有一个 cookie，直接返回
			if len(cookies) == 1 {
				result = append(result, cookies[0])
				addedNames[cookieName] = true
				continue
			}

			// 有多个同名 cookie，选择最匹配的（path 最长且匹配请求路径的）
			var bestMatch *http.Cookie
			bestMatchLen := -1

			for _, cookie := range cookies {
				cookiePath := cookie.Path
				if cookiePath == "" {
					cookiePath = "/"
				}

				// 检查 path 是否匹配
				if pathMatch(requestPath, cookiePath) {
					if len(cookiePath) > bestMatchLen {
						bestMatch = cookie
						bestMatchLen = len(cookiePath)
					}
				}
			}

			// 如果找到匹配的，添加到结果中
			if bestMatch != nil {
				result = append(result, bestMatch)
				addedNames[cookieName] = true
			}
		}
	}

	// 再添加没有 host 的全局 cookie（避免重复）
	if globalCookies, exists := p.hostMap[""]; exists {
		for cookieName, cookies := range globalCookies {
			// 如果该名称的 cookie 已经从 host 中添加了，跳过
			if addedNames[cookieName] {
				continue
			}

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

			// 有多个同名 cookie，选择最匹配的
			var bestMatch *http.Cookie
			bestMatchLen := -1

			for _, cookie := range cookies {
				cookiePath := cookie.Path
				if cookiePath == "" {
					cookiePath = "/"
				}

				if pathMatch(requestPath, cookiePath) {
					if len(cookiePath) > bestMatchLen {
						bestMatch = cookie
						bestMatchLen = len(cookiePath)
					}
				}
			}

			if bestMatch != nil {
				result = append(result, bestMatch)
			}
		}
	}

	return result
}

// pathMatch 检查请求路径是否匹配 cookie 的 path
// 根据 RFC 6265 的规则
func pathMatch(requestPath, cookiePath string) bool {
	if requestPath == cookiePath {
		return true
	}

	if strings.HasPrefix(requestPath, cookiePath) {
		// cookiePath 以 / 结尾，或者 requestPath 的下一个字符是 /
		if strings.HasSuffix(cookiePath, "/") {
			return true
		}
		if len(requestPath) > len(cookiePath) && requestPath[len(cookiePath)] == '/' {
			return true
		}
	}

	return false
}

func (p *PersistCookieJar) Serialize() string {
	p.RLock()
	defer p.RUnlock()

	json, err := sonic.Marshal(p.hostMap)

	utils.OkOrPanic(err)

	return string(json)
}

func (p *PersistCookieJar) Deserialize(content string) {
	p.Lock()
	defer p.Unlock()

	if content == "" {
		p.hostMap = make(hostMap)
		return
	}

	var hostMap hostMap
	err := sonic.Unmarshal([]byte(content), &hostMap)

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

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

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

	p.hostMap = hostMap
}

func (p *PersistCookieJar) SerializeBytes() []byte {
	p.RLock()
	defer p.RUnlock()

	ret := make([]byte, 0)
	retWriter := bytes.NewBuffer(ret)

	err := gob.NewEncoder(retWriter).
		Encode(p.hostMap)
	utils.OkOrPanic(err)

	return retWriter.Bytes()
}

func (p *PersistCookieJar) DeserializeBytes(content []byte) {
	p.Lock()
	defer p.Unlock()

	if len(content) == 0 {
		p.hostMap = make(hostMap)
		return
	}

	retReader := bytes.NewBuffer(content)

	var hostMap hostMap
	err := gob.NewDecoder(retReader).
		Decode(&hostMap)

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

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

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

	p.hostMap = hostMap
}

func (p *PersistCookieJar) DeserializeReader(reader io.Reader) {
	p.Lock()
	defer p.Unlock()

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

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

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

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

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

	p.hostMap = hostMap
}

// RemoveCookieInfo 将所有 host 的 cookie 移除 host 信息，存储为全局 cookie
func (p *PersistCookieJar) RemoveCookieInfo() {
	p.Lock()
	defer p.Unlock()

	// 确保全局 cookie map 存在
	if _, exists := p.hostMap[""]; !exists {
		p.hostMap[""] = make(cookieMap)
	}

	// 将所有 host 的 cookie 移动到全局
	for host, hostCookies := range p.hostMap {
		// 跳过全局 cookie 本身
		if host == "" {
			continue
		}

		// 将该 host 的所有 cookie 移动到全局
		for cookieName, cookies := range hostCookies {
			p.hostMap[""][cookieName] = cookies
		}

		// 删除该 host 的 cookie
		delete(p.hostMap, host)
	}
}