diff --git a/internal/turnstile/solver.go b/internal/turnstile/solver.go index cd2e588..a5982f1 100644 --- a/internal/turnstile/solver.go +++ b/internal/turnstile/solver.go @@ -5,6 +5,7 @@ package turnstile import ( "bytes" "context" + "crypto/rand" "encoding/json" "errors" "fmt" @@ -158,6 +159,64 @@ func (s *Solver) Invalidate(siteCode string) { } } +// guestCookieName is the cookie the upstream guest-token plugin keys a +// visitor's free daily tier to (verified against the live endpoints: the +// balance API echoes it as "gid" and accepts any freshly minted value). +const guestCookieName = "dsgt_gid" + +// guestIDAlphabet matches the 32-char ids the site itself issues. +const guestIDAlphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + +// randomGuestID mints a visitor id in the same shape the site uses. +func randomGuestID() string { + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return fmt.Sprintf("g%d", time.Now().UnixNano()) + } + for i := range b { + b[i] = guestIDAlphabet[int(b[i])%len(guestIDAlphabet)] + } + return string(b) +} + +// RotateGuest swaps the visitor cookie inside every cached session for the +// site. The upstream grants each fresh visitor id its own daily free tier — +// the same effect as opening a new private window — so a quota-exhausted +// site recovers without paying for a Turnstile re-solve. Returns the new id, +// or "" when the site has no live cached session to rotate. +func (s *Solver) RotateGuest(siteCode string) string { + gid := randomGuestID() + s.mu.Lock() + cores := make([]*core, 0, len(s.cores)) + for _, c := range s.cores { + cores = append(cores, c) + } + s.mu.Unlock() + rotated := false + for _, c := range cores { + c.mu.Lock() + if st, ok := c.cache[siteCode]; ok && time.Now().Before(st.expiresAt) { + st.cookies[guestCookieName] = gid + rotated = true + } + c.mu.Unlock() + } + if !rotated { + return "" + } + return gid +} + +// GuestID reports the visitor id cached for the site on one proxy route. +func (s *Solver) GuestID(siteCode, proxy string) (string, bool) { + ck := s.coreFor(proxy).load(siteCode) + if ck == nil { + return "", false + } + gid := ck[guestCookieName] + return gid, gid != "" +} + // ImportCookies stores cookies exported from a browser session that already // passed the Turnstile check, so chat works while [turnstile] is off or the // solver service is down. Cookies are bound to the IP/UA that obtained them, @@ -211,8 +270,8 @@ func ParseCookiePairs(raw string) map[string]string { // [turnstile] switch is off; only the solving step needs the switch. func (s *Solver) ApplyValidCookies(ctx context.Context, sess httpx.Session, site *config.Site, proxy string) error { c := s.coreFor(proxy) - if st := c.load(site.Code); st != nil { - sess.SetCookies(st.cookies) + if ck := c.load(site.Code); ck != nil { + sess.SetCookies(ck) return nil } if !s.Enabled() { @@ -221,8 +280,8 @@ func (s *Solver) ApplyValidCookies(ctx context.Context, sess httpx.Session, site l := c.lockFor(site.Code) l.Lock() defer l.Unlock() - if st := c.load(site.Code); st != nil { - sess.SetCookies(st.cookies) + if ck := c.load(site.Code); ck != nil { + sess.SetCookies(ck) return nil } return s.refreshLocked(ctx, sess, site, c) @@ -240,14 +299,20 @@ func (s *Solver) Refresh(ctx context.Context, sess httpx.Session, site *config.S return s.refreshLocked(ctx, sess, site, c) } -func (c *core) load(site string) *cookieState { +// load returns a copy of the cached cookies, or nil when absent or expired. +// Callers get a snapshot, so RotateGuest can swap entries concurrently. +func (c *core) load(site string) map[string]string { c.mu.Lock() defer c.mu.Unlock() st, ok := c.cache[site] if !ok || time.Now().After(st.expiresAt) { return nil } - return st + out := make(map[string]string, len(st.cookies)) + for k, v := range st.cookies { + out[k] = v + } + return out } func (s *Solver) refreshLocked(ctx context.Context, sess httpx.Session, site *config.Site, c *core) error { diff --git a/internal/turnstile/solver_test.go b/internal/turnstile/solver_test.go index 50b3b21..3888055 100644 --- a/internal/turnstile/solver_test.go +++ b/internal/turnstile/solver_test.go @@ -192,3 +192,50 @@ func TestHistoryIsCappedAndRecorded(t *testing.T) { t.Error("last record should be OK") } } + +func TestRotateGuestSwapsVisitorID(t *testing.T) { + cfg, err := config.Load(examplePath) + if err != nil { + t.Fatalf("load config: %v", err) + } + solver := New(cfg.Turnstile, httpx.DefaultUserAgent) + const oldGid = "OLDGIDOLDGIDOLDGIDOLDGIDOLDG12" + if _, _, err := solver.ImportCookies("de", "cf_clearance=abc; dsgt_gid="+oldGid, ""); err != nil { + t.Fatalf("import: %v", err) + } + gid, ok := solver.GuestID("de", "") + if !ok || gid != oldGid { + t.Fatalf("GuestID = %q %v, want %q", gid, ok, oldGid) + } + sess := &fakeSession{} + if err := solver.ApplyValidCookies(context.Background(), sess, cfg.Sites["de"], ""); err != nil { + t.Fatalf("apply: %v", err) + } + if sess.Cookies()["dsgt_gid"] != oldGid { + t.Fatalf("session gid = %q, want the cached %q", sess.Cookies()["dsgt_gid"], oldGid) + } + + rotated := solver.RotateGuest("de") + if rotated == "" { + t.Fatal("RotateGuest returned empty for a live session") + } + if rotated == oldGid { + t.Fatal("RotateGuest kept the old id") + } + if len(rotated) != 32 { + t.Errorf("rotated id length = %d, want 32", len(rotated)) + } + if got, _ := solver.GuestID("de", ""); got != rotated { + t.Errorf("GuestID after rotate = %q, want %q", got, rotated) + } + sess2 := &fakeSession{} + if err := solver.ApplyValidCookies(context.Background(), sess2, cfg.Sites["de"], ""); err != nil { + t.Fatalf("apply 2: %v", err) + } + if sess2.Cookies()["dsgt_gid"] != rotated { + t.Errorf("fresh session gid = %q, want rotated %q", sess2.Cookies()["dsgt_gid"], rotated) + } + if got := solver.RotateGuest("no-such-site"); got != "" { + t.Errorf("rotating an unknown site returned %q, want empty", got) + } +} diff --git a/internal/upstream/balance.go b/internal/upstream/balance.go new file mode 100644 index 0000000..8fed1f7 --- /dev/null +++ b/internal/upstream/balance.go @@ -0,0 +1,78 @@ +package upstream + +import ( + "context" + "encoding/json" + "strconv" + "time" + + "github.com/nyoungo/dsfree2api/internal/config" + "github.com/nyoungo/dsfree2api/internal/httpx" +) + +// guestCookieName mirrors the upstream guest-token plugin's visitor cookie; +// its value is the "gid" the balance endpoint reports on. +const guestCookieName = "dsgt_gid" + +// SiteBalance is the response of GET /wp-json/dsgt/v1/balance. The free-tier +// block only appears when the request carries a bot_id. +type SiteBalance struct { + Gid string `json:"gid"` + Balance float64 `json:"balance"` + Free struct { + Limit int `json:"limit"` + Used int `json:"used"` + Remaining int `json:"remaining"` + ResetPeriod string `json:"reset_period"` + } `json:"free"` +} + +// FetchBalance queries the site's guest balance endpoint for one visitor id. +// No nonce is required for guests; the id travels as the visitor cookie. +// botID selects the per-bot free tier (0 omits it, which also omits free). +func (c *Client) FetchBalance( + ctx context.Context, + site config.Site, + gid string, + botID int, + route Route, +) (*SiteBalance, error) { + if gid == "" { + return nil, errf("no guest id for site %s", site.Code) + } + sess, err := httpx.NewSession(httpx.Options{ProxyURL: route.Proxy, Timeout: 20 * time.Second}) + if err != nil { + return nil, err + } + defer sess.Close() + sess.SetCookies(map[string]string{guestCookieName: gid}) + + url := site.BaseURL + "/wp-json/dsgt/v1/balance?_t=" + strconv.FormatInt(time.Now().UnixMilli(), 10) + if botID > 0 { + url += "&bot_id=" + itoa(botID) + } + c.cfg.RLock() + ua := c.cfg.Upstream.UserAgent + c.cfg.RUnlock() + resp, err := sess.Do(ctx, httpx.Request{ + Method: "GET", + URL: url, + Header: map[string]string{ + "User-Agent": ua, + "Accept": "application/json", + "Referer": site.BaseURL + "/", + }, + Timeout: 20 * time.Second, + }) + if err != nil { + return nil, err + } + if resp.Status >= 400 { + return nil, errf("balance http %d: %s", resp.Status, truncate(string(resp.Body), 200)) + } + var out SiteBalance + if err := json.Unmarshal(resp.Body, &out); err != nil { + return nil, errf("balance bad json: %v", err) + } + return &out, nil +} diff --git a/internal/upstream/client.go b/internal/upstream/client.go index dcc9d6c..d068935 100644 --- a/internal/upstream/client.go +++ b/internal/upstream/client.go @@ -156,6 +156,43 @@ func (c *Client) sortQuotaCoolLast(ids []string) []string { return append(rest, cooled...) } +// clearQuotaCool lifts a site's cooldown once it serves again after a quota +// error — the fresh visitor identity restored its daily tier. +func (c *Client) clearQuotaCool(code string) { + c.stateMu.Lock() + _, ok := c.quotaCool[code] + if ok { + delete(c.quotaCool, code) + } + c.stateMu.Unlock() + if ok { + c.log.Info("site quota restored", "site", code) + } +} + +// probeBalanceAsync logs the site's current free tier after a quota +// recovery. Best effort — failures only show up at DEBUG. +func (c *Client) probeBalanceAsync(site config.Site, model config.Model, route Route) { + if c.ts == nil { + return + } + gid, ok := c.ts.GuestID(site.Code, route.Proxy) + if !ok { + return + } + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + bal, err := c.FetchBalance(ctx, site, gid, model.BotID, route) + if err != nil { + c.log.Debug("balance probe failed", "site", site.Code, "error", err) + return + } + c.log.Info("site balance", "site", site.Code, + "remaining", bal.Free.Remaining, "limit", bal.Free.Limit) + }() +} + // logUpstreamErr records one failed upstream attempt. Quota errors cool the // site down and repeats of an identical error inside a short window drop to // DEBUG, so a balance-starved site neither taxes every continuation round @@ -321,12 +358,17 @@ func (c *Client) chatWithRetries(ctx context.Context, modelID string, model conf routes := c.Routes() var lastErr error + sawQuota := false for attempt := 0; ; attempt++ { var failed []Route for i, route := range routes { allowSlow := i == 0 && len(routes) > 1 err := c.chatOnceSlowStart(ctx, site, modelID, model, prompt, route, allowSlow, info, yield) if err == nil { + if sawQuota { + c.clearQuotaCool(site.Code) + c.probeBalanceAsync(site, model, route) + } return nil } lastErr = err @@ -358,6 +400,9 @@ func (c *Client) chatWithRetries(ctx context.Context, modelID string, model conf } failed = append(failed, route) c.logUpstreamErr(site.Code, modelID, attempt, route, err) + if isQuota(err) { + sawQuota = true + } if i < len(routes)-1 { continue } @@ -384,6 +429,15 @@ func (c *Client) chatWithRetries(ctx context.Context, modelID string, model conf // first cycle retries cheaply and only the second one pays for // a Turnstile re-solve. forceCookie := isSession(lastErr) || (isQuota(lastErr) && attempt >= 1) + // The free daily tier is keyed to the visitor cookie: swapping + // in a brand-new visitor id restores it instantly, the same + // effect as a fresh private window. Only a re-solve (attempt + // >= 1) then remains as the slower fallback. + if isQuota(lastErr) && attempt == 0 && c.ts != nil { + if gid := c.ts.RotateGuest(site.Code); gid != "" { + c.log.Info("rotated visitor identity", "site", site.Code, "reason", "quota") + } + } for _, route := range failed { if err := c.refreshSession(ctx, site, modelID, route, forceCookie); err != nil { c.log.Warn("session refresh failed", "model", modelID, "error", err) diff --git a/internal/upstream/client_test.go b/internal/upstream/client_test.go index 5de4a4c..fb445b2 100644 --- a/internal/upstream/client_test.go +++ b/internal/upstream/client_test.go @@ -10,6 +10,8 @@ import ( "time" "github.com/nyoungo/dsfree2api/internal/config" + "github.com/nyoungo/dsfree2api/internal/httpx" + "github.com/nyoungo/dsfree2api/internal/turnstile" ) const examplePath = "../../config.example.toml" @@ -336,3 +338,40 @@ func TestCacheEmptyFailsOverToMirrorWithoutRefreshCycles(t *testing.T) { t.Fatal("mirror was not attempted after cache rejection") } } + +func TestQuotaRotationRecoversWithoutSolve(t *testing.T) { + cfg := testConfig(t) + cfg.Upstream.AutoRefresh = true + cfg.Upstream.RefreshRetries = 3 + solver := turnstile.New(cfg.Turnstile, httpx.DefaultUserAgent) + const oldGid = "OLDGIDOLDGIDOLDGIDOLDGIDOLDG12" + if _, _, err := solver.ImportCookies("de", "cf_clearance=abc; dsgt_gid="+oldGid, ""); err != nil { + t.Fatalf("seed cookies: %v", err) + } + c := New(cfg, solver, nil) + calls := 0 + c.chatOnceOverride = func(_ context.Context, site config.Site, modelID string, model config.Model, + _ string, _ Route, _ *ServeInfo, _ func(Event) error) error { + calls++ + if calls == 1 { + return newQuota("sse quota exhausted: test") + } + return nil + } + if err := c.Chat(context.Background(), "deepseek-v4-flash-de", "p", nil, func(Event) error { return nil }); err != nil { + t.Fatalf("Chat: %v", err) + } + if calls != 2 { + t.Fatalf("chat calls = %d, want 2 (quota then success)", calls) + } + gid, ok := solver.GuestID("de", "") + if !ok { + t.Fatal("guest id lost after rotation") + } + if gid == oldGid { + t.Fatal("visitor id was not rotated after the quota error") + } + if c.quotaCooled("de") { + t.Error("cooldown should be cleared once the site serves again") + } +}