diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..9733074a --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,13 @@ +{ + "[rust]": { + "editor.formatOnSave": true, + "editor.defaultFormatter": "rust-lang.rust-analyzer" + }, + "rust-analyzer.rustfmt.extraArgs": [ + "+nightly" + ], + "files.eol": "\n", + "files.insertFinalNewline": true, + "files.trimTrailingWhitespace": true, + "editor.formatOnSave": true +} diff --git a/Cargo.lock b/Cargo.lock index 62020d5f..ccdcb8ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2498,7 +2498,7 @@ checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" [[package]] name = "mcpmux" -version = "0.0.1" +version = "0.0.3" dependencies = [ "anyhow", "async-trait", @@ -2533,7 +2533,7 @@ dependencies = [ [[package]] name = "mcpmux-core" -version = "0.0.1" +version = "0.0.3" dependencies = [ "anyhow", "async-trait", @@ -2555,7 +2555,7 @@ dependencies = [ [[package]] name = "mcpmux-gateway" -version = "0.0.1" +version = "0.0.3" dependencies = [ "anyhow", "async-stream", @@ -2595,7 +2595,7 @@ dependencies = [ [[package]] name = "mcpmux-mcp" -version = "0.0.1" +version = "0.0.3" dependencies = [ "anyhow", "async-trait", @@ -2614,7 +2614,7 @@ dependencies = [ [[package]] name = "mcpmux-storage" -version = "0.0.1" +version = "0.0.3" dependencies = [ "anyhow", "async-trait", diff --git a/crates/mcpmux-gateway/src/pool/credential_store.rs b/crates/mcpmux-gateway/src/pool/credential_store.rs index 6a6c0dd6..7b18de35 100644 --- a/crates/mcpmux-gateway/src/pool/credential_store.rs +++ b/crates/mcpmux-gateway/src/pool/credential_store.rs @@ -13,7 +13,6 @@ use mcpmux_core::{ }; use oauth2::{basic::BasicTokenType, AccessToken, RefreshToken, TokenResponse}; use rmcp::transport::auth::{AuthError, CredentialStore, OAuthTokenResponse, StoredCredentials}; -use tokio::sync::RwLock; use tracing::{debug, warn}; use uuid::Uuid; @@ -21,14 +20,16 @@ use uuid::Uuid; /// /// This adapter bridges our encrypted database storage to rmcp's CredentialStore trait, /// allowing the SDK to handle token refresh automatically while we maintain persistent storage. +/// +/// IMPORTANT: This store does NOT cache credentials to ensure that expires_in is always +/// recalculated on each load(). RMCP calls load() before each request to check token expiry, +/// so we must return fresh expiration data for automatic token refresh to work correctly. pub struct DatabaseCredentialStore { space_id: Uuid, server_id: String, server_url: String, credential_repo: Arc, backend_oauth_repo: Arc, - /// Cached credentials for performance (SDK calls load() frequently) - cache: RwLock>, } impl DatabaseCredentialStore { @@ -45,7 +46,6 @@ impl DatabaseCredentialStore { server_url: server_url.into(), credential_repo, backend_oauth_repo, - cache: RwLock::new(None), } } @@ -170,17 +170,12 @@ impl CredentialStore for DatabaseCredentialStore { self.space_id, self.server_id ); - // Check cache first - { - let cache = self.cache.read().await; - if cache.is_some() { - debug!( - "[CredentialStore] Returning cached credentials for {}/{}", - self.space_id, self.server_id - ); - return Ok(cache.clone()); - } - } + // NOTE: We intentionally DO NOT use a cache here because expires_in + // must be recalculated on every load() call. RMCP's AuthClient calls + // load() before each request to check if the token is expired. + // If we cache the StoredCredentials with the OAuthTokenResponse, + // the expires_in Duration becomes stale and RMCP won't refresh + // expired tokens properly. // Load from database let registration = self @@ -240,12 +235,6 @@ impl CredentialStore for DatabaseCredentialStore { } }; - // Update cache - { - let mut cache = self.cache.write().await; - *cache = stored.clone(); - } - Ok(stored) } @@ -253,12 +242,6 @@ impl CredentialStore for DatabaseCredentialStore { // Save to database self.save_to_database(&credentials).await?; - // Update cache - { - let mut cache = self.cache.write().await; - *cache = Some(credentials); - } - Ok(()) } @@ -269,12 +252,6 @@ impl CredentialStore for DatabaseCredentialStore { .await .map_err(|e| AuthError::InternalError(format!("Failed to clear tokens: {}", e)))?; - // Clear cache - { - let mut cache = self.cache.write().await; - *cache = None; - } - debug!( "[CredentialStore] Cleared tokens for {}/{}", self.space_id, self.server_id @@ -312,6 +289,101 @@ fn build_token_response( #[cfg(test)] mod tests { use super::*; + use std::sync::Arc; + + // Mock implementations for testing + #[derive(Clone)] + struct MockCredentialRepo { + credential: Arc>>, + } + + impl MockCredentialRepo { + fn new() -> Self { + Self { + credential: Arc::new(tokio::sync::RwLock::new(None)), + } + } + + async fn set(&self, cred: Credential) { + *self.credential.write().await = Some(cred); + } + } + + #[async_trait] + impl CredentialRepository for MockCredentialRepo { + async fn get( + &self, + _space_id: &Uuid, + _server_id: &str, + ) -> anyhow::Result> { + Ok(self.credential.read().await.clone()) + } + + async fn save(&self, credential: &Credential) -> anyhow::Result<()> { + *self.credential.write().await = Some(credential.clone()); + Ok(()) + } + + async fn delete(&self, _space_id: &Uuid, _server_id: &str) -> anyhow::Result<()> { + *self.credential.write().await = None; + Ok(()) + } + + async fn clear_tokens(&self, _space_id: &Uuid, _server_id: &str) -> anyhow::Result { + let had_token = self.credential.read().await.is_some(); + *self.credential.write().await = None; + Ok(had_token) + } + + async fn list_for_space(&self, _space_id: &Uuid) -> anyhow::Result> { + Ok(vec![]) + } + } + + #[derive(Clone)] + struct MockOAuthRepo { + registration: Arc>>, + } + + impl MockOAuthRepo { + fn new() -> Self { + Self { + registration: Arc::new(tokio::sync::RwLock::new(None)), + } + } + + async fn set(&self, reg: OutboundOAuthRegistration) { + *self.registration.write().await = Some(reg); + } + } + + #[async_trait] + impl OutboundOAuthRepository for MockOAuthRepo { + async fn get( + &self, + _space_id: &Uuid, + _server_id: &str, + ) -> anyhow::Result> { + Ok(self.registration.read().await.clone()) + } + + async fn save(&self, registration: &OutboundOAuthRegistration) -> anyhow::Result<()> { + *self.registration.write().await = Some(registration.clone()); + Ok(()) + } + + async fn delete(&self, _space_id: &Uuid, _server_id: &str) -> anyhow::Result<()> { + *self.registration.write().await = None; + Ok(()) + } + + async fn list_for_space( + &self, + _space_id: &Uuid, + ) -> anyhow::Result> { + Ok(vec![]) + } + } #[test] fn test_build_token_response() { @@ -327,4 +399,176 @@ mod tests { Some("refresh456") ); } + + #[tokio::test] + async fn test_expires_in_recalculated_on_each_load() { + // This test verifies the critical fix: expires_in must be recalculated + // on each load() call, not cached with stale values + let space_id = Uuid::new_v4(); + let server_id = "test-server"; + let server_url = "https://test.example.com"; + + let cred_repo = Arc::new(MockCredentialRepo::new()); + let oauth_repo = Arc::new(MockOAuthRepo::new()); + + // Set up a registration + let registration = OutboundOAuthRegistration::new( + space_id, + server_id, + server_url, + "test-client-id", + "http://localhost:3000/callback".to_string(), + ); + oauth_repo.set(registration).await; + + // Set up a credential that expires in 10 seconds + let expires_at = Utc::now() + Duration::seconds(10); + let credential = Credential { + space_id, + server_id: server_id.to_string(), + value: CredentialValue::OAuth { + access_token: "token123".to_string(), + refresh_token: Some("refresh123".to_string()), + expires_at: Some(expires_at), + token_type: "Bearer".to_string(), + scope: None, + }, + created_at: Utc::now(), + updated_at: Utc::now(), + last_used: Some(Utc::now()), + }; + cred_repo.set(credential).await; + + let store = + DatabaseCredentialStore::new(space_id, server_id, server_url, cred_repo, oauth_repo); + + // First load - should have ~10 seconds + let stored1 = store.load().await.unwrap().unwrap(); + let token1 = stored1.token_response.as_ref().unwrap(); + let expires_in_1 = token1.expires_in().unwrap(); + + assert!(expires_in_1.as_secs() >= 9 && expires_in_1.as_secs() <= 10); + + // Wait 2 seconds + tokio::time::sleep(tokio::time::Duration::from_secs(2)).await; + + // Second load - should have ~8 seconds (recalculated, not cached) + let stored2 = store.load().await.unwrap().unwrap(); + let token2 = stored2.token_response.as_ref().unwrap(); + let expires_in_2 = token2.expires_in().unwrap(); + + // This is the critical assertion: expires_in should decrease because it's recalculated + assert!( + expires_in_2.as_secs() >= 7 && expires_in_2.as_secs() <= 8, + "Expected expires_in to decrease from ~10s to ~8s, but got {} seconds", + expires_in_2.as_secs() + ); + + // Verify it actually decreased + assert!( + expires_in_2 < expires_in_1, + "expires_in should decrease on subsequent loads (was {}, now {})", + expires_in_1.as_secs(), + expires_in_2.as_secs() + ); + } + + #[tokio::test] + async fn test_expired_token_detected() { + // Verify that an expired token is properly detected + let space_id = Uuid::new_v4(); + let server_id = "test-server"; + let server_url = "https://test.example.com"; + + let cred_repo = Arc::new(MockCredentialRepo::new()); + let oauth_repo = Arc::new(MockOAuthRepo::new()); + + // Set up registration + let registration = OutboundOAuthRegistration::new( + space_id, + server_id, + server_url, + "test-client-id", + "http://localhost:3000/callback".to_string(), + ); + oauth_repo.set(registration).await; + + // Set up a credential that already expired (5 seconds ago) + let expires_at = Utc::now() - Duration::seconds(5); + let credential = Credential { + space_id, + server_id: server_id.to_string(), + value: CredentialValue::OAuth { + access_token: "expired_token".to_string(), + refresh_token: Some("refresh123".to_string()), + expires_at: Some(expires_at), + token_type: "Bearer".to_string(), + scope: None, + }, + created_at: Utc::now(), + updated_at: Utc::now(), + last_used: Some(Utc::now()), + }; + cred_repo.set(credential).await; + + let store = + DatabaseCredentialStore::new(space_id, server_id, server_url, cred_repo, oauth_repo); + + // Load should return token with expires_in = 0 (expired) + let stored = store.load().await.unwrap().unwrap(); + let token = stored.token_response.as_ref().unwrap(); + let expires_in = token.expires_in().unwrap(); + + assert_eq!( + expires_in.as_secs(), + 0, + "Expired token should have expires_in = 0, got {} seconds", + expires_in.as_secs() + ); + } + + #[tokio::test] + async fn test_save_updates_database() { + // Verify that save() writes to database, not just cache + let space_id = Uuid::new_v4(); + let server_id = "test-server"; + let server_url = "https://test.example.com"; + + let cred_repo = Arc::new(MockCredentialRepo::new()); + let oauth_repo = Arc::new(MockOAuthRepo::new()); + + let store = DatabaseCredentialStore::new( + space_id, + server_id, + server_url, + Arc::clone(&cred_repo) as Arc, + Arc::clone(&oauth_repo) as Arc, + ); + + // Save new credentials + let token_response = build_token_response( + "new_token".to_string(), + Some("new_refresh".to_string()), + Some(std::time::Duration::from_secs(3600)), + ); + + let credentials = StoredCredentials { + client_id: "new-client-id".to_string(), + token_response: Some(token_response), + }; + + store.save(credentials).await.unwrap(); + + // Verify they were written to database by checking the mock repo directly + let saved_cred = cred_repo.get(&space_id, server_id).await.unwrap().unwrap(); + match saved_cred.value { + CredentialValue::OAuth { access_token, .. } => { + assert_eq!(access_token, "new_token"); + } + _ => panic!("Expected OAuth credential"), + } + + let saved_reg = oauth_repo.get(&space_id, server_id).await.unwrap().unwrap(); + assert_eq!(saved_reg.client_id, "new-client-id"); + } }