1use std::collections::HashMap;
9use std::sync::{
10 Arc, Mutex as StdMutex,
11 atomic::{AtomicU64, Ordering},
12};
13use std::time::{Duration, Instant};
14
15use base64::Engine as _;
16use meerkat_core::{AuthBindingRef, AuthProfile, BackendProfile, CredentialSourceSpec, Provider};
17use meerkat_llm_core::provider_runtime::binding::NormalizedBackendKind;
18use meerkat_llm_core::provider_runtime::catalog::ProviderRuntimeCatalog;
19use parking_lot::Mutex;
20use serde::{Deserialize, Serialize};
21
22use crate::auth_oauth::{OAuthEndpoints, OAuthTokenRequestFormat};
23use crate::auth_store::{PersistedAuthMode, credential_source_uses_persisted_store};
24
25pub use meerkat_core::oauth_identity::OAuthProviderIdentity;
31
32const DEFAULT_MAX_OUTSTANDING_FLOWS: usize = 1024;
33
34const ANTHROPIC_CLIENT_ID: &str = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
35const ANTHROPIC_AUTHORIZE_URL: &str = "https://claude.com/cai/oauth/authorize";
36const ANTHROPIC_CONSOLE_AUTHORIZE_URL: &str = "https://platform.claude.com/oauth/authorize";
37const ANTHROPIC_TOKEN_URL: &str = "https://platform.claude.com/v1/oauth/token";
38const ANTHROPIC_CONSOLE_SCOPES: &[&str] = &["org:create_api_key", "user:profile"];
39const ANTHROPIC_CLAUDE_AI_SCOPES: &[&str] = &[
40 "user:profile",
41 "user:inference",
42 "user:sessions:claude_code",
43 "user:mcp_servers",
44 "user:file_upload",
45];
46
47const OPENAI_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
54const OPENAI_AUTHORIZE_URL: &str = "https://auth.openai.com/oauth/authorize";
55const OPENAI_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
56const OPENAI_SCOPES: &[&str] = &[
57 "openid",
58 "profile",
59 "email",
60 "offline_access",
61 "api.connectors.read",
62 "api.connectors.invoke",
63];
64const OPENAI_ORIGINATOR: &str = "codex_cli_rs";
65
66const GOOGLE_CLIENT_ID: &str = concat!(
67 "6812558",
68 "09395-oo8ft2oprdrnp9e3aqf6av3hmdib135j",
69 ".apps.googleusercontent.com",
70);
71const GOOGLE_AUTHORIZE_URL: &str = "https://accounts.google.com/o/oauth2/v2/auth";
72const GOOGLE_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
73const GOOGLE_DEVICE_CODE_URL: &str = "https://oauth2.googleapis.com/device/code";
74const GOOGLE_SCOPES: &[&str] = &[
75 "https://www.googleapis.com/auth/cloud-platform",
76 "https://www.googleapis.com/auth/userinfo.email",
77 "https://www.googleapis.com/auth/userinfo.profile",
78];
79
80const TEST_OAUTH_ENDPOINT_OVERRIDE_ENV: &str = "MEERKAT_TEST_OAUTH_ENDPOINT_OVERRIDE";
81const TEST_OAUTH_BASE_URL_ENV: &str = "MEERKAT_TEST_OAUTH_BASE_URL";
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub struct OAuthProviderDeclaration {
94 pub client_id: &'static str,
96 pub authorize_endpoint: &'static str,
98 pub token_endpoint: &'static str,
100 pub scopes: &'static [&'static str],
102 pub backend_kind: NormalizedBackendKind,
104 pub extra_authorize_params: &'static [(&'static str, &'static str)],
108}
109
110const OPENAI_EXTRA_AUTHORIZE_PARAMS: &[(&str, &str)] = &[
111 ("id_token_add_organizations", "true"),
112 ("codex_cli_simplified_flow", "true"),
113 ("originator", OPENAI_ORIGINATOR),
114];
115
116const ANTHROPIC_CLAUDE_AI_EXTRA_AUTHORIZE_PARAMS: &[(&str, &str)] = &[("code", "true")];
117
118pub fn oauth_provider_declaration(id: OAuthProviderIdentity) -> OAuthProviderDeclaration {
129 let backend_kind = meerkat_llm_core::provider_runtime::binding::oauth_provider_backend_kind(id);
130 match id {
131 OAuthProviderIdentity::AnthropicClaudeAi => OAuthProviderDeclaration {
132 client_id: ANTHROPIC_CLIENT_ID,
133 authorize_endpoint: ANTHROPIC_AUTHORIZE_URL,
134 token_endpoint: ANTHROPIC_TOKEN_URL,
135 scopes: ANTHROPIC_CLAUDE_AI_SCOPES,
136 backend_kind,
137 extra_authorize_params: ANTHROPIC_CLAUDE_AI_EXTRA_AUTHORIZE_PARAMS,
138 },
139 OAuthProviderIdentity::AnthropicConsoleApiKey => OAuthProviderDeclaration {
140 client_id: ANTHROPIC_CLIENT_ID,
141 authorize_endpoint: ANTHROPIC_CONSOLE_AUTHORIZE_URL,
142 token_endpoint: ANTHROPIC_TOKEN_URL,
143 scopes: ANTHROPIC_CONSOLE_SCOPES,
144 backend_kind,
145 extra_authorize_params: &[],
146 },
147 OAuthProviderIdentity::OpenAiChatGpt => OAuthProviderDeclaration {
148 client_id: OPENAI_CLIENT_ID,
149 authorize_endpoint: OPENAI_AUTHORIZE_URL,
150 token_endpoint: OPENAI_TOKEN_URL,
151 scopes: OPENAI_SCOPES,
152 backend_kind,
153 extra_authorize_params: OPENAI_EXTRA_AUTHORIZE_PARAMS,
154 },
155 OAuthProviderIdentity::GoogleCodeAssist => OAuthProviderDeclaration {
156 client_id: GOOGLE_CLIENT_ID,
157 authorize_endpoint: GOOGLE_AUTHORIZE_URL,
158 token_endpoint: GOOGLE_TOKEN_URL,
159 scopes: GOOGLE_SCOPES,
160 backend_kind,
161 extra_authorize_params: &[],
162 },
163 }
164}
165
166pub fn oauth_provider_endpoints(
173 id: OAuthProviderIdentity,
174 redirect_uri: impl Into<String>,
175) -> OAuthEndpoints {
176 let declaration = oauth_provider_declaration(id);
177 let extra_authorize_params = declaration
178 .extra_authorize_params
179 .iter()
180 .map(|(key, value)| ((*key).to_string(), (*value).to_string()))
181 .collect();
182 let endpoints = match id {
183 OAuthProviderIdentity::AnthropicClaudeAi => OAuthEndpoints {
184 client_id: declaration.client_id.into(),
185 authorize_url: declaration.authorize_endpoint.into(),
186 token_url: declaration.token_endpoint.into(),
187 device_code_url: None,
188 redirect_uri: redirect_uri.into(),
189 scopes: strings(declaration.scopes),
190 extra_authorize_params,
191 token_request_format: OAuthTokenRequestFormat::Json,
192 include_state_in_token_exchange: true,
193 extra_token_params: Vec::new(),
194 refresh_scopes: strings(declaration.scopes),
195 extra_headers: Vec::new(),
196 },
197 OAuthProviderIdentity::AnthropicConsoleApiKey => OAuthEndpoints {
198 client_id: declaration.client_id.into(),
199 authorize_url: declaration.authorize_endpoint.into(),
200 token_url: declaration.token_endpoint.into(),
201 device_code_url: None,
202 redirect_uri: redirect_uri.into(),
203 scopes: strings(declaration.scopes),
204 extra_authorize_params,
205 token_request_format: OAuthTokenRequestFormat::Json,
206 include_state_in_token_exchange: true,
207 extra_token_params: Vec::new(),
208 refresh_scopes: strings(declaration.scopes),
209 extra_headers: Vec::new(),
210 },
211 OAuthProviderIdentity::OpenAiChatGpt => OAuthEndpoints {
212 client_id: declaration.client_id.into(),
213 authorize_url: declaration.authorize_endpoint.into(),
214 token_url: declaration.token_endpoint.into(),
215 device_code_url: None,
216 redirect_uri: redirect_uri.into(),
217 scopes: strings(declaration.scopes),
218 extra_authorize_params,
219 token_request_format: OAuthTokenRequestFormat::FormUrlEncoded,
220 include_state_in_token_exchange: false,
221 extra_token_params: Vec::new(),
222 refresh_scopes: Vec::new(),
223 extra_headers: Vec::new(),
224 },
225 OAuthProviderIdentity::GoogleCodeAssist => OAuthEndpoints {
226 client_id: declaration.client_id.into(),
227 authorize_url: declaration.authorize_endpoint.into(),
228 token_url: declaration.token_endpoint.into(),
229 device_code_url: Some(GOOGLE_DEVICE_CODE_URL.into()),
230 redirect_uri: redirect_uri.into(),
231 scopes: strings(declaration.scopes),
232 extra_authorize_params,
233 token_request_format: OAuthTokenRequestFormat::FormUrlEncoded,
234 include_state_in_token_exchange: false,
235 extra_token_params: Vec::new(),
236 refresh_scopes: Vec::new(),
237 extra_headers: Vec::new(),
238 },
239 };
240 apply_test_oauth_endpoint_override(id, endpoints)
241}
242
243#[derive(Debug, Clone)]
244pub struct OAuthProviderResolution {
245 pub identity: OAuthProviderIdentity,
246 pub provider: Provider,
247 pub endpoints: OAuthEndpoints,
248 pub auth_mode: PersistedAuthMode,
249 pub client_secret: Option<&'static str>,
250}
251
252#[derive(Debug, thiserror::Error, PartialEq, Eq)]
253#[error("Unknown provider '{provider}'. Supported: anthropic, openai, google.")]
254pub struct OAuthProviderResolutionError {
255 pub provider: String,
256}
257
258pub fn resolve_oauth_provider(
259 provider: &str,
260 redirect_uri: impl Into<String>,
261) -> Result<OAuthProviderResolution, OAuthProviderResolutionError> {
262 let identity = OAuthProviderIdentity::from_alias(provider).ok_or_else(|| {
263 OAuthProviderResolutionError {
264 provider: provider.to_string(),
265 }
266 })?;
267 Ok(OAuthProviderResolution {
268 identity,
269 provider: identity.provider(),
270 endpoints: oauth_provider_endpoints(identity, redirect_uri),
271 auth_mode: identity.auth_mode(),
272 client_secret: identity.client_secret(),
273 })
274}
275
276#[doc(hidden)]
280pub fn apply_test_oauth_endpoint_override(
281 identity: OAuthProviderIdentity,
282 mut endpoints: OAuthEndpoints,
283) -> OAuthEndpoints {
284 #[cfg(not(target_arch = "wasm32"))]
285 {
286 let enabled = std::env::var(TEST_OAUTH_ENDPOINT_OVERRIDE_ENV)
287 .map(|value| {
288 matches!(
289 value.as_str(),
290 "1" | "true" | "TRUE" | "yes" | "YES" | "on" | "ON"
291 )
292 })
293 .unwrap_or(false);
294 if !enabled {
295 return endpoints;
296 }
297 let Ok(base_url) = std::env::var(TEST_OAUTH_BASE_URL_ENV) else {
298 return endpoints;
299 };
300 let base_url = base_url.trim_end_matches('/');
301 let provider = identity.canonical_alias();
302 endpoints.authorize_url = format!("{base_url}/{provider}/authorize");
303 endpoints.token_url = format!("{base_url}/{provider}/token");
304 if endpoints.device_code_url.is_some() {
305 endpoints.device_code_url = Some(format!("{base_url}/{provider}/device/code"));
306 }
307 }
308 endpoints
309}
310
311#[derive(Debug, thiserror::Error, PartialEq, Eq)]
312pub enum OAuthTargetValidationError {
313 #[error("OAuth target backend provider mismatch: expected {expected:?}, got {actual:?}")]
314 BackendProviderMismatch {
315 expected: Provider,
316 actual: Provider,
317 },
318 #[error(
319 "OAuth target backend_kind '{backend_kind}' cannot store credential mode {expected_mode:?}; expected backend_kind {expected_backend_kind:?}"
320 )]
321 BackendKindMismatch {
322 backend_kind: String,
323 expected_backend_kind: NormalizedBackendKind,
324 expected_mode: PersistedAuthMode,
325 },
326 #[error("OAuth target provider mismatch: expected {expected:?}, got {actual:?}")]
327 ProviderMismatch {
328 expected: Provider,
329 actual: Provider,
330 },
331 #[error(
332 "OAuth target auth_method '{auth_method}' cannot store credential mode {expected_mode:?}"
333 )]
334 AuthMethodMismatch {
335 auth_method: String,
336 expected_mode: PersistedAuthMode,
337 },
338 #[error(
339 "OAuth target source '{source_kind}' cannot store OAuth credentials; expected source.kind = 'managed_store' or 'platform_default'"
340 )]
341 SourceMismatch { source_kind: &'static str },
342}
343
344pub fn validate_oauth_login_target(
345 auth_profile: &AuthProfile,
346 identity: OAuthProviderIdentity,
347) -> Result<(), OAuthTargetValidationError> {
348 validate_oauth_target_for_auth_mode(auth_profile, identity.provider(), identity.auth_mode())
349}
350
351pub fn validate_oauth_login_binding(
352 backend_profile: &BackendProfile,
353 auth_profile: &AuthProfile,
354 identity: OAuthProviderIdentity,
355) -> Result<(), OAuthTargetValidationError> {
356 validate_oauth_target_binding_for_auth_mode(
357 backend_profile,
358 auth_profile,
359 identity.provider(),
360 identity.auth_mode(),
361 meerkat_llm_core::provider_runtime::binding::oauth_provider_backend_kind(identity),
362 )
363}
364
365pub fn validate_oauth_target_for_auth_mode(
366 auth_profile: &AuthProfile,
367 expected_provider: Provider,
368 expected_mode: PersistedAuthMode,
369) -> Result<(), OAuthTargetValidationError> {
370 if auth_profile.provider != expected_provider {
371 return Err(OAuthTargetValidationError::ProviderMismatch {
372 expected: expected_provider,
373 actual: auth_profile.provider,
374 });
375 }
376 let actual_mode =
381 ProviderRuntimeCatalog::normalize_auth(auth_profile.provider, &auth_profile.auth_method)
382 .ok()
383 .and_then(|method| method.persisted_auth_mode());
384 match actual_mode {
385 Some(actual_mode) if actual_mode == expected_mode => {}
386 _ => {
387 return Err(OAuthTargetValidationError::AuthMethodMismatch {
388 auth_method: auth_profile.auth_method.clone(),
389 expected_mode,
390 });
391 }
392 }
393 if !oauth_source_can_store_flow_credentials(&auth_profile.source) {
394 return Err(OAuthTargetValidationError::SourceMismatch {
395 source_kind: source_kind_label(&auth_profile.source),
396 });
397 }
398 Ok(())
399}
400
401pub fn validate_oauth_target_binding_for_auth_mode(
402 backend_profile: &BackendProfile,
403 auth_profile: &AuthProfile,
404 expected_provider: Provider,
405 expected_mode: PersistedAuthMode,
406 expected_backend_kind: NormalizedBackendKind,
407) -> Result<(), OAuthTargetValidationError> {
408 validate_oauth_target_for_auth_mode(auth_profile, expected_provider, expected_mode)?;
409 if backend_profile.provider != expected_provider {
410 return Err(OAuthTargetValidationError::BackendProviderMismatch {
411 expected: expected_provider,
412 actual: backend_profile.provider,
413 });
414 }
415 let actual_backend_kind =
418 ProviderRuntimeCatalog::normalize_backend(expected_provider, &backend_profile.backend_kind)
419 .ok();
420 if actual_backend_kind != Some(expected_backend_kind) {
421 return Err(OAuthTargetValidationError::BackendKindMismatch {
422 backend_kind: backend_profile.backend_kind.clone(),
423 expected_backend_kind,
424 expected_mode,
425 });
426 }
427 Ok(())
428}
429
430fn oauth_source_can_store_flow_credentials(source: &CredentialSourceSpec) -> bool {
431 credential_source_uses_persisted_store(source)
432}
433
434fn source_kind_label(source: &CredentialSourceSpec) -> &'static str {
435 match source {
436 CredentialSourceSpec::InlineSecret { .. } => "inline_secret",
437 CredentialSourceSpec::ManagedStore => "managed_store",
438 CredentialSourceSpec::Env { .. } => "env",
439 CredentialSourceSpec::ExternalResolver { .. } => "external_resolver",
440 CredentialSourceSpec::PlatformDefault => "platform_default",
441 CredentialSourceSpec::Command { .. } => "command",
442 CredentialSourceSpec::FileDescriptor { .. } => "file_descriptor",
443 }
444}
445
446#[derive(Debug, Clone, PartialEq, Eq)]
447pub struct OAuthFlowRecord {
448 pub target: AuthBindingRef,
449 pub provider: OAuthProviderIdentity,
450 pub redirect_uri: String,
451 pub pkce_verifier: String,
452 pub created_at: Instant,
453}
454
455#[derive(Debug, Clone, PartialEq, Eq)]
456pub struct OAuthDeviceFlowRecord {
457 pub target: AuthBindingRef,
458 pub provider: OAuthProviderIdentity,
459 pub device_code: String,
460 pub created_at: Instant,
461 pub expires_at: Instant,
462}
463
464#[derive(Debug, Clone, Default, PartialEq, Eq)]
465pub struct OAuthPrunedFlows {
466 pub browser: Vec<(String, AuthBindingRef)>,
467 pub device: Vec<(String, AuthBindingRef)>,
468}
469
470impl OAuthPrunedFlows {
471 fn from_expired(
472 browser: Vec<(String, OAuthFlowRecord)>,
473 device: Vec<OAuthDeviceFlowRecord>,
474 ) -> Self {
475 Self {
476 browser: browser
477 .into_iter()
478 .map(|(flow_id, record)| (flow_id, record.target))
479 .collect(),
480 device: device
481 .into_iter()
482 .map(|record| (record.device_code, record.target))
483 .collect(),
484 }
485 }
486}
487
488#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
489pub struct OAuthFlowRegistrySnapshot {
490 #[serde(default)]
491 pub browser: Vec<PersistedOAuthBrowserFlow>,
492 #[serde(default)]
493 pub device: Vec<PersistedOAuthDeviceFlow>,
494}
495
496#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
497pub struct PersistedOAuthBrowserFlow {
498 pub state: String,
499 pub target: AuthBindingRef,
500 pub provider: OAuthProviderIdentity,
501 pub redirect_uri: String,
502 pub pkce_verifier: String,
503 pub created_at_millis: u64,
504 pub expires_at_millis: u64,
505}
506
507#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
508pub struct PersistedOAuthDeviceFlow {
509 pub target: AuthBindingRef,
510 pub provider: OAuthProviderIdentity,
511 pub device_code: String,
512 pub created_at_millis: u64,
513 pub expires_at_millis: u64,
514}
515
516#[derive(Debug, Clone, PartialEq, Eq)]
517struct OAuthDeviceFlowState {
518 record: OAuthDeviceFlowRecord,
519 poll_lease: Option<OAuthDevicePollLeaseState>,
520}
521
522#[derive(Debug, Clone, Copy, PartialEq, Eq)]
523struct OAuthDevicePollLeaseState {
524 id: u64,
525}
526
527#[derive(Debug, thiserror::Error, PartialEq, Eq)]
528pub enum OAuthFlowError {
529 #[error("oauth state is missing or expired")]
530 Missing,
531 #[error("oauth registry projection payload missing after AuthMachine accepted {operation}")]
532 RegistryProjectionMissing { operation: &'static str },
533 #[error("oauth state provider mismatch: expected {expected}, got {actual}")]
534 ProviderMismatch {
535 expected: OAuthProviderIdentity,
536 actual: OAuthProviderIdentity,
537 },
538 #[error("oauth state redirect_uri mismatch")]
539 RedirectUriMismatch,
540 #[error("oauth state target mismatch: expected {expected:?}, got {actual:?}")]
541 TargetMismatch {
542 expected: Box<AuthBindingRef>,
543 actual: Box<AuthBindingRef>,
544 },
545 #[error("failed to generate oauth state token")]
546 StateGenerationFailed,
547 #[error("oauth device code poll is already in progress")]
548 DevicePollInProgress,
549 #[error("oauth device code is already admitted")]
550 DeviceCodeAlreadyAdmitted,
551 #[error("oauth device code expiry is out of range")]
552 DeviceExpiryOutOfRange,
553 #[error("oauth flow lifecycle transition rejected during {operation}: {detail}")]
554 LifecycleRejected {
555 operation: &'static str,
556 detail: String,
557 },
558 #[error("oauth flow durable persistence failed during {operation}: {detail}")]
559 PersistenceFailed {
560 operation: &'static str,
561 detail: String,
562 },
563}
564
565pub trait OAuthDevicePollLifecycle: Send + Sync {
566 fn device_flow_state_is_authmachine_owned(&self) -> bool {
567 false
568 }
569
570 fn finish_device_poll(
571 &self,
572 target: &AuthBindingRef,
573 device_code: &str,
574 ) -> Result<(), OAuthFlowError>;
575
576 fn consume_device_flow(
577 &self,
578 target: &AuthBindingRef,
579 device_code: &str,
580 provider: OAuthProviderIdentity,
581 ) -> Result<(), OAuthFlowError>;
582
583 fn expire_device_flow(
584 &self,
585 target: &AuthBindingRef,
586 device_code: &str,
587 ) -> Result<(), OAuthFlowError>;
588
589 fn restore_device_flow(&self, _record: &OAuthDeviceFlowRecord) -> Result<(), OAuthFlowError> {
590 Ok(())
591 }
592
593 fn device_flow_payloads_changed(&self) -> Result<(), OAuthFlowError> {
594 Ok(())
595 }
596
597 fn device_flow_payload_removed(
598 &self,
599 _record: &OAuthDeviceFlowRecord,
600 ) -> Result<(), OAuthFlowError> {
601 self.device_flow_payloads_changed()
602 }
603}
604
605pub struct OAuthDevicePollLease {
606 device_flows: Arc<Mutex<HashMap<String, OAuthDeviceFlowState>>>,
607 target: AuthBindingRef,
608 device_code: String,
609 provider: OAuthProviderIdentity,
610 lease_id: u64,
611 lifecycle: Option<Arc<dyn OAuthDevicePollLifecycle>>,
612 operation_lock: Option<Arc<StdMutex<()>>>,
613 active: bool,
614}
615
616impl std::fmt::Debug for OAuthDevicePollLease {
617 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
618 f.debug_struct("OAuthDevicePollLease")
619 .field("target", &self.target)
620 .field("device_code", &self.device_code)
621 .field("provider", &self.provider)
622 .field("lease_id", &self.lease_id)
623 .field("has_lifecycle", &self.lifecycle.is_some())
624 .field("has_operation_lock", &self.operation_lock.is_some())
625 .field("active", &self.active)
626 .finish()
627 }
628}
629
630impl OAuthDevicePollLease {
631 fn new(
632 device_flows: Arc<Mutex<HashMap<String, OAuthDeviceFlowState>>>,
633 target: AuthBindingRef,
634 device_code: String,
635 provider: OAuthProviderIdentity,
636 lease_id: u64,
637 ) -> Self {
638 Self {
639 device_flows,
640 target,
641 device_code,
642 provider,
643 lease_id,
644 lifecycle: None,
645 operation_lock: None,
646 active: true,
647 }
648 }
649
650 pub fn terminal_flow_state_is_authmachine_owned(&self) -> bool {
651 self.lifecycle
652 .as_ref()
653 .map(|lifecycle| lifecycle.device_flow_state_is_authmachine_owned())
654 .unwrap_or(false)
655 }
656
657 fn local_missing_error(&self, operation: &'static str) -> OAuthFlowError {
658 if self.terminal_flow_state_is_authmachine_owned() {
659 OAuthFlowError::RegistryProjectionMissing { operation }
660 } else {
661 OAuthFlowError::Missing
662 }
663 }
664
665 pub fn with_lifecycle(mut self, lifecycle: Arc<dyn OAuthDevicePollLifecycle>) -> Self {
666 self.lifecycle = Some(lifecycle);
667 self
668 }
669
670 pub fn with_operation_lock(mut self, operation_lock: Arc<StdMutex<()>>) -> Self {
671 self.operation_lock = Some(operation_lock);
672 self
673 }
674
675 pub fn finish(mut self) -> Result<(), OAuthFlowError> {
676 let _operation_guard = self.operation_lock.as_ref().map(|lock| {
677 lock.lock()
678 .unwrap_or_else(std::sync::PoisonError::into_inner)
679 });
680 let verify_result = {
681 let mut flows = self.device_flows.lock();
682 prune_expired_device_locked(&mut flows);
683 verify_device_poll_lease_locked(
684 &mut flows,
685 &self.device_code,
686 self.provider,
687 self.lease_id,
688 )
689 .map(|_| ())
690 };
691 if matches!(verify_result, Err(OAuthFlowError::Missing)) {
692 if self.terminal_flow_state_is_authmachine_owned()
693 && let Some(lifecycle) = &self.lifecycle
694 {
695 let _ = lifecycle.finish_device_poll(&self.target, &self.device_code);
696 }
697 return Err(self.local_missing_error("finish_oauth_device_poll"));
698 }
699 verify_result?;
700
701 if let Some(lifecycle) = &self.lifecycle {
702 lifecycle.finish_device_poll(&self.target, &self.device_code)?;
703 }
704
705 let result = {
706 let mut flows = self.device_flows.lock();
707 let result = release_device_poll_lease_locked(
708 &mut flows,
709 &self.device_code,
710 self.provider,
711 self.lease_id,
712 );
713 prune_expired_device_locked(&mut flows);
714 result
715 };
716 match result {
717 Ok(()) => {
718 self.active = false;
719 if let Some(lifecycle) = &self.lifecycle {
720 lifecycle.device_flow_payloads_changed()?;
721 }
722 Ok(())
723 }
724 Err(OAuthFlowError::Missing) => {
725 Err(self.local_missing_error("finish_oauth_device_poll"))
726 }
727 Err(err) => Err(err),
728 }
729 }
730
731 pub fn verify(&self) -> Result<OAuthDeviceFlowRecord, OAuthFlowError> {
732 let mut flows = self.device_flows.lock();
733 prune_expired_device_locked(&mut flows);
734 let result = verify_device_poll_lease_locked(
735 &mut flows,
736 &self.device_code,
737 self.provider,
738 self.lease_id,
739 );
740 prune_expired_device_locked(&mut flows);
741 if matches!(result, Err(OAuthFlowError::Missing)) {
742 return Err(self.local_missing_error("verify_oauth_device_poll"));
743 }
744 result
745 }
746
747 pub fn consume(mut self) -> Result<OAuthDeviceFlowRecord, OAuthFlowError> {
748 let _operation_guard = self.operation_lock.as_ref().map(|lock| {
749 lock.lock()
750 .unwrap_or_else(std::sync::PoisonError::into_inner)
751 });
752 let verified = {
753 let mut flows = self.device_flows.lock();
754 prune_expired_device_locked(&mut flows);
755 verify_device_poll_lease_locked(
756 &mut flows,
757 &self.device_code,
758 self.provider,
759 self.lease_id,
760 )
761 };
762 let verified = match verified {
763 Ok(verified) => verified,
764 Err(OAuthFlowError::Missing) => {
765 if self.terminal_flow_state_is_authmachine_owned()
766 && let Some(lifecycle) = &self.lifecycle
767 {
768 let _ = lifecycle.finish_device_poll(&self.target, &self.device_code);
769 }
770 return Err(self.local_missing_error("consume_oauth_device_flow"));
771 }
772 Err(err) => return Err(err),
773 };
774
775 if let Some(lifecycle) = &self.lifecycle {
776 lifecycle.consume_device_flow(&self.target, &self.device_code, self.provider)?;
777 }
778
779 let result = {
780 let mut flows = self.device_flows.lock();
781 let result = consume_device_poll_lease_locked(
782 &mut flows,
783 &self.device_code,
784 self.provider,
785 self.lease_id,
786 );
787 prune_expired_device_locked(&mut flows);
788 result
789 };
790 match result {
791 Ok(record) => {
792 self.active = false;
793 if let Some(lifecycle) = &self.lifecycle
794 && let Err(err) = lifecycle.device_flow_payload_removed(&record)
795 {
796 if matches!(
797 err,
798 OAuthFlowError::Missing | OAuthFlowError::RegistryProjectionMissing { .. }
799 ) {
800 return Err(err);
801 }
802 let _ = lifecycle.restore_device_flow(&verified);
803 let mut flows = self.device_flows.lock();
804 flows.insert(
805 verified.device_code.clone(),
806 OAuthDeviceFlowState {
807 record: verified,
808 poll_lease: None,
809 },
810 );
811 return Err(err);
812 }
813 Ok(record)
814 }
815 Err(err) => {
816 if matches!(err, OAuthFlowError::Missing)
817 && self.terminal_flow_state_is_authmachine_owned()
818 {
819 return Err(self.local_missing_error("consume_oauth_device_flow"));
820 }
821 if let Some(lifecycle) = &self.lifecycle {
822 let _ = lifecycle.restore_device_flow(&verified);
823 if matches!(err, OAuthFlowError::Missing) {
824 let _ = lifecycle.expire_device_flow(&self.target, &self.device_code);
825 }
826 }
827 Err(err)
828 }
829 }
830 }
831}
832
833impl Drop for OAuthDevicePollLease {
834 fn drop(&mut self) {
835 if !self.active {
836 return;
837 }
838 let mut flows = self.device_flows.lock();
839 prune_expired_device_locked(&mut flows);
840 let result = release_device_poll_lease_locked(
841 &mut flows,
842 &self.device_code,
843 self.provider,
844 self.lease_id,
845 );
846 prune_expired_device_locked(&mut flows);
847 if let Some(lifecycle) = &self.lifecycle {
848 if matches!(result, Err(OAuthFlowError::Missing)) {
849 if lifecycle.device_flow_state_is_authmachine_owned() {
850 let _ = lifecycle.finish_device_poll(&self.target, &self.device_code);
851 } else {
852 let _ = lifecycle.expire_device_flow(&self.target, &self.device_code);
853 }
854 } else {
855 let _ = lifecycle.finish_device_poll(&self.target, &self.device_code);
856 }
857 }
858 }
859}
860
861pub trait OAuthFlowAuthority: Send + Sync {
862 fn terminal_flow_state_is_authmachine_owned(&self) -> bool {
863 false
864 }
865
866 fn start(
867 &self,
868 target: AuthBindingRef,
869 provider: OAuthProviderIdentity,
870 redirect_uri: String,
871 pkce_verifier: String,
872 ) -> Result<String, OAuthFlowError>;
873
874 fn verify(
875 &self,
876 state: &str,
877 target: &AuthBindingRef,
878 provider: OAuthProviderIdentity,
879 redirect_uri: &str,
880 ) -> Result<OAuthFlowRecord, OAuthFlowError>;
881
882 fn consume(
883 &self,
884 state: &str,
885 target: &AuthBindingRef,
886 provider: OAuthProviderIdentity,
887 redirect_uri: &str,
888 ) -> Result<OAuthFlowRecord, OAuthFlowError>;
889
890 fn admit_device_code(
891 &self,
892 target: AuthBindingRef,
893 provider: OAuthProviderIdentity,
894 device_code: String,
895 expires_in: Duration,
896 ) -> Result<(), OAuthFlowError>;
897
898 fn verify_device_code(
899 &self,
900 device_code: &str,
901 target: &AuthBindingRef,
902 provider: OAuthProviderIdentity,
903 ) -> Result<OAuthDeviceFlowRecord, OAuthFlowError>;
904
905 fn begin_device_code_poll(
906 &self,
907 device_code: &str,
908 target: &AuthBindingRef,
909 provider: OAuthProviderIdentity,
910 ) -> Result<OAuthDevicePollLease, OAuthFlowError>;
911}
912
913#[derive(Debug)]
914pub struct OAuthFlowRegistry {
915 ttl: Duration,
916 max_outstanding: usize,
917 flows: Mutex<HashMap<String, OAuthFlowRecord>>,
918 device_flows: Arc<Mutex<HashMap<String, OAuthDeviceFlowState>>>,
919 next_device_poll_lease_id: AtomicU64,
920}
921
922impl OAuthFlowRegistry {
923 pub fn new(ttl: Duration) -> Self {
924 Self::new_with_capacity(ttl, DEFAULT_MAX_OUTSTANDING_FLOWS)
925 }
926
927 pub fn new_with_capacity(ttl: Duration, max_outstanding: usize) -> Self {
928 Self {
929 ttl,
930 max_outstanding: max_outstanding.max(1),
931 flows: Mutex::new(HashMap::new()),
932 device_flows: Arc::new(Mutex::new(HashMap::new())),
933 next_device_poll_lease_id: AtomicU64::new(1),
934 }
935 }
936
937 pub fn max_outstanding(&self) -> usize {
938 self.max_outstanding
939 }
940
941 pub fn ttl(&self) -> Duration {
942 self.ttl
943 }
944
945 pub fn new_state() -> Result<String, OAuthFlowError> {
946 new_state_token()
947 }
948
949 pub fn start(
950 &self,
951 target: AuthBindingRef,
952 provider: OAuthProviderIdentity,
953 redirect_uri: impl Into<String>,
954 pkce_verifier: impl Into<String>,
955 ) -> Result<String, OAuthFlowError> {
956 <Self as OAuthFlowAuthority>::start(
957 self,
958 target,
959 provider,
960 redirect_uri.into(),
961 pkce_verifier.into(),
962 )
963 }
964
965 pub fn verify(
966 &self,
967 state: &str,
968 target: &AuthBindingRef,
969 provider: OAuthProviderIdentity,
970 redirect_uri: &str,
971 ) -> Result<OAuthFlowRecord, OAuthFlowError> {
972 <Self as OAuthFlowAuthority>::verify(self, state, target, provider, redirect_uri)
973 }
974
975 pub fn consume(
976 &self,
977 state: &str,
978 target: &AuthBindingRef,
979 provider: OAuthProviderIdentity,
980 redirect_uri: &str,
981 ) -> Result<OAuthFlowRecord, OAuthFlowError> {
982 <Self as OAuthFlowAuthority>::consume(self, state, target, provider, redirect_uri)
983 }
984
985 pub fn admit_device_code(
986 &self,
987 target: AuthBindingRef,
988 provider: OAuthProviderIdentity,
989 device_code: impl Into<String>,
990 expires_in: Duration,
991 ) -> Result<(), OAuthFlowError> {
992 <Self as OAuthFlowAuthority>::admit_device_code(
993 self,
994 target,
995 provider,
996 device_code.into(),
997 expires_in,
998 )
999 }
1000
1001 pub fn verify_device_code(
1002 &self,
1003 device_code: &str,
1004 target: &AuthBindingRef,
1005 provider: OAuthProviderIdentity,
1006 ) -> Result<OAuthDeviceFlowRecord, OAuthFlowError> {
1007 <Self as OAuthFlowAuthority>::verify_device_code(self, device_code, target, provider)
1008 }
1009
1010 pub fn begin_device_code_poll(
1011 &self,
1012 device_code: &str,
1013 target: &AuthBindingRef,
1014 provider: OAuthProviderIdentity,
1015 ) -> Result<OAuthDevicePollLease, OAuthFlowError> {
1016 <Self as OAuthFlowAuthority>::begin_device_code_poll(self, device_code, target, provider)
1017 }
1018
1019 pub fn expire_device_code(
1020 &self,
1021 device_code: &str,
1022 target: &AuthBindingRef,
1023 provider: OAuthProviderIdentity,
1024 ) -> Result<(), OAuthFlowError> {
1025 let mut flows = self.device_flows.lock();
1026 prune_expired_device_locked(&mut flows);
1027 let Some(state) = flows.get(device_code) else {
1028 return Err(OAuthFlowError::Missing);
1029 };
1030 verify_device_record(&state.record, target, provider)?;
1031 flows.remove(device_code);
1032 Ok(())
1033 }
1034
1035 pub fn prune_expired_browser_flows(&self) -> Vec<(String, AuthBindingRef)> {
1036 let mut flows = self.flows.lock();
1037 take_expired_locked(&mut flows, self.ttl)
1038 .into_iter()
1039 .map(|(flow_id, record)| (flow_id, record.target))
1040 .collect()
1041 }
1042
1043 pub fn prune_expired_device_flows(&self) -> Vec<(String, AuthBindingRef)> {
1044 let mut flows = self.device_flows.lock();
1045 take_expired_device_locked(&mut flows)
1046 .into_iter()
1047 .map(|record| (record.device_code, record.target))
1048 .collect()
1049 }
1050
1051 pub fn retain_flows_with_lifecycle(
1052 &self,
1053 mut browser_active: impl FnMut(&AuthBindingRef, &str) -> bool,
1054 mut device_active: impl FnMut(&AuthBindingRef, &str) -> bool,
1055 ) -> OAuthPrunedFlows {
1056 let mut flows = self.flows.lock();
1057 let mut browser = Vec::new();
1058 flows.retain(|flow_id, record| {
1059 let keep = browser_active(&record.target, flow_id);
1060 if !keep {
1061 browser.push((flow_id.clone(), record.target.clone()));
1062 }
1063 keep
1064 });
1065
1066 let mut device_flows = self.device_flows.lock();
1067 let mut device = Vec::new();
1068 device_flows.retain(|device_code, state| {
1069 let keep = device_active(&state.record.target, device_code);
1070 if !keep {
1071 device.push((device_code.clone(), state.record.target.clone()));
1072 }
1073 keep
1074 });
1075 OAuthPrunedFlows { browser, device }
1076 }
1077
1078 pub fn snapshot_for_persistence(&self, now_millis: u64) -> OAuthFlowRegistrySnapshot {
1079 let flows = self.flows.lock();
1080 let browser = flows
1081 .iter()
1082 .filter_map(|(state, record)| {
1083 let elapsed = record.created_at.elapsed();
1084 if elapsed > self.ttl {
1085 return None;
1086 }
1087 let elapsed_millis = duration_millis_u64(elapsed);
1088 let ttl_millis = duration_millis_u64(self.ttl);
1089 Some(PersistedOAuthBrowserFlow {
1090 state: state.clone(),
1091 target: record.target.clone(),
1092 provider: record.provider,
1093 redirect_uri: record.redirect_uri.clone(),
1094 pkce_verifier: record.pkce_verifier.clone(),
1095 created_at_millis: now_millis.saturating_sub(elapsed_millis),
1096 expires_at_millis: now_millis
1097 .saturating_add(ttl_millis.saturating_sub(elapsed_millis)),
1098 })
1099 })
1100 .collect::<Vec<_>>();
1101 drop(flows);
1102
1103 let device_flows = self.device_flows.lock();
1104 let now = Instant::now();
1105 let device = device_flows
1106 .values()
1107 .filter_map(|state| {
1108 let remaining = state.record.expires_at.checked_duration_since(now)?;
1109 let created_elapsed = state.record.created_at.elapsed();
1110 Some(PersistedOAuthDeviceFlow {
1111 target: state.record.target.clone(),
1112 provider: state.record.provider,
1113 device_code: state.record.device_code.clone(),
1114 created_at_millis: now_millis
1115 .saturating_sub(duration_millis_u64(created_elapsed)),
1116 expires_at_millis: now_millis.saturating_add(duration_millis_u64(remaining)),
1117 })
1118 })
1119 .collect();
1120
1121 OAuthFlowRegistrySnapshot { browser, device }
1122 }
1123
1124 pub fn insert_restored_browser_flow(
1125 &self,
1126 state: String,
1127 target: AuthBindingRef,
1128 provider: OAuthProviderIdentity,
1129 redirect_uri: String,
1130 pkce_verifier: String,
1131 created_at: Instant,
1132 ) -> Result<(), OAuthFlowError> {
1133 let mut flows = self.flows.lock();
1134 flows.insert(
1135 state,
1136 OAuthFlowRecord {
1137 target,
1138 provider,
1139 redirect_uri,
1140 pkce_verifier,
1141 created_at,
1142 },
1143 );
1144 Ok(())
1145 }
1146
1147 pub fn insert_restored_device_flow(
1148 &self,
1149 target: AuthBindingRef,
1150 provider: OAuthProviderIdentity,
1151 device_code: String,
1152 created_at: Instant,
1153 expires_at: Instant,
1154 ) -> Result<(), OAuthFlowError> {
1155 let mut device_flows = self.device_flows.lock();
1156 if device_flows.contains_key(&device_code) {
1157 return Err(OAuthFlowError::DeviceCodeAlreadyAdmitted);
1158 }
1159 let record = OAuthDeviceFlowRecord {
1160 target,
1161 provider,
1162 device_code: device_code.clone(),
1163 created_at,
1164 expires_at,
1165 };
1166 device_flows.insert(
1167 device_code,
1168 OAuthDeviceFlowState {
1169 record,
1170 poll_lease: None,
1171 },
1172 );
1173 Ok(())
1174 }
1175
1176 pub fn start_with_pruned(
1177 &self,
1178 target: AuthBindingRef,
1179 provider: OAuthProviderIdentity,
1180 redirect_uri: String,
1181 pkce_verifier: String,
1182 ) -> Result<(String, OAuthPrunedFlows), OAuthFlowError> {
1183 let state = new_state_token()?;
1184 let record = OAuthFlowRecord {
1185 target,
1186 provider,
1187 redirect_uri,
1188 pkce_verifier,
1189 created_at: Instant::now(),
1190 };
1191 let mut flows = self.flows.lock();
1192 let mut device_flows = self.device_flows.lock();
1193 let expired_browser = take_expired_locked(&mut flows, self.ttl);
1194 let expired_device = take_expired_device_locked(&mut device_flows);
1195 flows.insert(state.clone(), record);
1196 Ok((
1197 state,
1198 OAuthPrunedFlows::from_expired(expired_browser, expired_device),
1199 ))
1200 }
1201
1202 pub fn insert_browser_flow_with_pruned(
1203 &self,
1204 state: String,
1205 target: AuthBindingRef,
1206 provider: OAuthProviderIdentity,
1207 redirect_uri: String,
1208 pkce_verifier: String,
1209 ) -> Result<OAuthPrunedFlows, OAuthFlowError> {
1210 let record = OAuthFlowRecord {
1211 target,
1212 provider,
1213 redirect_uri,
1214 pkce_verifier,
1215 created_at: Instant::now(),
1216 };
1217 let mut flows = self.flows.lock();
1218 let mut device_flows = self.device_flows.lock();
1219 let expired_browser = take_expired_locked(&mut flows, self.ttl);
1220 let expired_device = take_expired_device_locked(&mut device_flows);
1221 flows.insert(state, record);
1222 Ok(OAuthPrunedFlows::from_expired(
1223 expired_browser,
1224 expired_device,
1225 ))
1226 }
1227
1228 pub fn admit_device_code_with_pruned(
1229 &self,
1230 target: AuthBindingRef,
1231 provider: OAuthProviderIdentity,
1232 device_code: String,
1233 expires_in: Duration,
1234 ) -> Result<OAuthPrunedFlows, OAuthFlowError> {
1235 let mut flows = self.flows.lock();
1236 let mut device_flows = self.device_flows.lock();
1237 let expired_browser = take_expired_locked(&mut flows, self.ttl);
1238 let expired_device = take_expired_device_locked(&mut device_flows);
1239 if device_flows.contains_key(&device_code) {
1240 return Err(OAuthFlowError::DeviceCodeAlreadyAdmitted);
1241 }
1242 let now = Instant::now();
1243 let expires_at = now
1244 .checked_add(expires_in)
1245 .ok_or(OAuthFlowError::DeviceExpiryOutOfRange)?;
1246 let record = OAuthDeviceFlowRecord {
1247 target,
1248 provider,
1249 device_code: device_code.clone(),
1250 created_at: now,
1251 expires_at,
1252 };
1253 device_flows.insert(
1254 device_code,
1255 OAuthDeviceFlowState {
1256 record,
1257 poll_lease: None,
1258 },
1259 );
1260 Ok(OAuthPrunedFlows::from_expired(
1261 expired_browser,
1262 expired_device,
1263 ))
1264 }
1265}
1266
1267impl Default for OAuthFlowRegistry {
1268 fn default() -> Self {
1269 Self::new(Duration::from_secs(10 * 60))
1270 }
1271}
1272
1273impl OAuthFlowAuthority for OAuthFlowRegistry {
1274 fn start(
1275 &self,
1276 target: AuthBindingRef,
1277 provider: OAuthProviderIdentity,
1278 redirect_uri: String,
1279 pkce_verifier: String,
1280 ) -> Result<String, OAuthFlowError> {
1281 self.start_with_pruned(target, provider, redirect_uri, pkce_verifier)
1282 .map(|(state, _)| state)
1283 }
1284
1285 fn verify(
1286 &self,
1287 state: &str,
1288 target: &AuthBindingRef,
1289 provider: OAuthProviderIdentity,
1290 redirect_uri: &str,
1291 ) -> Result<OAuthFlowRecord, OAuthFlowError> {
1292 let mut flows = self.flows.lock();
1293 prune_expired_locked(&mut flows, self.ttl);
1294 let Some(record) = flows.get(state) else {
1295 return Err(OAuthFlowError::Missing);
1296 };
1297 verify_browser_record(record, target, provider, redirect_uri)?;
1298 Ok(record.clone())
1299 }
1300
1301 fn consume(
1302 &self,
1303 state: &str,
1304 target: &AuthBindingRef,
1305 provider: OAuthProviderIdentity,
1306 redirect_uri: &str,
1307 ) -> Result<OAuthFlowRecord, OAuthFlowError> {
1308 let mut flows = self.flows.lock();
1309 prune_expired_locked(&mut flows, self.ttl);
1310 let Some(record) = flows.get(state) else {
1311 return Err(OAuthFlowError::Missing);
1312 };
1313 verify_browser_record(record, target, provider, redirect_uri)?;
1314 flows.remove(state).ok_or(OAuthFlowError::Missing)
1315 }
1316
1317 fn admit_device_code(
1318 &self,
1319 target: AuthBindingRef,
1320 provider: OAuthProviderIdentity,
1321 device_code: String,
1322 expires_in: Duration,
1323 ) -> Result<(), OAuthFlowError> {
1324 self.admit_device_code_with_pruned(target, provider, device_code, expires_in)
1325 .map(|_| ())
1326 }
1327
1328 fn verify_device_code(
1329 &self,
1330 device_code: &str,
1331 target: &AuthBindingRef,
1332 provider: OAuthProviderIdentity,
1333 ) -> Result<OAuthDeviceFlowRecord, OAuthFlowError> {
1334 let mut flows = self.device_flows.lock();
1335 prune_expired_device_locked(&mut flows);
1336 let Some(record) = flows.get(device_code) else {
1337 return Err(OAuthFlowError::Missing);
1338 };
1339 verify_device_record(&record.record, target, provider)?;
1340 Ok(record.record.clone())
1341 }
1342
1343 fn begin_device_code_poll(
1344 &self,
1345 device_code: &str,
1346 target: &AuthBindingRef,
1347 provider: OAuthProviderIdentity,
1348 ) -> Result<OAuthDevicePollLease, OAuthFlowError> {
1349 let mut flows = self.device_flows.lock();
1350 prune_expired_device_locked(&mut flows);
1351 let Some(state) = flows.get_mut(device_code) else {
1352 return Err(OAuthFlowError::Missing);
1353 };
1354 verify_device_record(&state.record, target, provider)?;
1355 if state.poll_lease.is_some() {
1356 return Err(OAuthFlowError::DevicePollInProgress);
1357 }
1358 let lease_id = self
1359 .next_device_poll_lease_id
1360 .fetch_add(1, Ordering::Relaxed);
1361 state.poll_lease = Some(OAuthDevicePollLeaseState { id: lease_id });
1362 Ok(OAuthDevicePollLease::new(
1363 Arc::clone(&self.device_flows),
1364 target.clone(),
1365 device_code.to_string(),
1366 provider,
1367 lease_id,
1368 ))
1369 }
1370}
1371
1372fn new_state_token() -> Result<String, OAuthFlowError> {
1373 let mut bytes = [0_u8; 32];
1374 getrandom::fill(&mut bytes).map_err(|_| OAuthFlowError::StateGenerationFailed)?;
1375 Ok(format!(
1376 "st-{}",
1377 base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
1378 ))
1379}
1380
1381fn duration_millis_u64(duration: Duration) -> u64 {
1382 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
1383}
1384
1385fn prune_expired_locked(flows: &mut HashMap<String, OAuthFlowRecord>, ttl: Duration) {
1386 let _ = take_expired_locked(flows, ttl);
1387}
1388
1389fn take_expired_locked(
1390 flows: &mut HashMap<String, OAuthFlowRecord>,
1391 ttl: Duration,
1392) -> Vec<(String, OAuthFlowRecord)> {
1393 let expired = flows
1394 .iter()
1395 .filter(|(_, record)| record.created_at.elapsed() > ttl)
1396 .map(|(flow_id, _)| flow_id.clone())
1397 .collect::<Vec<_>>();
1398 expired
1399 .into_iter()
1400 .filter_map(|flow_id| flows.remove(&flow_id).map(|record| (flow_id, record)))
1401 .collect()
1402}
1403
1404fn release_device_poll_lease_locked(
1405 flows: &mut HashMap<String, OAuthDeviceFlowState>,
1406 device_code: &str,
1407 provider: OAuthProviderIdentity,
1408 lease_id: u64,
1409) -> Result<(), OAuthFlowError> {
1410 let Some(state) = flows.get_mut(device_code) else {
1411 return Err(OAuthFlowError::Missing);
1412 };
1413 if state.record.provider != provider {
1414 return Err(OAuthFlowError::ProviderMismatch {
1415 expected: state.record.provider,
1416 actual: provider,
1417 });
1418 }
1419 match state.poll_lease {
1420 Some(lease) if lease.id == lease_id => {
1421 state.poll_lease = None;
1422 Ok(())
1423 }
1424 Some(_) => Err(OAuthFlowError::DevicePollInProgress),
1425 None => Err(OAuthFlowError::Missing),
1426 }
1427}
1428
1429fn verify_browser_record(
1430 record: &OAuthFlowRecord,
1431 target: &AuthBindingRef,
1432 provider: OAuthProviderIdentity,
1433 redirect_uri: &str,
1434) -> Result<(), OAuthFlowError> {
1435 if &record.target != target {
1436 return Err(OAuthFlowError::TargetMismatch {
1437 expected: Box::new(record.target.clone()),
1438 actual: Box::new(target.clone()),
1439 });
1440 }
1441 if record.provider != provider {
1442 return Err(OAuthFlowError::ProviderMismatch {
1443 expected: record.provider,
1444 actual: provider,
1445 });
1446 }
1447 if record.redirect_uri != redirect_uri {
1448 return Err(OAuthFlowError::RedirectUriMismatch);
1449 }
1450 Ok(())
1451}
1452
1453fn verify_device_record(
1454 record: &OAuthDeviceFlowRecord,
1455 target: &AuthBindingRef,
1456 provider: OAuthProviderIdentity,
1457) -> Result<(), OAuthFlowError> {
1458 if &record.target != target {
1459 return Err(OAuthFlowError::TargetMismatch {
1460 expected: Box::new(record.target.clone()),
1461 actual: Box::new(target.clone()),
1462 });
1463 }
1464 if record.provider != provider {
1465 return Err(OAuthFlowError::ProviderMismatch {
1466 expected: record.provider,
1467 actual: provider,
1468 });
1469 }
1470 Ok(())
1471}
1472
1473fn verify_device_poll_lease_locked(
1474 flows: &mut HashMap<String, OAuthDeviceFlowState>,
1475 device_code: &str,
1476 provider: OAuthProviderIdentity,
1477 lease_id: u64,
1478) -> Result<OAuthDeviceFlowRecord, OAuthFlowError> {
1479 let Some(state) = flows.get(device_code) else {
1480 return Err(OAuthFlowError::Missing);
1481 };
1482 if state.record.provider != provider {
1483 return Err(OAuthFlowError::ProviderMismatch {
1484 expected: state.record.provider,
1485 actual: provider,
1486 });
1487 }
1488 match state.poll_lease {
1489 Some(lease) if lease.id == lease_id => {}
1490 Some(_) => return Err(OAuthFlowError::DevicePollInProgress),
1491 None => return Err(OAuthFlowError::Missing),
1492 }
1493 Ok(state.record.clone())
1494}
1495
1496fn consume_device_poll_lease_locked(
1497 flows: &mut HashMap<String, OAuthDeviceFlowState>,
1498 device_code: &str,
1499 provider: OAuthProviderIdentity,
1500 lease_id: u64,
1501) -> Result<OAuthDeviceFlowRecord, OAuthFlowError> {
1502 verify_device_poll_lease_locked(flows, device_code, provider, lease_id)?;
1503 flows
1504 .remove(device_code)
1505 .map(|state| state.record)
1506 .ok_or(OAuthFlowError::Missing)
1507}
1508
1509fn prune_expired_device_locked(flows: &mut HashMap<String, OAuthDeviceFlowState>) {
1510 let _ = take_expired_device_locked(flows);
1511}
1512
1513fn take_expired_device_locked(
1514 flows: &mut HashMap<String, OAuthDeviceFlowState>,
1515) -> Vec<OAuthDeviceFlowRecord> {
1516 let now = Instant::now();
1517 let expired = flows
1518 .iter()
1519 .filter(|(_, state)| state.record.expires_at < now)
1520 .map(|(device_code, _)| device_code.clone())
1521 .collect::<Vec<_>>();
1522 expired
1523 .into_iter()
1524 .filter_map(|device_code| flows.remove(&device_code).map(|state| state.record))
1525 .collect()
1526}
1527
1528fn strings(values: &[&str]) -> Vec<String> {
1529 values.iter().map(|value| (*value).to_string()).collect()
1530}
1531
1532#[cfg(test)]
1533#[allow(clippy::expect_used)]
1534mod tests {
1535 use super::*;
1536
1537 fn target() -> AuthBindingRef {
1538 AuthBindingRef {
1539 realm: meerkat_core::RealmId::parse("dev").expect("valid realm"),
1540 binding: meerkat_core::BindingId::parse("default_openai").expect("valid binding"),
1541 profile: None,
1542 origin: meerkat_core::BindingOrigin::Configured,
1543 }
1544 }
1545
1546 fn alternate_target() -> AuthBindingRef {
1547 AuthBindingRef {
1548 realm: meerkat_core::RealmId::parse("dev").expect("valid realm"),
1549 binding: meerkat_core::BindingId::parse("alternate_openai").expect("valid binding"),
1550 profile: None,
1551 origin: meerkat_core::BindingOrigin::Configured,
1552 }
1553 }
1554
1555 #[test]
1556 fn oauth_provider_identity_survives_serde_round_trip() {
1557 for identity in [
1561 OAuthProviderIdentity::AnthropicClaudeAi,
1562 OAuthProviderIdentity::AnthropicConsoleApiKey,
1563 OAuthProviderIdentity::OpenAiChatGpt,
1564 OAuthProviderIdentity::GoogleCodeAssist,
1565 ] {
1566 let json = serde_json::to_string(&identity).expect("serialize identity");
1567 let restored: OAuthProviderIdentity =
1568 serde_json::from_str(&json).expect("deserialize identity");
1569 assert_eq!(identity, restored);
1570 }
1571 }
1572
1573 #[test]
1574 fn persisted_browser_flow_round_trips_typed_provider() {
1575 let persisted = PersistedOAuthBrowserFlow {
1576 state: "state-token".to_string(),
1577 target: target(),
1578 provider: OAuthProviderIdentity::AnthropicClaudeAi,
1579 redirect_uri: "https://example/callback".to_string(),
1580 pkce_verifier: "verifier".to_string(),
1581 created_at_millis: 1_000,
1582 expires_at_millis: 61_000,
1583 };
1584 let json = serde_json::to_string(&persisted).expect("serialize browser flow");
1585 let restored: PersistedOAuthBrowserFlow =
1586 serde_json::from_str(&json).expect("deserialize browser flow");
1587 assert_eq!(restored.provider, OAuthProviderIdentity::AnthropicClaudeAi);
1588 assert_eq!(persisted, restored);
1589 }
1590
1591 #[test]
1592 fn persisted_device_flow_round_trips_typed_provider() {
1593 let persisted = PersistedOAuthDeviceFlow {
1594 target: target(),
1595 provider: OAuthProviderIdentity::GoogleCodeAssist,
1596 device_code: "device-code".to_string(),
1597 created_at_millis: 1_000,
1598 expires_at_millis: 61_000,
1599 };
1600 let json = serde_json::to_string(&persisted).expect("serialize device flow");
1601 let restored: PersistedOAuthDeviceFlow =
1602 serde_json::from_str(&json).expect("deserialize device flow");
1603 assert_eq!(restored.provider, OAuthProviderIdentity::GoogleCodeAssist);
1604 assert_eq!(persisted, restored);
1605 }
1606
1607 #[test]
1608 fn from_alias_rejects_unknown_alias() {
1609 assert!(OAuthProviderIdentity::from_alias("unknown").is_none());
1610 assert!(OAuthProviderIdentity::from_alias("").is_none());
1611 }
1612
1613 #[test]
1614 fn oauth_state_pkce_round_trip() {
1615 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1616 let state = registry
1617 .start(
1618 target(),
1619 OAuthProviderIdentity::OpenAiChatGpt,
1620 "http://127.0.0.1/callback",
1621 "verifier",
1622 )
1623 .expect("state generation succeeds");
1624 let record = registry.consume(
1625 &state,
1626 &target(),
1627 OAuthProviderIdentity::OpenAiChatGpt,
1628 "http://127.0.0.1/callback",
1629 );
1630 assert!(
1631 record.is_ok(),
1632 "state should resolve once: {:?}",
1633 record.err()
1634 );
1635 if let Ok(record) = record {
1636 assert_eq!(record.pkce_verifier, "verifier");
1637 }
1638 assert!(matches!(
1639 registry.consume(
1640 &state,
1641 &target(),
1642 OAuthProviderIdentity::OpenAiChatGpt,
1643 "http://127.0.0.1/callback"
1644 ),
1645 Err(OAuthFlowError::Missing)
1646 ));
1647 }
1648
1649 #[test]
1650 fn oauth_state_rejects_mismatch() {
1651 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1652 let state = registry
1653 .start(
1654 target(),
1655 OAuthProviderIdentity::OpenAiChatGpt,
1656 "http://127.0.0.1/callback",
1657 "verifier",
1658 )
1659 .expect("state generation succeeds");
1660 assert!(matches!(
1661 registry.consume(
1662 &state,
1663 &target(),
1664 OAuthProviderIdentity::AnthropicClaudeAi,
1665 "http://127.0.0.1/callback"
1666 ),
1667 Err(OAuthFlowError::ProviderMismatch { .. })
1668 ));
1669 }
1670
1671 #[test]
1672 fn oauth_state_provider_mismatch_does_not_consume_state() {
1673 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1674 let state = registry
1675 .start(
1676 target(),
1677 OAuthProviderIdentity::OpenAiChatGpt,
1678 "http://127.0.0.1/callback",
1679 "verifier",
1680 )
1681 .expect("state generation succeeds");
1682 assert!(matches!(
1683 registry.consume(
1684 &state,
1685 &target(),
1686 OAuthProviderIdentity::AnthropicClaudeAi,
1687 "http://127.0.0.1/callback"
1688 ),
1689 Err(OAuthFlowError::ProviderMismatch { .. })
1690 ));
1691
1692 let record = registry
1693 .consume(
1694 &state,
1695 &target(),
1696 OAuthProviderIdentity::OpenAiChatGpt,
1697 "http://127.0.0.1/callback",
1698 )
1699 .expect("provider mismatch must leave state retryable");
1700 assert_eq!(record.pkce_verifier, "verifier");
1701 }
1702
1703 #[test]
1704 fn oauth_state_rejects_redirect_uri_mismatch() {
1705 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1706 let state = registry
1707 .start(
1708 target(),
1709 OAuthProviderIdentity::OpenAiChatGpt,
1710 "http://127.0.0.1/callback",
1711 "verifier",
1712 )
1713 .expect("state generation succeeds");
1714 assert!(matches!(
1715 registry.consume(
1716 &state,
1717 &target(),
1718 OAuthProviderIdentity::OpenAiChatGpt,
1719 "http://127.0.0.1/other"
1720 ),
1721 Err(OAuthFlowError::RedirectUriMismatch)
1722 ));
1723 }
1724
1725 #[test]
1726 fn oauth_state_redirect_uri_mismatch_does_not_consume_state() {
1727 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1728 let state = registry
1729 .start(
1730 target(),
1731 OAuthProviderIdentity::OpenAiChatGpt,
1732 "http://127.0.0.1/callback",
1733 "verifier",
1734 )
1735 .expect("state generation succeeds");
1736 assert!(matches!(
1737 registry.consume(
1738 &state,
1739 &target(),
1740 OAuthProviderIdentity::OpenAiChatGpt,
1741 "http://127.0.0.1/other"
1742 ),
1743 Err(OAuthFlowError::RedirectUriMismatch)
1744 ));
1745
1746 let record = registry
1747 .consume(
1748 &state,
1749 &target(),
1750 OAuthProviderIdentity::OpenAiChatGpt,
1751 "http://127.0.0.1/callback",
1752 )
1753 .expect("redirect mismatch must leave state retryable");
1754 assert_eq!(record.pkce_verifier, "verifier");
1755 }
1756
1757 #[test]
1758 fn oauth_state_target_mismatch_does_not_consume_state() {
1759 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1760 let state = registry
1761 .start(
1762 target(),
1763 OAuthProviderIdentity::OpenAiChatGpt,
1764 "http://127.0.0.1/callback",
1765 "verifier",
1766 )
1767 .expect("state generation succeeds");
1768 assert!(matches!(
1769 registry.consume(
1770 &state,
1771 &alternate_target(),
1772 OAuthProviderIdentity::OpenAiChatGpt,
1773 "http://127.0.0.1/callback"
1774 ),
1775 Err(OAuthFlowError::TargetMismatch { .. })
1776 ));
1777
1778 let record = registry
1779 .consume(
1780 &state,
1781 &target(),
1782 OAuthProviderIdentity::OpenAiChatGpt,
1783 "http://127.0.0.1/callback",
1784 )
1785 .expect("target mismatch must leave state retryable");
1786 assert_eq!(record.pkce_verifier, "verifier");
1787 }
1788
1789 #[test]
1790 fn oauth_state_verify_does_not_consume_state() {
1791 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1792 let state = registry
1793 .start(
1794 target(),
1795 OAuthProviderIdentity::OpenAiChatGpt,
1796 "http://127.0.0.1/callback",
1797 "verifier",
1798 )
1799 .expect("state generation succeeds");
1800
1801 let verified = registry
1802 .verify(
1803 &state,
1804 &target(),
1805 OAuthProviderIdentity::OpenAiChatGpt,
1806 "http://127.0.0.1/callback",
1807 )
1808 .expect("verify should read state before terminal commit");
1809
1810 assert_eq!(verified.pkce_verifier, "verifier");
1811 registry
1812 .consume(
1813 &state,
1814 &target(),
1815 OAuthProviderIdentity::OpenAiChatGpt,
1816 "http://127.0.0.1/callback",
1817 )
1818 .expect("verified browser state remains available for terminal consume");
1819 }
1820
1821 #[test]
1822 fn oauth_state_random_tokens_are_urlsafe_and_unique() {
1823 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1824 let first = registry
1825 .start(
1826 target(),
1827 OAuthProviderIdentity::OpenAiChatGpt,
1828 "http://127.0.0.1/callback",
1829 "verifier-a",
1830 )
1831 .expect("state generation succeeds");
1832 let second = registry
1833 .start(
1834 target(),
1835 OAuthProviderIdentity::OpenAiChatGpt,
1836 "http://127.0.0.1/callback",
1837 "verifier-b",
1838 )
1839 .expect("state generation succeeds");
1840 assert_ne!(first, second);
1841 assert!(first.starts_with("st-"));
1842 assert!(
1843 first[3..]
1844 .chars()
1845 .all(|ch| { ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' })
1846 );
1847 }
1848
1849 #[test]
1850 fn oauth_state_expired_records_are_pruned() {
1851 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1852 registry.flows.lock().insert(
1853 "st-old".to_string(),
1854 OAuthFlowRecord {
1855 target: target(),
1856 provider: OAuthProviderIdentity::OpenAiChatGpt,
1857 redirect_uri: "http://127.0.0.1/callback".to_string(),
1858 pkce_verifier: "verifier".to_string(),
1859 created_at: Instant::now()
1860 .checked_sub(Duration::from_secs(61))
1861 .expect("test duration is representable"),
1862 },
1863 );
1864 assert!(matches!(
1865 registry.consume(
1866 "st-old",
1867 &target(),
1868 OAuthProviderIdentity::OpenAiChatGpt,
1869 "http://127.0.0.1/callback"
1870 ),
1871 Err(OAuthFlowError::Missing)
1872 ));
1873 }
1874
1875 #[test]
1876 fn oauth_state_cannot_cross_login_lifecycle_authorities() {
1877 let admitting_authority = OAuthFlowRegistry::new(Duration::from_secs(60));
1878 let unrelated_authority = OAuthFlowRegistry::new(Duration::from_secs(60));
1879 let state = admitting_authority
1880 .start(
1881 target(),
1882 OAuthProviderIdentity::OpenAiChatGpt,
1883 "http://127.0.0.1/callback",
1884 "verifier",
1885 )
1886 .expect("state generation succeeds");
1887
1888 assert!(matches!(
1889 unrelated_authority.consume(
1890 &state,
1891 &target(),
1892 OAuthProviderIdentity::OpenAiChatGpt,
1893 "http://127.0.0.1/callback"
1894 ),
1895 Err(OAuthFlowError::Missing)
1896 ));
1897
1898 let record = admitting_authority
1899 .consume(
1900 &state,
1901 &target(),
1902 OAuthProviderIdentity::OpenAiChatGpt,
1903 "http://127.0.0.1/callback",
1904 )
1905 .expect("admitting authority owns the flow");
1906 assert_eq!(record.pkce_verifier, "verifier");
1907 }
1908
1909 #[test]
1910 fn oauth_device_flow_is_retained_until_terminal_consume() {
1911 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1912 registry
1913 .admit_device_code(
1914 target(),
1915 OAuthProviderIdentity::GoogleCodeAssist,
1916 "device-code",
1917 Duration::from_secs(600),
1918 )
1919 .expect("device code admitted");
1920
1921 let observed = registry
1922 .verify_device_code(
1923 "device-code",
1924 &target(),
1925 OAuthProviderIdentity::GoogleCodeAssist,
1926 )
1927 .expect("pending device flow remains visible");
1928 assert_eq!(observed.device_code, "device-code");
1929
1930 let poll = registry
1931 .begin_device_code_poll(
1932 "device-code",
1933 &target(),
1934 OAuthProviderIdentity::GoogleCodeAssist,
1935 )
1936 .expect("poll begins");
1937 let consumed = poll.consume().expect("terminal device flow consumes");
1938 assert_eq!(consumed.provider, OAuthProviderIdentity::GoogleCodeAssist);
1939 assert!(matches!(
1940 registry.verify_device_code(
1941 "device-code",
1942 &target(),
1943 OAuthProviderIdentity::GoogleCodeAssist
1944 ),
1945 Err(OAuthFlowError::Missing)
1946 ));
1947 }
1948
1949 #[test]
1950 fn oauth_device_poll_verify_keeps_terminal_consume_available() {
1951 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1952 registry
1953 .admit_device_code(
1954 target(),
1955 OAuthProviderIdentity::GoogleCodeAssist,
1956 "device-code",
1957 Duration::from_secs(600),
1958 )
1959 .expect("device code admitted");
1960 let poll = registry
1961 .begin_device_code_poll(
1962 "device-code",
1963 &target(),
1964 OAuthProviderIdentity::GoogleCodeAssist,
1965 )
1966 .expect("poll begins");
1967
1968 let verified = poll
1969 .verify()
1970 .expect("terminal preflight verifies the current lease");
1971
1972 assert_eq!(verified.device_code, "device-code");
1973 assert!(matches!(
1974 registry.begin_device_code_poll(
1975 "device-code",
1976 &target(),
1977 OAuthProviderIdentity::GoogleCodeAssist
1978 ),
1979 Err(OAuthFlowError::DevicePollInProgress)
1980 ));
1981 let consumed = poll.consume().expect("verified lease still consumes");
1982 assert_eq!(consumed.device_code, "device-code");
1983 }
1984
1985 #[test]
1986 fn oauth_device_admission_does_not_replace_active_poll_lease() {
1987 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
1988 registry
1989 .admit_device_code(
1990 target(),
1991 OAuthProviderIdentity::GoogleCodeAssist,
1992 "device-code",
1993 Duration::from_secs(600),
1994 )
1995 .expect("device code admitted");
1996 let poll = registry
1997 .begin_device_code_poll(
1998 "device-code",
1999 &target(),
2000 OAuthProviderIdentity::GoogleCodeAssist,
2001 )
2002 .expect("poll begins");
2003
2004 let duplicate = registry.admit_device_code(
2005 target(),
2006 OAuthProviderIdentity::GoogleCodeAssist,
2007 "device-code",
2008 Duration::from_secs(600),
2009 );
2010
2011 assert_eq!(duplicate, Err(OAuthFlowError::DeviceCodeAlreadyAdmitted));
2012 let consumed = poll
2013 .consume()
2014 .expect("duplicate admission must not replace the active poll lease");
2015 assert_eq!(consumed.device_code, "device-code");
2016 }
2017
2018 #[test]
2019 fn oauth_device_flow_rejects_provider_mismatch() {
2020 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2021 registry
2022 .admit_device_code(
2023 target(),
2024 OAuthProviderIdentity::GoogleCodeAssist,
2025 "device-code",
2026 Duration::from_secs(600),
2027 )
2028 .expect("device code admitted");
2029
2030 assert!(matches!(
2031 registry.verify_device_code(
2032 "device-code",
2033 &target(),
2034 OAuthProviderIdentity::OpenAiChatGpt
2035 ),
2036 Err(OAuthFlowError::ProviderMismatch { .. })
2037 ));
2038 let poll = registry
2039 .begin_device_code_poll(
2040 "device-code",
2041 &target(),
2042 OAuthProviderIdentity::GoogleCodeAssist,
2043 )
2044 .expect("correct provider poll begins after mismatch");
2045 assert!(poll.consume().is_ok());
2046 }
2047
2048 #[test]
2049 fn oauth_device_flow_rejects_target_mismatch_without_consuming() {
2050 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2051 registry
2052 .admit_device_code(
2053 target(),
2054 OAuthProviderIdentity::GoogleCodeAssist,
2055 "device-code",
2056 Duration::from_secs(600),
2057 )
2058 .expect("device code admitted");
2059
2060 assert!(matches!(
2061 registry.verify_device_code(
2062 "device-code",
2063 &alternate_target(),
2064 OAuthProviderIdentity::GoogleCodeAssist
2065 ),
2066 Err(OAuthFlowError::TargetMismatch { .. })
2067 ));
2068 let poll = registry
2069 .begin_device_code_poll(
2070 "device-code",
2071 &target(),
2072 OAuthProviderIdentity::GoogleCodeAssist,
2073 )
2074 .expect("correct target poll begins after mismatch");
2075 assert!(poll.consume().is_ok());
2076 }
2077
2078 #[test]
2079 fn oauth_device_flow_expired_records_are_pruned() {
2080 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2081 registry.device_flows.lock().insert(
2082 "device-code".to_string(),
2083 OAuthDeviceFlowState {
2084 record: OAuthDeviceFlowRecord {
2085 target: target(),
2086 provider: OAuthProviderIdentity::GoogleCodeAssist,
2087 device_code: "device-code".to_string(),
2088 created_at: Instant::now()
2089 .checked_sub(Duration::from_secs(601))
2090 .expect("test duration is representable"),
2091 expires_at: Instant::now()
2092 .checked_sub(Duration::from_secs(1))
2093 .expect("test duration is representable"),
2094 },
2095 poll_lease: None,
2096 },
2097 );
2098
2099 assert!(matches!(
2100 registry.verify_device_code(
2101 "device-code",
2102 &target(),
2103 OAuthProviderIdentity::GoogleCodeAssist
2104 ),
2105 Err(OAuthFlowError::Missing)
2106 ));
2107 }
2108
2109 #[test]
2110 fn oauth_device_flow_rejects_unrepresentable_expiry() {
2111 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2112 let err = registry
2113 .admit_device_code(
2114 target(),
2115 OAuthProviderIdentity::GoogleCodeAssist,
2116 "device-code",
2117 Duration::MAX,
2118 )
2119 .expect_err("unrepresentable device expiry should be rejected");
2120
2121 assert_eq!(err, OAuthFlowError::DeviceExpiryOutOfRange);
2122 assert!(matches!(
2123 registry.verify_device_code(
2124 "device-code",
2125 &target(),
2126 OAuthProviderIdentity::GoogleCodeAssist
2127 ),
2128 Err(OAuthFlowError::Missing)
2129 ));
2130 }
2131
2132 #[test]
2133 fn oauth_device_terminal_consume_rejects_local_expiry_boundary() {
2134 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2135 registry
2136 .admit_device_code(
2137 target(),
2138 OAuthProviderIdentity::GoogleCodeAssist,
2139 "device-code",
2140 Duration::from_secs(600),
2141 )
2142 .expect("device code admitted");
2143 let poll = registry
2144 .begin_device_code_poll(
2145 "device-code",
2146 &target(),
2147 OAuthProviderIdentity::GoogleCodeAssist,
2148 )
2149 .expect("poll begins before local expiry boundary");
2150 {
2151 let mut flows = registry.device_flows.lock();
2152 flows
2153 .get_mut("device-code")
2154 .expect("device flow exists")
2155 .record
2156 .expires_at = Instant::now()
2157 .checked_sub(Duration::from_secs(1))
2158 .expect("test duration is representable");
2159 }
2160
2161 assert!(matches!(poll.consume(), Err(OAuthFlowError::Missing)));
2162 assert!(matches!(
2163 registry.verify_device_code(
2164 "device-code",
2165 &target(),
2166 OAuthProviderIdentity::GoogleCodeAssist
2167 ),
2168 Err(OAuthFlowError::Missing)
2169 ));
2170 }
2171
2172 #[test]
2173 fn oauth_device_terminal_consume_rejects_intervening_prune_while_polling() {
2174 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2175 registry
2176 .admit_device_code(
2177 target(),
2178 OAuthProviderIdentity::GoogleCodeAssist,
2179 "device-code",
2180 Duration::from_secs(600),
2181 )
2182 .expect("device code admitted");
2183 let poll = registry
2184 .begin_device_code_poll(
2185 "device-code",
2186 &target(),
2187 OAuthProviderIdentity::GoogleCodeAssist,
2188 )
2189 .expect("poll begins");
2190 assert!(matches!(
2191 registry.begin_device_code_poll(
2192 "device-code",
2193 &target(),
2194 OAuthProviderIdentity::GoogleCodeAssist
2195 ),
2196 Err(OAuthFlowError::DevicePollInProgress)
2197 ));
2198 {
2199 let mut flows = registry.device_flows.lock();
2200 flows
2201 .get_mut("device-code")
2202 .expect("device flow exists")
2203 .record
2204 .expires_at = Instant::now()
2205 .checked_sub(Duration::from_secs(1))
2206 .expect("test duration is representable");
2207 }
2208
2209 registry
2210 .start(
2211 target(),
2212 OAuthProviderIdentity::OpenAiChatGpt,
2213 "http://127.0.0.1/callback",
2214 "verifier",
2215 )
2216 .expect("unrelated start prunes expired device flow");
2217 assert!(matches!(poll.consume(), Err(OAuthFlowError::Missing)));
2218 }
2219
2220 #[test]
2221 fn oauth_device_poll_drop_releases_in_progress_lifecycle() {
2222 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2223 registry
2224 .admit_device_code(
2225 target(),
2226 OAuthProviderIdentity::GoogleCodeAssist,
2227 "device-code",
2228 Duration::from_secs(600),
2229 )
2230 .expect("device code admitted");
2231
2232 let poll = registry
2233 .begin_device_code_poll(
2234 "device-code",
2235 &target(),
2236 OAuthProviderIdentity::GoogleCodeAssist,
2237 )
2238 .expect("poll begins");
2239 assert!(matches!(
2240 registry.begin_device_code_poll(
2241 "device-code",
2242 &target(),
2243 OAuthProviderIdentity::GoogleCodeAssist
2244 ),
2245 Err(OAuthFlowError::DevicePollInProgress)
2246 ));
2247
2248 drop(poll);
2249
2250 registry
2251 .begin_device_code_poll(
2252 "device-code",
2253 &target(),
2254 OAuthProviderIdentity::GoogleCodeAssist,
2255 )
2256 .expect("dropped poll lease releases the in-progress lifecycle");
2257 }
2258
2259 #[test]
2260 fn oauth_device_poll_drop_prunes_expired_in_progress_record() {
2261 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2262 registry
2263 .admit_device_code(
2264 target(),
2265 OAuthProviderIdentity::GoogleCodeAssist,
2266 "device-code",
2267 Duration::from_secs(600),
2268 )
2269 .expect("device code admitted");
2270 let poll = registry
2271 .begin_device_code_poll(
2272 "device-code",
2273 &target(),
2274 OAuthProviderIdentity::GoogleCodeAssist,
2275 )
2276 .expect("poll begins");
2277 {
2278 let mut flows = registry.device_flows.lock();
2279 flows
2280 .get_mut("device-code")
2281 .expect("device flow exists")
2282 .record
2283 .expires_at = Instant::now()
2284 .checked_sub(Duration::from_secs(1))
2285 .expect("test duration is representable");
2286 }
2287
2288 drop(poll);
2289
2290 assert!(matches!(
2291 registry.verify_device_code(
2292 "device-code",
2293 &target(),
2294 OAuthProviderIdentity::GoogleCodeAssist
2295 ),
2296 Err(OAuthFlowError::Missing)
2297 ));
2298 }
2299
2300 struct RejectConsumeLifecycle;
2301
2302 impl OAuthDevicePollLifecycle for RejectConsumeLifecycle {
2303 fn finish_device_poll(
2304 &self,
2305 _target: &AuthBindingRef,
2306 _device_code: &str,
2307 ) -> Result<(), OAuthFlowError> {
2308 Ok(())
2309 }
2310
2311 fn consume_device_flow(
2312 &self,
2313 _target: &AuthBindingRef,
2314 _device_code: &str,
2315 _provider: OAuthProviderIdentity,
2316 ) -> Result<(), OAuthFlowError> {
2317 Err(OAuthFlowError::LifecycleRejected {
2318 operation: "consume_oauth_device_flow",
2319 detail: "injected failure".to_string(),
2320 })
2321 }
2322
2323 fn expire_device_flow(
2324 &self,
2325 _target: &AuthBindingRef,
2326 _device_code: &str,
2327 ) -> Result<(), OAuthFlowError> {
2328 Ok(())
2329 }
2330 }
2331
2332 #[test]
2333 fn oauth_device_lifecycle_consume_failure_keeps_flow_retryable() {
2334 let registry = OAuthFlowRegistry::new(Duration::from_secs(60));
2335 registry
2336 .admit_device_code(
2337 target(),
2338 OAuthProviderIdentity::GoogleCodeAssist,
2339 "device-code",
2340 Duration::from_secs(600),
2341 )
2342 .expect("device code admitted");
2343 let poll = registry
2344 .begin_device_code_poll(
2345 "device-code",
2346 &target(),
2347 OAuthProviderIdentity::GoogleCodeAssist,
2348 )
2349 .expect("poll begins")
2350 .with_lifecycle(std::sync::Arc::new(RejectConsumeLifecycle));
2351
2352 assert!(matches!(
2353 poll.consume(),
2354 Err(OAuthFlowError::LifecycleRejected {
2355 operation: "consume_oauth_device_flow",
2356 ..
2357 })
2358 ));
2359
2360 let retry = registry
2361 .begin_device_code_poll(
2362 "device-code",
2363 &target(),
2364 OAuthProviderIdentity::GoogleCodeAssist,
2365 )
2366 .expect("failed lifecycle consume keeps device flow retryable");
2367 assert!(retry.consume().is_ok());
2368 }
2369
2370 #[test]
2371 fn oauth_provider_resolution_preserves_aliases() {
2372 let cases = [
2373 (
2374 "anthropic",
2375 OAuthProviderIdentity::AnthropicClaudeAi,
2376 Provider::Anthropic,
2377 PersistedAuthMode::ClaudeAiOauth,
2378 ),
2379 (
2380 "claude",
2381 OAuthProviderIdentity::AnthropicClaudeAi,
2382 Provider::Anthropic,
2383 PersistedAuthMode::ClaudeAiOauth,
2384 ),
2385 (
2386 "claude.ai",
2387 OAuthProviderIdentity::AnthropicClaudeAi,
2388 Provider::Anthropic,
2389 PersistedAuthMode::ClaudeAiOauth,
2390 ),
2391 (
2392 "openai",
2393 OAuthProviderIdentity::OpenAiChatGpt,
2394 Provider::OpenAI,
2395 PersistedAuthMode::ChatgptOauth,
2396 ),
2397 (
2398 "chatgpt",
2399 OAuthProviderIdentity::OpenAiChatGpt,
2400 Provider::OpenAI,
2401 PersistedAuthMode::ChatgptOauth,
2402 ),
2403 (
2404 "google",
2405 OAuthProviderIdentity::GoogleCodeAssist,
2406 Provider::Gemini,
2407 PersistedAuthMode::GoogleOauth,
2408 ),
2409 (
2410 "gemini",
2411 OAuthProviderIdentity::GoogleCodeAssist,
2412 Provider::Gemini,
2413 PersistedAuthMode::GoogleOauth,
2414 ),
2415 (
2416 "code_assist",
2417 OAuthProviderIdentity::GoogleCodeAssist,
2418 Provider::Gemini,
2419 PersistedAuthMode::GoogleOauth,
2420 ),
2421 ];
2422
2423 for (alias, identity, provider, auth_mode) in cases {
2424 let resolved =
2425 resolve_oauth_provider(alias, "http://127.0.0.1/callback").expect("alias resolves");
2426 assert_eq!(resolved.identity, identity);
2427 assert_eq!(resolved.provider, provider);
2428 assert_eq!(resolved.auth_mode, auth_mode);
2429 assert_eq!(resolved.endpoints.redirect_uri, "http://127.0.0.1/callback");
2430 }
2431 }
2432
2433 #[test]
2434 fn oauth_provider_resolution_exposes_google_device_secret() {
2435 let resolved =
2436 resolve_oauth_provider("code_assist", "").expect("google code assist resolves");
2437
2438 assert_eq!(resolved.identity, OAuthProviderIdentity::GoogleCodeAssist);
2439 assert!(resolved.endpoints.device_code_url.is_some());
2440 assert_eq!(
2441 resolved.client_secret,
2442 Some(meerkat_core::oauth_identity::GOOGLE_CLIENT_SECRET)
2443 );
2444 }
2445
2446 #[cfg(feature = "oauth")]
2447 #[test]
2448 fn openai_provider_resolution_matches_codex_authorize_contract() {
2449 let resolved = resolve_oauth_provider("openai", "http://localhost:1455/auth/callback")
2450 .expect("openai resolves");
2451 let pkce = crate::auth_oauth::PkcePair::generate_s256();
2452 let authorize_url = resolved
2453 .endpoints
2454 .authorize_url_with_pkce(&pkce.challenge, "state-abc");
2455
2456 assert_eq!(
2457 resolved.endpoints.redirect_uri,
2458 "http://localhost:1455/auth/callback"
2459 );
2460 assert_eq!(
2461 resolved.endpoints.token_request_format,
2462 OAuthTokenRequestFormat::FormUrlEncoded
2463 );
2464 assert!(
2465 authorize_url.contains("redirect_uri=http%3A%2F%2Flocalhost%3A1455%2Fauth%2Fcallback")
2466 );
2467 assert!(authorize_url.contains("id_token_add_organizations=true"));
2468 assert!(authorize_url.contains("codex_cli_simplified_flow=true"));
2469 assert!(authorize_url.contains("originator=codex_cli_rs"));
2470 }
2471
2472 #[test]
2473 fn anthropic_provider_resolution_matches_claude_code_token_contract() {
2474 let resolved = resolve_oauth_provider("anthropic", "http://localhost:1455/callback")
2475 .expect("anthropic resolves");
2476
2477 assert_eq!(
2478 resolved.endpoints.token_request_format,
2479 OAuthTokenRequestFormat::Json
2480 );
2481 assert!(resolved.endpoints.include_state_in_token_exchange);
2482 assert_eq!(
2483 resolved.endpoints.refresh_scopes,
2484 strings(ANTHROPIC_CLAUDE_AI_SCOPES)
2485 );
2486 assert!(
2487 resolved
2488 .endpoints
2489 .extra_authorize_params
2490 .contains(&("code".to_string(), "true".to_string()))
2491 );
2492 }
2493
2494 #[test]
2495 fn openai_declaration_is_the_single_owner_of_chatgpt_oauth_facts() {
2496 use meerkat_core::provider_matrix::openai::OpenAiBackendKind;
2497
2498 let declaration = oauth_provider_declaration(OAuthProviderIdentity::OpenAiChatGpt);
2503 assert_eq!(declaration.client_id, "app_EMoamEEZ73f0CkXaXp7hrann");
2504 assert_eq!(
2505 declaration.authorize_endpoint,
2506 "https://auth.openai.com/oauth/authorize"
2507 );
2508 assert_eq!(
2509 declaration.token_endpoint,
2510 "https://auth.openai.com/oauth/token"
2511 );
2512 assert_eq!(
2513 declaration.scopes,
2514 &[
2515 "openid",
2516 "profile",
2517 "email",
2518 "offline_access",
2519 "api.connectors.read",
2520 "api.connectors.invoke",
2521 ]
2522 );
2523 assert_eq!(
2524 declaration.backend_kind,
2525 NormalizedBackendKind::OpenAi(OpenAiBackendKind::ChatGptBackend)
2526 );
2527 assert_eq!(
2528 declaration.extra_authorize_params,
2529 &[
2530 ("id_token_add_organizations", "true"),
2531 ("codex_cli_simplified_flow", "true"),
2532 ("originator", "codex_cli_rs"),
2533 ]
2534 );
2535 }
2536
2537 #[test]
2538 fn endpoints_are_built_from_the_declaration() {
2539 for identity in [
2542 OAuthProviderIdentity::AnthropicClaudeAi,
2543 OAuthProviderIdentity::AnthropicConsoleApiKey,
2544 OAuthProviderIdentity::OpenAiChatGpt,
2545 OAuthProviderIdentity::GoogleCodeAssist,
2546 ] {
2547 let declaration = oauth_provider_declaration(identity);
2548 let endpoints = oauth_provider_endpoints(identity, "http://127.0.0.1:0/callback");
2549 assert_eq!(endpoints.client_id, declaration.client_id);
2550 assert_eq!(endpoints.authorize_url, declaration.authorize_endpoint);
2551 assert_eq!(endpoints.token_url, declaration.token_endpoint);
2552 assert_eq!(endpoints.scopes, strings(declaration.scopes));
2553 let expected_params: Vec<(String, String)> = declaration
2554 .extra_authorize_params
2555 .iter()
2556 .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
2557 .collect();
2558 assert_eq!(endpoints.extra_authorize_params, expected_params);
2559 }
2560 }
2561}