jspkit/common/lxrod/cookie.go

package lxrod

import (
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"time"

	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/proto"
)

// 将 p 中所有 cookie 提取到 jar 中。cookie 的 domain 设置为 u。
// TODO: 双向绑定
func SetCookieFor(p *rod.Page, u *url.URL) (http.CookieJar, error) {
	jar, _ := cookiejar.New(nil)
	cookies, err := p.Cookies(nil)
	if err != nil {
		return nil, err
	}
	cs := make([]*http.Cookie, len(cookies))
	for i, c := range cookies {
		cs[i] = CookieFrom(c)
	}
	jar.SetCookies(u, cs)
	return jar, nil
}

// 将 rod 的 cookie 格式转换为 go http 的 cookie。
// go 1.19 以后的新字段不支持。
func CookieFrom(c *proto.NetworkCookie) *http.Cookie {
	var t time.Time
	if c.Expires != -1 {
		t = c.Expires.Time()
	}

	return &http.Cookie{
		Name:     c.Name,
		Value:    c.Value,
		Path:     c.Path,
		Domain:   c.Domain,
		Expires:  t,
		Secure:   c.Secure,
		HttpOnly: c.HTTPOnly,
		SameSite: SameSite(c.SameSite),
	}
}

// 从 rod cookie 中提取创建 go cookie 所用的 SameSite 数据
func SameSite(s proto.NetworkCookieSameSite) http.SameSite {
	switch s {
	case proto.NetworkCookieSameSiteStrict:
		return http.SameSiteStrictMode
	case proto.NetworkCookieSameSiteLax:
		return http.SameSiteLaxMode
	case proto.NetworkCookieSameSiteNone:
		return http.SameSiteNoneMode
	default:
		return http.SameSiteDefaultMode
	}
}