use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use async_trait::async_trait;
use rust_mcp_sdk::auth::{
AuthInfo, AuthProvider as SdkAuthProvider, AuthenticationError, OauthEndpoint,
};
use rust_mcp_sdk::mcp_http::{GenericBody, McpAppState};
use rust_mcp_sdk::mcp_server::error::TransportServerError;
use crate::server::auth::ApiKeyConfig;
const API_KEY_TTL_SECS: u64 = 87_600 * 60 * 60;
const API_KEY_TOKEN_ID: &str = "api-key";
pub struct ApiKeyAuthProvider {
config: ApiKeyConfig,
}
impl ApiKeyAuthProvider {
#[must_use]
pub fn new(config: ApiKeyConfig) -> Self {
Self { config }
}
}
#[async_trait]
impl SdkAuthProvider for ApiKeyAuthProvider {
async fn verify_token(&self, access_token: String) -> Result<AuthInfo, AuthenticationError> {
if !self.config.enabled {
return Err(AuthenticationError::InvalidToken {
description: "API key authentication is disabled",
});
}
if access_token.trim().is_empty() {
return Err(AuthenticationError::InvalidToken {
description: "Empty API key",
});
}
if self.config.is_valid_key(&access_token) {
Ok(AuthInfo {
token_unique_id: API_KEY_TOKEN_ID.to_string(),
client_id: None,
user_id: None,
scopes: None,
expires_at: Some(SystemTime::now() + Duration::from_secs(API_KEY_TTL_SECS)),
audience: None,
extra: None,
})
} else {
Err(AuthenticationError::InvalidToken {
description: "Invalid API key",
})
}
}
fn auth_endpoints(&self) -> Option<&HashMap<String, OauthEndpoint>> {
None
}
async fn handle_request(
&self,
_request: http::Request<&str>,
_state: Arc<McpAppState>,
) -> Result<http::Response<GenericBody>, TransportServerError> {
Err(TransportServerError::HttpError(
"API-key authentication exposes no OAuth endpoints".to_string(),
))
}
fn protected_resource_metadata_url(&self) -> Option<&str> {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
fn config_with_key() -> (ApiKeyConfig, String) {
let generated = ApiKeyConfig::default()
.generate_key()
.expect("failed to generate API key");
(
ApiKeyConfig {
enabled: true,
keys: vec![generated.hash],
..Default::default()
},
generated.key,
)
}
#[tokio::test]
async fn verify_token_accepts_valid_key_with_future_expiry() {
let (config, key) = config_with_key();
let provider = ApiKeyAuthProvider::new(config);
let info = provider
.verify_token(key)
.await
.expect("valid key should be accepted");
let expires_at = info.expires_at.expect("expires_at must be Some");
assert!(expires_at > SystemTime::now());
assert_eq!(info.token_unique_id, API_KEY_TOKEN_ID);
}
#[tokio::test]
async fn verify_token_rejects_invalid_key() {
let (config, _key) = config_with_key();
let provider = ApiKeyAuthProvider::new(config);
let err = provider
.verify_token("not-a-valid-key".to_string())
.await
.expect_err("invalid key should be rejected");
assert!(matches!(err, AuthenticationError::InvalidToken { .. }));
}
#[tokio::test]
async fn verify_token_rejects_empty_token() {
let (config, _key) = config_with_key();
let provider = ApiKeyAuthProvider::new(config);
for token in ["", " "] {
let err = provider
.verify_token(token.to_string())
.await
.expect_err("empty token should be rejected");
assert!(matches!(err, AuthenticationError::InvalidToken { .. }));
}
}
#[tokio::test]
async fn verify_token_rejects_valid_key_when_disabled() {
let (mut config, key) = config_with_key();
config.enabled = false;
let provider = ApiKeyAuthProvider::new(config);
let err = provider
.verify_token(key)
.await
.expect_err("a disabled provider must reject every token");
assert!(matches!(err, AuthenticationError::InvalidToken { .. }));
}
#[test]
fn exposes_no_oauth_surface() {
let (config, _key) = config_with_key();
let provider = ApiKeyAuthProvider::new(config);
assert!(provider.auth_endpoints().is_none());
assert!(provider.protected_resource_metadata_url().is_none());
}
}