1use core::future::Future;
28use core::pin::Pin;
29use std::collections::HashMap;
30use std::sync::{Arc, Mutex, PoisonError};
31
32use serde_json::{Value, json};
33use url::Url;
34
35use crate::application::client::bearer::{BearerAuth, BearerError};
36use crate::application::oauth::callback::{CallbackError, authorize_callback};
37use crate::application::oauth::client::{AuthorizationSecrets, BeginError, begin_authorization};
38use crate::application::oauth::environment::{EnvironmentError, Fetched, OAuthEnvironment};
39use crate::application::oauth::flow::{
40 initial_scope, may_step_up, step_up_scope, with_offline_access,
41};
42use crate::application::oauth::metadata::{AuthorizationServer, MetadataError, ProtectedResource};
43use crate::application::oauth::redeem::{
44 RedeemError, RefreshRequest, TokenEndpoint as _, redeem_with_scope,
45};
46use crate::application::oauth::registration::ClientMetadata;
47use crate::domain::ProtocolVersion;
48use crate::domain::mcp::authorization::{AuthorizationModel, authorization_model};
49use crate::domain::mcp::resource_metadata::ResourceUri;
50use crate::domain::oauth::challenge::{BearerChallenge, parse_bearer};
51use crate::domain::oauth::discovery::{
52 authorization_server_metadata_urls, canonical_resource, protected_resource_metadata_urls,
53};
54use crate::domain::oauth::pkce::ChallengeMethod;
55use crate::domain::oauth::registration::{
56 ClientId, ClientSecret, RegisteredClient, TokenAuthMethod,
57};
58use crate::domain::oauth::resource::ResourceParameters;
59
60#[derive(Debug, thiserror::Error)]
62pub enum AuthorizeError {
63 #[error("{what} could not be read: {reason}")]
65 Discovery {
66 what: &'static str,
68 reason: String,
70 },
71 #[error("the resource metadata is for `{found}`, not `{expected}`")]
76 ResourceMismatch {
77 expected: String,
79 found: String,
81 },
82 #[error("the authorization server does not advertise PKCE with S256")]
87 NoPkce,
88 #[error("no client registration is possible with this authorization server")]
91 NoRegistration,
92 #[error("dynamic client registration failed: {0}")]
94 Registration(String),
95 #[error(
97 "the pre-registered credentials belong to `{registered}`, but the server now names `{named}`"
98 )]
99 CredentialsForAnotherIssuer {
100 registered: String,
102 named: String,
104 },
105 #[error(transparent)]
107 Begin(#[from] BeginError),
108 #[error(transparent)]
110 Callback(#[from] CallbackError),
111 #[error(transparent)]
113 Redeem(#[from] RedeemError),
114 #[error(transparent)]
116 Environment(#[from] EnvironmentError),
117 #[error(transparent)]
119 Metadata(#[from] MetadataError),
120 #[error("the server still refused after {0} authorization attempts")]
122 AttemptsExhausted(u32),
123}
124
125#[derive(Debug, Clone)]
127pub struct PreRegistered {
128 issuer: Option<String>,
129 client_id: ClientId,
130 client_secret: Option<ClientSecret>,
131}
132
133impl PreRegistered {
134 #[must_use]
136 pub fn new(client_id: &str, client_secret: Option<&str>) -> Self {
137 Self {
138 issuer: None,
139 client_id: ClientId::opaque(client_id),
140 client_secret: client_secret.map(ClientSecret::issued),
141 }
142 }
143
144 #[must_use]
151 pub fn issued_by(mut self, issuer: impl Into<String>) -> Self {
152 self.issuer = Some(issuer.into());
153 self
154 }
155}
156
157#[derive(Debug, Clone)]
159pub struct OAuthConfig {
160 server: Url,
161 model: AuthorizationModel,
162 redirect_uri: String,
163 client_name: String,
164 client_metadata_url: Option<String>,
165 pre_registered: Option<PreRegistered>,
166 refresh_token: bool,
167}
168
169impl OAuthConfig {
170 pub fn new(server: &str, version: ProtocolVersion) -> Result<Self, url::ParseError> {
176 Ok(Self {
177 server: Url::parse(server)?,
178 model: authorization_model(version),
179 redirect_uri: "http://localhost:3000/callback".to_owned(),
180 client_name: "protocol-droid".to_owned(),
181 client_metadata_url: None,
182 pre_registered: None,
183 refresh_token: false,
184 })
185 }
186
187 #[must_use]
189 pub fn redirect_to(mut self, redirect_uri: impl Into<String>) -> Self {
190 self.redirect_uri = redirect_uri.into();
191 self
192 }
193
194 #[must_use]
196 pub fn named(mut self, client_name: impl Into<String>) -> Self {
197 self.client_name = client_name.into();
198 self
199 }
200
201 #[must_use]
204 pub fn with_metadata_document(mut self, url: impl Into<String>) -> Self {
205 self.client_metadata_url = Some(url.into());
206 self
207 }
208
209 #[must_use]
211 pub fn with_pre_registered(mut self, credentials: PreRegistered) -> Self {
212 self.pre_registered = Some(credentials);
213 self
214 }
215
216 #[must_use]
221 pub const fn wanting_refresh_tokens(mut self) -> Self {
222 self.refresh_token = true;
223 self
224 }
225}
226
227#[derive(Clone)]
230struct Refreshable {
231 token: String,
232 token_endpoint: String,
233 client: RegisteredClient,
234 resource: ResourceParameters,
235}
236
237impl core::fmt::Debug for Refreshable {
238 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
240 f.debug_struct("Refreshable").finish_non_exhaustive()
241 }
242}
243
244#[derive(Default)]
245struct State {
246 token: Option<String>,
247 granted: Option<String>,
248 refresh: Option<Refreshable>,
249 clients: HashMap<String, RegisteredClient>,
250 pre_registered: Option<PreRegistered>,
251 attempts: u32,
252}
253
254pub struct OAuthAuthorizer<E, S> {
256 config: OAuthConfig,
257 environment: E,
258 secrets: S,
259 state: Mutex<State>,
260}
261
262impl<E, S> core::fmt::Debug for OAuthAuthorizer<E, S> {
263 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
265 f.debug_struct("OAuthAuthorizer")
266 .field("server", &self.config.server.as_str())
267 .finish_non_exhaustive()
268 }
269}
270
271impl<E: OAuthEnvironment, S: AuthorizationSecrets + Send + Sync + 'static> OAuthAuthorizer<E, S> {
272 #[must_use]
275 pub fn new(config: OAuthConfig, environment: E, secrets: S) -> Arc<Self> {
276 let pre_registered = config.pre_registered.clone();
277 Arc::new(Self {
278 config,
279 environment,
280 secrets,
281 state: Mutex::new(State {
282 pre_registered,
283 ..State::default()
284 }),
285 })
286 }
287
288 fn state(&self) -> std::sync::MutexGuard<'_, State> {
289 self.state.lock().unwrap_or_else(PoisonError::into_inner)
290 }
291
292 fn record(&self, token: String, granted: Option<String>, refresh: Option<Refreshable>) {
294 let mut state = self.state();
295 state.token = Some(token);
296 state.granted = granted;
297 state.refresh = refresh;
298 }
299
300 async fn try_refresh(&self) -> bool {
306 let Some(held) = self.state().refresh.clone() else {
307 return false;
308 };
309 let endpoint = self.environment.token_endpoint(&held.token_endpoint);
310 let request = RefreshRequest::new(&held.token, &held.client, &held.resource, None);
311 let outcome = endpoint.refresh(request).await;
312 let mut state = self.state();
313 if let Ok(tokens) = outcome {
314 let rotated = tokens
317 .refresh_token()
318 .map_or_else(|| held.token.clone(), |token| token.as_str().to_owned());
319 state.token = Some(tokens.access_token().as_str().to_owned());
320 state.refresh = Some(Refreshable {
321 token: rotated,
322 ..held
323 });
324 true
325 } else {
326 state.refresh = None;
327 false
328 }
329 }
330
331 fn held_client(&self, server: &AuthorizationServer) -> Option<RegisteredClient> {
333 self.state().clients.get(&server.issuer).cloned()
334 }
335
336 fn bind_pre_registered(
339 &self,
340 server: &AuthorizationServer,
341 ) -> Result<Option<RegisteredClient>, AuthorizeError> {
342 let mut pending = self.state().pre_registered.take();
343 let Some(credentials) = pending.take() else {
344 return Ok(None);
345 };
346 if let Some(registered) = &credentials.issuer
347 && *registered != server.issuer
348 {
349 self.state().pre_registered = Some(credentials.clone());
351 return Err(AuthorizeError::CredentialsForAnotherIssuer {
352 registered: registered.clone(),
353 named: server.issuer.clone(),
354 });
355 }
356 let method = auth_method(credentials.client_secret.is_some(), server);
357 let Ok(client) = RegisteredClient::dynamic(
358 &server.issuer,
359 credentials.client_id,
360 credentials.client_secret,
361 ) else {
362 return Ok(None);
363 };
364 let client = client.with_auth_method(method);
365 self.keep_client(&server.issuer, &client);
366 Ok(Some(client))
367 }
368
369 fn keep_client(&self, issuer: &str, client: &RegisteredClient) {
371 self.state()
372 .clients
373 .insert(issuer.to_owned(), client.clone());
374 }
375
376 pub async fn authorize(&self, challenge: &BearerChallenge) -> Result<(), AuthorizeError> {
382 let resource = self.protected_resource(challenge).await?;
383 let issuer = resource
384 .authorization_servers
385 .first()
386 .cloned()
387 .ok_or_else(|| AuthorizeError::Discovery {
388 what: "the protected resource metadata",
389 reason: "it lists no authorization server".to_owned(),
390 })?;
391 let server = self.authorization_server(&issuer).await?;
392
393 let advertises_s256 = server
394 .code_challenge_methods_supported
395 .as_ref()
396 .is_some_and(|methods| methods.iter().any(|method| method == "S256"));
397 if !advertises_s256 {
398 return Err(AuthorizeError::NoPkce);
399 }
400
401 let client = self.client_for(&server).await?;
402 let scope = self.scope_for(challenge, &resource, &server);
403
404 let resource_uri =
409 ResourceUri::new(&resource.resource).map_err(|error| AuthorizeError::Discovery {
410 what: "the resource identifier in the protected resource metadata",
411 reason: error.to_string(),
412 })?;
413 let mut pending = begin_authorization(
414 &self.config.model,
415 &client,
416 &self.config.redirect_uri,
417 ResourceParameters::single(&resource_uri),
418 &self.secrets,
419 )?
420 .with_iss_advertised(server.authorization_response_iss_parameter_supported);
421 if let Some(scope) = &scope {
422 pending = pending.with_scope(scope);
423 }
424 let url =
425 pending.authorization_url(&server.authorization_endpoint, ChallengeMethod::S256)?;
426
427 let callback = self.environment.authorize(url.as_str()).await?;
428 let code = authorize_callback(&self.config.model, &pending, &callback)?;
429 let endpoint = self.environment.token_endpoint(&server.token_endpoint);
430 let tokens = redeem_with_scope(code, &client, &endpoint, scope.as_deref()).await?;
431
432 let granted = tokens.scope().map(ToOwned::to_owned).or(scope);
433 let refresh = tokens.refresh_token().map(|token| Refreshable {
434 token: token.as_str().to_owned(),
435 token_endpoint: server.token_endpoint.clone(),
436 client: client.clone(),
437 resource: ResourceParameters::single(&resource_uri),
438 });
439 self.record(tokens.access_token().as_str().to_owned(), granted, refresh);
440 Ok(())
441 }
442
443 fn scope_for(
445 &self,
446 challenge: &BearerChallenge,
447 resource: &ProtectedResource,
448 server: &AuthorizationServer,
449 ) -> Option<String> {
450 let previous = self.state().granted.clone();
451 let scope = if previous.is_some() {
452 step_up_scope(previous.as_deref(), challenge)
453 } else {
454 initial_scope(challenge, resource.scopes_supported.as_deref())
455 };
456 if self.config.refresh_token {
457 with_offline_access(scope, &server.scopes_supported)
458 } else {
459 scope
460 }
461 }
462
463 async fn get_document(&self, url: &str) -> Result<Fetched, EnvironmentError> {
464 self.environment.get(url).await
465 }
466
467 async fn protected_resource(
470 &self,
471 challenge: &BearerChallenge,
472 ) -> Result<ProtectedResource, AuthorizeError> {
473 let urls = challenge.resource_metadata.as_ref().map_or_else(
474 || protected_resource_metadata_urls(&self.config.server),
475 |url| vec![url.clone()],
476 );
477 let mut last = "no location answered".to_owned();
478 for url in urls {
479 let fetched = match self.get_document(&url).await {
480 Ok(fetched) if fetched.is_success() => fetched,
481 Ok(fetched) => {
482 last = format!("{url} answered {}", fetched.status);
483 continue;
484 }
485 Err(error) => {
486 last = error.to_string();
487 continue;
488 }
489 };
490 let document = ProtectedResource::parse(&fetched.body)?;
491 return self.check_resource(document);
492 }
493 Err(AuthorizeError::Discovery {
494 what: "the protected resource metadata",
495 reason: last,
496 })
497 }
498
499 fn check_resource(
505 &self,
506 document: ProtectedResource,
507 ) -> Result<ProtectedResource, AuthorizeError> {
508 let expected = canonical_resource(&self.config.server);
509 let found = Url::parse(&document.resource).map_or_else(
510 |_| document.resource.clone(),
511 |url| canonical_resource(&url),
512 );
513 let covers = expected == found
514 || expected
515 .strip_prefix(&found)
516 .is_some_and(|rest| rest.starts_with('/'));
517 if covers {
518 Ok(document)
519 } else {
520 Err(AuthorizeError::ResourceMismatch { expected, found })
521 }
522 }
523
524 async fn authorization_server(
528 &self,
529 issuer: &str,
530 ) -> Result<AuthorizationServer, AuthorizeError> {
531 let parsed = Url::parse(issuer).map_err(|error| AuthorizeError::Discovery {
532 what: "the authorization server's issuer",
533 reason: error.to_string(),
534 })?;
535 let mut last = "no location answered".to_owned();
536 for url in authorization_server_metadata_urls(&parsed) {
537 let Ok(fetched) = self.get_document(&url).await else {
538 continue;
539 };
540 if !fetched.is_success() {
541 last = format!("{url} answered {}", fetched.status);
542 continue;
543 }
544 match AuthorizationServer::parse(&fetched.body, issuer) {
545 Ok(server) => return Ok(server),
546 Err(error) => last = error.to_string(),
549 }
550 }
551 Err(AuthorizeError::Discovery {
552 what: "the authorization server metadata",
553 reason: last,
554 })
555 }
556
557 async fn client_for(
565 &self,
566 server: &AuthorizationServer,
567 ) -> Result<RegisteredClient, AuthorizeError> {
568 if let Some(held) = self.held_client(server) {
569 return Ok(held);
570 }
571 if let Some(bound) = self.bind_pre_registered(server)? {
572 return Ok(bound);
573 }
574 let model = &self.config.model;
575 if server.client_id_metadata_document_supported
576 && model.client_registration.supports_metadata_document()
577 && let Some(url) = &self.config.client_metadata_url
578 && let Ok(id) = ClientId::metadata_url(url)
579 {
580 let client = RegisteredClient::from_metadata_document(&server.issuer, id);
581 self.keep_client(&server.issuer, &client);
582 return Ok(client);
583 }
584 if model.client_registration.supports_dynamic()
585 && let Some(endpoint) = &server.registration_endpoint
586 {
587 let client = self.register_dynamic(endpoint, server).await?;
588 self.keep_client(&server.issuer, &client);
589 return Ok(client);
590 }
591 Err(AuthorizeError::NoRegistration)
592 }
593
594 async fn register_dynamic(
599 &self,
600 endpoint: &str,
601 server: &AuthorizationServer,
602 ) -> Result<RegisteredClient, AuthorizeError> {
603 let mut grants = vec!["authorization_code"];
604 if self.config.refresh_token {
605 grants.push("refresh_token");
606 }
607 let metadata = ClientMetadata::new(vec![self.config.redirect_uri.clone()])
608 .map_err(|error| AuthorizeError::Registration(error.to_string()))?
609 .with_client_name(&self.config.client_name)
610 .with_application_type("native")
611 .with_field("grant_types", json!(grants))
612 .and_then(|metadata| metadata.with_field("response_types", json!(["code"])))
613 .map_err(|error| AuthorizeError::Registration(error.to_string()))?;
614 let metadata = match preferred_method(server) {
615 Some(method) => metadata.with_token_endpoint_auth_method(method),
616 None => metadata,
617 };
618 let body = serde_json::to_string(&metadata)
619 .map_err(|error| AuthorizeError::Registration(error.to_string()))?;
620 let fetched = self.environment.post_json(endpoint, &body).await?;
621 if !fetched.is_success() {
622 return Err(AuthorizeError::Registration(format!(
623 "the registration endpoint answered {}",
624 fetched.status
625 )));
626 }
627 let response: Value = serde_json::from_slice(&fetched.body)
628 .map_err(|error| AuthorizeError::Registration(error.to_string()))?;
629 let id = response
630 .get("client_id")
631 .and_then(Value::as_str)
632 .ok_or_else(|| {
633 AuthorizeError::Registration("the response has no client_id".to_owned())
634 })?;
635 let secret = response
636 .get("client_secret")
637 .and_then(Value::as_str)
638 .map(ClientSecret::issued);
639 let method = response
640 .get("token_endpoint_auth_method")
641 .and_then(Value::as_str)
642 .and_then(TokenAuthMethod::parse)
643 .unwrap_or_else(|| auth_method(secret.is_some(), server));
644 let client = RegisteredClient::dynamic(&server.issuer, ClientId::opaque(id), secret)
645 .map_err(|error| AuthorizeError::Registration(error.to_string()))?;
646 Ok(client.with_auth_method(method))
647 }
648}
649
650fn auth_method(has_secret: bool, server: &AuthorizationServer) -> TokenAuthMethod {
656 if !has_secret {
657 return TokenAuthMethod::None;
658 }
659 let offered = &server.token_endpoint_auth_methods_supported;
660 if offered.is_empty() || offered.iter().any(|method| method == "client_secret_basic") {
661 TokenAuthMethod::ClientSecretBasic
662 } else if offered.iter().any(|method| method == "client_secret_post") {
663 TokenAuthMethod::ClientSecretPost
664 } else {
665 TokenAuthMethod::None
666 }
667}
668
669fn preferred_method(server: &AuthorizationServer) -> Option<&'static str> {
671 let offered = &server.token_endpoint_auth_methods_supported;
672 ["client_secret_basic", "client_secret_post", "none"]
673 .into_iter()
674 .find(|method| offered.iter().any(|candidate| candidate == method))
675}
676
677impl<E: OAuthEnvironment, S: AuthorizationSecrets + Send + Sync + 'static> BearerAuth
678 for OAuthAuthorizer<E, S>
679{
680 fn bearer(&self) -> Option<String> {
681 self.state().token.clone()
682 }
683
684 fn challenged<'a>(
685 &'a self,
686 status: u16,
687 www_authenticate: Option<&'a str>,
688 ) -> Pin<Box<dyn Future<Output = Result<(), BearerError>> + Send + 'a>> {
689 Box::pin(async move {
690 let attempts = {
691 let mut state = self.state();
692 state.attempts = state.attempts.saturating_add(1);
693 state.attempts
694 };
695 if !may_step_up(attempts.saturating_sub(1)) {
699 return Err(BearerError(
700 AuthorizeError::AttemptsExhausted(attempts.saturating_sub(1)).to_string(),
701 ));
702 }
703 let challenge = www_authenticate.and_then(parse_bearer).unwrap_or_default();
704 let _ = status;
705 if !challenge.is_insufficient_scope() && self.try_refresh().await {
708 return Ok(());
709 }
710 self.authorize(&challenge)
711 .await
712 .map_err(|error| BearerError(error.to_string()))
713 })
714 }
715
716 fn accepted(&self) {
717 self.state().attempts = 0;
718 }
719}
720
721#[cfg(test)]
722mod tests;