Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 66 additions & 3 deletions crates/rmcp/src/transport/auth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,9 @@ const DEFAULT_APPLICATION_TYPE: &str = "native";
#[non_exhaustive]
pub struct StoredCredentials {
pub client_id: String,
/// Client authentication material required for token refresh after restart.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_secret: Option<ClientSecret>,
pub token_response: Option<OAuthTokenResponse>,
#[serde(default)]
pub granted_scopes: Vec<String>,
Expand All @@ -237,6 +240,10 @@ impl std::fmt::Debug for StoredCredentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StoredCredentials")
.field("client_id", &self.client_id)
.field(
"client_secret",
&self.client_secret.as_ref().map(|_| "[REDACTED]"),
)
.field(
"token_response",
&self.token_response.as_ref().map(|_| "[REDACTED]"),
Expand All @@ -258,13 +265,20 @@ impl StoredCredentials {
) -> Self {
Self {
client_id,
client_secret: None,
token_response,
granted_scopes,
token_received_at,
issuer: None,
}
}

/// Retain client authentication with the token grant. Empty secrets denote public clients.
pub fn with_client_secret(mut self, secret: Option<ClientSecret>) -> Self {
self.client_secret = secret.filter(|value| !value.secret().is_empty());
self
}

pub fn with_issuer(mut self, issuer: Option<String>) -> Self {
self.issuer = issuer;
self
Expand Down Expand Up @@ -1105,6 +1119,7 @@ pub struct AuthorizationManager {
refresh_redirect_policy: OAuthHttpRedirectPolicy,
metadata: Option<AuthorizationMetadata>,
oauth_client: Option<OAuthClient>,
client_secret: Option<ClientSecret>,
credential_store: Arc<dyn CredentialStore>,
state_store: Arc<dyn StateStore>,
base_url: Url,
Expand Down Expand Up @@ -1352,6 +1367,7 @@ impl AuthorizationManager {
refresh_redirect_policy,
metadata: None,
oauth_client: None,
client_secret: None,
credential_store: Arc::new(InMemoryCredentialStore::new()),
state_store: Arc::new(InMemoryStateStore::new()),
base_url,
Expand Down Expand Up @@ -1464,7 +1480,9 @@ impl AuthorizationManager {
}
}

self.configure_client_id(&stored.client_id)?;
let mut config = OAuthClientConfig::new(&stored.client_id, self.base_url.to_string());
config.client_secret = stored.client_secret.map(|secret| secret.secret().clone());
self.configure_client(config)?;
Comment on lines +1483 to +1485

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was dropping it here intentional, or would threading the secret through that path keep both restore branches in step?

return Ok(true);
}
Ok(false)
Expand Down Expand Up @@ -1612,6 +1630,11 @@ impl AuthorizationManager {
Ok((client_id.to_string(), token_response))
}

/// Return the secret needed to persist this client registration.
pub fn client_secret(&self) -> Option<&ClientSecret> {
self.client_secret.as_ref()
}

/// configure oauth2 client with client credentials
pub fn configure_client(&mut self, config: OAuthClientConfig) -> Result<(), AuthError> {
if self.metadata.is_none() {
Expand Down Expand Up @@ -1640,8 +1663,12 @@ impl AuthorizationManager {
.set_token_uri(token_url)
.set_redirect_uri(redirect_url);

if let Some(secret) = config.client_secret {
client_builder = client_builder.set_client_secret(ClientSecret::new(secret));
let client_secret = config
.client_secret
.filter(|secret| !secret.is_empty())
.map(ClientSecret::new);
if let Some(secret) = &client_secret {
client_builder = client_builder.set_client_secret(secret.clone());
}

let uses_secret_post = metadata
Expand All @@ -1662,6 +1689,7 @@ impl AuthorizationManager {
}

self.oauth_client = Some(client_builder);
self.client_secret = client_secret;
Ok(())
}
/// validate authorization server metadata before starting authorization.
Expand Down Expand Up @@ -2162,6 +2190,7 @@ impl AuthorizationManager {
let client_id = oauth_client.client_id().to_string();
let stored = StoredCredentials {
client_id,
client_secret: self.client_secret.clone(),
token_response: Some(token_result.clone()),
granted_scopes,
token_received_at: Some(Self::now_epoch_secs()),
Expand Down Expand Up @@ -2319,6 +2348,7 @@ impl AuthorizationManager {
let client_id = oauth_client.client_id().to_string();
let stored = StoredCredentials {
client_id,
client_secret: self.client_secret.clone(),
token_response: Some(token_result.clone()),
granted_scopes,
token_received_at: Some(Self::now_epoch_secs()),
Expand Down Expand Up @@ -3143,6 +3173,7 @@ impl AuthorizationManager {
let client_id = config.client_id().to_string();
let stored = StoredCredentials {
client_id,
client_secret: None,
token_response: Some(token_result.clone()),
granted_scopes,
token_received_at: Some(Self::now_epoch_secs()),
Expand Down Expand Up @@ -3264,6 +3295,7 @@ impl AuthorizationManager {

let stored = StoredCredentials {
client_id: client_id.clone(),
client_secret: None,
token_response: Some(token_result.clone()),
granted_scopes,
token_received_at: Some(Self::now_epoch_secs()),
Expand Down Expand Up @@ -3669,6 +3701,15 @@ impl OAuthState {
}
}

/// Return registration authentication material for external credential storage.
pub fn client_secret(&self) -> Option<&ClientSecret> {
match self {
Self::Unauthorized(manager) | Self::Authorized(manager) => manager.client_secret(),
Self::Session(session) => session.auth_manager.client_secret(),
Self::AuthorizedHttpClient(client) => client.auth_manager.client_secret(),
}
}

/// 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(
Expand Down Expand Up @@ -3697,6 +3738,7 @@ impl OAuthState {

let stored = StoredCredentials {
client_id: client_id.to_string(),
client_secret: None,

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How is a consumer that uses get_credentials and set_credentials instead of a CredentialStore expected to restore the secret on the next run?

token_response: Some(credentials),
granted_scopes,
token_received_at: Some(AuthorizationManager::now_epoch_secs()),
Expand Down Expand Up @@ -6258,6 +6300,7 @@ mod tests {
);
let creds = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(token_response),
granted_scopes: vec![],
token_received_at: None,
Expand Down Expand Up @@ -6453,6 +6496,7 @@ mod tests {
store
.save(StoredCredentials {
client_id: "dcr-client".to_string(),
client_secret: None,
token_response: Some(make_token_response("old-token", Some(3600))),
granted_scopes: vec![],
token_received_at: Some(AuthorizationManager::now_epoch_secs()),
Expand Down Expand Up @@ -7391,6 +7435,7 @@ mod tests {
let manager = AuthorizationManager::new("http://localhost").await.unwrap();
let stored = StoredCredentials {
client_id: "test".to_string(),
client_secret: None,
token_response: Some(make_token_response("my-access-token", Some(3600))),
granted_scopes: vec![],
token_received_at: Some(AuthorizationManager::now_epoch_secs()),
Expand All @@ -7409,6 +7454,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response("stale-token", Some(3600))),
granted_scopes: vec![],
token_received_at: Some(AuthorizationManager::now_epoch_secs() - 7200),
Expand All @@ -7428,6 +7474,7 @@ mod tests {
let manager = AuthorizationManager::new("http://localhost").await.unwrap();
let stored = StoredCredentials {
client_id: "test".to_string(),
client_secret: None,
token_response: Some(make_token_response("no-expiry-token", None)),
granted_scopes: vec![],
token_received_at: None,
Expand All @@ -7446,6 +7493,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response("almost-expired", Some(3600))),
granted_scopes: vec![],
token_received_at: Some(AuthorizationManager::now_epoch_secs() - 3590),
Expand All @@ -7465,6 +7513,7 @@ mod tests {
let manager = AuthorizationManager::new("http://localhost").await.unwrap();
let stored = StoredCredentials {
client_id: "test".to_string(),
client_secret: None,
token_response: Some(make_token_response("stale-token", Some(3600))),
granted_scopes: vec![],
token_received_at: Some(AuthorizationManager::now_epoch_secs() - 7200),
Expand Down Expand Up @@ -7878,6 +7927,7 @@ mod tests {
.credential_store
.save(StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response_with_refresh(
"old-token",
"my-refresh-token",
Expand Down Expand Up @@ -7910,6 +7960,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: None,
granted_scopes: vec![],
token_received_at: None,
Expand All @@ -7931,6 +7982,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response("old-token", Some(3600))),
granted_scopes: vec![],
token_received_at: Some(AuthorizationManager::now_epoch_secs()),
Expand Down Expand Up @@ -8019,6 +8071,7 @@ mod tests {
.credential_store
.save(StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response_with_refresh(
"old-token",
"my-refresh-token",
Expand Down Expand Up @@ -8225,6 +8278,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response_with_refresh(
"old-token",
"my-refresh-token",
Expand Down Expand Up @@ -8263,6 +8317,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response_with_refresh(
"old-token",
"my-refresh-token",
Expand Down Expand Up @@ -8301,6 +8356,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response_with_refresh(
"old-token",
"my-refresh-token",
Expand Down Expand Up @@ -8339,6 +8395,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response_with_refresh(
"old-token",
"my-refresh-token",
Expand Down Expand Up @@ -8378,6 +8435,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response_with_refresh(
"old-token",
"my-refresh-token",
Expand Down Expand Up @@ -8442,6 +8500,7 @@ mod tests {

let stored = StoredCredentials {
client_id: "my-client".to_string(),
client_secret: None,
token_response: Some(make_token_response_with_refresh(
"old-token",
"my-refresh-token",
Expand Down Expand Up @@ -8805,3 +8864,7 @@ mod tests {
assert!(store.lock.try_lock().is_ok());
}
}

#[cfg(test)]
#[path = "auth/client_secret_tests.rs"]
mod client_secret_tests;
Loading
Loading