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    token_expired_at(creds, now_millis())
183}
184
185/// The expiry rule against a given clock reading: a token is expired from
186/// its `expires` millisecond on.
187fn token_expired_at(creds: &OAuthCredentials, now_ms: u64) -> bool {
188    now_ms >= creds.expires
189}
190
191// ── Helpers ─────────────────────────────────────────────────────────────────
192
193pub(crate) fn now_millis() -> u64 {
194    crate::epoch_millis()
195}
196
197/// Try to parse manual input as a redirect URL or code#state shorthand.
198///
199/// Returns `(Some(code), Some(state))` only when both components are present
200/// and unambiguous. The `None` state path has been intentionally removed:
201/// previously the caller defaulted a missing state to the expected value
202/// (`parsed_state.unwrap_or(manual_state)`), which made the downstream
203/// `result.state != state` check compare a value against itself — CSRF
204/// protection silently nullified. By never returning `Some(code)` without
205/// `Some(state)`, the CSRF guard is guaranteed to do real work.
206fn parse_manual_input(input: &str) -> (Option<String>, Option<String>) {
207    let trimmed = input.trim();
208
209    // Try as full URL (e.g. the redirect from the browser)
210    if let Ok(url) = url::Url::parse(trimmed) {
211        let code = url
212            .query_pairs()
213            .find(|(k, _)| k == "code")
214            .map(|(_, v)| v.to_string());
215        let state = url
216            .query_pairs()
217            .find(|(k, _)| k == "state")
218            .map(|(_, v)| v.to_string());
219        if code.is_some() {
220            return (code, state);
221        }
222    }
223
224    // Try as "code#state" format (Claude Code manual flow)
225    if trimmed.contains('#') {
226        let parts: Vec<&str> = trimmed.splitn(2, '#').collect();
227        if parts.len() == 2 && !parts[0].is_empty() && !parts[1].is_empty() {
228            return (Some(parts[0].to_string()), Some(parts[1].to_string()));
229        }
230    }
231
232    // Bare code with no state: REJECTED. Accepting it and defaulting the state
233    // to `manual_state` would cause the CSRF check to compare a value to itself.
234    // The caller (`manual_paste_to_callback`) will return `None` for this input.
235    (None, None)
236}
237
238/// Validate user-pasted OAuth authorization input for the Anthropic manual
239/// fallback flow.
240///
241/// Returns `Some(CallbackResult)` only if the input contains BOTH a `code`
242/// and a `state`. Previously the caller defaulted a missing `state` to the
243/// original CSRF token (`parsed_state.unwrap_or(manual_state)`), which
244/// silently bypassed the CSRF check — the state comparison became
245/// `manual_state == state` (always true). By requiring an explicit state
246/// here, a bare code paste or a URL without `state` is rejected outright.
247fn manual_paste_to_callback(input: &str) -> Option<CallbackResult> {
248    let (code, state) = parse_manual_input(input);
249    Some(CallbackResult {
250        code: code?,
251        state: state?,
252    })
253}
254
255// ── High-level login flow ───────────────────────────────────────────────────
256
257/// Run the full OAuth login flow into the default Anthropic slot.
258pub async fn login() -> std::result::Result<OAuthCredentials, String> {
259    login_into(Some(OAuthProviderId::Anthropic.as_str())).await
260}
261
262/// Run the full OAuth login flow. With `Some(key)` the credential is persisted
263/// ONLY into that slot (`anthropic` or `anthropic@<label>`); with `None` it
264/// is returned unsaved so the caller can run pre-persistence checks.
265pub async fn login_into(
266    persist_key: Option<&str>,
267) -> std::result::Result<OAuthCredentials, String> {
268    let port = CALLBACK_PORT;
269
270    // 1. Generate PKCE
271    let verifier = generate_code_verifier();
272    let challenge = generate_code_challenge(&verifier);
273    let state = generate_state();
274
275    // 2. Start callback server
276    let (rx, server_handle) = start_callback_server(state.clone(), port).await?;
277
278    // 3. Build URL and open browser
279    let auth_url = build_auth_url(&challenge, &state, port);
280
281    eprintln!("\n\x1b[1mOpening browser to sign in...\x1b[0m\n");
282
283    if let Err(e) = open_browser(&auth_url) {
284        eprintln!("Could not open browser automatically: {}", e);
285    }
286
287    eprintln!("\x1b[2mIf the browser didn't open, visit this URL:\x1b[0m");
288    eprintln!("\x1b[36m{}\x1b[0m\n", auth_url);
289
290    // Also provide manual paste option
291    let (manual_tx, manual_rx) = oneshot::channel::<CallbackResult>();
292    let stdin_task = tokio::spawn(async move {
293        eprintln!("\x1b[2mOr paste the redirect URL here (must include `state`):\x1b[0m");
294
295        let mut line = String::new();
296        let result = tokio::task::spawn_blocking(move || {
297            std::io::stdin().read_line(&mut line).ok();
298            line.trim().to_string()
299        })
300        .await;
301
302        if let Ok(input) = result {
303            match manual_paste_to_callback(&input) {
304                Some(callback) => {
305                    let _ = manual_tx.send(callback);
306                }
307                None if !input.is_empty() => {
308                    eprintln!(
309                        "\x1b[31m✗ Pasted input did not contain both `code` and `state`.\x1b[0m"
310                    );
311                    eprintln!(
312                        "\x1b[2m  Paste the full redirect URL (e.g. http://localhost:53692/callback?code=…&state=…).\x1b[0m"
313                    );
314                }
315                None => {}
316            }
317        }
318    });
319
320    // 4. Wait for either callback or manual input
321    let result = tokio::select! {
322        callback = rx => {
323            match callback {
324                Ok(CallbackOutcome::Authorized(result)) => result,
325                Ok(CallbackOutcome::Denied { error, description }) => return Err(format!("OAuth denied: {}{}", error, description.map(|d| format!(": {d}")).unwrap_or_default())),
326                Ok(CallbackOutcome::Invalid) => return Err("Invalid OAuth callback".to_string()),
327                Err(_) => return Err("Callback channel closed".to_string()),
328            }
329        }
330        manual = manual_rx => {
331            match manual {
332                Ok(result) => result,
333                Err(_) => return Err("Manual input channel closed".to_string()),
334            }
335        }
336    };
337
338    stdin_task.abort();
339
340    // 5. Verify state
341    if result.state != state {
342        server_handle.shutdown().await;
343        return Err("OAuth state mismatch — possible CSRF attack".to_string());
344    }
345
346    eprintln!("\n\x1b[1mExchanging code for tokens...\x1b[0m");
347
348    // 6. Exchange code for tokens
349    let creds = exchange_code_for_tokens(&result.code, &result.state, &verifier, port).await?;
350
351    // 7. Shut down callback server
352    server_handle.shutdown().await;
353
354    // 8. Save to auth.json — exactly one slot (or none: caller persists).
355    if let Some(key) = persist_key {
356        save_provider_auth(key, &creds)?;
357    }
358
359    Ok(creds)
360}
361
362// ── Anthropic-specific CSRF regression tests ─────────────────────────────────
363//
364// See also: openai_codex::tests for the parallel codex test suite.
365// These mirror the codex tests — same invariants, same naming pattern.
366
367/// Pinned experimental Copilot models base URL for broker catalog proxy.
368pub(crate) fn github_copilot_models_base_url() -> &'static str {
369    "https://api.githubcopilot.com"
370}
371
372/// Headers for the experimental Copilot models GET (no secrets).
373pub(crate) fn github_copilot_models_request_headers() -> &'static [(&'static str, &'static str)] {
374    &[
375        ("User-Agent", "SynapsCLI/0.6.0"),
376        ("Accept", "application/json"),
377        ("Editor-Version", "vscode/1.107.0"),
378        ("Editor-Plugin-Version", "copilot-chat/0.35.0"),
379        ("Copilot-Integration-Id", "vscode-chat"),
380        ("X-Github-Api-Version", "2025-10-01"),
381    ]
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
388
389    #[test]
390    fn test_generate_code_verifier() {
391        let verifier = generate_code_verifier();
392        assert!(!verifier.is_empty(), "Code verifier should not be empty");
393        assert!(
394            verifier.len() > 20,
395            "Code verifier should be longer than 20 characters"
396        );
397        let verifier2 = generate_code_verifier();
398        assert_ne!(
399            verifier, verifier2,
400            "Two calls should produce different verifiers"
401        );
402    }
403
404    #[test]
405    fn test_generate_code_challenge() {
406        let verifier = "test_verifier_123";
407        let challenge = generate_code_challenge(verifier);
408        assert!(!challenge.is_empty(), "Code challenge should not be empty");
409        let challenge2 = generate_code_challenge(verifier);
410        assert_eq!(
411            challenge, challenge2,
412            "Same verifier should produce same challenge"
413        );
414        let different_challenge = generate_code_challenge("different_verifier_456");
415        assert_ne!(
416            challenge, different_challenge,
417            "Different verifiers should produce different challenges"
418        );
419    }
420
421    #[test]
422    fn test_generate_state() {
423        let state = generate_state();
424        assert!(!state.is_empty(), "State should not be empty");
425        let state2 = generate_state();
426        assert_ne!(state, state2, "Two calls should produce different states");
427    }
428
429    #[test]
430    fn test_build_auth_url() {
431        let challenge = "test_challenge";
432        let state = "test_state";
433        let port = 8080;
434        let url = build_auth_url(challenge, state, port);
435        assert!(url.contains("claude.ai/oauth/authorize"));
436        assert!(url.contains("client_id=9d1c250a-e61b-44d9-88ed-5944d1962f5e"));
437        assert!(url.contains(&format!("code_challenge={}", challenge)));
438        assert!(url.contains(&format!("state={}", state)));
439        assert!(url.contains("localhost"));
440        assert!(url.contains(&port.to_string()));
441        assert!(url.contains("redirect_uri="));
442    }
443
444    #[test]
445    fn test_is_token_expired() {
446        let expired_creds = OAuthCredentials {
447            auth_type: "oauth".to_string(),
448            refresh: "test_refresh".to_string(),
449            access: "test_access".to_string(),
450            expires: 0,
451            account_id: None,
452        };
453        assert!(is_token_expired(&expired_creds));
454
455        let future_time = now_millis() + 3600000;
456        let fresh_creds = OAuthCredentials {
457            auth_type: "oauth".to_string(),
458            refresh: "test_refresh".to_string(),
459            access: "test_access".to_string(),
460            expires: future_time,
461            account_id: None,
462        };
463        assert!(!is_token_expired(&fresh_creds));
464        assert_eq!(fresh_creds.auth_type, "oauth");
465    }
466
467    #[test]
468    fn test_pkce_challenge_sha256() {
469        let verifier = "test_verifier_string";
470        let challenge = generate_code_challenge(verifier);
471
472        use sha2::{Digest, Sha256};
473        let mut hasher = Sha256::new();
474        hasher.update(verifier.as_bytes());
475        let hash = hasher.finalize();
476        let expected = URL_SAFE_NO_PAD.encode(hash);
477
478        assert_eq!(challenge, expected);
479    }
480
481    #[test]
482    fn test_code_verifier_length() {
483        let verifier = generate_code_verifier();
484        assert_eq!(verifier.len(), 43);
485    }
486
487    #[test]
488    fn test_state_length() {
489        let state = generate_state();
490        assert_eq!(state.len(), 43);
491    }
492
493    #[test]
494    fn test_build_auth_url_required_params() {
495        let url = build_auth_url("test_challenge", "test_state", 8080);
496        assert!(url.contains("response_type=code"));
497        assert!(url.contains("code_challenge_method=S256"));
498        assert!(url.contains("scope="));
499        assert!(url.contains("redirect_uri="));
500        assert!(url.contains("8080"));
501    }
502
503    #[test]
504    fn test_is_token_expired_edge_cases() {
505        // Pin the clock: against the live one, "1 ms in the future" has
506        // already passed by the time a loaded CI runner checks it.
507        let current_time = now_millis();
508
509        let exactly_now_creds = OAuthCredentials {
510            auth_type: "oauth".to_string(),
511            refresh: "test_refresh".to_string(),
512            access: "test_access".to_string(),
513            expires: current_time,
514            account_id: None,
515        };
516        assert!(token_expired_at(&exactly_now_creds, current_time));
517
518        let one_ms_future_creds = OAuthCredentials {
519            auth_type: "oauth".to_string(),
520            refresh: "test_refresh".to_string(),
521            access: "test_access".to_string(),
522            expires: current_time + 1,
523            account_id: None,
524        };
525        assert!(!token_expired_at(&one_ms_future_creds, current_time));
526        assert!(token_expired_at(&one_ms_future_creds, current_time + 1));
527    }
528
529    #[test]
530    fn test_auth_file_path() {
531        let path = auth_file_path();
532        let path_str = path.to_string_lossy();
533        assert!(path_str.ends_with("auth.json"));
534    }
535
536    #[test]
537    fn test_oauth_credentials_serialization_roundtrip() {
538        let original_creds = OAuthCredentials {
539            auth_type: "oauth".to_string(),
540            refresh: "test_refresh_token".to_string(),
541            access: "test_access_token".to_string(),
542            expires: 1234567890,
543            account_id: None,
544        };
545
546        let json = serde_json::to_string(&original_creds).expect("Should serialize");
547        let deserialized_creds: OAuthCredentials =
548            serde_json::from_str(&json).expect("Should deserialize");
549
550        assert_eq!(original_creds.auth_type, deserialized_creds.auth_type);
551        assert_eq!(original_creds.refresh, deserialized_creds.refresh);
552        assert_eq!(original_creds.access, deserialized_creds.access);
553        assert_eq!(original_creds.expires, deserialized_creds.expires);
554    }
555
556    // ── manual_paste_to_callback: CSRF guard on the Anthropic manual paste path ──
557    //
558    // Mirrors the regression suite in openai_codex::tests (manual_paste_to_callback).
559    // Pre-fix: parse_manual_input returned (Some(code), None) for bare codes,
560    // and the caller did `parsed_state.unwrap_or(manual_state)`, making the
561    // CSRF check `manual_state == state` — always true — bypass.
562    // Fix: removed bare-code fallback from parse_manual_input + extracted
563    // manual_paste_to_callback that requires both code AND state.
564
565    #[test]
566    fn anthropic_manual_paste_accepts_full_redirect_url() {
567        let result = manual_paste_to_callback("http://localhost:53692/callback?code=abc&state=xyz")
568            .expect("URL with code+state must be accepted");
569        assert_eq!(result.code, "abc");
570        assert_eq!(result.state, "xyz");
571    }
572
573    #[test]
574    fn anthropic_manual_paste_rejects_bare_code() {
575        // CSRF regression: a bare code with no state was previously defaulted
576        // to `manual_state`, making the CSRF check trivially true.
577        assert!(
578            manual_paste_to_callback("abc123_bare_code").is_none(),
579            "bare code with no state must be rejected — was the root of the CSRF bypass"
580        );
581    }
582
583    #[test]
584    fn anthropic_manual_paste_rejects_url_without_state() {
585        assert!(
586            manual_paste_to_callback("http://localhost:53692/callback?code=abc").is_none(),
587            "URL missing `state` must be rejected — would bypass CSRF check"
588        );
589    }
590
591    #[test]
592    fn anthropic_manual_paste_accepts_code_hash_state() {
593        // "code#state" is the Claude Code shorthand; both components explicit.
594        let result = manual_paste_to_callback("mycode#mystate")
595            .expect("code#state shorthand must be accepted");
596        assert_eq!(result.code, "mycode");
597        assert_eq!(result.state, "mystate");
598    }
599
600    #[test]
601    fn anthropic_manual_paste_rejects_empty_input() {
602        assert!(manual_paste_to_callback("").is_none());
603        assert!(manual_paste_to_callback("   ").is_none());
604    }
605
606    #[test]
607    fn anthropic_manual_paste_rejects_url_with_only_state() {
608        assert!(
609            manual_paste_to_callback("http://localhost:53692/callback?state=xyz").is_none(),
610            "URL missing `code` must be rejected"
611        );
612    }
613}