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 raw code.
92fn parse_manual_input(input: &str) -> (Option<String>, Option<String>) {
93    let trimmed = input.trim();
94
95    // Try as full URL
96    if let Ok(url) = url::Url::parse(trimmed) {
97        let code = url.query_pairs().find(|(k, _)| k == "code").map(|(_, v)| v.to_string());
98        let state = url.query_pairs().find(|(k, _)| k == "state").map(|(_, v)| v.to_string());
99        if code.is_some() {
100            return (code, state);
101        }
102    }
103
104    // Try as "code#state" format (Claude Code manual flow)
105    if trimmed.contains('#') {
106        let parts: Vec<&str> = trimmed.splitn(2, '#').collect();
107        if parts.len() == 2 && !parts[0].is_empty() && !parts[1].is_empty() {
108            return (Some(parts[0].to_string()), Some(parts[1].to_string()));
109        }
110    }
111
112    // Treat as raw code
113    if !trimmed.is_empty() {
114        return (Some(trimmed.to_string()), None);
115    }
116
117    (None, None)
118}
119
120// ── High-level login flow ───────────────────────────────────────────────────
121
122/// Run the full OAuth login flow. Returns saved credentials.
123pub async fn login() -> std::result::Result<OAuthCredentials, String> {
124    let port = CALLBACK_PORT;
125
126    // 1. Generate PKCE
127    let verifier = generate_code_verifier();
128    let challenge = generate_code_challenge(&verifier);
129    let state = generate_state();
130
131    // 2. Start callback server
132    let (rx, server_handle) = start_callback_server(state.clone(), port).await?;
133
134    // 3. Build URL and open browser
135    let auth_url = build_auth_url(&challenge, &state, port);
136
137    eprintln!("\n\x1b[1mOpening browser to sign in...\x1b[0m\n");
138
139    if let Err(e) = open_browser(&auth_url) {
140        eprintln!("Could not open browser automatically: {}", e);
141    }
142
143    eprintln!("\x1b[2mIf the browser didn't open, visit this URL:\x1b[0m");
144    eprintln!("\x1b[36m{}\x1b[0m\n", auth_url);
145
146    // Also provide manual paste option
147    let (manual_tx, manual_rx) = oneshot::channel::<CallbackResult>();
148    let manual_state = state.clone();
149    let stdin_task = tokio::spawn(async move {
150        eprintln!("\x1b[2mOr paste the authorization code here:\x1b[0m");
151
152        let mut line = String::new();
153        let result = tokio::task::spawn_blocking(move || {
154            std::io::stdin().read_line(&mut line).ok();
155            line.trim().to_string()
156        })
157        .await;
158
159        if let Ok(input) = result {
160            if !input.is_empty() {
161                let (code, parsed_state) = parse_manual_input(&input);
162                if let Some(code) = code {
163                    let _ = manual_tx.send(CallbackResult {
164                        code,
165                        state: parsed_state.unwrap_or(manual_state),
166                    });
167                }
168            }
169        }
170    });
171
172    // 4. Wait for either callback or manual input
173    let result = tokio::select! {
174        callback = rx => {
175            match callback {
176                Ok(result) => result,
177                Err(_) => return Err("Callback channel closed".to_string()),
178            }
179        }
180        manual = manual_rx => {
181            match manual {
182                Ok(result) => result,
183                Err(_) => return Err("Manual input channel closed".to_string()),
184            }
185        }
186    };
187
188    stdin_task.abort();
189
190    // 5. Verify state
191    if result.state != state {
192        server_handle.shutdown().await;
193        return Err("OAuth state mismatch — possible CSRF attack".to_string());
194    }
195
196    eprintln!("\n\x1b[1mExchanging code for tokens...\x1b[0m");
197
198    // 6. Exchange code for tokens
199    let creds = exchange_code_for_tokens(&result.code, &result.state, &verifier, port).await?;
200
201    // 7. Shut down callback server
202    server_handle.shutdown().await;
203
204    // 8. Save to auth.json
205    save_auth(&creds)?;
206
207    Ok(creds)
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
214
215    #[test]
216    fn test_generate_code_verifier() {
217        let verifier = generate_code_verifier();
218        assert!(!verifier.is_empty(), "Code verifier should not be empty");
219        assert!(verifier.len() > 20, "Code verifier should be longer than 20 characters");
220        let verifier2 = generate_code_verifier();
221        assert_ne!(verifier, verifier2, "Two calls should produce different verifiers");
222    }
223
224    #[test]
225    fn test_generate_code_challenge() {
226        let verifier = "test_verifier_123";
227        let challenge = generate_code_challenge(verifier);
228        assert!(!challenge.is_empty(), "Code challenge should not be empty");
229        let challenge2 = generate_code_challenge(verifier);
230        assert_eq!(challenge, challenge2, "Same verifier should produce same challenge");
231        let different_challenge = generate_code_challenge("different_verifier_456");
232        assert_ne!(challenge, different_challenge, "Different verifiers should produce different challenges");
233    }
234
235    #[test]
236    fn test_generate_state() {
237        let state = generate_state();
238        assert!(!state.is_empty(), "State should not be empty");
239        let state2 = generate_state();
240        assert_ne!(state, state2, "Two calls should produce different states");
241    }
242
243    #[test]
244    fn test_build_auth_url() {
245        let challenge = "test_challenge";
246        let state = "test_state";
247        let port = 8080;
248        let url = build_auth_url(challenge, state, port);
249        assert!(url.contains("claude.ai/oauth/authorize"));
250        assert!(url.contains("client_id=9d1c250a-e61b-44d9-88ed-5944d1962f5e"));
251        assert!(url.contains(&format!("code_challenge={}", challenge)));
252        assert!(url.contains(&format!("state={}", state)));
253        assert!(url.contains("localhost"));
254        assert!(url.contains(&port.to_string()));
255        assert!(url.contains("redirect_uri="));
256    }
257
258    #[test]
259    fn test_is_token_expired() {
260        let expired_creds = OAuthCredentials {
261            auth_type: "oauth".to_string(),
262            refresh: "test_refresh".to_string(),
263            access: "test_access".to_string(),
264            expires: 0,
265            account_id: None,
266        };
267        assert!(is_token_expired(&expired_creds));
268
269        let future_time = now_millis() + 3600000;
270        let fresh_creds = OAuthCredentials {
271            auth_type: "oauth".to_string(),
272            refresh: "test_refresh".to_string(),
273            access: "test_access".to_string(),
274            expires: future_time,
275            account_id: None,
276        };
277        assert!(!is_token_expired(&fresh_creds));
278        assert_eq!(fresh_creds.auth_type, "oauth");
279    }
280
281    #[test]
282    fn test_pkce_challenge_sha256() {
283        let verifier = "test_verifier_string";
284        let challenge = generate_code_challenge(verifier);
285        
286        use sha2::{Sha256, Digest};
287        let mut hasher = Sha256::new();
288        hasher.update(verifier.as_bytes());
289        let hash = hasher.finalize();
290        let expected = URL_SAFE_NO_PAD.encode(hash);
291        
292        assert_eq!(challenge, expected);
293    }
294
295    #[test]
296    fn test_code_verifier_length() {
297        let verifier = generate_code_verifier();
298        assert_eq!(verifier.len(), 43);
299    }
300
301    #[test]
302    fn test_state_length() {
303        let state = generate_state();
304        assert_eq!(state.len(), 43);
305    }
306
307    #[test]
308    fn test_build_auth_url_required_params() {
309        let url = build_auth_url("test_challenge", "test_state", 8080);
310        assert!(url.contains("response_type=code"));
311        assert!(url.contains("code_challenge_method=S256"));
312        assert!(url.contains("scope="));
313        assert!(url.contains("redirect_uri="));
314        assert!(url.contains("8080"));
315    }
316
317    #[test]
318    fn test_is_token_expired_edge_cases() {
319        let current_time = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_millis() as u64;
320        
321        let exactly_now_creds = OAuthCredentials {
322            auth_type: "oauth".to_string(),
323            refresh: "test_refresh".to_string(),
324            access: "test_access".to_string(),
325            expires: current_time,
326            account_id: None,
327        };
328        assert!(is_token_expired(&exactly_now_creds));
329        
330        let one_ms_future_creds = OAuthCredentials {
331            auth_type: "oauth".to_string(),
332            refresh: "test_refresh".to_string(),
333            access: "test_access".to_string(),
334            expires: current_time + 1,
335            account_id: None,
336        };
337        assert!(!is_token_expired(&one_ms_future_creds));
338    }
339
340    #[test]
341    fn test_auth_file_path() {
342        let path = auth_file_path();
343        let path_str = path.to_string_lossy();
344        assert!(path_str.ends_with("auth.json"));
345    }
346
347    #[test]
348    fn test_oauth_credentials_serialization_roundtrip() {
349        let original_creds = OAuthCredentials {
350            auth_type: "oauth".to_string(),
351            refresh: "test_refresh_token".to_string(),
352            access: "test_access_token".to_string(),
353            expires: 1234567890,
354            account_id: None,
355        };
356        
357        let json = serde_json::to_string(&original_creds).expect("Should serialize");
358        let deserialized_creds: OAuthCredentials = serde_json::from_str(&json).expect("Should deserialize");
359        
360        assert_eq!(original_creds.auth_type, deserialized_creds.auth_type);
361        assert_eq!(original_creds.refresh, deserialized_creds.refresh);
362        assert_eq!(original_creds.access, deserialized_creds.access);
363        assert_eq!(original_creds.expires, deserialized_creds.expires);
364    }
365}
agent_core/core/auth/mod.rs

agent_core/core/auth/
mod.rs
