agent_core/core/auth/
mod.rs

1//! OAuth 2.0 Authorization Code + PKCE flow for Anthropic (Claude Pro/Max).
2//!
3//! Implements the same flow as Claude Code and Pi coding agent:
4//! 1. Generate PKCE verifier + challenge
5//! 2. Start localhost callback server
6//! 3. Open browser to claude.ai/oauth/authorize
7//! 4. Capture redirect with auth code
8//! 5. Exchange code for access + refresh tokens
9//! 6. Save to ~/.pi/agent/auth.json (shared with Pi)
10
11use serde::{Deserialize, Serialize};
12use tokio::sync::oneshot;
13
14mod pkce;
15mod callback;
16mod token;
17mod storage;
18mod browser;
19mod openai_codex;
20mod credential_source;
21
22// ── Re-exports ──────────────────────────────────────────────────────────────────
23
24pub use pkce::{generate_code_verifier, generate_code_challenge, generate_state, build_auth_url};
25pub use callback::{CallbackServerHandle, start_callback_server};
26pub use token::{exchange_code_for_tokens, refresh_token, ensure_fresh_token, ensure_fresh_provider_token};
27pub use storage::{auth_file_path, load_auth, load_provider_auth, save_auth, save_provider_auth};
28pub use browser::open_browser;
29pub use openai_codex::{extract_account_id as extract_codex_account_id, login as login_openai_codex};
30pub use credential_source::{CredentialSource, BrokerToken, BrokerClient, TokenCache, TokenFetcher, resolve_remote, resolve_remote_token, resolve_access_token, is_expired_with_margin, DEFAULT_MARGIN_MS};
31
32// ── Constants (match Claude Code / Pi) ──────────────────────────────────────
33
34pub(super) const CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
35pub(super) const AUTHORIZE_URL: &str = "https://claude.ai/oauth/authorize";
36pub(super) const TOKEN_URL: &str = "https://platform.claude.com/v1/oauth/token";
37pub(super) const CALLBACK_HOST: &str = "127.0.0.1";
38pub(super) const CALLBACK_PORT: u16 = 53692;
39pub(super) const SCOPES: &str = "org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload";
40
41// ── Types ───────────────────────────────────────────────────────────────────
42
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct OAuthCredentials {
45    #[serde(rename = "type")]
46    pub auth_type: String,
47    pub refresh: String,
48    pub access: String,
49    pub expires: u64,
50    #[serde(rename = "accountId", skip_serializing_if = "Option::is_none")]
51    pub account_id: Option<String>,
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct AuthFile {
56    pub anthropic: OAuthCredentials,
57    #[serde(rename = "openai-codex", default, skip_serializing_if = "Option::is_none")]
58    pub openai_codex: Option<OAuthCredentials>,
59}
60
61#[derive(Debug, Deserialize)]
62pub(crate) struct TokenResponse {
63    pub(crate) access_token: String,
64    pub(crate) refresh_token: String,
65    pub(crate) expires_in: u64,
66}
67
68/// Result from the OAuth callback.
69#[derive(Debug, Clone)]
70pub struct CallbackResult {
71    pub code: String,
72    pub state: String,
73}
74
75/// Check if the current token is expired.
76///
77/// Note the effective ~5-minute safety margin lives in the stored value, not
78/// here: `refresh_token`/`exchange_code_for_tokens` bake a 5-minute buffer into
79/// `expires` (`now + expires_in*1000 - 5min`). So a bare `now >= expires` check
80/// already fires ~5 minutes before the credential actually dies at the provider.
81pub fn is_token_expired(creds: &OAuthCredentials) -> bool {
82    now_millis() >= creds.expires
83}
84
85// ── Helpers ─────────────────────────────────────────────────────────────────
86
87pub(crate) fn now_millis() -> u64 {
88    crate::epoch_millis()
89}
90
91/// Try to parse manual input as a redirect URL or code#state shorthand.
92///
93/// Returns `(Some(code), Some(state))` only when both components are present
94/// and unambiguous. The `None` state path has been intentionally removed:
95/// previously the caller defaulted a missing state to the expected value
96/// (`parsed_state.unwrap_or(manual_state)`), which made the downstream
97/// `result.state != state` check compare a value against itself — CSRF
98/// protection silently nullified. By never returning `Some(code)` without
99/// `Some(state)`, the CSRF guard is guaranteed to do real work.
100fn parse_manual_input(input: &str) -> (Option<String>, Option<String>) {
101    let trimmed = input.trim();
102
103    // Try as full URL (e.g. the redirect from the browser)
104    if let Ok(url) = url::Url::parse(trimmed) {
105        let code = url.query_pairs().find(|(k, _)| k == "code").map(|(_, v)| v.to_string());
106        let state = url.query_pairs().find(|(k, _)| k == "state").map(|(_, v)| v.to_string());
107        if code.is_some() {
108            return (code, state);
109        }
110    }
111
112    // Try as "code#state" format (Claude Code manual flow)
113    if trimmed.contains('#') {
114        let parts: Vec<&str> = trimmed.splitn(2, '#').collect();
115        if parts.len() == 2 && !parts[0].is_empty() && !parts[1].is_empty() {
116            return (Some(parts[0].to_string()), Some(parts[1].to_string()));
117        }
118    }
119
120    // Bare code with no state: REJECTED. Accepting it and defaulting the state
121    // to `manual_state` would cause the CSRF check to compare a value to itself.
122    // The caller (`manual_paste_to_callback`) will return `None` for this input.
123    (None, None)
124}
125
126/// Validate user-pasted OAuth authorization input for the Anthropic manual
127/// fallback flow.
128///
129/// Returns `Some(CallbackResult)` only if the input contains BOTH a `code`
130/// and a `state`. Previously the caller defaulted a missing `state` to the
131/// original CSRF token (`parsed_state.unwrap_or(manual_state)`), which
132/// silently bypassed the CSRF check — the state comparison became
133/// `manual_state == state` (always true). By requiring an explicit state
134/// here, a bare code paste or a URL without `state` is rejected outright.
135fn manual_paste_to_callback(input: &str) -> Option<CallbackResult> {
136    let (code, state) = parse_manual_input(input);
137    Some(CallbackResult { code: code?, state: state? })
138}
139
140// ── High-level login flow ───────────────────────────────────────────────────
141
142/// Run the full OAuth login flow. Returns saved credentials.
143pub async fn login() -> std::result::Result<OAuthCredentials, String> {
144    let port = CALLBACK_PORT;
145
146    // 1. Generate PKCE
147    let verifier = generate_code_verifier();
148    let challenge = generate_code_challenge(&verifier);
149    let state = generate_state();
150
151    // 2. Start callback server
152    let (rx, server_handle) = start_callback_server(state.clone(), port).await?;
153
154    // 3. Build URL and open browser
155    let auth_url = build_auth_url(&challenge, &state, port);
156
157    eprintln!("\n\x1b[1mOpening browser to sign in...\x1b[0m\n");
158
159    if let Err(e) = open_browser(&auth_url) {
160        eprintln!("Could not open browser automatically: {}", e);
161    }
162
163    eprintln!("\x1b[2mIf the browser didn't open, visit this URL:\x1b[0m");
164    eprintln!("\x1b[36m{}\x1b[0m\n", auth_url);
165
166    // Also provide manual paste option
167    let (manual_tx, manual_rx) = oneshot::channel::<CallbackResult>();
168    let stdin_task = tokio::spawn(async move {
169        eprintln!("\x1b[2mOr paste the redirect URL here (must include `state`):\x1b[0m");
170
171        let mut line = String::new();
172        let result = tokio::task::spawn_blocking(move || {
173            std::io::stdin().read_line(&mut line).ok();
174            line.trim().to_string()
175        })
176        .await;
177
178        if let Ok(input) = result {
179            match manual_paste_to_callback(&input) {
180                Some(callback) => {
181                    let _ = manual_tx.send(callback);
182                }
183                None if !input.is_empty() => {
184                    eprintln!(
185                        "\x1b[31m✗ Pasted input did not contain both `code` and `state`.\x1b[0m"
186                    );
187                    eprintln!(
188                        "\x1b[2m  Paste the full redirect URL (e.g. http://localhost:53692/callback?code=…&state=…).\x1b[0m"
189                    );
190                }
191                None => {}
192            }
193        }
194    });
195
196    // 4. Wait for either callback or manual input
197    let result = tokio::select! {
198        callback = rx => {
199            match callback {
200                Ok(result) => result,
201                Err(_) => return Err("Callback channel closed".to_string()),
202            }
203        }
204        manual = manual_rx => {
205            match manual {
206                Ok(result) => result,
207                Err(_) => return Err("Manual input channel closed".to_string()),
208            }
209        }
210    };
211
212    stdin_task.abort();
213
214    // 5. Verify state
215    if result.state != state {
216        server_handle.shutdown().await;
217        return Err("OAuth state mismatch — possible CSRF attack".to_string());
218    }
219
220    eprintln!("\n\x1b[1mExchanging code for tokens...\x1b[0m");
221
222    // 6. Exchange code for tokens
223    let creds = exchange_code_for_tokens(&result.code, &result.state, &verifier, port).await?;
224
225    // 7. Shut down callback server
226    server_handle.shutdown().await;
227
228    // 8. Save to auth.json
229    save_auth(&creds)?;
230
231    Ok(creds)
232}
233
234// ── Anthropic-specific CSRF regression tests ─────────────────────────────────
235//
236// See also: openai_codex::tests for the parallel codex test suite.
237// These mirror the codex tests — same invariants, same naming pattern.
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
243
244    #[test]
245    fn test_generate_code_verifier() {
246        let verifier = generate_code_verifier();
247        assert!(!verifier.is_empty(), "Code verifier should not be empty");
248        assert!(verifier.len() > 20, "Code verifier should be longer than 20 characters");
249        let verifier2 = generate_code_verifier();
250        assert_ne!(verifier, verifier2, "Two calls should produce different verifiers");
251    }
252
253    #[test]
254    fn test_generate_code_challenge() {
255        let verifier = "test_verifier_123";
256        let challenge = generate_code_challenge(verifier);
257        assert!(!challenge.is_empty(), "Code challenge should not be empty");
258        let challenge2 = generate_code_challenge(verifier);
259        assert_eq!(challenge, challenge2, "Same verifier should produce same challenge");
260        let different_challenge = generate_code_challenge("different_verifier_456");
261        assert_ne!(challenge, different_challenge, "Different verifiers should produce different challenges");
262    }
263
264    #[test]
265    fn test_generate_state() {
266        let state = generate_state();
267        assert!(!state.is_empty(), "State should not be empty");
268        let state2 = generate_state();
269        assert_ne!(state, state2, "Two calls should produce different states");
270    }
271
272    #[test]
273    fn test_build_auth_url() {
274        let challenge = "test_challenge";
275        let state = "test_state";
276        let port = 8080;
277        let url = build_auth_url(challenge, state, port);
278        assert!(url.contains("claude.ai/oauth/authorize"));
279        assert!(url.contains("client_id=9d1c250a-e61b-44d9-88ed-5944d1962f5e"));
280        assert!(url.contains(&format!("code_challenge={}", challenge)));
281        assert!(url.contains(&format!("state={}", state)));
282        assert!(url.contains("localhost"));
283        assert!(url.contains(&port.to_string()));
284        assert!(url.contains("redirect_uri="));
285    }
286
287    #[test]
288    fn test_is_token_expired() {
289        let expired_creds = OAuthCredentials {
290            auth_type: "oauth".to_string(),
291            refresh: "test_refresh".to_string(),
292            access: "test_access".to_string(),
293            expires: 0,
294            account_id: None,
295        };
296        assert!(is_token_expired(&expired_creds));
297
298        let future_time = now_millis() + 3600000;
299        let fresh_creds = OAuthCredentials {
300            auth_type: "oauth".to_string(),
301            refresh: "test_refresh".to_string(),
302            access: "test_access".to_string(),
303            expires: future_time,
304            account_id: None,
305        };
306        assert!(!is_token_expired(&fresh_creds));
307        assert_eq!(fresh_creds.auth_type, "oauth");
308    }
309
310    #[test]
311    fn test_pkce_challenge_sha256() {
312        let verifier = "test_verifier_string";
313        let challenge = generate_code_challenge(verifier);
314        
315        use sha2::{Sha256, Digest};
316        let mut hasher = Sha256::new();
317        hasher.update(verifier.as_bytes());
318        let hash = hasher.finalize();
319        let expected = URL_SAFE_NO_PAD.encode(hash);
320        
321        assert_eq!(challenge, expected);
322    }
323
324    #[test]
325    fn test_code_verifier_length() {
326        let verifier = generate_code_verifier();
327        assert_eq!(verifier.len(), 43);
328    }
329
330    #[test]
331    fn test_state_length() {
332        let state = generate_state();
333        assert_eq!(state.len(), 43);
334    }
335
336    #[test]
337    fn test_build_auth_url_required_params() {
338        let url = build_auth_url("test_challenge", "test_state", 8080);
339        assert!(url.contains("response_type=code"));
340        assert!(url.contains("code_challenge_method=S256"));
341        assert!(url.contains("scope="));
342        assert!(url.contains("redirect_uri="));
343        assert!(url.contains("8080"));
344    }
345
346    #[test]
347    fn test_is_token_expired_edge_cases() {
348        let current_time = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_millis() as u64;
349        
350        let exactly_now_creds = OAuthCredentials {
351            auth_type: "oauth".to_string(),
352            refresh: "test_refresh".to_string(),
353            access: "test_access".to_string(),
354            expires: current_time,
355            account_id: None,
356        };
357        assert!(is_token_expired(&exactly_now_creds));
358        
359        let one_ms_future_creds = OAuthCredentials {
360            auth_type: "oauth".to_string(),
361            refresh: "test_refresh".to_string(),
362            access: "test_access".to_string(),
363            expires: current_time + 1,
364            account_id: None,
365        };
366        assert!(!is_token_expired(&one_ms_future_creds));
367    }
368
369    #[test]
370    fn test_auth_file_path() {
371        let path = auth_file_path();
372        let path_str = path.to_string_lossy();
373        assert!(path_str.ends_with("auth.json"));
374    }
375
376    #[test]
377    fn test_oauth_credentials_serialization_roundtrip() {
378        let original_creds = OAuthCredentials {
379            auth_type: "oauth".to_string(),
380            refresh: "test_refresh_token".to_string(),
381            access: "test_access_token".to_string(),
382            expires: 1234567890,
383            account_id: None,
384        };
385        
386        let json = serde_json::to_string(&original_creds).expect("Should serialize");
387        let deserialized_creds: OAuthCredentials = serde_json::from_str(&json).expect("Should deserialize");
388        
389        assert_eq!(original_creds.auth_type, deserialized_creds.auth_type);
390        assert_eq!(original_creds.refresh, deserialized_creds.refresh);
391        assert_eq!(original_creds.access, deserialized_creds.access);
392        assert_eq!(original_creds.expires, deserialized_creds.expires);
393    }
394
395    // ── manual_paste_to_callback: CSRF guard on the Anthropic manual paste path ──
396    //
397    // Mirrors the regression suite in openai_codex::tests (manual_paste_to_callback).
398    // Pre-fix: parse_manual_input returned (Some(code), None) for bare codes,
399    // and the caller did `parsed_state.unwrap_or(manual_state)`, making the
400    // CSRF check `manual_state == state` — always true — bypass.
401    // Fix: removed bare-code fallback from parse_manual_input + extracted
402    // manual_paste_to_callback that requires both code AND state.
403
404    #[test]
405    fn anthropic_manual_paste_accepts_full_redirect_url() {
406        let result = manual_paste_to_callback(
407            "http://localhost:53692/callback?code=abc&state=xyz",
408        )
409        .expect("URL with code+state must be accepted");
410        assert_eq!(result.code, "abc");
411        assert_eq!(result.state, "xyz");
412    }
413
414    #[test]
415    fn anthropic_manual_paste_rejects_bare_code() {
416        // CSRF regression: a bare code with no state was previously defaulted
417        // to `manual_state`, making the CSRF check trivially true.
418        assert!(
419            manual_paste_to_callback("abc123_bare_code").is_none(),
420            "bare code with no state must be rejected — was the root of the CSRF bypass"
421        );
422    }
423
424    #[test]
425    fn anthropic_manual_paste_rejects_url_without_state() {
426        assert!(
427            manual_paste_to_callback("http://localhost:53692/callback?code=abc").is_none(),
428            "URL missing `state` must be rejected — would bypass CSRF check"
429        );
430    }
431
432    #[test]
433    fn anthropic_manual_paste_accepts_code_hash_state() {
434        // "code#state" is the Claude Code shorthand; both components explicit.
435        let result = manual_paste_to_callback("mycode#mystate")
436            .expect("code#state shorthand must be accepted");
437        assert_eq!(result.code, "mycode");
438        assert_eq!(result.state, "mystate");
439    }
440
441    #[test]
442    fn anthropic_manual_paste_rejects_empty_input() {
443        assert!(manual_paste_to_callback("").is_none());
444        assert!(manual_paste_to_callback("   ").is_none());
445    }
446
447    #[test]
448    fn anthropic_manual_paste_rejects_url_with_only_state() {
449        assert!(
450            manual_paste_to_callback("http://localhost:53692/callback?state=xyz").is_none(),
451            "URL missing `code` must be rejected"
452        );
453    }
454}
agent_core/core/auth/mod.rs

agent_core/core/auth/
mod.rs
