feat(gemini-web): Implement proactive PSIDTS cookie rotation

This commit is contained in:
hkfires
2025-09-22 21:54:52 +08:00
parent ed87dda0a6
commit 22a69333a0
5 changed files with 86 additions and 41 deletions

View File

@@ -772,7 +772,18 @@ func (cs *ChatSession) RCID() string {
} }
return "" return ""
} }
func (cs *ChatSession) setCID(v string) {
if len(cs.metadata) < 1 {
cs.metadata = normalizeMeta(cs.metadata)
}
cs.metadata[0] = v
}
func (cs *ChatSession) setRID(v string) {
if len(cs.metadata) < 2 {
cs.metadata = normalizeMeta(cs.metadata)
}
cs.metadata[1] = v
}
func (cs *ChatSession) setRCID(v string) { func (cs *ChatSession) setRCID(v string) {
if len(cs.metadata) < 3 { if len(cs.metadata) < 3 {
cs.metadata = normalizeMeta(cs.metadata) cs.metadata = normalizeMeta(cs.metadata)

View File

@@ -13,7 +13,6 @@ import (
"path/filepath" "path/filepath"
"regexp" "regexp"
"strings" "strings"
"sync"
"time" "time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces" "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
@@ -494,12 +493,6 @@ type FileStreamingLogWriter struct {
// statusWritten indicates whether the response status has been written. // statusWritten indicates whether the response status has been written.
statusWritten bool statusWritten bool
// mu protects concurrent access to the writer state.
mu sync.RWMutex
// closed indicates whether the streaming writer has been closed.
closed bool
} }
// WriteChunkAsync writes a response chunk asynchronously (non-blocking). // WriteChunkAsync writes a response chunk asynchronously (non-blocking).
@@ -507,10 +500,7 @@ type FileStreamingLogWriter struct {
// Parameters: // Parameters:
// - chunk: The response chunk to write // - chunk: The response chunk to write
func (w *FileStreamingLogWriter) WriteChunkAsync(chunk []byte) { func (w *FileStreamingLogWriter) WriteChunkAsync(chunk []byte) {
w.mu.RLock() if w.chunkChan == nil {
defer w.mu.RUnlock()
if w.chunkChan == nil || w.closed {
return return
} }
@@ -535,9 +525,6 @@ func (w *FileStreamingLogWriter) WriteChunkAsync(chunk []byte) {
// Returns: // Returns:
// - error: An error if writing fails, nil otherwise // - error: An error if writing fails, nil otherwise
func (w *FileStreamingLogWriter) WriteStatus(status int, headers map[string][]string) error { func (w *FileStreamingLogWriter) WriteStatus(status int, headers map[string][]string) error {
w.mu.Lock()
defer w.mu.Unlock()
if w.file == nil || w.statusWritten { if w.file == nil || w.statusWritten {
return nil return nil
} }
@@ -566,38 +553,21 @@ func (w *FileStreamingLogWriter) WriteStatus(status int, headers map[string][]st
// Returns: // Returns:
// - error: An error if closing fails, nil otherwise // - error: An error if closing fails, nil otherwise
func (w *FileStreamingLogWriter) Close() error { func (w *FileStreamingLogWriter) Close() error {
w.mu.Lock() if w.chunkChan != nil {
if w.closed { close(w.chunkChan)
w.mu.Unlock()
return nil
}
w.closed = true
chunkChan := w.chunkChan
closeChan := w.closeChan
file := w.file
w.mu.Unlock()
if chunkChan != nil {
close(chunkChan)
} }
// Wait for async writer to finish // Wait for async writer to finish
if closeChan != nil { if w.closeChan != nil {
<-closeChan <-w.closeChan
w.chunkChan = nil
} }
var err error if w.file != nil {
if file != nil { return w.file.Close()
err = file.Close()
} }
w.mu.Lock() return nil
w.chunkChan = nil
w.closeChan = nil
w.file = nil
w.mu.Unlock()
return err
} }
// asyncWriter runs in a goroutine to handle async chunk writing. // asyncWriter runs in a goroutine to handle async chunk writing.

View File

@@ -5,6 +5,8 @@ import (
"context" "context"
"errors" "errors"
"fmt" "fmt"
"net/http"
"net/url"
"path/filepath" "path/filepath"
"strings" "strings"
"sync" "sync"
@@ -141,6 +143,7 @@ func (s *geminiWebState) ensureClient() error {
if s.cfg != nil && s.cfg.GeminiWeb.TokenRefreshSeconds > 0 { if s.cfg != nil && s.cfg.GeminiWeb.TokenRefreshSeconds > 0 {
refresh = s.cfg.GeminiWeb.TokenRefreshSeconds refresh = s.cfg.GeminiWeb.TokenRefreshSeconds
} }
// Use explicit refresh; background auto-refresh disabled here
if err := s.client.Init(float64(timeout), false, 300, false, float64(refresh), false); err != nil { if err := s.client.Init(float64(timeout), false, 300, false, float64(refresh), false); err != nil {
s.client = nil s.client = nil
return err return err
@@ -166,9 +169,12 @@ func (s *geminiWebState) refresh(ctx context.Context) error {
if s.cfg != nil && s.cfg.GeminiWeb.TokenRefreshSeconds > 0 { if s.cfg != nil && s.cfg.GeminiWeb.TokenRefreshSeconds > 0 {
refresh = s.cfg.GeminiWeb.TokenRefreshSeconds refresh = s.cfg.GeminiWeb.TokenRefreshSeconds
} }
// Use explicit refresh; background auto-refresh disabled here
if err := s.client.Init(float64(timeout), false, 300, false, float64(refresh), false); err != nil { if err := s.client.Init(float64(timeout), false, 300, false, float64(refresh), false); err != nil {
return err return err
} }
// Attempt rotation proactively to persist new TS sooner
_ = s.tryRotatePSIDTS(proxyURL)
s.lastRefresh = time.Now() s.lastRefresh = time.Now()
return nil return nil
} }
@@ -195,6 +201,59 @@ func (s *geminiWebState) tokenSnapshot() *gemini.GeminiWebTokenStorage {
return &c return &c
} }
// tryRotatePSIDTS performs a best-effort rotation of __Secure-1PSIDTS using
// the public RotateCookies endpoint. On success it updates both the in-memory
// token and the live client's cookie jar so that subsequent requests adopt the
// new value. Any error is ignored by the caller to avoid disrupting refresh.
func (s *geminiWebState) tryRotatePSIDTS(proxy string) error {
cookies := map[string]string{
"__Secure-1PSID": s.token.Secure1PSID,
"__Secure-1PSIDTS": s.token.Secure1PSIDTS,
}
tr := &http.Transport{}
if proxy != "" {
if pu, err := url.Parse(proxy); err == nil {
tr.Proxy = http.ProxyURL(pu)
}
}
client := &http.Client{Transport: tr, Timeout: 60 * time.Second}
req, _ := http.NewRequest(http.MethodPost, geminiwebapi.EndpointRotateCookies, bytes.NewReader([]byte("[000,\"-0000000000000000000\"]")))
for k, vs := range geminiwebapi.HeadersRotateCookies {
for _, v := range vs {
req.Header.Add(k, v)
}
}
for k, v := range cookies {
req.AddCookie(&http.Cookie{Name: k, Value: v})
}
resp, err := client.Do(req)
if err != nil {
return err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
for _, c := range resp.Cookies() {
if c == nil {
continue
}
if c.Name == "__Secure-1PSIDTS" && c.Value != "" && c.Value != s.token.Secure1PSIDTS {
s.tokenMu.Lock()
s.token.Secure1PSIDTS = c.Value
s.tokenDirty = true
if s.client != nil && s.client.Cookies != nil {
s.client.Cookies["__Secure-1PSIDTS"] = c.Value
}
s.tokenMu.Unlock()
break
}
}
}
return nil
}
func (s *geminiWebState) ShouldRefresh(now time.Time, _ *cliproxyauth.Auth) bool { func (s *geminiWebState) ShouldRefresh(now time.Time, _ *cliproxyauth.Auth) bool {
interval := s.refreshInterval interval := s.refreshInterval
if interval <= 0 { if interval <= 0 {

View File

@@ -861,7 +861,11 @@ func (m *Manager) refreshAuth(ctx context.Context, id string) {
if updated == nil { if updated == nil {
updated = cloned updated = cloned
} }
updated.Runtime = auth.Runtime // Preserve runtime created by the executor during Refresh.
// If executor didn't set one, fall back to the previous runtime.
if updated.Runtime == nil {
updated.Runtime = auth.Runtime
}
updated.LastRefreshedAt = now updated.LastRefreshedAt = now
updated.NextRefreshAfter = time.Time{} updated.NextRefreshAfter = time.Time{}
updated.LastError = nil updated.LastError = nil

View File

@@ -133,6 +133,7 @@ var defaultAuthenticatorFactories = map[string]func() clipauth.Authenticator{
"qwen": func() clipauth.Authenticator { return clipauth.NewQwenAuthenticator() }, "qwen": func() clipauth.Authenticator { return clipauth.NewQwenAuthenticator() },
"gemini": func() clipauth.Authenticator { return clipauth.NewGeminiAuthenticator() }, "gemini": func() clipauth.Authenticator { return clipauth.NewGeminiAuthenticator() },
"gemini-cli": func() clipauth.Authenticator { return clipauth.NewGeminiAuthenticator() }, "gemini-cli": func() clipauth.Authenticator { return clipauth.NewGeminiAuthenticator() },
"gemini-web": func() clipauth.Authenticator { return clipauth.NewGeminiWebAuthenticator() },
} }
var expireKeys = [...]string{"expired", "expire", "expires_at", "expiresAt", "expiry", "expires"} var expireKeys = [...]string{"expired", "expire", "expires_at", "expiresAt", "expiry", "expires"}