vct_core/quota/
refresh.rs

1//! Shared OAuth token-refresh primitives used by the Claude / Codex quota
2//! fetchers.
3//!
4//! The provider-specific request/response shapes live in each fetcher; this
5//! module owns the cross-cutting concerns:
6//!
7//! - a per-provider refresh **backoff** so a revoked token cannot make the
8//!   worker hammer the token endpoint every tick (`RefreshCooldown`),
9//! - **write-back that preserves unknown fields** by mutating a whole
10//!   `serde_json::Value` and re-checking the file mtime just before the
11//!   atomic write (TOCTOU guard against a concurrently-running official CLI),
12//! - the near-expiry check and the RFC3339 timestamp formats the provider
13//!   CLIs write.
14//!
15//! No function here ever logs a request/response body or a token; callers only
16//! surface the HTTP status.
17
18use anyhow::{Context, Result};
19use reqwest::StatusCode;
20use reqwest::blocking::RequestBuilder;
21use serde_json::Value;
22use std::path::Path;
23use std::time::SystemTime;
24
25/// Cooldown after a refresh failure, so a revoked/rotated-away refresh token
26/// cannot make the 10s worker retry the token endpoint every tick (B1).
27pub const REFRESH_COOLDOWN_SECS: i64 = 300;
28
29/// Skew applied to proactive expiry checks, to absorb a slightly-fast clock.
30pub const EXPIRY_SKEW_SECS: i64 = 60;
31
32/// Per-provider refresh backoff keyed on the credential file's mtime.
33///
34/// After a refresh failure the worker arms the cooldown; while it is active a
35/// refresh is skipped (the panel shows the login hint) — *unless* the
36/// credential file's mtime changes, which means the user re-logged in or the
37/// official CLI rotated the token, so a retry is worthwhile immediately.
38#[derive(Default)]
39pub struct RefreshCooldown {
40    until: i64,
41    mtime: Option<SystemTime>,
42}
43
44impl RefreshCooldown {
45    /// Whether a refresh should be skipped right now.
46    pub fn active(&self, now: i64, cur_mtime: Option<SystemTime>) -> bool {
47        now < self.until && self.mtime == cur_mtime
48    }
49
50    /// Arms the cooldown after a refresh failure.
51    pub fn arm(&mut self, now: i64, mtime: Option<SystemTime>) {
52        self.until = now + REFRESH_COOLDOWN_SECS;
53        self.mtime = mtime;
54    }
55
56    /// Clears the cooldown after a success.
57    pub fn clear(&mut self) {
58        self.until = 0;
59        self.mtime = None;
60    }
61}
62
63/// Returns a file's modification time, or `None` if it cannot be read.
64pub fn file_mtime(path: &Path) -> Option<SystemTime> {
65    std::fs::metadata(path).and_then(|m| m.modified()).ok()
66}
67
68/// True when `expires_at_secs` is known and within `skew_secs` of `now_secs`.
69///
70/// An unknown expiry (`None`) returns `false` — we never refresh blindly, which
71/// would spin the token endpoint (see B1).
72pub fn is_expiring(expires_at_secs: Option<i64>, now_secs: i64, skew_secs: i64) -> bool {
73    expires_at_secs.is_some_and(|exp| exp - now_secs <= skew_secs)
74}
75
76/// Sends a refresh request and returns `(status, body)` without ever logging
77/// the body (which echoes tokens). The caller checks the status.
78///
79/// # Errors
80///
81/// Returns an error only if the request could not be sent (network).
82pub fn send_refresh(req: RequestBuilder) -> Result<(StatusCode, String)> {
83    let resp = req.send().context("token refresh request failed")?;
84    let status = resp.status();
85    let text = resp.text().unwrap_or_default();
86    Ok((status, text))
87}
88
89/// Reads a credential file, applies `mutate` to only the fields it touches, and
90/// writes the whole `Value` back atomically (pretty), preserving every other
91/// key.
92///
93/// Re-checks the file's mtime against `expected_mtime` just before writing and
94/// **aborts the write** (returns `Ok(false)`) if it changed — a concurrently
95/// running official CLI may have rotated the token, and clobbering it would log
96/// that CLI out. Returns `Ok(true)` when the write happened.
97///
98/// # Errors
99///
100/// Returns an error if the file cannot be read or parsed, or the write fails.
101pub fn update_json_file_in_place(
102    path: &Path,
103    expected_mtime: Option<SystemTime>,
104    mutate: impl FnOnce(&mut Value) -> Result<()>,
105) -> Result<bool> {
106    // TOCTOU guard: bail before reading if the file changed under us.
107    if file_mtime(path) != expected_mtime {
108        return Ok(false);
109    }
110    let body = std::fs::read_to_string(path)
111        .with_context(|| format!("Failed to read {}", path.display()))?;
112    let mut root: Value = serde_json::from_str(&body).context("Failed to parse credential file")?;
113    mutate(&mut root)?;
114    // Re-check right before writing: a concurrent official CLI may have rotated
115    // the credential during the read/mutate above. This shrinks the TOCTOU
116    // window to the serialize+rename below (fully closing it would need file
117    // locking).
118    if file_mtime(path) != expected_mtime {
119        return Ok(false);
120    }
121    crate::utils::write_json_atomic_pretty(path, &root)?;
122    Ok(true)
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn is_expiring_unknown_is_never_expiring() {
131        assert!(!is_expiring(None, 1000, 60));
132    }
133
134    #[test]
135    fn is_expiring_boundaries() {
136        // Expires in 61s, skew 60 → not expiring yet.
137        assert!(!is_expiring(Some(1061), 1000, 60));
138        // Expires in exactly 60s → expiring.
139        assert!(is_expiring(Some(1060), 1000, 60));
140        // Already expired → expiring.
141        assert!(is_expiring(Some(900), 1000, 60));
142    }
143
144    #[test]
145    fn cooldown_respects_mtime_change() {
146        let mt0 = SystemTime::UNIX_EPOCH;
147        let mt1 = SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(5);
148        let mut cd = RefreshCooldown::default();
149        cd.arm(1000, Some(mt0));
150        // Within cooldown, same mtime → active (skip refresh).
151        assert!(cd.active(1100, Some(mt0)));
152        // Within cooldown but mtime changed → not active (retry).
153        assert!(!cd.active(1100, Some(mt1)));
154        // Past cooldown → not active.
155        assert!(!cd.active(2000, Some(mt0)));
156    }
157
158    #[test]
159    fn update_in_place_preserves_unknown_and_aborts_on_mtime_change() {
160        use std::io::Write;
161        let dir = tempfile::tempdir().unwrap();
162        let path = dir.path().join("auth.json");
163        let mut f = std::fs::File::create(&path).unwrap();
164        write!(
165            f,
166            r#"{{"auth_mode":"chatgpt","OPENAI_API_KEY":null,"tokens":{{"access_token":"old","refresh_token":"oldr"}},"extra":123}}"#
167        )
168        .unwrap();
169        drop(f);
170
171        let mtime = file_mtime(&path);
172        let wrote = update_json_file_in_place(&path, mtime, |root| {
173            let t = root
174                .get_mut("tokens")
175                .and_then(|v| v.as_object_mut())
176                .unwrap();
177            t.insert("access_token".into(), serde_json::json!("new"));
178            t.insert("refresh_token".into(), serde_json::json!("newr"));
179            Ok(())
180        })
181        .unwrap();
182        assert!(wrote);
183
184        let v: Value = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
185        assert_eq!(v["tokens"]["access_token"], "new");
186        assert_eq!(v["tokens"]["refresh_token"], "newr");
187        // Unknown fields survive.
188        assert_eq!(v["auth_mode"], "chatgpt");
189        assert_eq!(v["extra"], 123);
190        assert!(v.get("OPENAI_API_KEY").is_some());
191
192        // A stale expected_mtime aborts the write.
193        let wrote2 = update_json_file_in_place(&path, Some(SystemTime::UNIX_EPOCH), |root| {
194            root["tokens"]["access_token"] = serde_json::json!("should-not-write");
195            Ok(())
196        })
197        .unwrap();
198        assert!(!wrote2);
199        let v2: Value = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
200        assert_eq!(v2["tokens"]["access_token"], "new");
201    }
202
203    #[test]
204    fn claude_write_back_preserves_design_oauth_and_ms_expiry() {
205        use std::io::Write;
206        let dir = tempfile::tempdir().unwrap();
207        let path = dir.path().join(".credentials.json");
208        let mut f = std::fs::File::create(&path).unwrap();
209        write!(
210            f,
211            r#"{{"claudeAiOauth":{{"accessToken":"old","refreshToken":"oldr","expiresAt":1,"scopes":["user:inference"],"subscriptionType":"max","rateLimitTier":"t"}},"designOauth":{{"accessToken":"design"}}}}"#
212        )
213        .unwrap();
214        drop(f);
215
216        let mtime = file_mtime(&path);
217        let wrote = update_json_file_in_place(&path, mtime, |root| {
218            let o = root
219                .get_mut("claudeAiOauth")
220                .and_then(|v| v.as_object_mut())
221                .unwrap();
222            o.insert("accessToken".into(), serde_json::json!("newacc"));
223            o.insert("refreshToken".into(), serde_json::json!("newref"));
224            o.insert("expiresAt".into(), serde_json::json!(1783108188604i64));
225            Ok(())
226        })
227        .unwrap();
228        assert!(wrote);
229
230        let v: Value = serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
231        assert_eq!(v["claudeAiOauth"]["accessToken"], "newacc");
232        assert_eq!(v["claudeAiOauth"]["refreshToken"], "newref");
233        // expiresAt stays a NUMBER (ms), not a string.
234        assert_eq!(v["claudeAiOauth"]["expiresAt"], 1783108188604i64);
235        assert!(v["claudeAiOauth"]["expiresAt"].is_number());
236        // Preserved siblings.
237        assert_eq!(v["claudeAiOauth"]["subscriptionType"], "max");
238        assert_eq!(v["claudeAiOauth"]["rateLimitTier"], "t");
239        assert_eq!(v["designOauth"]["accessToken"], "design");
240    }
241}
vct_core/quota/refresh.rs

vct_core/quota/
refresh.rs
