several oauth fixes: capture extra fields, allow for caching of client credentials

Author: the-rooster

crates/rmcp/src/transport/auth.rs

@@ -1,6 +1,5 @@
 use std::{
-    sync::Arc,
-    time::{Duration, Instant},
+    collections::HashMap, sync::Arc, time::{Duration, Instant}
 };
 
 use oauth2::{
@@ -70,6 +69,9 @@ pub struct AuthorizationMetadata {
     pub issuer: Option<String>,
     pub jwks_uri: Option<String>,
     pub scopes_supported: Option<Vec<String>>,
+    // allow additional fields
+    #[serde(flatten)]
+    pub additional_fields: HashMap<String, serde_json::Value>,
 }
 
 /// oauth2 client config
@@ -124,9 +126,12 @@ pub struct ClientRegistrationRequest {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ClientRegistrationResponse {
     pub client_id: String,
-    pub client_secret: String,
+    pub client_secret: Option<String>,
     pub client_name: String,
     pub redirect_uris: Vec<String>,
+    // allow additional fields
+    #[serde(flatten)]
+    pub additional_fields: HashMap<String, serde_json::Value>,
 }
 
 impl AuthorizationManager {
@@ -191,10 +196,21 @@ impl AuthorizationManager {
                 issuer: None,
                 jwks_uri: None,
                 scopes_supported: None,
+                additional_fields:HashMap::new(),
             })
         }
     }
 
+    pub async fn get_credentials(&self) -> Result<(String, Option<OAuthTokenResponse>), AuthError> {
+        let credentials = self.credentials.read().await;
+        let client_id = self
+            .oauth_client
+            .as_ref()
+            .ok_or_else(|| AuthError::InternalError("OAuth client not configured".to_string()))?
+            .client_id();
+        Ok((client_id.to_string(), credentials.clone()))
+    }
+
     /// configure oauth2 client with client credentials
     pub fn configure_client(&mut self, config: OAuthClientConfig) -> Result<(), AuthError> {
         if self.metadata.is_none() {
@@ -287,6 +303,8 @@ impl AuthorizationManager {
                 status, error_text
             )));
         }
+
+
         debug!("registration response: {:?}", response);
         let reg_response = match response.json::<ClientRegistrationResponse>().await {
             Ok(response) => response,
@@ -301,7 +319,7 @@ impl AuthorizationManager {
 
         let config = OAuthClientConfig {
             client_id: reg_response.client_id,
-            client_secret: Some(reg_response.client_secret),
+            client_secret: reg_response.client_secret,
             redirect_uri: redirect_uri.to_string(),
             scopes: vec![],
         };
@@ -310,6 +328,23 @@ impl AuthorizationManager {
         Ok(config)
     }
 
+
+    /// use provided client id to configure oauth2 client instead of dynamic registration
+    /// this is useful when you have a stored client id from previous registration
+    pub fn configure_client_id(
+        &mut self,
+        client_id: &str,
+    ) -> Result<(), AuthError> {
+        let config = OAuthClientConfig {
+            client_id: client_id.to_string(),
+            client_secret: None,
+            scopes: vec![],
+            redirect_uri: self.base_url.to_string(),
+        };
+        self.configure_client(config)
+    }
+    
+
     /// generate authorization url
     pub async fn get_authorization_url(&self, scopes: &[&str]) -> Result<String, AuthError> {
         let oauth_client = self
@@ -513,6 +548,12 @@ impl AuthorizationSession {
         })
     }
 
+
+    /// get client_id and credentials
+    pub async fn get_credentials(&self) -> Result<(String, Option<OAuthTokenResponse>), AuthError> {
+        self.auth_manager.get_credentials().await
+    }
+
     /// get authorization url
     pub fn get_authorization_url(&self) -> &str {
         &self.auth_url
@@ -590,9 +631,52 @@ impl OAuthState {
         if let Some(client) = client {
             manager.with_client(client)?;
         }
+
         Ok(OAuthState::Unauthorized(manager))
     }
 
+    /// Get client_id and OAuth credentials
+    pub async fn get_credentials(&self) -> Result<(String, Option<OAuthTokenResponse>), AuthError> {
+        // return client_id and credentials
+        match self {
+            OAuthState::Unauthorized(manager) => manager.get_credentials().await,
+            OAuthState::Session(session) => session.get_credentials().await,
+            OAuthState::Authorized(manager) => manager.get_credentials().await,
+            OAuthState::AuthorizedHttpClient(client) => client.auth_manager.get_credentials().await,
+        }
+            
+    }
+
+
+    /// Manually set credentials and move into authorized state
+    /// Useful if you're caching credentials externally and wish to reuse them
+    pub async fn set_credentials(
+        &mut self,
+        client_id: &str,
+        credentials: OAuthTokenResponse,
+    ) -> Result<(), AuthError> {
+        if let OAuthState::Unauthorized(manager) = self {
+            let mut manager = std::mem::replace(manager, AuthorizationManager::new("http://localhost").await?);
+            
+            // write credentials
+            *manager.credentials.write().await = Some(credentials);
+
+            // discover metadata
+            let metadata = manager.discover_metadata().await?;
+            manager.metadata = Some(metadata);
+                        
+            // set client id and secret
+            manager.configure_client_id(client_id)?;
+
+            *self = OAuthState::Authorized(manager);
+            Ok(())
+        } else {
+            Err(AuthError::InternalError(
+                "Cannot set credentials in this state".to_string(),
+            ))
+        }
+    }
+
     /// start authorization
     pub async fn start_authorization(
         &mut self,

examples/servers/src/mcp_oauth_server.rs

@@ -525,6 +525,7 @@ async fn oauth_authorization_server() -> impl IntoResponse {
         registration_endpoint: format!("http://{}/oauth/register", BIND_ADDRESS),
         issuer: Some(BIND_ADDRESS.to_string()),
         jwks_uri: Some(format!("http://{}/oauth/jwks", BIND_ADDRESS)),
+        additional_fields: HashMap::new(),
     };
     debug!("metadata: {:?}", metadata);
     (StatusCode::OK, Json(metadata))
@@ -567,9 +568,10 @@ async fn oauth_register(
     // return client information
     let response = ClientRegistrationResponse {
         client_id,
-        client_secret,
+        client_secret: Some(client_secret),
         client_name: req.client_name,
         redirect_uris: req.redirect_uris,
+        additional_fields: HashMap::new(),
     };
 
     (StatusCode::CREATED, Json(response)).into_response()