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)
		}

		// 查找是否存在相同 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
	if cookiesList, exists := p.hostMap[host]; exists {
		result := make([]*http.Cookie, 0, len(cookiesList))
		requestPath := u.Path
		if requestPath == "" {
			requestPath = "/"
		}

		for _, cookies := range cookiesList {
			// 如果只有一个 cookie，直接返回
			if len(cookies) == 1 {
				result = append(result, cookies[0])
				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)
			}
		}
		return result
	}
	return nil
}

// 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)

	utils.OkOrPanic(err)

	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)
	utils.OkOrPanic(err)

	p.hostMap = hostMap
}

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

	var hostMap hostMap
	err := gob.NewDecoder(reader).
		Decode(&hostMap)
	utils.OkOrPanic(err)

	p.hostMap = hostMap
}