1use std::collections::HashMap;
38use std::fmt;
39use std::sync::Arc;
40use std::time::{Duration, Instant};
41
42use async_trait::async_trait;
43use tokio::sync::{Mutex, RwLock, oneshot};
44
45use super::oauth::{
46 OAuthBearerChallenge, OAuthClientError, OAuthTokenEndpointAuthMethod, TokenProvider,
47};
48
49fn generate_code_verifier() -> String {
55 use base64::Engine;
56 let mut bytes = [0u8; 32];
57 getrandom::fill(&mut bytes).expect("getrandom failed");
58 base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
59}
60
61fn compute_code_challenge(verifier: &str) -> String {
63 use base64::Engine;
64 use sha2::{Digest, Sha256};
65 let hash = Sha256::digest(verifier.as_bytes());
66 base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hash)
67}
68
69fn generate_state() -> String {
71 use base64::Engine;
72 let mut bytes = [0u8; 16];
73 getrandom::fill(&mut bytes).expect("getrandom failed");
74 base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
75}
76
77#[derive(Debug, Clone, serde::Deserialize)]
86pub struct OAuthAuthorizationServerMetadata {
87 pub issuer: String,
89 pub authorization_endpoint: String,
91 pub token_endpoint: String,
93 pub registration_endpoint: Option<String>,
95 #[serde(default)]
97 pub client_id_metadata_document_supported: bool,
98 #[serde(default)]
102 pub authorization_response_iss_parameter_supported: bool,
103 #[serde(default)]
105 pub code_challenge_methods_supported: Vec<String>,
106 #[serde(default)]
108 pub token_endpoint_auth_methods_supported: Vec<String>,
109 #[serde(default)]
111 pub grant_types_supported: Vec<String>,
112 #[serde(default)]
114 pub scopes_supported: Vec<String>,
115}
116
117#[derive(Debug, Clone, serde::Deserialize)]
122pub struct OAuthProtectedResourceMetadata {
123 pub resource: String,
125 #[serde(default)]
127 pub authorization_servers: Vec<String>,
128 #[serde(default)]
130 pub scopes_supported: Vec<String>,
131}
132
133#[derive(Debug, Clone)]
135pub struct OAuthAuthorizationDiscovery {
136 pub resource: String,
138 pub protected_resource_metadata: OAuthProtectedResourceMetadata,
140 pub authorization_servers: Vec<OAuthAuthorizationServerMetadata>,
143 pub challenge: Option<OAuthBearerChallenge>,
145}
146
147impl OAuthAuthorizationDiscovery {
148 pub fn authorization_server(&self, issuer: &str) -> Option<&OAuthAuthorizationServerMetadata> {
150 self.authorization_servers
151 .iter()
152 .find(|metadata| metadata.issuer == issuer)
153 }
154}
155
156pub async fn discover_oauth_authorization_server(
159 server_url: &str,
160 client: &reqwest::Client,
161) -> Result<OAuthAuthorizationServerMetadata, OAuthClientError> {
162 let discovery = discover_oauth_authorization(server_url, None, client).await?;
163 discovery
164 .authorization_servers
165 .into_iter()
166 .next()
167 .ok_or_else(|| OAuthClientError::Discovery("no authorization server discovered".into()))
168}
169
170pub async fn probe_oauth_bearer_challenge(
172 resource_url: &str,
173 client: &reqwest::Client,
174) -> Result<Option<OAuthBearerChallenge>, OAuthClientError> {
175 let response = client
176 .get(resource_url)
177 .send()
178 .await
179 .map_err(|e| OAuthClientError::Discovery(e.to_string()))?;
180 Ok(response
181 .headers()
182 .get_all(reqwest::header::WWW_AUTHENTICATE)
183 .iter()
184 .filter_map(|value| value.to_str().ok())
185 .find_map(OAuthBearerChallenge::from_www_authenticate))
186}
187
188pub async fn discover_oauth_authorization(
196 server_url: &str,
197 challenge: Option<OAuthBearerChallenge>,
198 client: &reqwest::Client,
199) -> Result<OAuthAuthorizationDiscovery, OAuthClientError> {
200 let challenge = match challenge {
201 Some(challenge) => Some(challenge),
202 None => probe_oauth_bearer_challenge(server_url, client).await?,
203 };
204
205 let challenge_metadata_url = challenge
206 .as_ref()
207 .and_then(|challenge| challenge.resource_metadata.as_ref())
208 .cloned();
209 let (metadata_url, protected_resource_metadata) = if let Some(url) = challenge_metadata_url {
210 let metadata = fetch_json::<OAuthProtectedResourceMetadata>(client, &url).await?;
211 (url, metadata)
212 } else {
213 let mut discovered = None;
214 for url in protected_resource_metadata_urls(server_url)? {
215 match fetch_json::<OAuthProtectedResourceMetadata>(client, &url).await {
216 Ok(metadata) => {
217 discovered = Some((url, metadata));
218 break;
219 }
220 Err(error) => {
221 tracing::debug!(%url, %error, "OAuth protected-resource metadata candidate failed")
222 }
223 }
224 }
225 discovered.ok_or_else(|| {
226 OAuthClientError::Discovery(format!(
227 "could not discover Protected Resource Metadata for `{server_url}`"
228 ))
229 })?
230 };
231 validate_resource_identifier(server_url, &protected_resource_metadata.resource)?;
232 if protected_resource_metadata.authorization_servers.is_empty() {
233 return Err(OAuthClientError::Discovery(format!(
234 "protected resource metadata at `{metadata_url}` omitted authorization_servers"
235 )));
236 }
237
238 let resource = protected_resource_metadata.resource.clone();
239 let issuers = protected_resource_metadata.authorization_servers.clone();
240
241 let mut authorization_servers = Vec::new();
242 let mut last_error = None;
243 for issuer in issuers {
244 match discover_authorization_server_from_issuer(&issuer, client).await {
245 Ok(metadata) => authorization_servers.push(metadata),
246 Err(error) => last_error = Some(error),
247 }
248 }
249 if authorization_servers.is_empty() {
250 return Err(last_error.unwrap_or_else(|| {
251 OAuthClientError::Discovery(
252 "protected resource advertised no usable authorization server".into(),
253 )
254 }));
255 }
256
257 Ok(OAuthAuthorizationDiscovery {
258 resource,
259 protected_resource_metadata,
260 authorization_servers,
261 challenge,
262 })
263}
264
265async fn discover_authorization_server_from_issuer(
266 issuer: &str,
267 client: &reqwest::Client,
268) -> Result<OAuthAuthorizationServerMetadata, OAuthClientError> {
269 let mut last_error = None;
270 for url in authorization_server_metadata_urls(issuer)? {
271 match fetch_json::<OAuthAuthorizationServerMetadata>(client, &url).await {
272 Ok(metadata) => {
273 validate_metadata_issuer(&metadata, issuer)?;
274 return Ok(metadata);
275 }
276 Err(error) => last_error = Some(error),
277 }
278 }
279 Err(last_error.unwrap_or_else(|| {
280 OAuthClientError::Discovery(format!(
281 "could not discover authorization server metadata for `{issuer}`"
282 ))
283 }))
284}
285
286async fn fetch_json<T: serde::de::DeserializeOwned>(
287 client: &reqwest::Client,
288 url: &str,
289) -> Result<T, OAuthClientError> {
290 client
291 .get(url)
292 .send()
293 .await
294 .map_err(|error| OAuthClientError::Discovery(error.to_string()))?
295 .error_for_status()
296 .map_err(|error| OAuthClientError::Discovery(error.to_string()))?
297 .json()
298 .await
299 .map_err(|error| OAuthClientError::Discovery(error.to_string()))
300}
301
302fn protected_resource_metadata_urls(server_url: &str) -> Result<Vec<String>, OAuthClientError> {
303 let parsed = reqwest::Url::parse(server_url)
304 .map_err(|error| OAuthClientError::Discovery(error.to_string()))?;
305 let origin = parsed.origin().ascii_serialization();
306 let path = parsed.path().trim_end_matches('/');
307 let mut urls = Vec::new();
308 if !path.is_empty() {
309 push_unique(
310 &mut urls,
311 format!("{origin}/.well-known/oauth-protected-resource{path}"),
312 );
313 }
314 push_unique(
315 &mut urls,
316 format!("{origin}/.well-known/oauth-protected-resource"),
317 );
318 Ok(urls)
319}
320
321fn authorization_server_metadata_urls(issuer: &str) -> Result<Vec<String>, OAuthClientError> {
322 let parsed = reqwest::Url::parse(issuer)
323 .map_err(|error| OAuthClientError::Discovery(error.to_string()))?;
324 let origin = parsed.origin().ascii_serialization();
325 let path = parsed.path().trim_end_matches('/');
326 let trimmed = issuer.trim_end_matches('/');
327 let mut urls = Vec::new();
328 if !path.is_empty() {
329 push_unique(
330 &mut urls,
331 format!("{origin}/.well-known/oauth-authorization-server{path}"),
332 );
333 push_unique(
334 &mut urls,
335 format!("{origin}/.well-known/openid-configuration{path}"),
336 );
337 push_unique(
338 &mut urls,
339 format!("{trimmed}/.well-known/openid-configuration"),
340 );
341 } else {
342 push_unique(
343 &mut urls,
344 format!("{origin}/.well-known/oauth-authorization-server"),
345 );
346 push_unique(
347 &mut urls,
348 format!("{origin}/.well-known/openid-configuration"),
349 );
350 }
351 Ok(urls)
352}
353
354fn validate_resource_identifier(server_url: &str, resource: &str) -> Result<(), OAuthClientError> {
355 let server = reqwest::Url::parse(server_url)
356 .map_err(|error| OAuthClientError::Discovery(error.to_string()))?;
357 let metadata = reqwest::Url::parse(resource)
358 .map_err(|error| OAuthClientError::Discovery(error.to_string()))?;
359 let matches = metadata.fragment().is_none()
360 && server.scheme() == metadata.scheme()
361 && server.host_str() == metadata.host_str()
362 && server.port_or_known_default() == metadata.port_or_known_default()
363 && server.path() == metadata.path()
364 && server.query() == metadata.query();
365 if matches {
366 Ok(())
367 } else {
368 Err(OAuthClientError::Discovery(format!(
369 "protected resource metadata mismatch: expected `{server_url}`, got `{resource}`"
370 )))
371 }
372}
373
374fn push_unique(values: &mut Vec<String>, value: String) {
375 if !values.contains(&value) {
376 values.push(value);
377 }
378}
379
380fn validate_metadata_issuer(
386 metadata: &OAuthAuthorizationServerMetadata,
387 expected: &str,
388) -> Result<(), OAuthClientError> {
389 if metadata.issuer == expected {
390 Ok(())
391 } else {
392 Err(OAuthClientError::Discovery(format!(
393 "authorization server metadata issuer mismatch: expected `{expected}`, got `{}`",
394 metadata.issuer
395 )))
396 }
397}
398
399#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
405#[serde(rename_all = "snake_case")]
406#[non_exhaustive]
407pub enum OAuthClientRegistrationMethod {
408 PreRegistered,
410 ClientIdMetadataDocument,
412 Dynamic,
414}
415
416#[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
426pub struct OAuthClientRegistration {
427 client_id: String,
428 client_secret: Option<String>,
429 method: OAuthClientRegistrationMethod,
430 bound_issuer: Option<String>,
431}
432
433impl fmt::Debug for OAuthClientRegistration {
434 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
435 f.debug_struct("OAuthClientRegistration")
436 .field("client_id", &self.client_id)
437 .field(
438 "client_secret",
439 &self.client_secret.as_ref().map(|_| "[REDACTED]"),
440 )
441 .field("method", &self.method)
442 .field("bound_issuer", &self.bound_issuer)
443 .finish()
444 }
445}
446
447impl OAuthClientRegistration {
448 pub fn pre_registered(
450 issuer: impl Into<String>,
451 client_id: impl Into<String>,
452 client_secret: Option<String>,
453 ) -> Self {
454 Self {
455 client_id: client_id.into(),
456 client_secret,
457 method: OAuthClientRegistrationMethod::PreRegistered,
458 bound_issuer: Some(issuer.into()),
459 }
460 }
461
462 pub fn dynamically_registered(
471 issuer: impl Into<String>,
472 client_id: impl Into<String>,
473 client_secret: Option<String>,
474 ) -> Self {
475 Self {
476 client_id: client_id.into(),
477 client_secret,
478 method: OAuthClientRegistrationMethod::Dynamic,
479 bound_issuer: Some(issuer.into()),
480 }
481 }
482
483 pub fn client_id_metadata_document(client_id: impl Into<String>) -> Self {
485 Self {
486 client_id: client_id.into(),
487 client_secret: None,
488 method: OAuthClientRegistrationMethod::ClientIdMetadataDocument,
489 bound_issuer: None,
490 }
491 }
492
493 pub fn client_id(&self) -> &str {
495 &self.client_id
496 }
497
498 pub fn client_secret(&self) -> Option<&str> {
500 self.client_secret.as_deref()
501 }
502
503 pub fn method(&self) -> OAuthClientRegistrationMethod {
505 self.method
506 }
507
508 pub fn bound_issuer(&self) -> Option<&str> {
512 self.bound_issuer.as_deref()
513 }
514}
515
516#[async_trait]
526pub trait OAuthClientRegistrationStore: Send + Sync {
527 async fn load(&self, issuer: &str)
529 -> Result<Option<OAuthClientRegistration>, OAuthClientError>;
530
531 async fn save(
533 &self,
534 issuer: &str,
535 registration: &OAuthClientRegistration,
536 ) -> Result<(), OAuthClientError>;
537
538 async fn remove(&self, issuer: &str) -> Result<(), OAuthClientError>;
540}
541
542#[derive(Clone, Default)]
548pub struct MemoryOAuthClientRegistrationStore {
549 registrations: Arc<RwLock<HashMap<String, OAuthClientRegistration>>>,
550}
551
552impl fmt::Debug for MemoryOAuthClientRegistrationStore {
553 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
554 f.debug_struct("MemoryOAuthClientRegistrationStore")
555 .finish_non_exhaustive()
556 }
557}
558
559impl MemoryOAuthClientRegistrationStore {
560 pub fn new() -> Self {
562 Self::default()
563 }
564
565 pub async fn len(&self) -> usize {
567 self.registrations.read().await.len()
568 }
569
570 pub async fn is_empty(&self) -> bool {
572 self.registrations.read().await.is_empty()
573 }
574}
575
576#[async_trait]
577impl OAuthClientRegistrationStore for MemoryOAuthClientRegistrationStore {
578 async fn load(
579 &self,
580 issuer: &str,
581 ) -> Result<Option<OAuthClientRegistration>, OAuthClientError> {
582 Ok(self.registrations.read().await.get(issuer).cloned())
583 }
584
585 async fn save(
586 &self,
587 issuer: &str,
588 registration: &OAuthClientRegistration,
589 ) -> Result<(), OAuthClientError> {
590 self.registrations
591 .write()
592 .await
593 .insert(issuer.to_string(), registration.clone());
594 Ok(())
595 }
596
597 async fn remove(&self, issuer: &str) -> Result<(), OAuthClientError> {
598 self.registrations.write().await.remove(issuer);
599 Ok(())
600 }
601}
602
603#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
605#[serde(rename_all = "lowercase")]
606pub enum OAuthApplicationType {
607 Native,
609 Web,
611}
612
613#[derive(Debug, Clone, serde::Serialize)]
618#[non_exhaustive]
619pub struct OAuthDynamicClientRegistration {
620 pub client_name: String,
622 pub application_type: OAuthApplicationType,
624 pub redirect_uris: Vec<String>,
626 pub grant_types: Vec<String>,
628 pub response_types: Vec<String>,
630 pub token_endpoint_auth_method: String,
632}
633
634impl OAuthDynamicClientRegistration {
635 pub fn native(
637 client_name: impl Into<String>,
638 redirect_uris: impl IntoIterator<Item = impl Into<String>>,
639 ) -> Self {
640 Self {
641 client_name: client_name.into(),
642 application_type: OAuthApplicationType::Native,
643 redirect_uris: redirect_uris.into_iter().map(Into::into).collect(),
644 grant_types: vec!["authorization_code".to_string()],
645 response_types: vec!["code".to_string()],
646 token_endpoint_auth_method: "none".to_string(),
647 }
648 }
649
650 pub fn web(
652 client_name: impl Into<String>,
653 redirect_uris: impl IntoIterator<Item = impl Into<String>>,
654 ) -> Self {
655 Self {
656 application_type: OAuthApplicationType::Web,
657 ..Self::native(client_name, redirect_uris)
658 }
659 }
660
661 pub fn grant_types(mut self, grant_types: impl IntoIterator<Item = impl Into<String>>) -> Self {
663 self.grant_types = grant_types.into_iter().map(Into::into).collect();
664 self
665 }
666
667 pub fn token_endpoint_auth_method(mut self, method: impl Into<String>) -> Self {
669 self.token_endpoint_auth_method = method.into();
670 self
671 }
672}
673
674#[derive(Debug, Clone, Default)]
680pub struct OAuthClientRegistrationOptions {
681 pub pre_registered: Option<OAuthClientRegistration>,
683 pub client_id_metadata_document: Option<String>,
685 pub dynamic_registration: Option<OAuthDynamicClientRegistration>,
687}
688
689impl OAuthClientRegistrationOptions {
690 pub fn new() -> Self {
692 Self::default()
693 }
694
695 pub fn with_pre_registered(mut self, registration: OAuthClientRegistration) -> Self {
697 self.pre_registered = Some(registration);
698 self
699 }
700
701 pub fn with_client_id_metadata_document(mut self, client_id: impl Into<String>) -> Self {
703 self.client_id_metadata_document = Some(client_id.into());
704 self
705 }
706
707 pub fn with_dynamic_registration(
709 mut self,
710 registration: OAuthDynamicClientRegistration,
711 ) -> Self {
712 self.dynamic_registration = Some(registration);
713 self
714 }
715}
716
717#[derive(Debug, serde::Deserialize)]
718struct DynamicClientRegistrationResponse {
719 client_id: String,
720 client_secret: Option<String>,
721}
722
723pub async fn resolve_oauth_client_registration(
730 client: &reqwest::Client,
731 metadata: &OAuthAuthorizationServerMetadata,
732 options: &OAuthClientRegistrationOptions,
733) -> Result<OAuthClientRegistration, OAuthClientError> {
734 resolve_oauth_client_registration_inner(client, metadata, options, None).await
735}
736
737pub async fn resolve_oauth_client_registration_with_store(
751 client: &reqwest::Client,
752 metadata: &OAuthAuthorizationServerMetadata,
753 options: &OAuthClientRegistrationOptions,
754 store: &dyn OAuthClientRegistrationStore,
755) -> Result<OAuthClientRegistration, OAuthClientError> {
756 resolve_oauth_client_registration_inner(client, metadata, options, Some(store)).await
757}
758
759async fn resolve_oauth_client_registration_inner(
760 client: &reqwest::Client,
761 metadata: &OAuthAuthorizationServerMetadata,
762 options: &OAuthClientRegistrationOptions,
763 store: Option<&dyn OAuthClientRegistrationStore>,
764) -> Result<OAuthClientRegistration, OAuthClientError> {
765 if let Some(registration) = &options.pre_registered {
766 if registration.method != OAuthClientRegistrationMethod::PreRegistered {
767 return Err(OAuthClientError::BuildError(
768 "pre_registered must contain pre-registered credentials".to_string(),
769 ));
770 }
771 if registration.bound_issuer() != Some(metadata.issuer.as_str()) {
772 return Err(OAuthClientError::BuildError(format!(
773 "pre-registered credentials are bound to issuer {:?}, not `{}`",
774 registration.bound_issuer(),
775 metadata.issuer
776 )));
777 }
778 return Ok(registration.clone());
779 }
780
781 if metadata.client_id_metadata_document_supported
782 && let Some(client_id) = &options.client_id_metadata_document
783 {
784 validate_client_id_metadata_document_url(client_id)?;
785 return Ok(OAuthClientRegistration {
786 client_id: client_id.clone(),
787 client_secret: None,
788 method: OAuthClientRegistrationMethod::ClientIdMetadataDocument,
789 bound_issuer: None,
790 });
791 }
792
793 if options.dynamic_registration.is_some()
794 && let Some(store) = store
795 && let Some(registration) = store.load(&metadata.issuer).await?
796 {
797 validate_stored_dynamic_registration(®istration, &metadata.issuer)?;
798 return Ok(registration);
799 }
800
801 if let (Some(endpoint), Some(request)) = (
802 metadata.registration_endpoint.as_deref(),
803 options.dynamic_registration.as_ref(),
804 ) {
805 if request.redirect_uris.is_empty() {
806 return Err(OAuthClientError::BuildError(
807 "dynamic registration requires at least one redirect URI".to_string(),
808 ));
809 }
810
811 let response = client
812 .post(endpoint)
813 .json(request)
814 .send()
815 .await
816 .map_err(|error| OAuthClientError::Registration(error.to_string()))?;
817 let status = response.status();
818 if !status.is_success() {
819 let body: String = response
820 .text()
821 .await
822 .unwrap_or_default()
823 .chars()
824 .take(1024)
825 .collect();
826 return Err(OAuthClientError::Registration(format!(
827 "dynamic client registration failed with {status}: {body}"
828 )));
829 }
830 let response: DynamicClientRegistrationResponse = response
831 .json()
832 .await
833 .map_err(|error| OAuthClientError::Registration(error.to_string()))?;
834 let registration = OAuthClientRegistration::dynamically_registered(
835 metadata.issuer.clone(),
836 response.client_id,
837 response.client_secret,
838 );
839 if let Some(store) = store {
840 store.save(&metadata.issuer, ®istration).await?;
841 }
842 return Ok(registration);
843 }
844
845 Err(OAuthClientError::BuildError(
846 "authorization server supports none of the configured client registration mechanisms; \
847 prompt the user for pre-registered client information"
848 .to_string(),
849 ))
850}
851
852fn validate_stored_dynamic_registration(
853 registration: &OAuthClientRegistration,
854 issuer: &str,
855) -> Result<(), OAuthClientError> {
856 if registration.method() != OAuthClientRegistrationMethod::Dynamic {
857 return Err(OAuthClientError::CredentialStore(format!(
858 "stored registration for issuer `{issuer}` uses {:?}, expected dynamic registration",
859 registration.method()
860 )));
861 }
862 if registration.bound_issuer() != Some(issuer) {
863 return Err(OAuthClientError::CredentialStore(format!(
864 "stored registration is bound to issuer {:?}, not `{issuer}`",
865 registration.bound_issuer()
866 )));
867 }
868 Ok(())
869}
870
871fn validate_client_id_metadata_document_url(client_id: &str) -> Result<(), OAuthClientError> {
872 let url = reqwest::Url::parse(client_id).map_err(|error| {
873 OAuthClientError::BuildError(format!(
874 "invalid Client ID Metadata Document URL `{client_id}`: {error}"
875 ))
876 })?;
877 if url.scheme() != "https" || url.path() == "/" {
878 return Err(OAuthClientError::BuildError(format!(
879 "Client ID Metadata Document URL `{client_id}` must use HTTPS and contain a path"
880 )));
881 }
882 Ok(())
883}
884
885#[derive(Debug, Clone, serde::Deserialize)]
891struct TokenResponse {
892 access_token: String,
893 #[allow(dead_code)]
894 token_type: String,
895 expires_in: Option<u64>,
896 refresh_token: Option<String>,
897 #[allow(dead_code)]
898 scope: Option<String>,
899}
900
901#[derive(Debug, Clone)]
903struct CachedAuthCodeToken {
904 access_token: String,
905 refresh_token: Option<String>,
906 expires_at: Instant,
907}
908
909#[derive(Clone)]
918pub struct OAuthAuthorizationCode {
919 inner: Arc<OAuthAuthCodeInner>,
920}
921
922struct OAuthAuthCodeInner {
923 authorization_url: String,
925 token_endpoint: String,
927 client_id: String,
929 client_secret: Option<String>,
931 token_endpoint_auth_method: OAuthTokenEndpointAuthMethod,
933 resource: String,
935 code_verifier: String,
937 state: String,
939 redirect_uri: String,
941 scopes: Option<String>,
943 refresh_buffer: Duration,
945 client: reqwest::Client,
947 cache: RwLock<Option<CachedAuthCodeToken>>,
949 callback_rx: Mutex<Option<oneshot::Receiver<Result<CallbackResult, String>>>>,
951 _callback_task: tokio::task::JoinHandle<()>,
953 expected_issuer: Option<String>,
957 iss_required: bool,
961}
962
963#[derive(Debug)]
964struct CallbackResult {
965 code: String,
966 #[allow(dead_code)]
967 state: String,
968 iss: Option<String>,
971}
972
973impl fmt::Debug for OAuthAuthorizationCode {
974 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
975 f.debug_struct("OAuthAuthorizationCode")
976 .field("client_id", &self.inner.client_id)
977 .field("token_endpoint", &self.inner.token_endpoint)
978 .field("redirect_uri", &self.inner.redirect_uri)
979 .finish()
980 }
981}
982
983impl OAuthAuthorizationCode {
984 pub async fn start(server_url: &str, scopes: &[&str]) -> Result<Self, OAuthClientError> {
993 Self::start_with_config(server_url, scopes, OAuthAuthCodeConfig::default()).await
994 }
995
996 pub async fn start_with_config(
998 server_url: &str,
999 scopes: &[&str],
1000 mut config: OAuthAuthCodeConfig,
1001 ) -> Result<Self, OAuthClientError> {
1002 let client = config.http_client.take().unwrap_or_default();
1003
1004 let discovery =
1008 discover_oauth_authorization(server_url, config.challenge.take(), &client).await?;
1009 let metadata = match config.preferred_authorization_server.take() {
1010 Some(issuer) => discovery
1011 .authorization_server(&issuer)
1012 .cloned()
1013 .ok_or_else(|| {
1014 OAuthClientError::Discovery(format!(
1015 "preferred authorization server `{issuer}` was not advertised by the resource"
1016 ))
1017 })?,
1018 None => discovery
1019 .authorization_servers
1020 .first()
1021 .cloned()
1022 .ok_or_else(|| OAuthClientError::Discovery("no authorization server discovered".into()))?,
1023 };
1024 require_s256(&metadata)?;
1025 let token_endpoint_auth_method = OAuthTokenEndpointAuthMethod::select(
1026 &metadata.token_endpoint_auth_methods_supported,
1027 config.client_secret.is_some(),
1028 )?;
1029
1030 let code_verifier = generate_code_verifier();
1032 let code_challenge = compute_code_challenge(&code_verifier);
1033 let state = generate_state();
1034
1035 let callback_port = config.callback_port.unwrap_or(0);
1037 let (callback_tx, callback_rx) = oneshot::channel();
1038 let listener = tokio::net::TcpListener::bind(format!("127.0.0.1:{}", callback_port))
1039 .await
1040 .map_err(|e| OAuthClientError::BuildError(format!("Callback server bind: {}", e)))?;
1041 let actual_port = listener
1042 .local_addr()
1043 .map_err(|e| OAuthClientError::BuildError(format!("Get local addr: {}", e)))?
1044 .port();
1045 let redirect_uri = format!("http://127.0.0.1:{}/callback", actual_port);
1046
1047 let expected_state = state.clone();
1048 let callback_task = tokio::spawn(async move {
1049 run_callback_server(listener, callback_tx, expected_state).await;
1050 });
1051
1052 let scope_str = if !scopes.is_empty() {
1054 Some(scopes.join(" "))
1055 } else if let Some(challenge) = &discovery.challenge
1056 && !challenge.scopes.is_empty()
1057 {
1058 Some(challenge.scopes.join(" "))
1059 } else {
1060 (!discovery
1061 .protected_resource_metadata
1062 .scopes_supported
1063 .is_empty())
1064 .then(|| {
1065 discovery
1066 .protected_resource_metadata
1067 .scopes_supported
1068 .join(" ")
1069 })
1070 };
1071
1072 let client_id = config.client_id.unwrap_or_else(|| "tower-mcp".to_string());
1073 let mut auth_url = reqwest::Url::parse(&metadata.authorization_endpoint)
1074 .map_err(|error| OAuthClientError::Discovery(error.to_string()))?;
1075 {
1076 let mut query = auth_url.query_pairs_mut();
1077 query
1078 .append_pair("response_type", "code")
1079 .append_pair("client_id", &client_id)
1080 .append_pair("redirect_uri", &redirect_uri)
1081 .append_pair("state", &state)
1082 .append_pair("code_challenge", &code_challenge)
1083 .append_pair("code_challenge_method", "S256");
1084 if let Some(scopes) = &scope_str {
1085 query.append_pair("scope", scopes);
1086 }
1087 query.append_pair("resource", &discovery.resource);
1088 }
1089
1090 Ok(Self {
1091 inner: Arc::new(OAuthAuthCodeInner {
1092 authorization_url: auth_url.into(),
1093 token_endpoint: metadata.token_endpoint,
1094 client_id,
1095 client_secret: config.client_secret,
1096 token_endpoint_auth_method,
1097 resource: discovery.resource,
1098 code_verifier,
1099 state,
1100 redirect_uri,
1101 scopes: scope_str,
1102 refresh_buffer: config.refresh_buffer,
1103 client,
1104 cache: RwLock::new(None),
1105 callback_rx: Mutex::new(Some(callback_rx)),
1106 _callback_task: callback_task,
1107 expected_issuer: Some(metadata.issuer),
1108 iss_required: metadata.authorization_response_iss_parameter_supported,
1109 }),
1110 })
1111 }
1112
1113 pub fn authorization_url(&self) -> &str {
1115 &self.inner.authorization_url
1116 }
1117
1118 pub async fn wait_for_callback(&self) -> Result<(), OAuthClientError> {
1123 self.wait_for_callback_with_timeout(Duration::from_secs(300))
1124 .await
1125 }
1126
1127 pub async fn wait_for_callback_with_timeout(
1129 &self,
1130 timeout: Duration,
1131 ) -> Result<(), OAuthClientError> {
1132 let rx = self.inner.callback_rx.lock().await.take().ok_or_else(|| {
1133 OAuthClientError::InvalidResponse("Callback already consumed".to_string())
1134 })?;
1135
1136 let result = tokio::time::timeout(timeout, rx)
1137 .await
1138 .map_err(|_| {
1139 OAuthClientError::TokenRequest("Timed out waiting for OAuth callback".to_string())
1140 })?
1141 .map_err(|_| OAuthClientError::TokenRequest("Callback cancelled".to_string()))?
1142 .map_err(|e| OAuthClientError::TokenRequest(format!("Callback error: {}", e)))?;
1143
1144 if result.state != self.inner.state {
1146 return Err(OAuthClientError::InvalidResponse(
1147 "CSRF state mismatch".to_string(),
1148 ));
1149 }
1150
1151 validate_iss(
1155 result.iss.as_deref(),
1156 self.inner.expected_issuer.as_deref(),
1157 self.inner.iss_required,
1158 )
1159 .map_err(OAuthClientError::InvalidResponse)?;
1160
1161 let token = self.exchange_code(&result.code).await?;
1163 *self.inner.cache.write().await = Some(token);
1164
1165 Ok(())
1166 }
1167
1168 async fn exchange_code(&self, code: &str) -> Result<CachedAuthCodeToken, OAuthClientError> {
1170 let response = send_token_request(
1171 &self.inner.client,
1172 &self.inner.token_endpoint,
1173 vec![
1174 ("grant_type", "authorization_code".to_string()),
1175 ("code", code.to_string()),
1176 ("redirect_uri", self.inner.redirect_uri.clone()),
1177 ("code_verifier", self.inner.code_verifier.clone()),
1178 ("resource", self.inner.resource.clone()),
1179 ],
1180 self.inner.token_endpoint_auth_method,
1181 &self.inner.client_id,
1182 self.inner.client_secret.as_deref(),
1183 )
1184 .await?;
1185
1186 if !response.status().is_success() {
1187 let status = response.status();
1188 let body = response.text().await.unwrap_or_default();
1189 return Err(OAuthClientError::TokenRequest(format!(
1190 "HTTP {}: {}",
1191 status, body
1192 )));
1193 }
1194
1195 let token_response: TokenResponse = response
1196 .json()
1197 .await
1198 .map_err(|e| OAuthClientError::InvalidResponse(e.to_string()))?;
1199
1200 Ok(to_cached_token(token_response))
1201 }
1202
1203 async fn refresh_token(
1205 &self,
1206 refresh_token: &str,
1207 ) -> Result<CachedAuthCodeToken, OAuthClientError> {
1208 let mut params = vec![
1209 ("grant_type", "refresh_token".to_string()),
1210 ("refresh_token", refresh_token.to_string()),
1211 ("resource", self.inner.resource.clone()),
1212 ];
1213 if let Some(ref scopes) = self.inner.scopes {
1214 params.push(("scope", scopes.clone()));
1215 }
1216
1217 let response = send_token_request(
1218 &self.inner.client,
1219 &self.inner.token_endpoint,
1220 params,
1221 self.inner.token_endpoint_auth_method,
1222 &self.inner.client_id,
1223 self.inner.client_secret.as_deref(),
1224 )
1225 .await
1226 .map_err(|error| OAuthClientError::TokenRequest(format!("Refresh failed: {error}")))?;
1227
1228 if !response.status().is_success() {
1229 let status = response.status();
1230 let body = response.text().await.unwrap_or_default();
1231 return Err(OAuthClientError::TokenRequest(format!(
1232 "Refresh HTTP {}: {}",
1233 status, body
1234 )));
1235 }
1236
1237 let mut token_response: TokenResponse = response
1238 .json()
1239 .await
1240 .map_err(|e| OAuthClientError::InvalidResponse(e.to_string()))?;
1241
1242 if token_response.refresh_token.is_none() {
1244 token_response.refresh_token = Some(refresh_token.to_string());
1245 }
1246
1247 Ok(to_cached_token(token_response))
1248 }
1249}
1250
1251fn require_s256(metadata: &OAuthAuthorizationServerMetadata) -> Result<(), OAuthClientError> {
1252 if metadata
1253 .code_challenge_methods_supported
1254 .iter()
1255 .any(|method| method == "S256")
1256 {
1257 Ok(())
1258 } else {
1259 Err(OAuthClientError::Discovery(format!(
1260 "authorization server `{}` does not advertise PKCE S256 support",
1261 metadata.issuer
1262 )))
1263 }
1264}
1265
1266async fn send_token_request(
1267 client: &reqwest::Client,
1268 token_endpoint: &str,
1269 mut params: Vec<(&'static str, String)>,
1270 method: OAuthTokenEndpointAuthMethod,
1271 client_id: &str,
1272 client_secret: Option<&str>,
1273) -> Result<reqwest::Response, OAuthClientError> {
1274 let mut request = client.post(token_endpoint);
1275 match method {
1276 OAuthTokenEndpointAuthMethod::None => {
1277 params.push(("client_id", client_id.to_string()));
1278 }
1279 OAuthTokenEndpointAuthMethod::ClientSecretBasic => {
1280 let secret = client_secret.ok_or_else(|| {
1281 OAuthClientError::BuildError(
1282 "client_secret_basic requires a client secret".to_string(),
1283 )
1284 })?;
1285 request = request.basic_auth(client_id, Some(secret));
1286 }
1287 OAuthTokenEndpointAuthMethod::ClientSecretPost => {
1288 let secret = client_secret.ok_or_else(|| {
1289 OAuthClientError::BuildError(
1290 "client_secret_post requires a client secret".to_string(),
1291 )
1292 })?;
1293 params.push(("client_id", client_id.to_string()));
1294 params.push(("client_secret", secret.to_string()));
1295 }
1296 OAuthTokenEndpointAuthMethod::PrivateKeyJwt => {
1297 return Err(OAuthClientError::BuildError(
1298 "private_key_jwt requires OAuthAuthorizationFlow with a client assertion signer"
1299 .to_string(),
1300 ));
1301 }
1302 }
1303 request
1304 .form(¶ms)
1305 .send()
1306 .await
1307 .map_err(|error| OAuthClientError::TokenRequest(error.to_string()))
1308}
1309
1310fn to_cached_token(response: TokenResponse) -> CachedAuthCodeToken {
1311 let expires_in = Duration::from_secs(response.expires_in.unwrap_or(3600));
1312 CachedAuthCodeToken {
1313 access_token: response.access_token,
1314 refresh_token: response.refresh_token,
1315 expires_at: Instant::now() + expires_in,
1316 }
1317}
1318
1319fn is_token_valid(token: &CachedAuthCodeToken, buffer: Duration) -> bool {
1320 token
1321 .expires_at
1322 .checked_sub(buffer)
1323 .is_some_and(|effective| Instant::now() < effective)
1324}
1325
1326#[async_trait]
1327impl TokenProvider for OAuthAuthorizationCode {
1328 async fn get_token(&self) -> Result<String, OAuthClientError> {
1329 {
1331 let cache = self.inner.cache.read().await;
1332 if let Some(ref token) = *cache
1333 && is_token_valid(token, self.inner.refresh_buffer)
1334 {
1335 return Ok(token.access_token.clone());
1336 }
1337 }
1338
1339 let mut cache = self.inner.cache.write().await;
1341
1342 if let Some(ref token) = *cache
1344 && is_token_valid(token, self.inner.refresh_buffer)
1345 {
1346 return Ok(token.access_token.clone());
1347 }
1348
1349 if let Some(ref token) = *cache
1351 && let Some(ref refresh) = token.refresh_token
1352 {
1353 tracing::debug!("Refreshing OAuth access token");
1354 match self.refresh_token(refresh).await {
1355 Ok(new_token) => {
1356 let access = new_token.access_token.clone();
1357 *cache = Some(new_token);
1358 return Ok(access);
1359 }
1360 Err(e) => {
1361 tracing::warn!(error = %e, "Token refresh failed");
1362 }
1364 }
1365 }
1366
1367 Err(OAuthClientError::TokenRequest(
1368 "No valid token available. Call wait_for_callback() to authenticate.".to_string(),
1369 ))
1370 }
1371}
1372
1373pub struct OAuthAuthCodeConfig {
1379 pub client_id: Option<String>,
1381 pub client_secret: Option<String>,
1383 pub callback_port: Option<u16>,
1385 pub refresh_buffer: Duration,
1387 pub http_client: Option<reqwest::Client>,
1389 pub challenge: Option<OAuthBearerChallenge>,
1394 pub preferred_authorization_server: Option<String>,
1398}
1399
1400impl Default for OAuthAuthCodeConfig {
1401 fn default() -> Self {
1402 Self {
1403 client_id: None,
1404 client_secret: None,
1405 callback_port: None,
1406 refresh_buffer: Duration::from_secs(30),
1407 http_client: None,
1408 challenge: None,
1409 preferred_authorization_server: None,
1410 }
1411 }
1412}
1413
1414async fn run_callback_server(
1420 listener: tokio::net::TcpListener,
1421 tx: oneshot::Sender<Result<CallbackResult, String>>,
1422 expected_state: String,
1423) {
1424 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1425
1426 let mut tx = Some(tx);
1427
1428 let Ok((mut stream, _)) = listener.accept().await else {
1430 if let Some(tx) = tx.take() {
1431 let _ = tx.send(Err("Callback server accept failed".to_string()));
1432 }
1433 return;
1434 };
1435
1436 let mut buf = vec![0u8; 4096];
1437 let n = match stream.read(&mut buf).await {
1438 Ok(n) => n,
1439 Err(e) => {
1440 if let Some(tx) = tx.take() {
1441 let _ = tx.send(Err(format!("Read error: {}", e)));
1442 }
1443 return;
1444 }
1445 };
1446
1447 let request = String::from_utf8_lossy(&buf[..n]);
1448
1449 let result = if let Some(path) = request.lines().next().and_then(|line| {
1451 let parts: Vec<&str> = line.split_whitespace().collect();
1452 if parts.len() >= 2 {
1453 Some(parts[1])
1454 } else {
1455 None
1456 }
1457 }) {
1458 parse_callback_query(path, &expected_state)
1459 } else {
1460 Err("Invalid HTTP request".to_string())
1461 };
1462
1463 let (status, body) = match &result {
1465 Ok(_) => (
1466 "200 OK",
1467 "Authorization successful. You can close this tab.",
1468 ),
1469 Err(e) => ("400 Bad Request", e.as_str()),
1470 };
1471
1472 let response = format!(
1473 "HTTP/1.1 {}\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
1474 status,
1475 body.len(),
1476 body
1477 );
1478 let _ = stream.write_all(response.as_bytes()).await;
1479 let _ = stream.flush().await;
1480
1481 if let Some(tx) = tx.take() {
1482 let _ = tx.send(result);
1483 }
1484}
1485
1486fn parse_callback_query(path: &str, expected_state: &str) -> Result<CallbackResult, String> {
1488 let query = path
1489 .split('?')
1490 .nth(1)
1491 .ok_or_else(|| "No query parameters in callback".to_string())?;
1492
1493 let mut code = None;
1494 let mut state = None;
1495 let mut error = None;
1496 let mut iss = None;
1497
1498 for param in query.split('&') {
1499 let mut parts = param.splitn(2, '=');
1500 let key = parts.next().unwrap_or("");
1501 let value = parts.next().unwrap_or("");
1502 let decoded = urlencoding::decode(value).unwrap_or_default().to_string();
1503
1504 match key {
1505 "code" => code = Some(decoded),
1506 "state" => state = Some(decoded),
1507 "error" => error = Some(decoded),
1508 "error_description" if error.is_none() => error = Some(decoded),
1509 "iss" => iss = Some(decoded),
1511 _ => {}
1512 }
1513 }
1514
1515 if let Some(err) = error {
1516 return Err(format!("OAuth error: {}", err));
1517 }
1518
1519 let code = code.ok_or_else(|| "Missing 'code' parameter".to_string())?;
1520 let state = state.ok_or_else(|| "Missing 'state' parameter".to_string())?;
1521
1522 if state != expected_state {
1523 return Err("CSRF state mismatch".to_string());
1524 }
1525
1526 Ok(CallbackResult { code, state, iss })
1527}
1528
1529fn validate_iss(
1543 iss: Option<&str>,
1544 expected: Option<&str>,
1545 iss_required: bool,
1546) -> Result<(), String> {
1547 match (iss, expected, iss_required) {
1548 (Some(received), Some(want), _) => {
1549 if received == want {
1550 Ok(())
1551 } else {
1552 Err(format!(
1553 "Issuer mismatch (SEP-2468): expected `{}`, got `{}`",
1554 want, received
1555 ))
1556 }
1557 }
1558 (Some(_received), None, _) => {
1559 Ok(())
1564 }
1565 (None, _, true) => Err(
1566 "Authorization response missing `iss` (SEP-2468): the AS advertises \
1567 authorization_response_iss_parameter_supported but did not include iss"
1568 .to_string(),
1569 ),
1570 (None, _, false) => Ok(()),
1571 }
1572}
1573
1574#[cfg(test)]
1575mod tests {
1576 use super::*;
1577
1578 async fn spawn_discovery_server() -> (String, tokio::task::JoinHandle<Vec<String>>) {
1579 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1580
1581 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1582 let base = format!("http://{}", listener.local_addr().unwrap());
1583 let server_base = base.clone();
1584 let task = tokio::spawn(async move {
1585 let mut requests = Vec::new();
1586 for _ in 0..4 {
1587 let (mut stream, _) = listener.accept().await.unwrap();
1588 let mut bytes = Vec::new();
1589 loop {
1590 let mut chunk = [0_u8; 1024];
1591 let read = stream.read(&mut chunk).await.unwrap();
1592 assert!(read > 0);
1593 bytes.extend_from_slice(&chunk[..read]);
1594 if bytes.windows(4).any(|window| window == b"\r\n\r\n") {
1595 break;
1596 }
1597 }
1598 let request = String::from_utf8_lossy(&bytes);
1599 let path = request
1600 .lines()
1601 .next()
1602 .and_then(|line| line.split_ascii_whitespace().nth(1))
1603 .unwrap()
1604 .to_string();
1605 requests.push(path.clone());
1606
1607 let (status, extra_headers, body) = match path.as_str() {
1608 "/mcp" => (
1609 "401 Unauthorized",
1610 format!(
1611 "WWW-Authenticate: Bearer resource_metadata=\"{server_base}/metadata\", scope=\"challenge.scope\"\r\n"
1612 ),
1613 String::new(),
1614 ),
1615 "/metadata" => (
1616 "200 OK",
1617 String::new(),
1618 serde_json::json!({
1619 "resource": format!("{server_base}/mcp"),
1620 "authorization_servers": [
1621 format!("{server_base}/auth-a"),
1622 format!("{server_base}/auth-b")
1623 ],
1624 "scopes_supported": ["metadata.scope"]
1625 })
1626 .to_string(),
1627 ),
1628 "/.well-known/oauth-authorization-server/auth-a" => (
1629 "200 OK",
1630 String::new(),
1631 authorization_metadata_json(&server_base, "auth-a"),
1632 ),
1633 "/.well-known/oauth-authorization-server/auth-b" => (
1634 "200 OK",
1635 String::new(),
1636 authorization_metadata_json(&server_base, "auth-b"),
1637 ),
1638 other => panic!("unexpected request path: {other}"),
1639 };
1640 let response = format!(
1641 "HTTP/1.1 {status}\r\ncontent-type: application/json\r\n{extra_headers}content-length: {}\r\nconnection: close\r\n\r\n{body}",
1642 body.len()
1643 );
1644 stream.write_all(response.as_bytes()).await.unwrap();
1645 }
1646 requests
1647 });
1648 (base, task)
1649 }
1650
1651 fn authorization_metadata_json(base: &str, name: &str) -> String {
1652 let issuer = format!("{base}/{name}");
1653 serde_json::json!({
1654 "issuer": issuer,
1655 "authorization_endpoint": format!("{base}/{name}/authorize"),
1656 "token_endpoint": format!("{base}/{name}/token"),
1657 "code_challenge_methods_supported": ["S256"],
1658 "token_endpoint_auth_methods_supported": ["none"]
1659 })
1660 .to_string()
1661 }
1662
1663 #[test]
1664 fn test_pkce_code_verifier_length() {
1665 let verifier = generate_code_verifier();
1666 assert!(
1667 verifier.len() >= 43,
1668 "Verifier too short: {}",
1669 verifier.len()
1670 );
1671 assert!(
1672 verifier.len() <= 128,
1673 "Verifier too long: {}",
1674 verifier.len()
1675 );
1676 }
1677
1678 #[test]
1679 fn test_pkce_code_challenge_deterministic() {
1680 let challenge1 = compute_code_challenge("test-verifier");
1681 let challenge2 = compute_code_challenge("test-verifier");
1682 assert_eq!(challenge1, challenge2);
1683 }
1684
1685 #[test]
1686 fn test_pkce_code_challenge_differs_for_different_input() {
1687 let c1 = compute_code_challenge("verifier-a");
1688 let c2 = compute_code_challenge("verifier-b");
1689 assert_ne!(c1, c2);
1690 }
1691
1692 #[test]
1693 fn test_state_generation_unique() {
1694 let s1 = generate_state();
1695 let s2 = generate_state();
1696 assert_ne!(s1, s2);
1697 }
1698
1699 #[test]
1700 fn final_well_known_urls_are_path_aware() {
1701 assert_eq!(
1702 protected_resource_metadata_urls("https://mcp.example.com/team/mcp").unwrap(),
1703 vec![
1704 "https://mcp.example.com/.well-known/oauth-protected-resource/team/mcp",
1705 "https://mcp.example.com/.well-known/oauth-protected-resource",
1706 ]
1707 );
1708 let urls = authorization_server_metadata_urls("https://auth.example.com/tenant").unwrap();
1709 assert_eq!(
1710 urls[0],
1711 "https://auth.example.com/.well-known/oauth-authorization-server/tenant"
1712 );
1713 assert_eq!(
1714 urls[1],
1715 "https://auth.example.com/.well-known/openid-configuration/tenant"
1716 );
1717 assert!(
1718 urls.contains(
1719 &"https://auth.example.com/tenant/.well-known/openid-configuration".into()
1720 )
1721 );
1722 }
1723
1724 #[tokio::test]
1725 async fn public_discovery_honors_challenge_and_exposes_all_servers() {
1726 let (base, server) = spawn_discovery_server().await;
1727 let resource = format!("{base}/mcp");
1728 let discovery = discover_oauth_authorization(&resource, None, &reqwest::Client::new())
1729 .await
1730 .unwrap();
1731
1732 assert_eq!(discovery.resource, resource);
1733 assert_eq!(discovery.authorization_servers.len(), 2);
1734 assert!(
1735 discovery
1736 .authorization_server(&format!("{base}/auth-b"))
1737 .is_some()
1738 );
1739 assert_eq!(discovery.challenge.unwrap().scopes, vec!["challenge.scope"]);
1740 assert_eq!(
1741 server.await.unwrap(),
1742 vec![
1743 "/mcp",
1744 "/metadata",
1745 "/.well-known/oauth-authorization-server/auth-a",
1746 "/.well-known/oauth-authorization-server/auth-b",
1747 ]
1748 );
1749 }
1750
1751 #[tokio::test]
1752 async fn authorization_flow_selects_issuer_and_binds_resource() {
1753 let (base, server) = spawn_discovery_server().await;
1754 let resource = format!("{base}/mcp");
1755 let provider = OAuthAuthorizationCode::start_with_config(
1756 &resource,
1757 &[],
1758 OAuthAuthCodeConfig {
1759 client_id: Some("public-client".into()),
1760 preferred_authorization_server: Some(format!("{base}/auth-b")),
1761 ..OAuthAuthCodeConfig::default()
1762 },
1763 )
1764 .await
1765 .unwrap();
1766
1767 let authorization_url = reqwest::Url::parse(provider.authorization_url()).unwrap();
1768 assert_eq!(authorization_url.path(), "/auth-b/authorize");
1769 let parameters: HashMap<_, _> = authorization_url.query_pairs().into_owned().collect();
1770 assert_eq!(parameters.get("resource"), Some(&resource));
1771 assert_eq!(
1772 parameters.get("scope").map(String::as_str),
1773 Some("challenge.scope")
1774 );
1775 assert_eq!(
1776 parameters.get("code_challenge_method").map(String::as_str),
1777 Some("S256")
1778 );
1779 server.await.unwrap();
1780 }
1781
1782 #[tokio::test]
1783 async fn token_request_uses_metadata_selected_basic_auth_and_resource() {
1784 use tokio::io::{AsyncReadExt, AsyncWriteExt};
1785
1786 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1787 let endpoint = format!("http://{}/token", listener.local_addr().unwrap());
1788 let server = tokio::spawn(async move {
1789 let (mut stream, _) = listener.accept().await.unwrap();
1790 let mut bytes = Vec::new();
1791 let header_end = loop {
1792 let mut chunk = [0_u8; 1024];
1793 let read = stream.read(&mut chunk).await.unwrap();
1794 assert!(read > 0);
1795 bytes.extend_from_slice(&chunk[..read]);
1796 if let Some(index) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
1797 break index + 4;
1798 }
1799 };
1800 let headers = String::from_utf8_lossy(&bytes[..header_end]).to_string();
1801 let content_length = headers
1802 .lines()
1803 .find_map(|line| {
1804 let (name, value) = line.split_once(':')?;
1805 name.eq_ignore_ascii_case("content-length")
1806 .then(|| value.trim().parse::<usize>().unwrap())
1807 })
1808 .unwrap();
1809 while bytes.len() < header_end + content_length {
1810 let mut chunk = [0_u8; 1024];
1811 let read = stream.read(&mut chunk).await.unwrap();
1812 bytes.extend_from_slice(&chunk[..read]);
1813 }
1814 let body = String::from_utf8_lossy(&bytes[header_end..header_end + content_length])
1815 .to_string();
1816 let response_body = r#"{"access_token":"token","token_type":"Bearer"}"#;
1817 let response = format!(
1818 "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{response_body}",
1819 response_body.len()
1820 );
1821 stream.write_all(response.as_bytes()).await.unwrap();
1822 (headers, body)
1823 });
1824
1825 let response = send_token_request(
1826 &reqwest::Client::new(),
1827 &endpoint,
1828 vec![
1829 ("grant_type", "authorization_code".into()),
1830 ("resource", "https://mcp.example.com/team/mcp".into()),
1831 ],
1832 OAuthTokenEndpointAuthMethod::ClientSecretBasic,
1833 "client id",
1834 Some("client secret"),
1835 )
1836 .await
1837 .unwrap();
1838 assert!(response.status().is_success());
1839 let (headers, body) = server.await.unwrap();
1840 assert!(
1841 headers
1842 .to_ascii_lowercase()
1843 .contains("authorization: basic ")
1844 );
1845 assert!(body.contains("resource=https%3A%2F%2Fmcp.example.com%2Fteam%2Fmcp"));
1846 assert!(!body.contains("client_secret"));
1847 }
1848
1849 #[test]
1850 fn test_parse_callback_success() {
1851 let result = parse_callback_query("/callback?code=abc123&state=mystate", "mystate");
1852 let cb = result.unwrap();
1853 assert_eq!(cb.code, "abc123");
1854 assert_eq!(cb.state, "mystate");
1855 }
1856
1857 #[test]
1858 fn test_parse_callback_state_mismatch() {
1859 let result = parse_callback_query("/callback?code=abc123&state=wrong", "expected");
1860 assert!(result.is_err());
1861 assert!(result.unwrap_err().contains("CSRF"));
1862 }
1863
1864 #[test]
1865 fn test_parse_callback_error() {
1866 let result = parse_callback_query(
1867 "/callback?error=access_denied&error_description=User+denied+access",
1868 "state",
1869 );
1870 assert!(result.is_err());
1871 assert!(result.unwrap_err().contains("access_denied"));
1872 }
1873
1874 #[test]
1875 fn test_parse_callback_missing_code() {
1876 let result = parse_callback_query("/callback?state=mystate", "mystate");
1877 assert!(result.is_err());
1878 assert!(result.unwrap_err().contains("code"));
1879 }
1880
1881 #[test]
1886 fn parse_callback_extracts_iss_when_present() {
1887 let result = parse_callback_query(
1888 "/callback?code=abc&state=s&iss=https%3A%2F%2Fauth.example.com",
1889 "s",
1890 )
1891 .unwrap();
1892 assert_eq!(result.iss.as_deref(), Some("https://auth.example.com"));
1893 }
1894
1895 #[test]
1896 fn parse_callback_iss_is_none_when_absent() {
1897 let result = parse_callback_query("/callback?code=abc&state=s", "s").unwrap();
1898 assert!(result.iss.is_none());
1899 }
1900
1901 #[test]
1902 fn validate_iss_accepts_exact_match() {
1903 let expected = Some("https://auth.example.com");
1904 assert!(validate_iss(Some("https://auth.example.com"), expected, true).is_ok());
1905 assert!(validate_iss(Some("https://auth.example.com"), expected, false).is_ok());
1906 }
1907
1908 #[test]
1909 fn validate_iss_rejects_mismatch_regardless_of_required() {
1910 let expected = Some("https://auth.example.com");
1911 let bad = Some("https://evil.example.com");
1912 for required in [true, false] {
1913 let err = validate_iss(bad, expected, required).unwrap_err();
1914 assert!(
1915 err.contains("Issuer mismatch"),
1916 "should reject mismatch (required={required}), got: {err}"
1917 );
1918 }
1919 }
1920
1921 #[test]
1922 fn validate_iss_rejects_missing_when_as_advertises_support() {
1923 let err = validate_iss(None, Some("https://auth.example.com"), true).unwrap_err();
1925 assert!(err.contains("missing `iss`"), "got: {err}");
1926 }
1927
1928 #[test]
1929 fn validate_iss_accepts_missing_when_as_does_not_advertise_support() {
1930 assert!(validate_iss(None, Some("https://auth.example.com"), false).is_ok());
1933 }
1934
1935 #[test]
1936 fn validate_iss_accepts_when_no_expected_recorded() {
1937 assert!(validate_iss(Some("https://auth.example.com"), None, false).is_ok());
1940 assert!(validate_iss(None, None, false).is_ok());
1941 }
1942
1943 #[test]
1944 fn validate_metadata_issuer_accepts_exact_match() {
1945 let metadata = authorization_server_metadata("https://auth.example.com");
1946 assert!(validate_metadata_issuer(&metadata, "https://auth.example.com").is_ok());
1947 }
1948
1949 #[test]
1950 fn validate_metadata_issuer_rejects_trailing_slash_mismatch() {
1951 let metadata = authorization_server_metadata("https://auth.example.com/");
1952 let err = validate_metadata_issuer(&metadata, "https://auth.example.com").unwrap_err();
1953 assert!(err.to_string().contains("issuer mismatch"), "got: {err}");
1954 }
1955
1956 #[test]
1957 fn authorization_code_flow_requires_advertised_s256() {
1958 let mut metadata = authorization_server_metadata("https://auth.example.com");
1959 metadata.code_challenge_methods_supported.clear();
1960
1961 let error = require_s256(&metadata).unwrap_err();
1962 assert!(error.to_string().contains("PKCE S256"));
1963 }
1964
1965 #[tokio::test]
1966 async fn registration_prefers_pre_registered_credentials() {
1967 let mut metadata = authorization_server_metadata("https://auth.example.com");
1968 metadata.client_id_metadata_document_supported = true;
1969 metadata.registration_endpoint = Some("http://127.0.0.1:9/register".to_string());
1970 let pre_registered = OAuthClientRegistration::pre_registered(
1971 "https://auth.example.com",
1972 "configured-client",
1973 Some("secret".to_string()),
1974 );
1975 let options = OAuthClientRegistrationOptions::new()
1976 .with_pre_registered(pre_registered)
1977 .with_client_id_metadata_document("https://client.example.com/client.json")
1978 .with_dynamic_registration(OAuthDynamicClientRegistration::native(
1979 "test-client",
1980 ["http://127.0.0.1/callback"],
1981 ));
1982
1983 let registration =
1984 resolve_oauth_client_registration(&reqwest::Client::new(), &metadata, &options)
1985 .await
1986 .unwrap();
1987
1988 assert_eq!(
1989 registration.method(),
1990 OAuthClientRegistrationMethod::PreRegistered
1991 );
1992 assert_eq!(registration.client_id(), "configured-client");
1993 assert_eq!(registration.client_secret(), Some("secret"));
1994 assert_eq!(
1995 registration.bound_issuer(),
1996 Some("https://auth.example.com")
1997 );
1998 }
1999
2000 #[tokio::test]
2001 async fn registration_rejects_pre_registered_issuer_mismatch() {
2002 let metadata = authorization_server_metadata("https://new-auth.example.com");
2003 let options = OAuthClientRegistrationOptions::new().with_pre_registered(
2004 OAuthClientRegistration::pre_registered(
2005 "https://old-auth.example.com",
2006 "configured-client",
2007 None,
2008 ),
2009 );
2010
2011 let error = resolve_oauth_client_registration(&reqwest::Client::new(), &metadata, &options)
2012 .await
2013 .unwrap_err();
2014
2015 assert!(
2016 error.to_string().contains("bound to issuer"),
2017 "got: {error}"
2018 );
2019 }
2020
2021 #[tokio::test]
2022 async fn registration_prefers_cimd_over_dynamic_registration() {
2023 let mut metadata = authorization_server_metadata("https://auth.example.com");
2024 metadata.client_id_metadata_document_supported = true;
2025 metadata.registration_endpoint = Some("http://127.0.0.1:9/register".to_string());
2026 let options = OAuthClientRegistrationOptions::new()
2027 .with_client_id_metadata_document("https://client.example.com/client.json")
2028 .with_dynamic_registration(OAuthDynamicClientRegistration::native(
2029 "test-client",
2030 ["http://127.0.0.1/callback"],
2031 ));
2032
2033 let registration =
2034 resolve_oauth_client_registration(&reqwest::Client::new(), &metadata, &options)
2035 .await
2036 .unwrap();
2037
2038 assert_eq!(
2039 registration.method(),
2040 OAuthClientRegistrationMethod::ClientIdMetadataDocument
2041 );
2042 assert_eq!(
2043 registration.client_id(),
2044 "https://client.example.com/client.json"
2045 );
2046 assert_eq!(registration.bound_issuer(), None);
2047 }
2048
2049 #[tokio::test]
2050 async fn registration_rejects_invalid_cimd_url() {
2051 let mut metadata = authorization_server_metadata("https://auth.example.com");
2052 metadata.client_id_metadata_document_supported = true;
2053 let options = OAuthClientRegistrationOptions::new()
2054 .with_client_id_metadata_document("http://client.example.com/client.json");
2055
2056 let error = resolve_oauth_client_registration(&reqwest::Client::new(), &metadata, &options)
2057 .await
2058 .unwrap_err();
2059
2060 assert!(error.to_string().contains("must use HTTPS"), "got: {error}");
2061 }
2062
2063 #[tokio::test]
2064 async fn registration_falls_back_to_native_dcr_and_binds_issuer() {
2065 use tokio::io::{AsyncReadExt, AsyncWriteExt};
2066
2067 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2068 let address = listener.local_addr().unwrap();
2069 let (request_tx, request_rx) = oneshot::channel();
2070 let server = tokio::spawn(async move {
2071 let (mut stream, _) = listener.accept().await.unwrap();
2072 let mut bytes = Vec::new();
2073 let header_end = loop {
2074 let mut chunk = [0u8; 1024];
2075 let read = stream.read(&mut chunk).await.unwrap();
2076 assert!(read > 0);
2077 bytes.extend_from_slice(&chunk[..read]);
2078 if let Some(index) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
2079 break index + 4;
2080 }
2081 };
2082 let headers = String::from_utf8_lossy(&bytes[..header_end]);
2083 let content_length = headers
2084 .lines()
2085 .find_map(|line| {
2086 let (name, value) = line.split_once(':')?;
2087 name.eq_ignore_ascii_case("content-length")
2088 .then(|| value.trim().parse::<usize>().unwrap())
2089 })
2090 .unwrap();
2091 while bytes.len() < header_end + content_length {
2092 let mut chunk = [0u8; 1024];
2093 let read = stream.read(&mut chunk).await.unwrap();
2094 assert!(read > 0);
2095 bytes.extend_from_slice(&chunk[..read]);
2096 }
2097 let body: serde_json::Value =
2098 serde_json::from_slice(&bytes[header_end..header_end + content_length]).unwrap();
2099 request_tx.send(body).unwrap();
2100
2101 let response_body = r#"{"client_id":"dynamic-client","client_secret":"secret"}"#;
2102 let response = format!(
2103 "HTTP/1.1 201 Created\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
2104 response_body.len(),
2105 response_body
2106 );
2107 stream.write_all(response.as_bytes()).await.unwrap();
2108 });
2109
2110 let mut metadata = authorization_server_metadata("https://auth.example.com");
2111 metadata.registration_endpoint = Some(format!("http://{address}/register"));
2112 let options = OAuthClientRegistrationOptions::new().with_dynamic_registration(
2113 OAuthDynamicClientRegistration::native("test-client", ["http://127.0.0.1/callback"])
2114 .grant_types(["authorization_code", "refresh_token"])
2115 .token_endpoint_auth_method("client_secret_basic"),
2116 );
2117
2118 let registration =
2119 resolve_oauth_client_registration(&reqwest::Client::new(), &metadata, &options)
2120 .await
2121 .unwrap();
2122 let request = request_rx.await.unwrap();
2123 server.await.unwrap();
2124
2125 assert_eq!(
2126 registration.method(),
2127 OAuthClientRegistrationMethod::Dynamic
2128 );
2129 assert_eq!(registration.client_id(), "dynamic-client");
2130 assert_eq!(registration.client_secret(), Some("secret"));
2131 assert_eq!(
2132 registration.bound_issuer(),
2133 Some("https://auth.example.com")
2134 );
2135 assert_eq!(request["application_type"], "native");
2136 assert_eq!(request["grant_types"][1], "refresh_token");
2137 assert_eq!(request["token_endpoint_auth_method"], "client_secret_basic");
2138 }
2139
2140 #[test]
2141 fn registration_credentials_round_trip_for_persistent_stores() {
2142 let registration = OAuthClientRegistration::dynamically_registered(
2143 "https://auth.example.com",
2144 "dynamic-client",
2145 Some("stored-secret".to_string()),
2146 );
2147
2148 let json = serde_json::to_string(®istration).unwrap();
2149 let restored: OAuthClientRegistration = serde_json::from_str(&json).unwrap();
2150
2151 assert_eq!(restored, registration);
2152 assert!(!format!("{restored:?}").contains("stored-secret"));
2153 }
2154
2155 #[tokio::test]
2156 async fn stored_dynamic_registration_is_reused_for_exact_issuer() {
2157 let issuer = "https://auth.example.com";
2158 let registration = OAuthClientRegistration {
2159 client_id: "stored-client".to_string(),
2160 client_secret: Some("stored-secret".to_string()),
2161 method: OAuthClientRegistrationMethod::Dynamic,
2162 bound_issuer: Some(issuer.to_string()),
2163 };
2164 let store = MemoryOAuthClientRegistrationStore::new();
2165 store.save(issuer, ®istration).await.unwrap();
2166 let options = OAuthClientRegistrationOptions::new().with_dynamic_registration(
2167 OAuthDynamicClientRegistration::native("test-client", ["http://127.0.0.1/callback"]),
2168 );
2169 let metadata = authorization_server_metadata(issuer);
2170
2171 let resolved = resolve_oauth_client_registration_with_store(
2172 &reqwest::Client::new(),
2173 &metadata,
2174 &options,
2175 &store,
2176 )
2177 .await
2178 .unwrap();
2179
2180 assert_eq!(resolved, registration);
2181 assert_eq!(store.len().await, 1);
2182 }
2183
2184 #[tokio::test]
2185 async fn issuer_migration_registers_new_credentials_without_reusing_old() {
2186 let old_issuer = "https://old-auth.example.com";
2187 let new_issuer = "https://new-auth.example.com";
2188 let old_registration = OAuthClientRegistration {
2189 client_id: "old-client".to_string(),
2190 client_secret: Some("old-secret".to_string()),
2191 method: OAuthClientRegistrationMethod::Dynamic,
2192 bound_issuer: Some(old_issuer.to_string()),
2193 };
2194 let store = MemoryOAuthClientRegistrationStore::new();
2195 store.save(old_issuer, &old_registration).await.unwrap();
2196 let (registration_endpoint, registration_task) =
2197 dynamic_registration_endpoint("new-client", "new-secret").await;
2198 let mut metadata = authorization_server_metadata(new_issuer);
2199 metadata.registration_endpoint = Some(registration_endpoint);
2200 let options = OAuthClientRegistrationOptions::new().with_dynamic_registration(
2201 OAuthDynamicClientRegistration::native("test-client", ["http://127.0.0.1/callback"]),
2202 );
2203
2204 let resolved = resolve_oauth_client_registration_with_store(
2205 &reqwest::Client::new(),
2206 &metadata,
2207 &options,
2208 &store,
2209 )
2210 .await
2211 .unwrap();
2212 registration_task.await.unwrap();
2213
2214 assert_eq!(resolved.client_id(), "new-client");
2215 assert_eq!(resolved.bound_issuer(), Some(new_issuer));
2216 assert_eq!(store.len().await, 2);
2217 assert_eq!(
2218 store.load(old_issuer).await.unwrap(),
2219 Some(old_registration)
2220 );
2221 assert_eq!(
2222 store
2223 .load(new_issuer)
2224 .await
2225 .unwrap()
2226 .as_ref()
2227 .map(OAuthClientRegistration::client_id),
2228 Some("new-client")
2229 );
2230 }
2231
2232 #[tokio::test]
2233 async fn corrupted_store_binding_is_rejected_instead_of_reused() {
2234 let issuer = "https://new-auth.example.com";
2235 let registration = OAuthClientRegistration {
2236 client_id: "old-client".to_string(),
2237 client_secret: None,
2238 method: OAuthClientRegistrationMethod::Dynamic,
2239 bound_issuer: Some("https://old-auth.example.com".to_string()),
2240 };
2241 let store = MemoryOAuthClientRegistrationStore::new();
2242 store.save(issuer, ®istration).await.unwrap();
2243 let options = OAuthClientRegistrationOptions::new().with_dynamic_registration(
2244 OAuthDynamicClientRegistration::native("test-client", ["http://127.0.0.1/callback"]),
2245 );
2246
2247 let error = resolve_oauth_client_registration_with_store(
2248 &reqwest::Client::new(),
2249 &authorization_server_metadata(issuer),
2250 &options,
2251 &store,
2252 )
2253 .await
2254 .unwrap_err();
2255
2256 assert!(matches!(error, OAuthClientError::CredentialStore(_)));
2257 assert!(error.to_string().contains("old-auth.example.com"));
2258 }
2259
2260 #[tokio::test]
2261 async fn registration_reports_when_user_input_is_required() {
2262 let metadata = authorization_server_metadata("https://auth.example.com");
2263 let error = resolve_oauth_client_registration(
2264 &reqwest::Client::new(),
2265 &metadata,
2266 &OAuthClientRegistrationOptions::new(),
2267 )
2268 .await
2269 .unwrap_err();
2270
2271 assert!(
2272 error.to_string().contains("prompt the user"),
2273 "got: {error}"
2274 );
2275 }
2276
2277 fn authorization_server_metadata(issuer: &str) -> OAuthAuthorizationServerMetadata {
2278 OAuthAuthorizationServerMetadata {
2279 issuer: issuer.to_string(),
2280 authorization_endpoint: "https://auth.example.com/authorize".to_string(),
2281 token_endpoint: "https://auth.example.com/token".to_string(),
2282 registration_endpoint: None,
2283 client_id_metadata_document_supported: false,
2284 authorization_response_iss_parameter_supported: false,
2285 code_challenge_methods_supported: vec!["S256".to_string()],
2286 token_endpoint_auth_methods_supported: Vec::new(),
2287 grant_types_supported: Vec::new(),
2288 scopes_supported: Vec::new(),
2289 }
2290 }
2291
2292 async fn dynamic_registration_endpoint(
2293 client_id: &'static str,
2294 client_secret: &'static str,
2295 ) -> (String, tokio::task::JoinHandle<()>) {
2296 use tokio::io::{AsyncReadExt, AsyncWriteExt};
2297
2298 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2299 let address = listener.local_addr().unwrap();
2300 let task = tokio::spawn(async move {
2301 let (mut stream, _) = listener.accept().await.unwrap();
2302 let mut bytes = Vec::new();
2303 let header_end = loop {
2304 let mut chunk = [0u8; 1024];
2305 let read = stream.read(&mut chunk).await.unwrap();
2306 assert!(read > 0);
2307 bytes.extend_from_slice(&chunk[..read]);
2308 if let Some(index) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
2309 break index + 4;
2310 }
2311 };
2312 let headers = String::from_utf8_lossy(&bytes[..header_end]);
2313 let content_length = headers
2314 .lines()
2315 .find_map(|line| {
2316 let (name, value) = line.split_once(':')?;
2317 name.eq_ignore_ascii_case("content-length")
2318 .then(|| value.trim().parse::<usize>().unwrap())
2319 })
2320 .unwrap_or_default();
2321 while bytes.len() < header_end + content_length {
2322 let mut chunk = [0u8; 1024];
2323 let read = stream.read(&mut chunk).await.unwrap();
2324 assert!(read > 0);
2325 bytes.extend_from_slice(&chunk[..read]);
2326 }
2327
2328 let body = serde_json::json!({
2329 "client_id": client_id,
2330 "client_secret": client_secret,
2331 })
2332 .to_string();
2333 let response = format!(
2334 "HTTP/1.1 201 Created\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
2335 body.len(),
2336 body
2337 );
2338 stream.write_all(response.as_bytes()).await.unwrap();
2339 });
2340
2341 (format!("http://{address}/register"), task)
2342 }
2343
2344 #[test]
2345 fn test_token_validity_check() {
2346 let valid = CachedAuthCodeToken {
2347 access_token: "token".into(),
2348 refresh_token: None,
2349 expires_at: Instant::now() + Duration::from_secs(300),
2350 };
2351 assert!(is_token_valid(&valid, Duration::from_secs(30)));
2352
2353 let expiring = CachedAuthCodeToken {
2354 access_token: "token".into(),
2355 refresh_token: None,
2356 expires_at: Instant::now() + Duration::from_secs(10),
2357 };
2358 assert!(!is_token_valid(&expiring, Duration::from_secs(30)));
2359 }
2360}