Skip to main content

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
14pub mod account;
15pub mod anthropic_profile;
16pub mod auth_state;
17pub mod aws_bedrock;
18pub mod azure_openai;
19pub mod broker;
20pub mod broker_token;
21mod browser;
22mod callback;
23pub mod cloud;
24pub mod cloud_login;
25mod credential_source;
26pub mod github_copilot;
27pub mod google_gemini;
28pub mod google_vertex;
29pub mod identity;
30pub mod kimi_code;
31pub mod muse;
32mod openai_codex;
33mod pkce;
34pub mod provider;
35pub mod providers;
36pub mod quota_keeper;
37pub mod quota_policy;
38pub mod static_providers;
39pub mod auto_selection;
40pub use auto_selection::{AutoSelectionPolicy, PlanRow, SelectionPlan};
41pub use auth_state::{AuthStateQuery, AuthStateReport, AUTH_STATE_SCHEMA_VERSION};
42pub use quota_policy::{
43    AccountFailure, AccountFailureKind, ObservationFailure, SelectionFailure, SelectionFailureKind,
44};
45mod storage;
46mod token;
47pub mod usage;
48pub mod usage_reads;
49mod xai;
50
51// ── Re-exports ──────────────────────────────────────────────────────────────────
52
53pub use account::{
54    account_config_key, account_env_var, seat_fingerprint, Account, AccountLabel, AccountPolicy,
55    AccountSelector, AccountSummary, CredentialRef, SeatIdentity, AUTO_ACCOUNT_NAME,
56    DEFAULT_ACCOUNT_NAME,
57};
58pub use anthropic_profile::{
59    fetch_anthropic_profile, parse_anthropic_profile, AnthropicProfile, ANTHROPIC_PROFILE_URL,
60};
61pub use broker::{
62    broker_from_source, global_broker, global_broker_install_count, local_auth_state,
63    preflight_cloud_capability, set_global_broker, AccessToken, BrokerError, CredentialBroker,
64    CredentialKind, LocalBroker, PinnedToken, ProviderStatus, ProxyByteStream, ProxyMethod,
65    ProxyRequest, ProxyResponse, RemoteBroker, StaticKeyStatus, MAX_PROXY_REQUEST_BYTES,
66    MAX_PROXY_RESPONSE_BYTES, MAX_UPSTREAM_ERROR_BYTES, PROXY_REQUEST_TIMEOUT,
67};
68pub use broker_token::{
69    is_loopback_endpoint, load_or_create_broker_token, local_broker_token_path,
70    read_local_broker_token, unauthorized_hint, BrokerTokenOrigin, MachineToken,
71    LOCAL_BROKER_TOKEN_FILE,
72};
73pub use browser::open_browser;
74pub use callback::{
75    start_callback_server, start_callback_server_at, CallbackOutcome, CallbackServerHandle,
76};
77pub use cloud::{
78    AuthIdentity, AwsBedrockConfig, AzureOpenAiConfig, BrokerMessage, BrokerOperation, BrokerTool,
79    CloudProviderId, GoogleVertexConfig, InvokeOptions, InvokeRequest, MessageRole, ProviderId,
80};
81pub use credential_source::{
82    is_expired_with_margin, resolve_access_token, resolve_remote, resolve_remote_token,
83    BrokerClient, BrokerToken, CredentialSource, TokenCache, TokenFetcher, DEFAULT_MARGIN_MS,
84};
85pub use github_copilot::login as login_github_copilot;
86pub use identity::{
87    identity_http_client, resolve_seat, supports_seat_identity, ResolvedSeat, SeatResolution,
88};
89pub use kimi_code::login as login_kimi_code;
90pub use muse::login as login_muse;
91pub use openai_codex::{
92    extract_account_id as extract_codex_account_id, extract_email as extract_codex_email,
93    login as login_openai_codex,
94};
95pub use pkce::{build_auth_url, generate_code_challenge, generate_code_verifier, generate_state};
96pub use provider::{
97    BrokerCredentialStrategy, OAuthProviderDescriptor, OAuthProviderId, OAuthProviderRegistry,
98};
99pub use static_providers::{static_provider, StaticProviderSpec, LOCAL_PROVIDER_KEY};
100pub use storage::{
101    any_oauth_credential_present, auth_file_path, find_duplicate_identity, list_accounts,
102    list_accounts_detailed, list_all_accounts, load_auth, load_cloud_state, load_credential,
103    load_provider_auth, load_static_key, remove_credential, save_account_metadata, save_auth,
104    save_cloud_state, save_credential, save_credential_unless_duplicate, save_provider_auth,
105    save_static_key, set_slot_identity, AccountInventory, AccountMetadata, IdentityEvidence,
106    LoginPersistOutcome,
107};
108pub use token::{
109    ensure_fresh_credential, ensure_fresh_provider_token, ensure_fresh_token,
110    exchange_code_for_tokens, refresh_token, RefreshFailure, RefreshFailureKind,
111};
112pub use usage::{fetch_usage, UsageError, UsageSnapshot};
113pub use xai::login as login_xai;
114
115// ── Constants (match Claude Code / Pi) ──────────────────────────────────────
116
117pub(super) const CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
118pub(super) const AUTHORIZE_URL: &str = "https://claude.ai/oauth/authorize";
119pub(super) const TOKEN_URL: &str = "https://platform.claude.com/v1/oauth/token";
120pub(super) const CALLBACK_HOST: &str = "127.0.0.1";
121pub(super) const CALLBACK_PORT: u16 = 53692;
122pub(super) const SCOPES: &str = "org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload";
123
124// ── Types ───────────────────────────────────────────────────────────────────
125
126#[derive(Clone, Serialize, Deserialize)]
127pub struct OAuthCredentials {
128    #[serde(rename = "type")]
129    pub auth_type: String,
130    pub refresh: String,
131    pub access: String,
132    pub expires: u64,
133    #[serde(rename = "accountId", skip_serializing_if = "Option::is_none")]
134    pub account_id: Option<String>,
135}
136
137/// Redacting Debug — never print refresh/access tokens (review finding H1).
138impl std::fmt::Debug for OAuthCredentials {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        f.debug_struct("OAuthCredentials")
141            .field("auth_type", &self.auth_type)
142            .field("refresh", &"[REDACTED]")
143            .field("access", &"[REDACTED]")
144            .field("expires", &self.expires)
145            .field("account_id", &self.account_id)
146            .finish()
147    }
148}
149
150#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct AuthFile {
152    pub anthropic: OAuthCredentials,
153    #[serde(
154        rename = "openai-codex",
155        default,
156        skip_serializing_if = "Option::is_none"
157    )]
158    pub openai_codex: Option<OAuthCredentials>,
159}
160
161#[derive(Debug, Deserialize)]
162pub(crate) struct TokenResponse {
163    pub(crate) access_token: String,
164    pub(crate) refresh_token: String,
165    pub(crate) expires_in: u64,
166}
167
168/// Result from the OAuth callback.
169#[derive(Debug, Clone, PartialEq, Eq)]
170pub struct CallbackResult {
171    pub code: String,
172    pub state: String,
173}
174
175/// Check if the current token is expired.
176///
177/// Note the effective ~5-minute safety margin lives in the stored value, not
178/// here: `refresh_token`/`exchange_code_for_tokens` bake a 5-minute buffer into
179/// `expires` (`now + expires_in*1000 - 5min`). So a bare `now >= expires` check
180/// already fires ~5 minutes before the credential actually dies at the provider.
181pub fn is_token_expired(creds: &OAuthCredentials) -> bool {
182    now_millis() >= creds.expires
183}
184
185// ── Helpers ─────────────────────────────────────────────────────────────────
186
187pub(crate) fn now_millis() -> u64 {
188    crate::epoch_millis()
189}
190
191/// Try to parse manual input as a redirect URL or code#state shorthand.
192///
193/// Returns `(Some(code), Some(state))` only when both components are present
194/// and unambiguous. The `None` state path has been intentionally removed:
195/// previously the caller defaulted a missing state to the expected value
196/// (`parsed_state.unwrap_or(manual_state)`), which made the downstream
197/// `result.state != state` check compare a value against itself — CSRF
198/// protection silently nullified. By never returning `Some(code)` without
199/// `Some(state)`, the CSRF guard is guaranteed to do real work.
200fn parse_manual_input(input: &str) -> (Option<String>, Option<String>) {
201    let trimmed = input.trim();
202
203    // Try as full URL (e.g. the redirect from the browser)
204    if let Ok(url) = url::Url::parse(trimmed) {
205        let code = url
206            .query_pairs()
207            .find(|(k, _)| k == "code")
208            .map(|(_, v)| v.to_string());
209        let state = url
210            .query_pairs()
211            .find(|(k, _)| k == "state")
212            .map(|(_, v)| v.to_string());
213        if code.is_some() {
214            return (code, state);
215        }
216    }
217
218    // Try as "code#state" format (Claude Code manual flow)
219    if trimmed.contains('#') {
220        let parts: Vec<&str> = trimmed.splitn(2, '#').collect();
221        if parts.len() == 2 && !parts[0].is_empty() && !parts[1].is_empty() {
222            return (Some(parts[0].to_string()), Some(parts[1].to_string()));
223        }
224    }
225
226    // Bare code with no state: REJECTED. Accepting it and defaulting the state
227    // to `manual_state` would cause the CSRF check to compare a value to itself.
228    // The caller (`manual_paste_to_callback`) will return `None` for this input.
229    (None, None)
230}
231
232/// Validate user-pasted OAuth authorization input for the Anthropic manual
233/// fallback flow.
234///
235/// Returns `Some(CallbackResult)` only if the input contains BOTH a `code`
236/// and a `state`. Previously the caller defaulted a missing `state` to the
237/// original CSRF token (`parsed_state.unwrap_or(manual_state)`), which
238/// silently bypassed the CSRF check — the state comparison became
239/// `manual_state == state` (always true). By requiring an explicit state
240/// here, a bare code paste or a URL without `state` is rejected outright.
241fn manual_paste_to_callback(input: &str) -> Option<CallbackResult> {
242    let (code, state) = parse_manual_input(input);
243    Some(CallbackResult {
244        code: code?,
245        state: state?,
246    })
247}
248
249// ── High-level login flow ───────────────────────────────────────────────────
250
251/// Run the full OAuth login flow into the default Anthropic slot.
252pub async fn login() -> std::result::Result<OAuthCredentials, String> {
253    login_into(Some(OAuthProviderId::Anthropic.as_str())).await
254}
255
256/// Run the full OAuth login flow. With `Some(key)` the credential is persisted
257/// ONLY into that slot (`anthropic` or `anthropic@<label>`); with `None` it
258/// is returned unsaved so the caller can run pre-persistence checks.
259pub async fn login_into(
260    persist_key: Option<&str>,
261) -> std::result::Result<OAuthCredentials, String> {
262    let port = CALLBACK_PORT;
263
264    // 1. Generate PKCE
265    let verifier = generate_code_verifier();
266    let challenge = generate_code_challenge(&verifier);
267    let state = generate_state();
268
269    // 2. Start callback server
270    let (rx, server_handle) = start_callback_server(state.clone(), port).await?;
271
272    // 3. Build URL and open browser
273    let auth_url = build_auth_url(&challenge, &state, port);
274
275    eprintln!("\n\x1b[1mOpening browser to sign in...\x1b[0m\n");
276
277    if let Err(e) = open_browser(&auth_url) {
278        eprintln!("Could not open browser automatically: {}", e);
279    }
280
281    eprintln!("\x1b[2mIf the browser didn't open, visit this URL:\x1b[0m");
282    eprintln!("\x1b[36m{}\x1b[0m\n", auth_url);
283
284    // Also provide manual paste option
285    let (manual_tx, manual_rx) = oneshot::channel::<CallbackResult>();
286    let stdin_task = tokio::spawn(async move {
287        eprintln!("\x1b[2mOr paste the redirect URL here (must include `state`):\x1b[0m");
288
289        let mut line = String::new();
290        let result = tokio::task::spawn_blocking(move || {
291            std::io::stdin().read_line(&mut line).ok();
292            line.trim().to_string()
293        })
294        .await;
295
296        if let Ok(input) = result {
297            match manual_paste_to_callback(&input) {
298                Some(callback) => {
299                    let _ = manual_tx.send(callback);
300                }
301                None if !input.is_empty() => {
302                    eprintln!(
303                        "\x1b[31m✗ Pasted input did not contain both `code` and `state`.\x1b[0m"
304                    );
305                    eprintln!(
306                        "\x1b[2m  Paste the full redirect URL (e.g. http://localhost:53692/callback?code=…&state=…).\x1b[0m"
307                    );
308                }
309                None => {}
310            }
311        }
312    });
313
314    // 4. Wait for either callback or manual input
315    let result = tokio::select! {
316        callback = rx => {
317            match callback {
318                Ok(CallbackOutcome::Authorized(result)) => result,
319                Ok(CallbackOutcome::Denied { error, description }) => return Err(format!("OAuth denied: {}{}", error, description.map(|d| format!(": {d}")).unwrap_or_default())),
320                Ok(CallbackOutcome::Invalid) => return Err("Invalid OAuth callback".to_string()),
321                Err(_) => return Err("Callback channel closed".to_string()),
322            }
323        }
324        manual = manual_rx => {
325            match manual {
326                Ok(result) => result,
327                Err(_) => return Err("Manual input channel closed".to_string()),
328            }
329        }
330    };
331
332    stdin_task.abort();
333
334    // 5. Verify state
335    if result.state != state {
336        server_handle.shutdown().await;
337        return Err("OAuth state mismatch — possible CSRF attack".to_string());
338    }
339
340    eprintln!("\n\x1b[1mExchanging code for tokens...\x1b[0m");
341
342    // 6. Exchange code for tokens
343    let creds = exchange_code_for_tokens(&result.code, &result.state, &verifier, port).await?;
344
345    // 7. Shut down callback server
346    server_handle.shutdown().await;
347
348    // 8. Save to auth.json — exactly one slot (or none: caller persists).
349    if let Some(key) = persist_key {
350        save_provider_auth(key, &creds)?;
351    }
352
353    Ok(creds)
354}
355
356// ── Anthropic-specific CSRF regression tests ─────────────────────────────────
357//
358// See also: openai_codex::tests for the parallel codex test suite.
359// These mirror the codex tests — same invariants, same naming pattern.
360
361/// Pinned experimental Copilot models base URL for broker catalog proxy.
362pub(crate) fn github_copilot_models_base_url() -> &'static str {
363    "https://api.githubcopilot.com"
364}
365
366/// Headers for the experimental Copilot models GET (no secrets).
367pub(crate) fn github_copilot_models_request_headers() -> &'static [(&'static str, &'static str)] {
368    &[
369        ("User-Agent", "SynapsCLI/0.6.0"),
370        ("Accept", "application/json"),
371        ("Editor-Version", "vscode/1.107.0"),
372        ("Editor-Plugin-Version", "copilot-chat/0.35.0"),
373        ("Copilot-Integration-Id", "vscode-chat"),
374        ("X-Github-Api-Version", "2025-10-01"),
375    ]
376}
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
382
383    #[test]
384    fn test_generate_code_verifier() {
385        let verifier = generate_code_verifier();
386        assert!(!verifier.is_empty(), "Code verifier should not be empty");
387        assert!(
388            verifier.len() > 20,
389            "Code verifier should be longer than 20 characters"
390        );
391        let verifier2 = generate_code_verifier();
392        assert_ne!(
393            verifier, verifier2,
394            "Two calls should produce different verifiers"
395        );
396    }
397
398    #[test]
399    fn test_generate_code_challenge() {
400        let verifier = "test_verifier_123";
401        let challenge = generate_code_challenge(verifier);
402        assert!(!challenge.is_empty(), "Code challenge should not be empty");
403        let challenge2 = generate_code_challenge(verifier);
404        assert_eq!(
405            challenge, challenge2,
406            "Same verifier should produce same challenge"
407        );
408        let different_challenge = generate_code_challenge("different_verifier_456");
409        assert_ne!(
410            challenge, different_challenge,
411            "Different verifiers should produce different challenges"
412        );
413    }
414
415    #[test]
416    fn test_generate_state() {
417        let state = generate_state();
418        assert!(!state.is_empty(), "State should not be empty");
419        let state2 = generate_state();
420        assert_ne!(state, state2, "Two calls should produce different states");
421    }
422
423    #[test]
424    fn test_build_auth_url() {
425        let challenge = "test_challenge";
426        let state = "test_state";
427        let port = 8080;
428        let url = build_auth_url(challenge, state, port);
429        assert!(url.contains("claude.ai/oauth/authorize"));
430        assert!(url.contains("client_id=9d1c250a-e61b-44d9-88ed-5944d1962f5e"));
431        assert!(url.contains(&format!("code_challenge={}", challenge)));
432        assert!(url.contains(&format!("state={}", state)));
433        assert!(url.contains("localhost"));
434        assert!(url.contains(&port.to_string()));
435        assert!(url.contains("redirect_uri="));
436    }
437
438    #[test]
439    fn test_is_token_expired() {
440        let expired_creds = OAuthCredentials {
441            auth_type: "oauth".to_string(),
442            refresh: "test_refresh".to_string(),
443            access: "test_access".to_string(),
444            expires: 0,
445            account_id: None,
446        };
447        assert!(is_token_expired(&expired_creds));
448
449        let future_time = now_millis() + 3600000;
450        let fresh_creds = OAuthCredentials {
451            auth_type: "oauth".to_string(),
452            refresh: "test_refresh".to_string(),
453            access: "test_access".to_string(),
454            expires: future_time,
455            account_id: None,
456        };
457        assert!(!is_token_expired(&fresh_creds));
458        assert_eq!(fresh_creds.auth_type, "oauth");
459    }
460
461    #[test]
462    fn test_pkce_challenge_sha256() {
463        let verifier = "test_verifier_string";
464        let challenge = generate_code_challenge(verifier);
465
466        use sha2::{Digest, Sha256};
467        let mut hasher = Sha256::new();
468        hasher.update(verifier.as_bytes());
469        let hash = hasher.finalize();
470        let expected = URL_SAFE_NO_PAD.encode(hash);
471
472        assert_eq!(challenge, expected);
473    }
474
475    #[test]
476    fn test_code_verifier_length() {
477        let verifier = generate_code_verifier();
478        assert_eq!(verifier.len(), 43);
479    }
480
481    #[test]
482    fn test_state_length() {
483        let state = generate_state();
484        assert_eq!(state.len(), 43);
485    }
486
487    #[test]
488    fn test_build_auth_url_required_params() {
489        let url = build_auth_url("test_challenge", "test_state", 8080);
490        assert!(url.contains("response_type=code"));
491        assert!(url.contains("code_challenge_method=S256"));
492        assert!(url.contains("scope="));
493        assert!(url.contains("redirect_uri="));
494        assert!(url.contains("8080"));
495    }
496
497    #[test]
498    fn test_is_token_expired_edge_cases() {
499        let current_time = std::time::SystemTime::now()
500            .duration_since(std::time::UNIX_EPOCH)
501            .unwrap()
502            .as_millis() as u64;
503
504        let exactly_now_creds = OAuthCredentials {
505            auth_type: "oauth".to_string(),
506            refresh: "test_refresh".to_string(),
507            access: "test_access".to_string(),
508            expires: current_time,
509            account_id: None,
510        };
511        assert!(is_token_expired(&exactly_now_creds));
512
513        let one_ms_future_creds = OAuthCredentials {
514            auth_type: "oauth".to_string(),
515            refresh: "test_refresh".to_string(),
516            access: "test_access".to_string(),
517            expires: current_time + 1,
518            account_id: None,
519        };
520        assert!(!is_token_expired(&one_ms_future_creds));
521    }
522
523    #[test]
524    fn test_auth_file_path() {
525        let path = auth_file_path();
526        let path_str = path.to_string_lossy();
527        assert!(path_str.ends_with("auth.json"));
528    }
529
530    #[test]
531    fn test_oauth_credentials_serialization_roundtrip() {
532        let original_creds = OAuthCredentials {
533            auth_type: "oauth".to_string(),
534            refresh: "test_refresh_token".to_string(),
535            access: "test_access_token".to_string(),
536            expires: 1234567890,
537            account_id: None,
538        };
539
540        let json = serde_json::to_string(&original_creds).expect("Should serialize");
541        let deserialized_creds: OAuthCredentials =
542            serde_json::from_str(&json).expect("Should deserialize");
543
544        assert_eq!(original_creds.auth_type, deserialized_creds.auth_type);
545        assert_eq!(original_creds.refresh, deserialized_creds.refresh);
546        assert_eq!(original_creds.access, deserialized_creds.access);
547        assert_eq!(original_creds.expires, deserialized_creds.expires);
548    }
549
550    // ── manual_paste_to_callback: CSRF guard on the Anthropic manual paste path ──
551    //
552    // Mirrors the regression suite in openai_codex::tests (manual_paste_to_callback).
553    // Pre-fix: parse_manual_input returned (Some(code), None) for bare codes,
554    // and the caller did `parsed_state.unwrap_or(manual_state)`, making the
555    // CSRF check `manual_state == state` — always true — bypass.
556    // Fix: removed bare-code fallback from parse_manual_input + extracted
557    // manual_paste_to_callback that requires both code AND state.
558
559    #[test]
560    fn anthropic_manual_paste_accepts_full_redirect_url() {
561        let result = manual_paste_to_callback("http://localhost:53692/callback?code=abc&state=xyz")
562            .expect("URL with code+state must be accepted");
563        assert_eq!(result.code, "abc");
564        assert_eq!(result.state, "xyz");
565    }
566
567    #[test]
568    fn anthropic_manual_paste_rejects_bare_code() {
569        // CSRF regression: a bare code with no state was previously defaulted
570        // to `manual_state`, making the CSRF check trivially true.
571        assert!(
572            manual_paste_to_callback("abc123_bare_code").is_none(),
573            "bare code with no state must be rejected — was the root of the CSRF bypass"
574        );
575    }
576
577    #[test]
578    fn anthropic_manual_paste_rejects_url_without_state() {
579        assert!(
580            manual_paste_to_callback("http://localhost:53692/callback?code=abc").is_none(),
581            "URL missing `state` must be rejected — would bypass CSRF check"
582        );
583    }
584
585    #[test]
586    fn anthropic_manual_paste_accepts_code_hash_state() {
587        // "code#state" is the Claude Code shorthand; both components explicit.
588        let result = manual_paste_to_callback("mycode#mystate")
589            .expect("code#state shorthand must be accepted");
590        assert_eq!(result.code, "mycode");
591        assert_eq!(result.state, "mystate");
592    }
593
594    #[test]
595    fn anthropic_manual_paste_rejects_empty_input() {
596        assert!(manual_paste_to_callback("").is_none());
597        assert!(manual_paste_to_callback("   ").is_none());
598    }
599
600    #[test]
601    fn anthropic_manual_paste_rejects_url_with_only_state() {
602        assert!(
603            manual_paste_to_callback("http://localhost:53692/callback?state=xyz").is_none(),
604            "URL missing `code` must be rejected"
605        );
606    }
607}