likedream/internal/application/cookiejar/jar.go

package cookiejar

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

type hostMap map[string]cookieMap

type cookieMap map[string]*http.Cookie

// cookieKey 生成 cookie 的唯一标识 key，使用 name 和 path 的组合
func cookieKey(name, path string) string {
	return fmt.Sprintf("%s|%s", name, path)
}

// pathMatch 检查请求路径是否匹配 cookie 的 path
// 根据 RFC 6265 Section 5.1.4
func pathMatch(requestPath, cookiePath string) bool {
	// 如果路径完全相同
	if requestPath == cookiePath {
		return true
	}

	// 如果 cookie path 是请求 path 的前缀
	if strings.HasPrefix(requestPath, cookiePath) {
		// cookie path 以 / 结尾，或者请求 path 在 cookie path 后紧跟 /
		if strings.HasSuffix(cookiePath, "/") {
			return true
		}
		if len(requestPath) > len(cookiePath) && requestPath[len(cookiePath)] == '/' {
			return true
		}
	}

	return false
}

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 = make(hostMap)
	}

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

		// 如果 cookie 没有设置 Path，使用默认值 "/"
		cookiePath := cookie.Path
		if cookiePath == "" {
			cookiePath = "/"
		}

		// 使用 name 和 path 的组合作为 key
		key := cookieKey(cookie.Name, cookiePath)

		// 如果 MaxAge=0 或 MaxAge<0，表示删除 cookie
		// 只删除完全匹配的 cookie（name 和 path 都相同）
		if cookie.MaxAge < 0 || (cookie.MaxAge == 0 && cookie.Value == "") {
			delete(p.hostMap[host], key)
		} else {
			p.hostMap[host][key] = cookie
		}
	}

}

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

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

	if cookies, exists := p.hostMap[host]; exists {
		result := make([]*http.Cookie, 0, len(cookies))
		for _, cookie := range cookies {
			// 检查 cookie 的 path 是否匹配请求的 path
			cookiePath := cookie.Path
			if cookiePath == "" {
				cookiePath = "/"
			}

			if pathMatch(requestPath, cookiePath) {
				result = append(result, cookie)
			}
		}
		return result
	}
	return nil
}

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
}