diff --git a/client_config.example.json b/client_config.example.json index 65fa80f..e0f6245 100644 --- a/client_config.example.json +++ b/client_config.example.json @@ -8,6 +8,8 @@ "script_keys": [ {"id": "REPLACE_WITH_DEPLOYMENT_ID", "account": "acct-a"} ], + "_comment_quota_state_path": "Persists only Apps Script quota quarantines so restarting the client does not immediately retry an account already known to be over quota. Empty string disables persistence.", + "quota_state_path": ".goose-quota-state.json", "tunnel_key": "REPLACE_WITH_64_HEX_CHARACTER_RANDOM_KEY", "_comment_socks_auth": "Optional: require SOCKS5 username/password (RFC 1929). Both fields must be set together or both omitted.", diff --git a/internal/carrier/atomicfile.go b/internal/carrier/atomicfile.go new file mode 100644 index 0000000..605b090 --- /dev/null +++ b/internal/carrier/atomicfile.go @@ -0,0 +1,45 @@ +package carrier + +import ( + "fmt" + "os" + "path/filepath" +) + +func writeFileAtomic(path string, data []byte, perm os.FileMode) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o700); err != nil { + return err + } + tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".tmp-*") + if err != nil { + return err + } + tmpPath := tmp.Name() + cleanup := true + defer func() { + if cleanup { + _ = os.Remove(tmpPath) + } + }() + if err := tmp.Chmod(perm); err != nil { + _ = tmp.Close() + return err + } + if _, err := tmp.Write(data); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Sync(); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + if err := replaceFileAtomic(tmpPath, path); err != nil { + return fmt.Errorf("replace %s: %w", path, err) + } + cleanup = false + return syncParentDir(dir) +} diff --git a/internal/carrier/atomicfile_unix.go b/internal/carrier/atomicfile_unix.go new file mode 100644 index 0000000..b85c67d --- /dev/null +++ b/internal/carrier/atomicfile_unix.go @@ -0,0 +1,18 @@ +//go:build !windows + +package carrier + +import "os" + +func replaceFileAtomic(tmpPath, path string) error { + return os.Rename(tmpPath, path) +} + +func syncParentDir(dir string) error { + d, err := os.Open(dir) + if err != nil { + return nil + } + defer d.Close() + return d.Sync() +} diff --git a/internal/carrier/atomicfile_windows.go b/internal/carrier/atomicfile_windows.go new file mode 100644 index 0000000..5dfceb8 --- /dev/null +++ b/internal/carrier/atomicfile_windows.go @@ -0,0 +1,46 @@ +//go:build windows + +package carrier + +import ( + "os" + "syscall" + "unsafe" +) + +const ( + moveFileReplaceExisting = 0x1 + moveFileWriteThrough = 0x8 +) + +var ( + kernel32 = syscall.NewLazyDLL("kernel32.dll") + moveFileExW = kernel32.NewProc("MoveFileExW") +) + +func replaceFileAtomic(tmpPath, path string) error { + oldName, err := syscall.UTF16PtrFromString(tmpPath) + if err != nil { + return err + } + newName, err := syscall.UTF16PtrFromString(path) + if err != nil { + return err + } + r1, _, callErr := moveFileExW.Call( + uintptr(unsafe.Pointer(oldName)), + uintptr(unsafe.Pointer(newName)), + uintptr(moveFileReplaceExisting|moveFileWriteThrough), + ) + if r1 != 0 { + return nil + } + if callErr != syscall.Errno(0) { + return callErr + } + return os.Rename(tmpPath, path) +} + +func syncParentDir(string) error { + return nil +} diff --git a/internal/carrier/client.go b/internal/carrier/client.go index d99c818..e014ff5 100644 --- a/internal/carrier/client.go +++ b/internal/carrier/client.go @@ -4,6 +4,9 @@ import ( "bytes" "context" "crypto/rand" + "crypto/sha256" + "encoding/hex" + "encoding/json" "errors" "fmt" "io" @@ -11,6 +14,7 @@ import ( "net" "net/http" "os" + "path/filepath" "sort" "strings" "sync" @@ -191,6 +195,11 @@ type Config struct { AESKeyHex string // 64-char hex, must match server DebugTiming bool // when true, log per-session TTFB and per-poll Apps Script RTT + // QuotaStatePath persists durable Apps Script quota quarantines across + // client restarts. Empty disables persistence. Only quota walls are stored; + // transient network/backoff state remains in memory. + QuotaStatePath string + // CoalesceStep / CoalesceMax enable adaptive uplink coalescing on kick(). // When CoalesceStep > 0 the first kick of a burst arms a step timer; each // subsequent kick within the window resets it, bounded by CoalesceMax from @@ -211,6 +220,7 @@ type relayEndpoint struct { account string // optional human-readable Google account label, "" = unlabeled blacklistedTill time.Time localNetworkOffline bool + quotaExhaustedUntil time.Time failCount int statsOK uint64 statsFail uint64 @@ -281,6 +291,8 @@ type Client struct { bucketCount int // distinct in-flight buckets; one per labeled account, plus one per unlabeled endpoint idleSlotsPerBucket int // resolved from Config.IdleSlotsPerBucket; max concurrent polls per bucket clientVersion string + quotaStatePath string + quotaStateMu sync.Mutex // clientID is a random 16-byte identifier minted once per process. It is // embedded in every encrypted batch so the server can route downstream @@ -426,7 +438,7 @@ func New(cfg Config) (*Client, error) { numWorkers, len(endpoints), idleSlotsPerBucket) } - return &Client{ + c := &Client{ cfg: cfg, aead: aead, httpClients: NewFrontedClients(cfg.Fronting, pollTimeout, endpoints[0].url), @@ -435,6 +447,7 @@ func New(cfg Config) (*Client, error) { bucketCount: bucketCount, idleSlotsPerBucket: idleSlotsPerBucket, clientVersion: cfg.ClientVersion, + quotaStatePath: strings.TrimSpace(cfg.QuotaStatePath), clientID: clientID, sessions: make(map[[frame.SessionIDLen]byte]*session.Session), inFlight: make(map[[frame.SessionIDLen]byte]bool), @@ -445,7 +458,13 @@ func New(cfg Config) (*Client, error) { coalesceStep: cfg.CoalesceStep, coalesceMax: cfg.CoalesceMax, recoveryProbeAddr: recoveryProbeAddress(cfg), - }, nil + } + if c.quotaStatePath != "" { + if err := c.loadQuotaState(c.quotaStatePath); err != nil { + log.Printf("[carrier] WARN: quota state load skipped: %v", err) + } + } + return c, nil } // NewSession creates a tunneled session for target ("host:port") and registers @@ -855,7 +874,11 @@ func (c *Client) pollOnce(ctx context.Context) bool { if isLikelyNonBatchRelayPayload(respBody) { errReason, errHard := classifyRelayErrorBody(respBody) if errHard { - c.markEndpointHardFailure(endpointIdx) + if strings.Contains(strings.ToLower(errReason), "quota") { + c.markEndpointQuotaExhausted(endpointIdx) + } else { + c.markEndpointHardFailure(endpointIdx) + } } else { c.markEndpointFailure(endpointIdx) } @@ -944,7 +967,7 @@ func (c *Client) pickRelayEndpoint() (int, string) { for i := 0; i < n; i++ { idx := (start + i) % n ep := &c.endpoints[idx] - if ep.blacklistedTill.After(now) { + if c.endpointUnavailableLocked(ep, now) { continue } c.nextEndpoint = (idx + 1) % n @@ -975,7 +998,7 @@ func (c *Client) pickIdleEndpoint() (int, string) { for i := 0; i < n; i++ { idx := (start + i) % n ep := &c.endpoints[idx] - if ep.blacklistedTill.After(now) { + if c.endpointUnavailableLocked(ep, now) { continue } if c.inFlightByBucket[ep.bucket] >= c.idleSlotsPerBucket { @@ -1005,6 +1028,10 @@ func (c *Client) releaseBucketSlot(idx int) { } } +func (c *Client) endpointUnavailableLocked(ep *relayEndpoint, now time.Time) bool { + return ep.blacklistedTill.After(now) || ep.quotaExhaustedUntil.After(now) +} + func (c *Client) resetLocalNetworkFailures() int { c.endpointMu.Lock() defer c.endpointMu.Unlock() @@ -1089,6 +1116,144 @@ func (c *Client) markEndpointHardFailure(endpointIdx int) { c.markEndpointFailureWith(endpointIdx, 5) } +func endpointQuotaKey(ep *relayEndpoint) string { + if ep.account != "" { + return "account:" + ep.account + } + sum := sha256.Sum256([]byte(ep.url)) + return "urlsha256:" + hex.EncodeToString(sum[:]) +} + +type quotaStateFile struct { + Version int `json:"version"` + Entries []quotaStateEntry `json:"entries"` +} + +type quotaStateEntry struct { + Key string `json:"key"` + Until time.Time `json:"until"` +} + +func (c *Client) loadQuotaState(path string) error { + b, err := os.ReadFile(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil + } + return err + } + var state quotaStateFile + if err := json.Unmarshal(b, &state); err != nil { + return err + } + byKey := make(map[string]time.Time, len(state.Entries)) + now := time.Now() + for _, entry := range state.Entries { + key := strings.TrimSpace(entry.Key) + if key == "" || !entry.Until.After(now) { + continue + } + if prev := byKey[key]; prev.IsZero() || entry.Until.After(prev) { + byKey[key] = entry.Until + } + } + if len(byKey) == 0 { + return nil + } + c.endpointMu.Lock() + defer c.endpointMu.Unlock() + for i := range c.endpoints { + ep := &c.endpoints[i] + until := byKey[endpointQuotaKey(ep)] + if until.IsZero() { + continue + } + ep.quotaExhaustedUntil = until + ep.blacklistedTill = until + ep.localNetworkOffline = false + } + return nil +} + +func (c *Client) saveQuotaState() error { + path := strings.TrimSpace(c.quotaStatePath) + if path == "" { + return nil + } + now := time.Now() + byKey := make(map[string]time.Time) + c.endpointMu.Lock() + for i := range c.endpoints { + ep := &c.endpoints[i] + c.touchDailyWindow(ep, now) + if !ep.quotaExhaustedUntil.After(now) { + continue + } + key := endpointQuotaKey(ep) + if prev := byKey[key]; prev.IsZero() || ep.quotaExhaustedUntil.After(prev) { + byKey[key] = ep.quotaExhaustedUntil + } + } + c.endpointMu.Unlock() + + state := quotaStateFile{Version: 1, Entries: make([]quotaStateEntry, 0, len(byKey))} + for key, until := range byKey { + state.Entries = append(state.Entries, quotaStateEntry{Key: key, Until: until}) + } + body, err := json.MarshalIndent(state, "", " ") + if err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return err + } + c.quotaStateMu.Lock() + defer c.quotaStateMu.Unlock() + return writeFileAtomic(path, append(body, '\n'), 0o600) +} + +func (c *Client) markEndpointQuotaExhausted(endpointIdx int) { + c.endpointMu.Lock() + if endpointIdx < 0 || endpointIdx >= len(c.endpoints) { + c.endpointMu.Unlock() + return + } + now := time.Now() + ep := &c.endpoints[endpointIdx] + c.touchDailyWindow(ep, now) + resetAt := ep.dailyResetAt + if resetAt.IsZero() { + resetAt = nextQuotaReset(now) + } + url := ep.url + account := ep.account + scopeAccount := account != "" + for i := range c.endpoints { + if scopeAccount { + if c.endpoints[i].account != account { + continue + } + } else if i != endpointIdx { + continue + } + peer := &c.endpoints[i] + c.touchDailyWindow(peer, now) + peer.quotaExhaustedUntil = resetAt + peer.blacklistedTill = resetAt + peer.localNetworkOffline = false + if i == endpointIdx { + peer.failCount++ + peer.statsFail++ + } + } + c.endpointMu.Unlock() + log.Printf("[carrier] endpoint %s account=%q quota exhausted until approx %s; rotating away", + shortScriptKey(url), account, resetAt.Format(time.RFC3339)) + if err := c.saveQuotaState(); err != nil { + log.Printf("[carrier] WARN: quota state save failed: %v", err) + } +} + // markEndpointFailureWith is the shared implementation. minFailCount is a floor // applied before incrementing so callers can skip the slow 3-48 s ramp for // failure classes known not to self-heal quickly (quota, auth, rate-limit). diff --git a/internal/carrier/quota_test.go b/internal/carrier/quota_test.go index e260b38..6db35af 100644 --- a/internal/carrier/quota_test.go +++ b/internal/carrier/quota_test.go @@ -1,6 +1,8 @@ package carrier import ( + "os" + "path/filepath" "strings" "testing" "time" @@ -134,3 +136,48 @@ func TestRecordScriptStatsFromBody_RecoveryAfterRedeploy(t *testing.T) { t.Fatalf("scriptCount=%d want 7", c.endpoints[0].scriptCount) } } + +func TestMarkEndpointQuotaExhaustedScopesToAccountAndPersists(t *testing.T) { + path := filepath.Join(t.TempDir(), "quota.json") + now := time.Now() + c := &Client{ + quotaStatePath: path, + endpoints: []relayEndpoint{ + {url: "u1", account: "A", dailyResetAt: now.Add(time.Hour)}, + {url: "u2", account: "A", dailyResetAt: now.Add(time.Hour)}, + {url: "u3", account: "B", dailyResetAt: now.Add(time.Hour)}, + }, + } + + c.markEndpointQuotaExhausted(0) + + if !c.endpoints[0].quotaExhaustedUntil.After(now) { + t.Fatalf("endpoint 0 was not quota-quarantined") + } + if !c.endpoints[1].quotaExhaustedUntil.Equal(c.endpoints[0].quotaExhaustedUntil) { + t.Fatalf("same-account endpoint was not quarantined with endpoint 0") + } + if !c.endpoints[2].quotaExhaustedUntil.IsZero() { + t.Fatalf("different account was quarantined: %v", c.endpoints[2].quotaExhaustedUntil) + } + if _, err := os.Stat(path); err != nil { + t.Fatalf("quota state was not persisted: %v", err) + } + + restored := &Client{ + endpoints: []relayEndpoint{ + {url: "u1", account: "A"}, + {url: "u2", account: "A"}, + {url: "u3", account: "B"}, + }, + } + if err := restored.loadQuotaState(path); err != nil { + t.Fatalf("loadQuotaState: %v", err) + } + if !restored.endpoints[0].blacklistedTill.After(now) || !restored.endpoints[1].blacklistedTill.After(now) { + t.Fatalf("restored same-account quota quarantine = %#v", restored.endpoints) + } + if !restored.endpoints[2].blacklistedTill.IsZero() { + t.Fatalf("restored different account blacklist = %v", restored.endpoints[2].blacklistedTill) + } +} diff --git a/internal/carrier/stats.go b/internal/carrier/stats.go index 4cd9e98..3e49705 100644 --- a/internal/carrier/stats.go +++ b/internal/carrier/stats.go @@ -96,6 +96,10 @@ func (c *Client) endpointStatsLine() string { remaining := time.Until(ep.blacklistedTill).Round(time.Second) part = fmt.Sprintf("%s bl=%s", part, remaining) } + if ep.quotaExhaustedUntil.After(now) { + remaining := time.Until(ep.quotaExhaustedUntil).Round(time.Second) + part = fmt.Sprintf("%s quota_reset=%s", part, remaining) + } parts = append(parts, part) } return strings.Join(parts, " | ") diff --git a/internal/config/client.go b/internal/config/client.go index 6e56d24..ff1b447 100644 --- a/internal/config/client.go +++ b/internal/config/client.go @@ -11,6 +11,7 @@ import ( "net" "net/url" "os" + "path/filepath" "strconv" "strings" ) @@ -35,6 +36,11 @@ type Client struct { // empty string = unlabeled. Always the same length as ScriptURLs. ScriptAccounts []string + // QuotaStatePath persists durable Apps Script quota quarantines so a + // restart does not immediately hammer an account already known to be over + // its daily UrlFetch quota. Empty disables persistence. + QuotaStatePath string + // Adaptive uplink coalescing. When CoalesceStepMs > 0, the carrier waits // up to that many ms for more TX ops to arrive before sending, resetting // the timer on each new op. Bursts collapse into a single poll. Off by @@ -87,6 +93,8 @@ type clientFile struct { // is spending its time. Off by default. DebugTiming bool `json:"debug_timing"` + QuotaStatePath *string `json:"quota_state_path"` + // Optional SOCKS5 RFC 1929 credentials. When set, clients must supply // these credentials or the connection is rejected. Both must be non-empty // together — setting only one is an error. @@ -385,6 +393,14 @@ func LoadClient(path string) (*Client, error) { return nil, fmt.Errorf("socks_user and socks_pass must both be set or both be empty in %s", path) } + quotaStatePath := ".goose-quota-state.json" + if f.QuotaStatePath != nil { + quotaStatePath = strings.TrimSpace(*f.QuotaStatePath) + } + if quotaStatePath != "" && !filepath.IsAbs(quotaStatePath) { + quotaStatePath = filepath.Join(filepath.Dir(path), quotaStatePath) + } + if f.CoalesceStepMs < 0 { return nil, fmt.Errorf("coalesce_step_ms must be >= 0 in %s (got %d)", path, f.CoalesceStepMs) } @@ -398,19 +414,20 @@ func LoadClient(path string) (*Client, error) { } c := Client{ - ListenAddr: net.JoinHostPort(listenHost, strconv.Itoa(listenPort)), - GoogleIP: googleIP, - SNIHosts: sniHosts, - ScriptURLs: scriptURLs, - ScriptAccounts: scriptAccounts, - UseFronting: useFronting, - AESKeyHex: key, - DebugTiming: f.DebugTiming, - SocksUser: socksUser, - SocksPass: socksPass, - CoalesceStepMs: f.CoalesceStepMs, - CoalesceMaxMs: coalesceMax, - IdleSlotsPerBucket: f.IdleSlotsPerBucket, + ListenAddr: net.JoinHostPort(listenHost, strconv.Itoa(listenPort)), + GoogleIP: googleIP, + SNIHosts: sniHosts, + ScriptURLs: scriptURLs, + ScriptAccounts: scriptAccounts, + UseFronting: useFronting, + AESKeyHex: key, + DebugTiming: f.DebugTiming, + QuotaStatePath: quotaStatePath, + SocksUser: socksUser, + SocksPass: socksPass, + CoalesceStepMs: f.CoalesceStepMs, + CoalesceMaxMs: coalesceMax, + IdleSlotsPerBucket: f.IdleSlotsPerBucket, } return &c, nil } diff --git a/internal/config/client_test.go b/internal/config/client_test.go new file mode 100644 index 0000000..69c741e --- /dev/null +++ b/internal/config/client_test.go @@ -0,0 +1,54 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestLoadClient_QuotaStatePathDefaultsBesideConfig(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "client_config.json") + key := strings.Repeat("a", 64) + deploymentID := "AKfycb" + strings.Repeat("x", 60) + body := `{ + "script_keys": [{"id": "` + deploymentID + `", "account": "acct-a"}], + "tunnel_key": "` + key + `" +}` + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + t.Fatal(err) + } + + cfg, err := LoadClient(path) + if err != nil { + t.Fatalf("LoadClient: %v", err) + } + want := filepath.Join(dir, ".goose-quota-state.json") + if cfg.QuotaStatePath != want { + t.Fatalf("QuotaStatePath = %q, want %q", cfg.QuotaStatePath, want) + } +} + +func TestLoadClient_QuotaStatePathCanBeDisabled(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "client_config.json") + key := strings.Repeat("b", 64) + deploymentID := "AKfycb" + strings.Repeat("y", 60) + body := `{ + "script_keys": [{"id": "` + deploymentID + `", "account": "acct-a"}], + "quota_state_path": "", + "tunnel_key": "` + key + `" +}` + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + t.Fatal(err) + } + + cfg, err := LoadClient(path) + if err != nil { + t.Fatalf("LoadClient: %v", err) + } + if cfg.QuotaStatePath != "" { + t.Fatalf("QuotaStatePath = %q, want disabled empty path", cfg.QuotaStatePath) + } +}