diff --git a/Cargo.lock b/Cargo.lock index 6615855dc..7bca17d4f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -147,6 +147,9 @@ dependencies = [ "age", "apple-native-keyring-store", "async-trait", + "aws-config", + "aws-credential-types", + "aws-sigv4", "dbus-secret-service-keyring-store", "dirs", "futures", @@ -225,8 +228,6 @@ dependencies = [ "aws-config", "aws-credential-types", "aws-sdk-bedrockruntime", - "aws-sigv4", - "aws-smithy-runtime-api", "aws-smithy-types", "axum", "base64 0.23.1", @@ -334,6 +335,8 @@ version = "0.7.13" dependencies = [ "aether-auth", "async-stream", + "aws-credential-types", + "aws-sigv4", "axum", "chrono", "futures", @@ -341,6 +344,7 @@ dependencies = [ "rmcp", "serde", "serde_json", + "sse-stream", "thiserror 2.0.21", "tokio", "tokio-stream", diff --git a/Cargo.toml b/Cargo.toml index 4918cd7d3..0c12c9912 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,7 +21,6 @@ aws-sdk-bedrockruntime = { version = "1.133.0", default-features = false, featur # Capped below 1.5.0: it makes DeferredSignerSender::send generic, which breaks # type inference in aws-runtime 1.7.4 (latest). Lift once aws-runtime ships a fix. aws-smithy-types = ">=1.4.9, <1.9.0" -aws-smithy-runtime-api = "1.14.0" aws-credential-types = "1.3.0" aws-sigv4 = "1.5.1" @@ -51,6 +50,7 @@ tokio-tungstenite = { version = "0.30.0", features = ["rustls-tls-webpki-roots"] # MCP and API clients rmcp = { version = "^3.5.1", default-features = false } +sse-stream = "0.2.6" async-openai = { version = "^0.42.0", features = ["byot", "chat-completion", "responses"] } reqwest = { version = "^0.13.4", default-features = false, features = ["json", "query", "rustls", "http2", "stream"] } oauth2 = "5.0" diff --git a/crates/aether-auth/Cargo.toml b/crates/aether-auth/Cargo.toml index aff8955fc..abe61a51a 100644 --- a/crates/aether-auth/Cargo.toml +++ b/crates/aether-auth/Cargo.toml @@ -2,7 +2,7 @@ name = "aether-auth" version = "0.1.17" edition = "2024" -description = "OAuth credential storage and authorization flows for the Aether AI agent framework" +description = "OAuth credential storage, authorization flows and AWS SigV4 signing for the Aether AI agent framework" license = "MIT" repository = "https://github.com/contextbridge/aether" readme = "README.md" @@ -22,6 +22,7 @@ keyring = [ "dep:dbus-secret-service-keyring-store", ] mcp = ["dep:rmcp"] +sigv4 = ["dep:aws-config", "dep:aws-credential-types", "dep:aws-sigv4"] [dependencies] age = { workspace = true } @@ -43,6 +44,11 @@ oauth2 = { workspace = true } reqwest = { workspace = true } url = { workspace = true } +# sigv4 feature +aws-config = { workspace = true, optional = true } +aws-credential-types = { workspace = true, optional = true } +aws-sigv4 = { workspace = true, optional = true } + [target.'cfg(target_os = "macos")'.dependencies] apple-native-keyring-store = { workspace = true, optional = true } diff --git a/crates/aether-auth/src/lib.rs b/crates/aether-auth/src/lib.rs index 313c33a06..fc5ddc0af 100644 --- a/crates/aether-auth/src/lib.rs +++ b/crates/aether-auth/src/lib.rs @@ -25,3 +25,8 @@ pub mod mcp; pub use mcp::{ McpCredentialStore, OAuthClientRegistration, OAuthFlowOptions, create_auth_manager_from_store, perform_oauth_flow, }; + +#[cfg(feature = "sigv4")] +mod sigv4; +#[cfg(feature = "sigv4")] +pub use sigv4::{SigV4Config, SigV4Credentials, SigV4Error, SigV4Signer}; diff --git a/crates/aether-auth/src/sigv4.rs b/crates/aether-auth/src/sigv4.rs new file mode 100644 index 000000000..2d29189ab --- /dev/null +++ b/crates/aether-auth/src/sigv4.rs @@ -0,0 +1,129 @@ +use aws_config::Region; +use aws_config::default_provider::credentials::DefaultCredentialsChain; +use aws_credential_types::Credentials; +use aws_credential_types::provider::{ProvideCredentials, SharedCredentialsProvider}; +use aws_sigv4::http_request::{SignableBody, SignableRequest, SigningSettings, sign}; +use aws_sigv4::sign::v4; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use std::error::Error; +use std::sync::Arc; +use std::time::{Duration, SystemTime}; +use tokio::sync::{Mutex, OnceCell}; + +#[derive(Debug, Clone)] +pub struct SigV4Signer { + state: Arc, +} + +#[derive(Debug, Clone)] +pub struct SigV4Config { + pub service: String, + pub region: String, + pub credentials: SigV4Credentials, + pub refresh_before_expiry: Duration, +} + +#[derive(Debug, Clone)] +pub enum SigV4Credentials { + DefaultChain { profile: Option }, + Provider(SharedCredentialsProvider), +} + +#[derive(Debug, thiserror::Error)] +pub enum SigV4Error { + #[error("Could not resolve AWS credentials: {0}")] + Credentials(String), + + #[error("SigV4 signing failed: {0}")] + Signing(String), +} + +impl SigV4Signer { + pub fn new(config: SigV4Config) -> Self { + Self { state: Arc::new(SignerState { config, default_chain: OnceCell::new(), cached: Mutex::default() }) } + } + + pub fn config(&self) -> &SigV4Config { + &self.state.config + } + + pub async fn headers(&self, method: &str, url: &str, body: &[u8]) -> Result { + let identity = self.credentials().await?.into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(&self.config().region) + .name(&self.config().service) + .time(SystemTime::now()) + .settings(SigningSettings::default()) + .build() + .map_err(signing_error)? + .into(); + + let signable = + SignableRequest::new(method, url, std::iter::empty(), SignableBody::Bytes(body)).map_err(signing_error)?; + + let (instructions, _signature) = sign(signable, ¶ms).map_err(signing_error)?.into_parts(); + instructions.headers().map(|(name, value)| sensitive_header(name, value)).collect() + } + + async fn credentials(&self) -> Result { + let mut cached = self.state.cached.lock().await; + if let Some(credentials) = cached.as_ref().filter(|credentials| self.is_fresh(credentials)) { + return Ok(credentials.clone()); + } + let credentials = self + .provider() + .await + .provide_credentials() + .await + .map_err(|error| SigV4Error::Credentials(error_chain(&error)))?; + + *cached = Some(credentials.clone()); + Ok(credentials) + } + + async fn provider(&self) -> &SharedCredentialsProvider { + let config = self.config(); + match &config.credentials { + SigV4Credentials::Provider(provider) => provider, + SigV4Credentials::DefaultChain { profile } => { + self.state.default_chain.get_or_init(|| default_chain(&config.region, profile.as_deref())).await + } + } + } + + fn is_fresh(&self, credentials: &Credentials) -> bool { + credentials.expiry().is_none_or(|expiry| expiry > SystemTime::now() + self.config().refresh_before_expiry) + } +} + +#[derive(Debug)] +struct SignerState { + config: SigV4Config, + default_chain: OnceCell, + cached: Mutex>, +} + +async fn default_chain(region: &str, profile: Option<&str>) -> SharedCredentialsProvider { + let chain = DefaultCredentialsChain::builder().region(Region::new(region.to_string())); + let chain = match profile { + Some(profile) => chain.profile_name(profile), + None => chain, + }; + SharedCredentialsProvider::new(chain.build().await) +} + +fn sensitive_header(name: &str, value: &str) -> Result<(HeaderName, HeaderValue), SigV4Error> { + let invalid = || SigV4Error::Signing(format!("invalid signed header `{name}`")); + let mut value = HeaderValue::from_str(value).map_err(|_| invalid())?; + value.set_sensitive(true); + Ok((HeaderName::try_from(name).map_err(|_| invalid())?, value)) +} + +fn signing_error(error: impl Error + 'static) -> SigV4Error { + SigV4Error::Signing(error_chain(&error)) +} + +fn error_chain(error: &(dyn Error + 'static)) -> String { + std::iter::successors(Some(error), |&error| error.source()).map(ToString::to_string).collect::>().join(": ") +} diff --git a/crates/aether-auth/tests/sigv4.rs b/crates/aether-auth/tests/sigv4.rs new file mode 100644 index 000000000..3090d5c39 --- /dev/null +++ b/crates/aether-auth/tests/sigv4.rs @@ -0,0 +1,92 @@ +#![cfg(feature = "sigv4")] + +use aether_auth::{SigV4Config, SigV4Credentials, SigV4Error, SigV4Signer}; +use aws_credential_types::Credentials; +use aws_credential_types::credential_fn::provide_credentials_fn; +use aws_credential_types::provider::error::CredentialsError; +use aws_credential_types::provider::{ProvideCredentials, SharedCredentialsProvider}; +use reqwest::header::HeaderValue; +use std::future::ready; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, SystemTime}; + +#[tokio::test] +async fn credentials_are_reused_while_they_are_fresh() { + let signer = signer(rotating(Duration::from_hours(1))); + let first = signing_key(&signer).await; + let second = signing_key(&signer).await; + assert_eq!([first, second], ["AKID1", "AKID1"]); +} + +#[tokio::test] +async fn credentials_expiring_within_the_refresh_window_are_refreshed() { + let signer = SigV4Signer::new(SigV4Config { + refresh_before_expiry: Duration::from_hours(2), + ..config(rotating(Duration::from_hours(1))) + }); + + let first = signing_key(&signer).await; + let second = signing_key(&signer).await; + assert_eq!([first, second], ["AKID1", "AKID2"]); +} + +#[tokio::test] +async fn clones_share_resolved_credentials() { + let signer = signer(rotating(Duration::from_hours(1))); + let original = signing_key(&signer).await; + let clone = signing_key(&signer.clone()).await; + assert_eq!([original, clone], ["AKID1", "AKID1"]); +} + +#[tokio::test] +async fn signed_headers_carry_the_session_token_and_are_sensitive() { + let credentials = Credentials::new("AKID", "secret", Some("session".to_string()), None, "test"); + let signer = signer(credentials); + let headers = signer.headers("POST", "https://example.com/mcp", b"{}").await.unwrap(); + + assert_eq!(headers["x-amz-security-token"], "session"); + assert!(headers.contains_key("x-amz-date")); + assert!(headers.values().all(HeaderValue::is_sensitive), "{headers:?}"); +} + +#[tokio::test] +async fn credential_failures_explain_why() { + let signer = signer(missing("the SSO session has expired")); + let error = signer.headers("GET", "https://example.com/mcp", &[]).await.unwrap_err(); + + assert!(matches!(error, SigV4Error::Credentials(_))); + assert!(error.to_string().contains("the SSO session has expired"), "{error}"); +} + +fn rotating(lifetime: Duration) -> impl ProvideCredentials + 'static { + let issued = AtomicUsize::new(0); + provide_credentials_fn(move || { + let issued = issued.fetch_add(1, Ordering::SeqCst) + 1; + let expiry = SystemTime::now() + lifetime; + ready(Ok(Credentials::new(format!("AKID{issued}"), "secret", None, Some(expiry), "sts"))) + }) +} + +fn missing(reason: &'static str) -> impl ProvideCredentials + 'static { + provide_credentials_fn(move || ready(Err(CredentialsError::provider_error(reason)))) +} + +fn signer(credentials: impl ProvideCredentials + 'static) -> SigV4Signer { + SigV4Signer::new(config(credentials)) +} + +fn config(credentials: impl ProvideCredentials + 'static) -> SigV4Config { + SigV4Config { + service: "execute-api".to_string(), + region: "us-east-1".to_string(), + credentials: SigV4Credentials::Provider(SharedCredentialsProvider::new(credentials)), + refresh_before_expiry: Duration::from_mins(5), + } +} + +async fn signing_key(signer: &SigV4Signer) -> String { + let headers = signer.headers("POST", "https://example.com/mcp", b"{}").await.unwrap(); + let authorization = headers["authorization"].to_str().unwrap(); + let credential = authorization.split("Credential=").nth(1).expect("a credential scope"); + credential.split('/').next().unwrap().to_string() +} diff --git a/crates/aether-cli/src/acp/protocol/mcp.rs b/crates/aether-cli/src/acp/protocol/mcp.rs index ca3e0f4f7..059725b34 100644 --- a/crates/aether-cli/src/acp/protocol/mcp.rs +++ b/crates/aether-cli/src/acp/protocol/mcp.rs @@ -28,7 +28,7 @@ fn try_map_mcp_server(server: McpServer, cwd: &Path) -> Option { cwd: Some(cwd.to_path_buf()), }, ), - Http(http) => (http.name, http_transport(http.url, header_map(&http.headers)?, None)), + Http(http) => (http.name, http_transport(http.url, header_map(&http.headers)?)), _ => return None, }; Some(ServerSpec { name, transport }) @@ -42,6 +42,7 @@ fn header_map(headers: &[HttpHeader]) -> Option { mod tests { use super::*; use agent_client_protocol::schema::v2 as acp; + use mcp_utils::client::HttpAuth; fn map(servers: Vec) -> Vec { map_acp_mcp_servers(servers, Path::new("/workspace")) @@ -81,11 +82,11 @@ mod tests { assert_eq!(configs.len(), 1); match &configs[0].transport { - Transport::Http { url, headers, oauth } => { + Transport::Http { url, headers, auth } => { assert_eq!(configs[0].name, "http-server"); assert_eq!(url, "https://example.com/mcp"); assert_eq!(headers["authorization"], "Bearer token123"); - assert!(oauth.is_none(), "an Authorization header replaces OAuth"); + assert!(matches!(auth, HttpAuth::None), "an Authorization header replaces OAuth"); } other => panic!("Expected Http, got {other:?}"), } @@ -129,9 +130,9 @@ mod tests { let configs = map(servers); let Transport::Stdio { args, env, .. } = &configs[0].transport else { panic!("expected stdio") }; assert!(args.is_empty() && env.is_empty()); - let Transport::Http { headers, oauth, .. } = &configs[1].transport else { panic!("expected http") }; + let Transport::Http { headers, auth, .. } = &configs[1].transport else { panic!("expected http") }; assert!(headers.is_empty()); - assert!(oauth.is_some(), "servers without credentials can sign in with OAuth"); + assert!(matches!(auth, HttpAuth::OAuth(_)), "servers without credentials can sign in with OAuth"); } #[test] diff --git a/crates/aether-core/Cargo.toml b/crates/aether-core/Cargo.toml index eba968025..a1745a6c2 100644 --- a/crates/aether-core/Cargo.toml +++ b/crates/aether-core/Cargo.toml @@ -40,7 +40,7 @@ tracing = { workspace = true } thiserror = { workspace = true } glob = { workspace = true } http = { workspace = true } -aether-auth = { path = "../aether-auth", version = "0.1.17" } +aether-auth = { path = "../aether-auth", version = "0.1.17", features = ["sigv4"] } mcp_utils = { package = "aether-mcp-utils", path = "../mcp-utils", version = "0.7.13" } llm = { package = "aether-llm", path = "../llm", version = "0.9.13", features = ["codex"] } utils = { package = "aether-utils", path = "../utils", version = "0.2.19" } diff --git a/crates/aether-core/src/docs/mcp_server_config.md b/crates/aether-core/src/docs/mcp_server_config.md index 015b170c1..8fe120e27 100644 --- a/crates/aether-core/src/docs/mcp_server_config.md +++ b/crates/aether-core/src/docs/mcp_server_config.md @@ -35,8 +35,18 @@ The `oauth` object is optional. Set `clientMetadataUrl` for a custom CIMD, or `clientId` for a pre-registered public client; a configured `clientId` takes priority. `callbackPort` defaults to 3118 and must exactly match the client's registered redirect URI. If the authorization server does not advertise CIMD, -Aether falls back to deprecated Dynamic Client Registration. An explicit -`Authorization` header bypasses OAuth entirely. +Aether falls back to deprecated Dynamic Client Registration. An +`Authorization` header on its own turns off OAuth. + +A server can configure one of `oauth`, `sigv4`, or an `Authorization` header. For example: + +```json +{ + "type": "http", + "url": "https://abc123.execute-api.us-east-1.amazonaws.com/mcp", + "sigv4": { "service": "execute-api", "region": "us-east-1", "profile": "dev" } +} +``` A remote server using a bearer token: diff --git a/crates/aether-core/src/mcp/config.rs b/crates/aether-core/src/mcp/config.rs index 411e3036d..3c15e96cb 100644 --- a/crates/aether-core/src/mcp/config.rs +++ b/crates/aether-core/src/mcp/config.rs @@ -1,12 +1,13 @@ use super::ToolExposure; -use aether_auth::OAuthClientRegistration; +use aether_auth::{OAuthClientRegistration, SigV4Config, SigV4Credentials, SigV4Signer}; use http::header::{AUTHORIZATION, HeaderMap, HeaderName, HeaderValue, InvalidHeaderName, InvalidHeaderValue}; -use mcp_utils::client::{OAuthSettings, Transport}; +use mcp_utils::client::{HttpAuth, OAuthSettings, Transport}; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use std::collections::{BTreeMap, HashMap}; use std::num::NonZeroU16; use std::path::Path; +use std::time::Duration; use utils::variables::{VarError, Vars}; #[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, PartialEq)] @@ -53,6 +54,15 @@ pub struct McpOAuthConfig { pub callback_port: Option, } +#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct McpSigV4Config { + pub service: String, + pub region: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub profile: Option, +} + #[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] #[serde(deny_unknown_fields)] pub struct RemoteServerConfig { @@ -67,6 +77,9 @@ pub struct RemoteServerConfig { #[serde(default, skip_serializing_if = "Option::is_none")] pub oauth: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sigv4: Option, + #[serde(rename = "deferTools", default, skip_serializing_if = "is_model_visible")] pub defer_tools: ToolExposure, } @@ -133,11 +146,14 @@ pub enum ParseError { #[source] source: InvalidHeaderValue, }, + + #[error("A server can set at most one of `oauth`, `sigv4` or an `Authorization` header")] + ConflictingAuth, } -pub fn http_transport(url: String, headers: HeaderMap, oauth: Option) -> Transport { - let oauth = (!headers.contains_key(AUTHORIZATION)).then(|| oauth.unwrap_or_default().into_settings()); - Transport::Http { url, headers, oauth } +pub fn http_transport(url: String, headers: HeaderMap) -> Transport { + let auth = default_auth(&headers); + Transport::Http { url, headers, auth } } impl McpConfig { @@ -224,6 +240,18 @@ impl McpOAuthConfig { } } +impl McpSigV4Config { + fn into_signer(self, vars: &Vars) -> Result { + let profile = self.profile.map(|profile| vars.expand(&profile)).transpose()?; + Ok(SigV4Signer::new(SigV4Config { + service: vars.expand(&self.service)?, + region: vars.expand(&self.region)?, + credentials: SigV4Credentials::DefaultChain { profile }, + refresh_before_expiry: Duration::from_mins(5), + })) + } +} + impl RemoteServerConfig { pub fn into_transport(self, vars: &Vars) -> Result { let mut headers = HeaderMap::with_capacity(self.headers.len()); @@ -238,8 +266,14 @@ impl RemoteServerConfig { header_value.set_sensitive(true); headers.insert(header_name, header_value); } - let oauth = self.oauth.map(|oauth| expand_oauth(vars, oauth)).transpose()?; - Ok(http_transport(vars.expand(&self.url)?, headers, oauth)) + let has_authorization = headers.contains_key(AUTHORIZATION); + let auth = match (self.oauth, self.sigv4) { + (Some(oauth), None) if !has_authorization => HttpAuth::OAuth(expand_oauth(vars, oauth)?.into_settings()), + (None, Some(sigv4)) if !has_authorization => HttpAuth::SigV4(sigv4.into_signer(vars)?), + (None, None) => default_auth(&headers), + _ => return Err(ParseError::ConflictingAuth), + }; + Ok(Transport::Http { url: vars.expand(&self.url)?, headers, auth }) } } @@ -254,6 +288,14 @@ fn expand_all(vars: &Vars, values: &[String]) -> Result, VarError> { values.iter().map(|value| vars.expand(value)).collect() } +fn default_auth(headers: &HeaderMap) -> HttpAuth { + if headers.contains_key(AUTHORIZATION) { + HttpAuth::None + } else { + HttpAuth::OAuth(McpOAuthConfig::default().into_settings()) + } +} + fn expand_oauth(vars: &Vars, oauth: McpOAuthConfig) -> Result { let expand = |value: Option| value.map(|value| vars.expand(&value)).transpose(); Ok(McpOAuthConfig { diff --git a/crates/aether-core/src/mcp/mod.rs b/crates/aether-core/src/mcp/mod.rs index ef0d6cf8b..e39529c14 100644 --- a/crates/aether-core/src/mcp/mod.rs +++ b/crates/aether-core/src/mcp/mod.rs @@ -7,8 +7,8 @@ pub(crate) mod tool_bridge; mod tool_filter; pub use config::{ - InMemoryServerConfig, InMemoryType, McpConfig, McpOAuthConfig, McpServerConfig, ParseError, RemoteServerConfig, - RemoteType, StdioServerConfig, StdioType, http_transport, + InMemoryServerConfig, InMemoryType, McpConfig, McpOAuthConfig, McpServerConfig, McpSigV4Config, ParseError, + RemoteServerConfig, RemoteType, StdioServerConfig, StdioType, http_transport, }; pub(crate) use deferred_tools::PROGRESSIVE_DISCOVERY_INSTRUCTIONS; pub use deferred_tools::{ diff --git a/crates/aether-core/tests/mcp/config_parser_tests.rs b/crates/aether-core/tests/mcp/config_parser_tests.rs index 8616dbea5..ec398535f 100644 --- a/crates/aether-core/tests/mcp/config_parser_tests.rs +++ b/crates/aether-core/tests/mcp/config_parser_tests.rs @@ -1,12 +1,10 @@ -use aether_auth::OAuthClientRegistration; +use aether_auth::{OAuthClientRegistration, SigV4Credentials}; use aether_core::mcp::{ - McpConfig, McpOAuthConfig, McpServerConfig, ParseError, StdioServerConfig, ToolExposure, ToolFilter, ToolMatcher, - http_transport, + McpConfig, McpServerConfig, ParseError, StdioServerConfig, ToolExposure, ToolFilter, ToolMatcher, http_transport, }; use http::header::{AUTHORIZATION, HeaderMap, HeaderValue}; -use mcp_utils::client::{OAuthSettings, Transport}; +use mcp_utils::client::{HttpAuth, OAuthSettings, Transport}; use mcp_utils::model::{Tool, ToolAnnotations}; -use std::num::NonZeroU16; use std::path::{Path, PathBuf}; use std::sync::Arc; use tempfile::tempdir; @@ -266,7 +264,7 @@ fn invalid_http_headers_fail_resolution() { #[test] fn http_servers_sign_in_with_aether_client_metadata_by_default() { - let oauth = http_oauth(None, HeaderMap::new()).expect("OAuth is offered"); + let oauth = http_oauth(remote(r#""headers":{}"#).unwrap()).expect("OAuth is offered"); assert!(matches!( &oauth.registration, @@ -277,13 +275,9 @@ fn http_servers_sign_in_with_aether_client_metadata_by_default() { #[test] fn http_servers_prefer_a_preregistered_client_id_over_client_metadata() { - let config = McpOAuthConfig { - client_id: Some("registered".to_string()), - client_metadata_url: Some("https://client.example/oauth/client.json".to_string()), - callback_port: NonZeroU16::new(4000), - }; + let oauth = r#""oauth":{"clientId":"registered","clientMetadataUrl":"https://client.example/oauth/client.json","callbackPort":4000}"#; - let oauth = http_oauth(Some(config), HeaderMap::new()).expect("OAuth is offered"); + let oauth = http_oauth(remote(oauth).unwrap()).expect("OAuth is offered"); assert!(matches!(&oauth.registration, OAuthClientRegistration::PreRegistered(id) if id == "registered")); assert_eq!(oauth.callback_port, 4000); @@ -293,7 +287,8 @@ fn http_servers_prefer_a_preregistered_client_id_over_client_metadata() { fn http_servers_with_an_authorization_header_never_sign_in() { let headers = HeaderMap::from_iter([(AUTHORIZATION, HeaderValue::from_static("Bearer static"))]); - assert!(http_oauth(None, headers).is_none()); + assert!(http_oauth(http_transport("https://example.com/mcp".to_string(), headers)).is_none()); + assert!(http_oauth(remote(r#""headers":{"Authorization":"Bearer static"}"#).unwrap()).is_none()); } #[test] @@ -303,7 +298,7 @@ fn oauth_settings_expand_variables() { let json = format!( r#"{{"servers":{{"linear":{{"type":"http","url":"https://example.com/mcp","oauth":{settings}}}}}}}"# ); - let Transport::Http { oauth: Some(oauth), .. } = transport(&json, &vars) else { + let Transport::Http { auth: HttpAuth::OAuth(oauth), .. } = transport(&json, &vars) else { panic!("expected HTTP transport with OAuth settings") }; oauth @@ -320,6 +315,64 @@ fn oauth_settings_expand_variables() { )); } +#[test] +fn sigv4_servers_sign_requests_for_the_configured_service_region_and_profile() { + let vars = vars(&[("AWS_REGION", "us-west-2"), ("AWS_PROFILE", "gateway")]); + let json = r#"{"servers":{"gateway":{"type":"http","url":"https://abc.execute-api.us-west-2.amazonaws.com/mcp","sigv4":{"service":"execute-api","region":"${AWS_REGION}","profile":"$AWS_PROFILE"}}}}"#; + + let Transport::Http { auth: HttpAuth::SigV4(signer), .. } = transport(json, &vars) else { + panic!("expected HTTP transport with a SigV4 signer") + }; + + let config = signer.config(); + assert_eq!(config.service, "execute-api"); + assert_eq!(config.region, "us-west-2"); + assert!( + matches!(&config.credentials, SigV4Credentials::DefaultChain { profile: Some(profile) } if profile == "gateway") + ); +} + +#[test] +fn sigv4_servers_without_a_profile_use_the_default_credential_chain() { + let json = r#"{"servers":{"gateway":{"type":"http","url":"https://example.com/mcp","sigv4":{"service":"lambda","region":"us-east-1"}}}}"#; + + let Transport::Http { auth: HttpAuth::SigV4(signer), .. } = transport(json, &vars(&[])) else { + panic!("expected HTTP transport with a SigV4 signer") + }; + + assert!(matches!(signer.config().credentials, SigV4Credentials::DefaultChain { profile: None })); +} + +#[test] +fn sigv4_requires_a_service_and_region_and_rejects_unknown_fields() { + for sigv4 in [ + r#"{"region":"us-east-1"}"#, + r#"{"service":"execute-api"}"#, + r#"{"service":"execute-api","region":"us-east-1","servce":"lambda"}"#, + ] { + assert!(remote(&format!(r#""sigv4":{sigv4}"#)).is_err(), "invalid sigv4 config was accepted: {sigv4}"); + } +} + +#[test] +fn a_server_sets_at_most_one_of_oauth_sigv4_or_an_authorization_header() { + let oauth = r#""oauth":{"clientId":"id"}"#; + let sigv4 = r#""sigv4":{"service":"execute-api","region":"us-east-1"}"#; + let authorization = r#""headers":{"Authorization":"Bearer static"}"#; + let lowercase_authorization = r#""headers":{"authorization":"Bearer static"}"#; + + for server in [ + format!("{oauth},{sigv4}"), + format!("{authorization},{oauth}"), + format!("{authorization},{sigv4}"), + format!("{lowercase_authorization},{sigv4}"), + format!("{authorization},{oauth},{sigv4}"), + ] { + let result = remote(&server); + assert!(matches!(result, Err(ParseError::ConflictingAuth)), "{server} gave {result:?}"); + } +} + fn vars(pairs: &[(&str, &str)]) -> Vars { pairs.iter().fold(Vars::new().with_env_lookup(|_| None), |vars, (name, value)| vars.with(*name, *value)) } @@ -357,11 +410,21 @@ fn command(command: &str, exposure: ToolExposure) -> (String, ToolExposure) { (command.to_string(), exposure) } -fn http_oauth(oauth: Option, headers: HeaderMap) -> Option { - let Transport::Http { oauth, .. } = http_transport("https://example.com/mcp".to_string(), headers, oauth) else { - panic!("expected http") +fn http_oauth(transport: Transport) -> Option { + match transport { + Transport::Http { auth: HttpAuth::OAuth(oauth), .. } => Some(oauth), + Transport::Http { .. } => None, + other => panic!("expected http, got {other:?}"), + } +} + +fn remote(server: &str) -> Result { + let json = format!(r#"{{"servers":{{"remote":{{"type":"http","url":"https://example.com/mcp",{server}}}}}}}"#); + let McpServerConfig::Remote(config) = McpConfig::from_json(&json)?.servers.pop_first().expect("one server").1 + else { + panic!("expected a remote server") }; - oauth + config.into_transport(&vars(&[])) } fn deferred(allow: &[&str], deny: &[&str]) -> ToolExposure { diff --git a/crates/llm/Cargo.toml b/crates/llm/Cargo.toml index 3fe2a0f47..68ccec2a6 100644 --- a/crates/llm/Cargo.toml +++ b/crates/llm/Cargo.toml @@ -26,11 +26,11 @@ workspace = true [features] default = [] bedrock = [ + "dep:aether-auth", + "aether-auth/sigv4", "dep:aws-config", "dep:aws-credential-types", "dep:aws-sdk-bedrockruntime", - "dep:aws-sigv4", - "dep:aws-smithy-runtime-api", "dep:aws-smithy-types", "dep:base64", "dep:time", @@ -43,8 +43,6 @@ async-stream = { workspace = true } aws-config = { workspace = true, optional = true } aws-credential-types = { workspace = true, optional = true } aws-sdk-bedrockruntime = { workspace = true, optional = true } -aws-sigv4 = { workspace = true, optional = true } -aws-smithy-runtime-api = { workspace = true, optional = true } aws-smithy-types = { workspace = true, optional = true } base64 = { workspace = true, optional = true } chrono = { workspace = true } diff --git a/crates/llm/src/providers/bedrock/mantle/auth.rs b/crates/llm/src/providers/bedrock/mantle/auth.rs index 46ad748f4..718495b87 100644 --- a/crates/llm/src/providers/bedrock/mantle/auth.rs +++ b/crates/llm/src/providers/bedrock/mantle/auth.rs @@ -1,10 +1,7 @@ -use std::time::SystemTime; - -use aws_credential_types::provider::{ProvideCredentials, SharedCredentialsProvider}; -use aws_sigv4::http_request::{SignableBody, SignableRequest, SigningSettings, sign}; -use aws_sigv4::sign::v4; -use aws_smithy_runtime_api::client::identity::Identity; -use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderName, HeaderValue}; +use aether_auth::{SigV4Config, SigV4Credentials, SigV4Signer}; +use aws_credential_types::provider::SharedCredentialsProvider; +use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderValue}; +use std::time::Duration; use crate::{LlmError, Result}; @@ -17,7 +14,7 @@ pub enum MantleAuth { /// `AWS_BEARER_TOKEN_BEDROCK`, or a token supplied programmatically. BearerToken(String), /// `SigV4` over credentials resolved from the AWS credential chain. - SigV4 { credentials: SharedCredentialsProvider, region: String }, + SigV4(SigV4Signer), /// Send the request unsigned. None, } @@ -26,6 +23,15 @@ pub const BEARER_TOKEN_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK"; const SIGNING_SERVICE: &str = "bedrock"; impl MantleAuth { + pub fn sigv4(credentials: SharedCredentialsProvider, region: &str) -> Self { + Self::SigV4(SigV4Signer::new(SigV4Config { + service: SIGNING_SERVICE.to_string(), + region: region.to_string(), + credentials: SigV4Credentials::Provider(credentials), + refresh_before_expiry: Duration::from_mins(5), + })) + } + /// Read a bearer token from the environment, ignoring an empty value. pub fn bearer_token_from_env() -> Option { std::env::var(BEARER_TOKEN_ENV_VAR).ok().filter(|token| !token.is_empty()) @@ -41,7 +47,9 @@ impl MantleAuth { headers.insert(AUTHORIZATION, value); Ok(headers) } - Self::SigV4 { credentials, region } => sigv4_headers(credentials, region, method, url, body).await, + Self::SigV4(signer) => { + signer.headers(method, url, body).await.map_err(|e| LlmError::ProviderRequest(e.to_string())) + } Self::None => Ok(HeaderMap::new()), } } @@ -51,50 +59,8 @@ impl std::fmt::Display for MantleAuth { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str(match self { Self::BearerToken(_) => "bearer token", - Self::SigV4 { .. } => "SigV4", + Self::SigV4(_) => "SigV4", Self::None => "unsigned", }) } } - -async fn sigv4_headers( - credentials: &SharedCredentialsProvider, - region: &str, - method: &str, - url: &str, - body: &[u8], -) -> Result { - let credentials = credentials - .provide_credentials() - .await - .map_err(|e| LlmError::ProviderRequest(format!("Could not resolve AWS credentials: {e}")))?; - let identity = Identity::from(credentials); - - let params = v4::SigningParams::builder() - .identity(&identity) - .region(region) - .name(SIGNING_SERVICE) - .time(SystemTime::now()) - .settings(SigningSettings::default()) - .build() - .map_err(|e| LlmError::ProviderRequest(format!("Could not build SigV4 signing params: {e}")))? - .into(); - - let signable = SignableRequest::new(method, url, std::iter::empty(), SignableBody::Bytes(body)) - .map_err(|e| LlmError::ProviderRequest(format!("Could not build signable request: {e}")))?; - - let (instructions, _signature) = sign(signable, ¶ms) - .map_err(|e| LlmError::ProviderRequest(format!("SigV4 signing failed: {e}")))? - .into_parts(); - - let mut headers = HeaderMap::new(); - for (name, value) in instructions.headers() { - let name = HeaderName::try_from(name) - .map_err(|e| LlmError::ProviderRequest(format!("Invalid signed header name: {e}")))?; - let mut value = HeaderValue::from_str(value) - .map_err(|e| LlmError::ProviderRequest(format!("Invalid signed header value: {e}")))?; - value.set_sensitive(true); - headers.insert(name, value); - } - Ok(headers) -} diff --git a/crates/llm/src/providers/bedrock/provider.rs b/crates/llm/src/providers/bedrock/provider.rs index 636326d11..469149f74 100644 --- a/crates/llm/src/providers/bedrock/provider.rs +++ b/crates/llm/src/providers/bedrock/provider.rs @@ -295,7 +295,7 @@ fn mantle_auth(credentials: Option, region: &str) -> return MantleAuth::BearerToken(token); } match credentials { - Some(credentials) => MantleAuth::SigV4 { credentials, region: region.to_string() }, + Some(credentials) => MantleAuth::sigv4(credentials, region), None => MantleAuth::None, } } diff --git a/crates/mcp-utils/Cargo.toml b/crates/mcp-utils/Cargo.toml index e54294c4d..51c6a8df8 100644 --- a/crates/mcp-utils/Cargo.toml +++ b/crates/mcp-utils/Cargo.toml @@ -29,11 +29,13 @@ testing = ["client", "dep:axum", "rmcp/transport-streamable-http-server"] client = [ "dep:aether-auth", "aether-auth/mcp", + "aether-auth/sigv4", "dep:async-stream", "dep:chrono", "dep:reqwest", "dep:serde", "dep:serde_json", + "dep:sse-stream", "dep:tokio-stream", "rmcp/transport-streamable-http-client-reqwest", "rmcp/transport-child-process", @@ -56,6 +58,7 @@ chrono = { workspace = true, optional = true } reqwest = { workspace = true, optional = true } serde = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } +sse-stream = { workspace = true, optional = true } tokio-stream = { workspace = true, features = ["sync"], optional = true } [[test]] @@ -76,9 +79,15 @@ required-features = ["testing"] [[test]] name = "http_headers" -required-features = ["client"] +required-features = ["testing"] + +[[test]] +name = "sigv4" +required-features = ["testing"] [dev-dependencies] +aws-credential-types = { workspace = true } +aws-sigv4 = { workspace = true } axum = { workspace = true } rmcp = { workspace = true, features = ["macros", "transport-streamable-http-server"] } tokio = { workspace = true, features = ["rt-multi-thread", "test-util"] } diff --git a/crates/mcp-utils/src/client/mod.rs b/crates/mcp-utils/src/client/mod.rs index 6c6a68972..aa66c4bdb 100644 --- a/crates/mcp-utils/src/client/mod.rs +++ b/crates/mcp-utils/src/client/mod.rs @@ -6,6 +6,7 @@ mod handler; mod mcp_client; mod oauth; mod options; +mod sigv4; mod task; mod tool_call; mod transport; @@ -21,4 +22,4 @@ pub use reqwest::header::HeaderMap; pub use task::TaskErrorReason; pub use tokio_util::sync::CancellationToken; pub use tool_call::{ToolCall, ToolCallError, ToolCallEvent, ToolCallOptions}; -pub use transport::Transport; +pub use transport::{HttpAuth, Transport}; diff --git a/crates/mcp-utils/src/client/oauth.rs b/crates/mcp-utils/src/client/oauth.rs index b3dc00b56..904f49520 100644 --- a/crates/mcp-utils/src/client/oauth.rs +++ b/crates/mcp-utils/src/client/oauth.rs @@ -1,7 +1,7 @@ use super::error::{ClientError, Result}; use super::{ClientOptions, Elicitation, elicitation::ElicitationRequest}; use aether_auth::{ - OAuthClientRegistration, OAuthCredentialStorage, OAuthError, OAuthFlowOptions, OAuthHandler, accept_oauth_callback, + OAuthClientRegistration, OAuthError, OAuthFlowOptions, OAuthHandler, accept_oauth_callback, create_auth_manager_from_store, perform_oauth_flow, }; use futures::future::BoxFuture; @@ -23,8 +23,9 @@ pub(super) async fn restore( server: &str, url: &str, oauth: &OAuthSettings, - store: &Arc, + options: &ClientOptions, ) -> Option> { + let store = options.oauth_store.as_ref()?; let restored = create_auth_manager_from_store( server, url, @@ -34,7 +35,10 @@ pub(super) async fn restore( ) .await; match restored { - Ok(manager) => manager.map(|manager| AuthClient::new(reqwest::Client::default(), manager)), + Ok(manager) => manager.map(|manager| { + tracing::debug!(server, "Using stored OAuth credentials"); + AuthClient::new(reqwest::Client::default(), manager) + }), Err(error) => { tracing::warn!( server, diff --git a/crates/mcp-utils/src/client/sigv4.rs b/crates/mcp-utils/src/client/sigv4.rs new file mode 100644 index 000000000..1f1e85548 --- /dev/null +++ b/crates/mcp-utils/src/client/sigv4.rs @@ -0,0 +1,125 @@ +use aether_auth::SigV4Signer; +use futures::stream::BoxStream; +use reqwest::header::{HeaderName, HeaderValue}; +use rmcp::model::ClientJsonRpcMessage; +use rmcp::transport::auth::AuthError; +use rmcp::transport::streamable_http_client::{ + SseError, StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, +}; +use sse_stream::Sse; +use std::collections::HashMap; +use std::sync::Arc; + +#[derive(Clone)] +pub(super) struct SigV4HttpClient { + http: reqwest::Client, + signer: SigV4Signer, +} + +impl SigV4HttpClient { + pub(super) fn new(signer: SigV4Signer) -> Self { + Self { http: reqwest::Client::default(), signer } + } + + async fn sign_message( + &self, + uri: &str, + message: &ClientJsonRpcMessage, + headers: HashMap, + ) -> Result, StreamableHttpError> { + self.sign("POST", uri, &serde_json::to_vec(message)?, headers).await + } + + async fn sign( + &self, + method: &str, + uri: &str, + body: &[u8], + mut headers: HashMap, + ) -> Result, StreamableHttpError> { + let signed = self + .signer + .headers(method, uri, body) + .await + .map_err(|error| StreamableHttpError::Auth(AuthError::AuthorizationFailed(error.to_string())))?; + + headers.extend(signed.into_iter().filter_map(|(name, value)| Some((name?, value)))); + Ok(headers) + } +} + +impl StreamableHttpClient for SigV4HttpClient { + type Error = reqwest::Error; + + async fn post_message( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_header: Option, + custom_headers: HashMap, + ) -> Result> { + let headers = self.sign_message(&uri, &message, custom_headers).await?; + self.http.post_message(uri, message, session_id, auth_header, headers).await + } + + async fn post_message_with_max_sse_event_size( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + auth_header: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result> { + let headers = self.sign_message(&uri, &message, custom_headers).await?; + self.http + .post_message_with_max_sse_event_size(uri, message, session_id, auth_header, headers, max_sse_event_size) + .await + } + + async fn delete_session( + &self, + uri: Arc, + session_id: Arc, + auth_header: Option, + custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + let headers = self.sign("DELETE", &uri, &[], custom_headers).await?; + self.http.delete_session(uri, session_id, auth_header, headers).await + } + + async fn get_stream( + &self, + uri: Arc, + session_id: Option>, + last_event_id: Option, + auth_header: Option, + custom_headers: HashMap, + ) -> Result>, StreamableHttpError> { + let headers = self.sign("GET", &uri, &[], custom_headers).await?; + self.http.get_stream(uri, session_id, last_event_id, auth_header, headers).await + } + + async fn get_stream_with_max_sse_event_size( + &self, + uri: Arc, + session_id: Option>, + last_event_id: Option, + auth_header: Option, + custom_headers: HashMap, + max_sse_event_size: usize, + ) -> Result>, StreamableHttpError> { + let headers = self.sign("GET", &uri, &[], custom_headers).await?; + self.http + .get_stream_with_max_sse_event_size( + uri, + session_id, + last_event_id, + auth_header, + headers, + max_sse_event_size, + ) + .await + } +} diff --git a/crates/mcp-utils/src/client/transport.rs b/crates/mcp-utils/src/client/transport.rs index ee823bb07..f397d801b 100644 --- a/crates/mcp-utils/src/client/transport.rs +++ b/crates/mcp-utils/src/client/transport.rs @@ -1,10 +1,12 @@ use super::error::{ClientError, Result}; +use super::sigv4::SigV4HttpClient; use super::{ ClientOptions, handler::Handler, oauth::{self, OAuthSettings}, }; use crate::server::McpServer; +use aether_auth::SigV4Signer; use reqwest::header::HeaderMap; use rmcp::{ ClientLifecycleMode, RoleClient, @@ -29,16 +31,23 @@ use tokio::{ #[derive(Debug, Clone)] pub enum Transport { Stdio { command: String, args: Vec, env: HashMap, cwd: Option }, - Http { url: String, headers: HeaderMap, oauth: Option }, + Http { url: String, headers: HeaderMap, auth: HttpAuth }, InProcess(McpServer), Unix(PathBuf), } +#[derive(Debug, Clone)] +pub enum HttpAuth { + None, + OAuth(OAuthSettings), + SigV4(SigV4Signer), +} + pub(super) type ClientService = RunningService>; impl Transport { pub(crate) fn supports_oauth(&self, options: &ClientOptions) -> bool { - options.oauth_prompts().is_some() && matches!(self, Self::Http { oauth: Some(_), .. }) + options.oauth_prompts().is_some() && matches!(self, Self::Http { auth: HttpAuth::OAuth(_), .. }) } pub(super) async fn connect(self, handler: Handler, options: &ClientOptions) -> Result { @@ -46,7 +55,7 @@ impl Transport { let handler = Arc::new(handler); let connected = match self { Self::Stdio { command, args, env, cwd } => connect_stdio(handler, command, args, env, cwd).await, - Self::Http { url, headers, oauth } => connect_http(handler, url, headers, oauth, options).await, + Self::Http { url, headers, auth } => connect_http(handler, url, headers, auth, options).await, Self::InProcess(server) => serve(handler, server.serve_in_memory()).await, Self::Unix(path) => connect_unix(handler, path).await, }; @@ -65,7 +74,7 @@ impl Transport { challenge: Option, ) -> Result { let server = handler.server().to_string(); - let Self::Http { url, headers, oauth: Some(oauth) } = self else { + let Self::Http { url, headers, auth: HttpAuth::OAuth(oauth) } = self else { return Err(ClientError::OAuthUnavailable { server }); }; let auth_client = oauth::authorize(&server, &url, oauth, options, challenge).await?; @@ -115,22 +124,19 @@ async fn connect_http( handler: Arc, url: String, headers: HeaderMap, - oauth: Option, + auth: HttpAuth, options: &ClientOptions, ) -> Result { - let server = handler.server(); - let restored = match (options.oauth_store.as_ref(), &oauth) { - (Some(store), Some(oauth)) => oauth::restore(server, &url, oauth, store).await, - _ => None, - }; - let config = http_config(url, headers); - match restored { - Some(auth_client) => { - tracing::debug!(server, "Using stored OAuth credentials"); - serve(handler, StreamableHttpClientTransport::with_client(auth_client, config)).await + match auth { + HttpAuth::None => serve(handler, StreamableHttpClientTransport::from_config(config)).await, + HttpAuth::OAuth(oauth) => match oauth::restore(handler.server(), &config.uri, &oauth, options).await { + Some(auth_client) => serve(handler, StreamableHttpClientTransport::with_client(auth_client, config)).await, + None => serve(handler, StreamableHttpClientTransport::from_config(config)).await, + }, + HttpAuth::SigV4(signer) => { + serve(handler, StreamableHttpClientTransport::with_client(SigV4HttpClient::new(signer), config)).await } - None => serve(handler, StreamableHttpClientTransport::from_config(config)).await, } } diff --git a/crates/mcp-utils/src/testing.rs b/crates/mcp-utils/src/testing.rs index 77375a841..3cf8df8ca 100644 --- a/crates/mcp-utils/src/testing.rs +++ b/crates/mcp-utils/src/testing.rs @@ -4,6 +4,7 @@ pub use elicitation_script::{ pub use fake_auth_server::{ ACCESS_TOKEN, CLIENT_METADATA_URL, FakeAuthServer, OAuthPrompt, client_metadata, complete_authorization, }; +pub use fake_http_server::{FakeHttpServer, RecordedRequest}; pub use fake_mcp::{ CapturedTaskUpdate, CapturedToolCall, FakeMcpServer, FakeMcpState, FakeTool, FakeToolResponse, completed_task_payload, @@ -16,6 +17,7 @@ pub use gate::Gate; mod elicitation_script; mod fake_auth_server; +mod fake_http_server; mod fake_mcp; mod fake_task; mod gate; diff --git a/crates/mcp-utils/src/testing/fake_auth_server.rs b/crates/mcp-utils/src/testing/fake_auth_server.rs index 4bdd80ea2..2afb2a49d 100644 --- a/crates/mcp-utils/src/testing/fake_auth_server.rs +++ b/crates/mcp-utils/src/testing/fake_auth_server.rs @@ -1,28 +1,20 @@ -use super::FakeMcpServer; -use crate::client::{Elicitation, ElicitationRequest, OAuthSettings, Transport}; +use super::fake_http_server::origin; +use super::{FakeHttpServer, FakeMcpServer, RecordedRequest}; +use crate::client::{Elicitation, ElicitationRequest, HttpAuth, OAuthSettings, Transport}; use crate::gateway::ServerSpec; use aether_auth::OAuthClientRegistration; -use axum::extract::{Request, State}; use axum::http::{StatusCode, Uri, header::WWW_AUTHENTICATE}; -use axum::middleware::{Next, from_fn_with_state}; -use axum::response::{IntoResponse, Response}; +use axum::response::IntoResponse; use axum::routing::post; use axum::{Json, Router}; use reqwest::Url; use reqwest::header::{AUTHORIZATION, HeaderMap}; use rmcp::model::{ElicitRequestParams, ElicitResult, ElicitationAction}; -use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; -use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; use serde_json::{Value, json}; -use std::sync::{Arc, Mutex}; -use tokio::net::TcpListener; use tokio::sync::mpsc; -use tokio::task::JoinHandle; pub struct FakeAuthServer { - url: String, - auth: Arc, - task: JoinHandle<()>, + http: FakeHttpServer, } pub struct OAuthPrompt { @@ -34,33 +26,16 @@ pub const CLIENT_METADATA_URL: &str = "https://client.example/oauth/client-metad impl FakeAuthServer { pub async fn start(server: FakeMcpServer) -> Self { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind the fake authorization server"); - let origin = format!("http://{}", listener.local_addr().expect("a bound address")); - let auth = Arc::new(AuthState { origin: origin.clone(), bearers: Mutex::new(Vec::new()) }); - let mcp = StreamableHttpService::new( - move || Ok(server.clone()), - Arc::new(LocalSessionManager::default()), - StreamableHttpServerConfig::default(), - ); - let protected = - Router::new().nest_service("/mcp", mcp).layer(from_fn_with_state(Arc::clone(&auth), require_token)); - let app = Router::new() - .route("/token", post(token)) - .fallback(metadata) - .with_state(Arc::clone(&auth)) - .merge(protected); - let task = tokio::spawn(async move { - axum::serve(listener, app).await.expect("serve the fake authorization server"); - }); - Self { url: format!("{origin}/mcp"), auth, task } + let routes = Router::new().route("/token", post(token)).fallback(metadata); + Self { http: FakeHttpServer::serve(server, routes, require_token).await } } pub fn transport(&self) -> Transport { - self.transport_with(Some(client_metadata())) + self.transport_with(HttpAuth::OAuth(client_metadata())) } - pub fn transport_with(&self, oauth: Option) -> Transport { - Transport::Http { url: self.url.clone(), headers: HeaderMap::new(), oauth } + pub fn transport_with(&self, auth: HttpAuth) -> Transport { + self.http.transport(HeaderMap::new(), auth) } pub fn spec(&self, name: &str) -> ServerSpec { @@ -68,13 +43,11 @@ impl FakeAuthServer { } pub fn received_bearer(&self, token: &str) -> bool { - self.auth.bearers.lock().expect("bearers lock").iter().any(|bearer| bearer == token) - } -} - -impl Drop for FakeAuthServer { - fn drop(&mut self) { - self.task.abort(); + let bearer = format!("Bearer {token}"); + self.http + .requests() + .iter() + .any(|request| request.headers.get(AUTHORIZATION).is_some_and(|value| *value == bearer)) } } @@ -130,34 +103,22 @@ pub fn client_metadata() -> OAuthSettings { pub const ACCESS_TOKEN: &str = "access-token"; -struct AuthState { - origin: String, - bearers: Mutex>, -} - -async fn require_token(State(auth): State>, request: Request, next: Next) -> Response { - let bearer = request - .headers() - .get(AUTHORIZATION) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.strip_prefix("Bearer ")) - .map(str::to_string); - if let Some(bearer) = bearer { - auth.bearers.lock().expect("bearers lock").push(bearer.clone()); - if bearer == ACCESS_TOKEN { - return next.run(request).await; - } +fn require_token(request: &RecordedRequest) -> Result<(), impl IntoResponse + use<>> { + let bearer = format!("Bearer {ACCESS_TOKEN}"); + if request.headers.get(AUTHORIZATION).is_some_and(|value| *value == bearer) { + return Ok(()); } - let challenge = format!("Bearer resource_metadata=\"{}/.well-known/oauth-protected-resource/mcp\"", auth.origin); - (StatusCode::UNAUTHORIZED, [(WWW_AUTHENTICATE, challenge)]).into_response() + let challenge = + format!("Bearer resource_metadata=\"{}/.well-known/oauth-protected-resource/mcp\"", origin(&request.headers)); + Err((StatusCode::UNAUTHORIZED, [(WWW_AUTHENTICATE, challenge)])) } async fn token() -> Json { Json(json!({ "access_token": ACCESS_TOKEN, "token_type": "Bearer", "expires_in": 3600 })) } -async fn metadata(State(auth): State>, uri: Uri) -> Json { - let origin = &auth.origin; +async fn metadata(headers: HeaderMap, uri: Uri) -> Json { + let origin = origin(&headers); if uri.path().contains("oauth-protected-resource") { return Json(json!({ "resource": format!("{origin}/mcp"), "authorization_servers": [origin] })); } diff --git a/crates/mcp-utils/src/testing/fake_http_server.rs b/crates/mcp-utils/src/testing/fake_http_server.rs new file mode 100644 index 000000000..c82dc6c17 --- /dev/null +++ b/crates/mcp-utils/src/testing/fake_http_server.rs @@ -0,0 +1,125 @@ +use crate::client::{HeaderMap, HttpAuth, McpClient, Transport}; +use axum::Router; +use axum::body::{Body, Bytes, to_bytes}; +use axum::extract::{OriginalUri, Request, State}; +use axum::http::{Method, StatusCode, header::HOST}; +use axum::middleware::{Next, from_fn_with_state}; +use axum::response::{IntoResponse, Response}; +use rmcp::ServerHandler; +use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; +use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; +use std::sync::Arc; +use tokio::net::TcpListener; +use tokio::sync::watch; +use tokio::task::JoinHandle; + +/// Serves an MCP server over streamable HTTP at `/mcp` and records every request it receives there. +pub struct FakeHttpServer { + url: String, + requests: watch::Receiver>, + task: JoinHandle<()>, +} + +/// A request as the server received it. +#[derive(Debug, Clone)] +pub struct RecordedRequest { + pub method: Method, + pub url: String, + pub headers: HeaderMap, + pub body: Bytes, +} + +impl FakeHttpServer { + pub async fn start(server: impl ServerHandler + Clone) -> Self { + Self::start_guarded(server, |_| Ok::<_, StatusCode>(())).await + } + + /// Like [`Self::start`], but answers requests that `guard` rejects with its response. + pub async fn start_guarded( + server: impl ServerHandler + Clone, + guard: impl Fn(&RecordedRequest) -> Result<(), R> + Send + Sync + 'static, + ) -> Self { + Self::serve(server, Router::new(), guard).await + } + + pub fn transport(&self, headers: HeaderMap, auth: HttpAuth) -> Transport { + Transport::Http { url: self.url.clone(), headers, auth } + } + + /// Every request received so far, including rejected ones. + pub fn requests(&self) -> Vec { + self.requests.borrow().clone() + } + + /// Every request of a legacy session: waits for `client` to open its GET stream, then closes it so it sends DELETE. + pub async fn record_session(&self, client: McpClient) -> Vec { + let mut requests = self.requests.clone(); + let opened = requests.wait_for(|requests| requests.iter().any(|request| request.method == Method::GET)).await; + drop(opened.expect("the fake HTTP server stopped")); + client.close().await; + self.requests() + } + + /// Like [`Self::start_guarded`], and also serves `routes` beside the guarded `/mcp` endpoint. + pub(super) async fn serve( + server: impl ServerHandler + Clone, + routes: Router, + guard: impl Fn(&RecordedRequest) -> Result<(), R> + Send + Sync + 'static, + ) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind the fake HTTP server"); + let url = format!("http://{}/mcp", listener.local_addr().expect("a bound address")); + let (sender, requests) = watch::channel(Vec::new()); + let guard: Guard = Box::new(move |request| guard(request).err().map(IntoResponse::into_response)); + let mcp = StreamableHttpService::new( + move || Ok(server.clone()), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ); + let app = Router::new() + .nest_service("/mcp", mcp) + .layer(from_fn_with_state(Arc::new(Recorder { guard, requests: sender }), record)) + .merge(routes); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve the fake HTTP server"); + }); + Self { url, requests, task } + } +} + +impl Drop for FakeHttpServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +/// The origin a request was sent to, e.g. `http://127.0.0.1:4000`. +pub(super) fn origin(headers: &HeaderMap) -> String { + let host = headers.get(HOST).and_then(|host| host.to_str().ok()).unwrap_or_default(); + format!("http://{host}") +} + +/// Returns the response that rejects a request, or `None` to accept it. +type Guard = Box Option + Send + Sync>; + +struct Recorder { + guard: Guard, + requests: watch::Sender>, +} + +async fn record(State(recorder): State>, request: Request, next: Next) -> Response { + let (parts, body) = request.into_parts(); + let body = to_bytes(body, usize::MAX).await.expect("read the request body"); + let path = parts.extensions.get::().map_or(&parts.uri, |original| &original.0); + let received = RecordedRequest { + method: parts.method.clone(), + url: format!("{}{path}", origin(&parts.headers)), + headers: parts.headers.clone(), + body: body.clone(), + }; + let rejection = (recorder.guard)(&received); + recorder.requests.send_modify(|requests| requests.push(received)); + match rejection { + Some(rejection) => rejection, + None => next.run(Request::from_parts(parts, Body::from(body))).await, + } +} diff --git a/crates/mcp-utils/src/testing/fake_mcp.rs b/crates/mcp-utils/src/testing/fake_mcp.rs index d05d015e0..34b61a453 100644 --- a/crates/mcp-utils/src/testing/fake_mcp.rs +++ b/crates/mcp-utils/src/testing/fake_mcp.rs @@ -2,11 +2,11 @@ use rmcp::{ ErrorData as McpError, Peer, RoleServer, ServerHandler, model::{ CacheScope, CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, ClientCapabilities, - ContentBlock, CreateTaskResult, DetailedTask, DiscoverResult, ElicitResult, GetPromptRequestParams, - GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskResult, Implementation, ListPromptsResult, - ListToolsResult, PaginatedRequestParams, ProgressNotificationParam, ProgressToken, Prompt, PromptMessage, - PromptsCapability, ProtocolVersion, ResultType, Role, ServerCapabilities, ServerConfig, TaskPayload, Tool, - ToolAnnotations, UpdateTaskParams, + ContentBlock, CreateTaskResult, DetailedTask, DiscoverRequestMethod, DiscoverResult, ElicitResult, + GetPromptRequestParams, GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskResult, Implementation, + ListPromptsResult, ListToolsResult, PaginatedRequestParams, ProgressNotificationParam, ProgressToken, Prompt, + PromptMessage, PromptsCapability, ProtocolVersion, ResultType, Role, ServerCapabilities, ServerConfig, + TaskPayload, Tool, ToolAnnotations, UpdateTaskParams, }, service::RequestContext, }; @@ -122,6 +122,11 @@ impl FakeMcpServer { self } + pub fn with_legacy_protocol(self) -> Self { + self.state.lock().legacy_protocol = true; + self + } + pub fn state(&self) -> FakeMcpState { self.state.clone() } @@ -382,6 +387,9 @@ impl ServerHandler for FakeMcpServer { &self, context: RequestContext, ) -> impl Future> + Send + '_ { + if self.state.lock().legacy_protocol { + return std::future::ready(Err(McpError::method_not_found::())); + } self.state.lock().client_capabilities = context.meta.client_capabilities(); std::future::ready(Ok(DiscoverResult::from_server_info( ServerHandler::supported_protocol_versions(self).into_owned(), @@ -395,9 +403,12 @@ impl ServerHandler for FakeMcpServer { if !inner.prompts.is_empty() { capabilities.prompts = Some(PromptsCapability::default()); } - let config = ServerConfig::new(capabilities).with_server_info( + let mut config = ServerConfig::new(capabilities).with_server_info( Implementation::new("fake-mcp-server", "0.1.0").with_description("A fake MCP server for testing"), ); + if inner.legacy_protocol { + config = config.with_protocol_version(ProtocolVersion::V_2025_11_25); + } match &inner.instructions { Some(instructions) => config.with_instructions(instructions.clone()), None => config, @@ -559,6 +570,7 @@ struct FakeMcpStateInner { task_update_failures: usize, tool_list_requests: usize, tool_list_failures: usize, + legacy_protocol: bool, client: Option>, } diff --git a/crates/mcp-utils/tests/gateway.rs b/crates/mcp-utils/tests/gateway.rs index 01a179c1a..7392d1b8f 100644 --- a/crates/mcp-utils/tests/gateway.rs +++ b/crates/mcp-utils/tests/gateway.rs @@ -1,5 +1,7 @@ use futures::{FutureExt, StreamExt}; -use mcp_utils::client::{ClientError, ClientOptions, HeaderMap, OAuthSettings, ToolCall, ToolCallOptions, Transport}; +use mcp_utils::client::{ + ClientError, ClientOptions, HeaderMap, HttpAuth, OAuthSettings, ToolCall, ToolCallOptions, Transport, +}; use mcp_utils::gateway::{ConnectedServer, GatewayError, McpGateway, ServerSpec, ServerState}; use mcp_utils::model::{ContentBlock, GetPromptResult, Tool, ToolAnnotations}; use mcp_utils::server::{McpServer, UnixSocket}; @@ -282,7 +284,7 @@ fn stdio(name: &str, command: &str, args: &[&str]) -> ServerSpec { fn unreachable_http(name: &str) -> ServerSpec { let port = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port(); let url = format!("http://127.0.0.1:{port}/mcp"); - let transport = Transport::Http { url, headers: HeaderMap::new(), oauth: Some(OAuthSettings::default()) }; + let transport = Transport::Http { url, headers: HeaderMap::new(), auth: HttpAuth::OAuth(OAuthSettings::default()) }; ServerSpec { name: name.to_string(), transport } } diff --git a/crates/mcp-utils/tests/http_headers.rs b/crates/mcp-utils/tests/http_headers.rs index 1117bd831..1ba01f81f 100644 --- a/crates/mcp-utils/tests/http_headers.rs +++ b/crates/mcp-utils/tests/http_headers.rs @@ -1,113 +1,28 @@ -use axum::{ - Router, - extract::{Request, State}, - http::{HeaderMap, HeaderName, Method, header::AUTHORIZATION}, - middleware::{Next, from_fn_with_state}, - response::Response, -}; -use mcp_utils::client::{ClientOptions, McpClient, Transport}; -use mcp_utils::model::{DiscoverResult, ErrorCode, ErrorData, ProtocolVersion, ServerCapabilities, ServerConfig}; -use reqwest::header::{HeaderMap as ConfiguredHeaders, HeaderValue}; -use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; -use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; -use rmcp::{RoleServer, ServerHandler, service::RequestContext}; -use std::future::{Future, ready}; -use std::sync::Arc; -use tokio::{net::TcpListener, sync::mpsc, task::JoinHandle}; +use axum::http::{HeaderName, HeaderValue, Method, header::AUTHORIZATION}; +use mcp_utils::client::{ClientOptions, HeaderMap, HttpAuth, McpClient}; +use mcp_utils::testing::{FakeHttpServer, FakeMcpServer}; #[tokio::test] async fn every_request_of_a_session_sends_the_configured_headers() { - let mut server = SessionServer::start().await; - let transport = Transport::Http { url: server.url.clone(), headers: configured_headers(), oauth: None }; + let server = FakeHttpServer::start(FakeMcpServer::new().with_legacy_protocol()).await; + let transport = server.transport(configured_headers(), HttpAuth::None); let client = McpClient::connect("remote", transport, &ClientOptions::default()).await.unwrap(); - let mut requests = server.requests_until(Method::GET).await; - client.close().await; - requests.extend(server.received()); + let requests = server.record_session(client).await; for method in [Method::POST, Method::GET, Method::DELETE] { - let (_, headers) = - requests.iter().find(|(seen, _)| *seen == method).unwrap_or_else(|| panic!("the client sent no {method}")); - assert_eq!(headers.get_all(AUTHORIZATION).iter().count(), 1, "{method}"); - assert_eq!(headers[AUTHORIZATION], "Sentry-Bearer secret", "{method}"); - assert_eq!(headers["x-api-key"], "custom-secret", "{method}"); + let request = requests + .iter() + .find(|request| request.method == method) + .unwrap_or_else(|| panic!("the client sent no {method}")); + assert_eq!(request.headers.get_all(AUTHORIZATION).iter().count(), 1, "{method}"); + assert_eq!(request.headers[AUTHORIZATION], "Sentry-Bearer secret", "{method}"); + assert_eq!(request.headers["x-api-key"], "custom-secret", "{method}"); } } -struct SessionServer { - url: String, - requests: mpsc::UnboundedReceiver<(Method, HeaderMap)>, - task: JoinHandle<()>, -} - -impl SessionServer { - async fn start() -> Self { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let url = format!("http://{}/mcp", listener.local_addr().unwrap()); - let (sender, requests) = mpsc::unbounded_channel(); - let mcp = StreamableHttpService::new( - || Ok(LegacyServer), - Arc::new(LocalSessionManager::default()), - StreamableHttpServerConfig::default(), - ); - let app = Router::new().nest_service("/mcp", mcp).layer(from_fn_with_state(sender, record)); - let task = tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); - }); - Self { url, requests, task } - } - - async fn requests_until(&mut self, last: Method) -> Vec<(Method, HeaderMap)> { - let mut requests = Vec::new(); - loop { - let request = self.requests.recv().await.unwrap(); - let done = request.0 == last; - requests.push(request); - if done { - return requests; - } - } - } - - fn received(&mut self) -> Vec<(Method, HeaderMap)> { - std::iter::from_fn(|| self.requests.try_recv().ok()).collect() - } -} - -impl Drop for SessionServer { - fn drop(&mut self) { - self.task.abort(); - } -} - -#[derive(Clone)] -struct LegacyServer; - -impl ServerHandler for LegacyServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_protocol_version(ProtocolVersion::V_2025_11_25) - } - - fn discover( - &self, - _context: RequestContext, - ) -> impl Future> + Send + '_ { - ready(Err(ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "server/discover is not supported", None))) - } -} - -async fn record( - State(sender): State>, - request: Request, - next: Next, -) -> Response { - sender.send((request.method().clone(), request.headers().clone())).unwrap(); - next.run(request).await -} - -fn configured_headers() -> ConfiguredHeaders { - ConfiguredHeaders::from_iter([ +fn configured_headers() -> HeaderMap { + HeaderMap::from_iter([ (AUTHORIZATION, HeaderValue::from_static("Sentry-Bearer secret")), (HeaderName::from_static("x-api-key"), HeaderValue::from_static("custom-secret")), ]) diff --git a/crates/mcp-utils/tests/oauth.rs b/crates/mcp-utils/tests/oauth.rs index d926321a8..0b06cf2aa 100644 --- a/crates/mcp-utils/tests/oauth.rs +++ b/crates/mcp-utils/tests/oauth.rs @@ -1,6 +1,6 @@ use aether_auth::{OAuthClientRegistration, OAuthError}; use mcp_utils::client::{ - ClientError, ClientOptions, Elicitation, ElicitationSink, McpClient, OAuthSettings, Transport, + ClientError, ClientOptions, Elicitation, ElicitationSink, HttpAuth, McpClient, OAuthSettings, Transport, }; use mcp_utils::gateway::{GatewayError, McpGateway, ServerState}; use mcp_utils::model::{ElicitResult, ElicitationAction, ElicitationCapability, FormElicitationCapability}; @@ -31,8 +31,8 @@ async fn servers_without_oauth_settings_are_not_offered_oauth() { let server = math_auth_server().await; let (options, _host) = elicitation_host(); - let connected = McpClient::connect("linear", server.transport_with(None), &options).await; - let authorized = McpClient::authorize("linear", server.transport_with(None), &options, None).await; + let connected = McpClient::connect("linear", server.transport_with(HttpAuth::None), &options).await; + let authorized = McpClient::authorize("linear", server.transport_with(HttpAuth::None), &options, None).await; assert!(matches!(connected, Err(ClientError::Connect { server, .. }) if server == "linear")); assert!(matches!(authorized, Err(ClientError::OAuthUnavailable { server }) if server == "linear")); @@ -84,8 +84,9 @@ async fn authorization_presents_the_client_metadata_and_a_loopback_redirect() { async fn authorization_redirects_to_the_configured_callback_port() { let server = math_auth_server().await; let port = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port(); - let mut flow = - AuthFlow::start(server.transport_with(Some(OAuthSettings { callback_port: port, ..client_metadata() }))); + let mut flow = AuthFlow::start( + server.transport_with(HttpAuth::OAuth(OAuthSettings { callback_port: port, ..client_metadata() })), + ); let prompt = flow.prompt().await; let redirect_uri = prompt.query("redirect_uri"); @@ -102,7 +103,7 @@ async fn authorization_presents_a_preregistered_client_id() { registration: OAuthClientRegistration::PreRegistered("registered".to_string()), callback_port: 0, }; - let mut flow = AuthFlow::start(server.transport_with(Some(oauth))); + let mut flow = AuthFlow::start(server.transport_with(HttpAuth::OAuth(oauth))); let prompt = flow.prompt().await; let client_id = prompt.query("client_id"); @@ -242,7 +243,7 @@ async fn a_configured_callback_port_that_is_taken_fails_authorization() { let oauth = OAuthSettings { callback_port: taken.local_addr().unwrap().port(), ..client_metadata() }; let (options, _host) = elicitation_host(); - let error = McpClient::authorize("linear", server.transport_with(Some(oauth)), &options, None) + let error = McpClient::authorize("linear", server.transport_with(HttpAuth::OAuth(oauth)), &options, None) .await .expect_err("the callback port is in use"); diff --git a/crates/mcp-utils/tests/sigv4.rs b/crates/mcp-utils/tests/sigv4.rs new file mode 100644 index 000000000..2a1adefa1 --- /dev/null +++ b/crates/mcp-utils/tests/sigv4.rs @@ -0,0 +1,100 @@ +use aether_auth::{SigV4Config, SigV4Credentials, SigV4Signer}; +use aws_credential_types::Credentials; +use aws_credential_types::provider::SharedCredentialsProvider; +use aws_sigv4::http_request::{SignableBody, SignableRequest, SigningSettings, sign}; +use aws_sigv4::sign::v4; +use axum::http::{Method, StatusCode, header::AUTHORIZATION}; +use chrono::NaiveDateTime; +use mcp_utils::client::{ClientError, ClientOptions, HeaderMap, HttpAuth, McpClient, Transport}; +use mcp_utils::testing::{FakeHttpServer, FakeMcpServer, FakeTool, FakeToolResponse, RecordedRequest, call}; +use rmcp::ServerHandler; +use serde_json::json; +use std::time::{Duration, SystemTime}; + +#[tokio::test] +async fn tool_calls_are_signed_with_the_configured_credentials() { + let server = FakeMcpServer::new().with_tool(FakeTool::new("echo").responds(FakeToolResponse::text("ok"))); + let gateway = sigv4_gateway(server, credentials("secret")).await; + let client = McpClient::connect("gateway", transport(&gateway, credentials("secret")), &ClientOptions::default()) + .await + .unwrap(); + + let result = call(&client, "echo", json!({ "text": "héllo \"quoted\" & " })).await; + + assert!(result.is_ok(), "the signed tool call was rejected: {result:?}"); +} + +#[tokio::test] +async fn every_request_of_a_legacy_session_is_signed() { + let gateway = sigv4_gateway(FakeMcpServer::new().with_legacy_protocol(), credentials("secret")).await; + let client = McpClient::connect("gateway", transport(&gateway, credentials("secret")), &ClientOptions::default()) + .await + .unwrap(); + + let requests = gateway.record_session(client).await; + + for method in [Method::POST, Method::GET, Method::DELETE] { + assert!(requests.iter().any(|request| request.method == method), "the client sent no {method}"); + } + let unsigned = requests.iter().find(|request| !signed_by(&credentials("secret"), request)); + assert!(unsigned.is_none(), "the client sent an unsigned request: {unsigned:?}"); +} + +#[tokio::test] +async fn requests_signed_with_the_wrong_secret_are_rejected() { + let gateway = sigv4_gateway(FakeMcpServer::new(), credentials("secret")).await; + + let result = + McpClient::connect("gateway", transport(&gateway, credentials("wrong")), &ClientOptions::default()).await; + + assert!(matches!(result, Err(ClientError::Connect { .. })), "expected the connection to be rejected"); +} + +const SERVICE: &str = "execute-api"; +const REGION: &str = "us-east-1"; + +async fn sigv4_gateway(server: impl ServerHandler + Clone, credentials: Credentials) -> FakeHttpServer { + FakeHttpServer::start_guarded(server, move |request| { + signed_by(&credentials, request).then_some(()).ok_or(StatusCode::FORBIDDEN) + }) + .await +} + +fn transport(gateway: &FakeHttpServer, credentials: Credentials) -> Transport { + let signer = SigV4Signer::new(SigV4Config { + service: SERVICE.to_string(), + region: REGION.to_string(), + credentials: SigV4Credentials::Provider(SharedCredentialsProvider::new(credentials)), + refresh_before_expiry: Duration::from_mins(5), + }); + gateway.transport(HeaderMap::new(), HttpAuth::SigV4(signer)) +} + +fn signed_by(credentials: &Credentials, request: &RecordedRequest) -> bool { + let header = |name: &str| request.headers.get(name).and_then(|value| value.to_str().ok()); + header(AUTHORIZATION.as_str()) + .zip(header("x-amz-date")) + .is_some_and(|(authorization, date)| authorization == expected_authorization(credentials, request, date)) +} + +fn expected_authorization(credentials: &Credentials, request: &RecordedRequest, date: &str) -> String { + let time: SystemTime = NaiveDateTime::parse_from_str(date, "%Y%m%dT%H%M%SZ").unwrap().and_utc().into(); + let identity = credentials.clone().into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(REGION) + .name(SERVICE) + .time(time) + .settings(SigningSettings::default()) + .build() + .unwrap() + .into(); + let body = SignableBody::Bytes(&request.body); + let signable = SignableRequest::new(request.method.as_str(), &request.url, std::iter::empty(), body).unwrap(); + let (instructions, _) = sign(signable, ¶ms).unwrap().into_parts(); + instructions.headers().find(|(name, _)| name.eq_ignore_ascii_case("authorization")).unwrap().1.to_string() +} + +fn credentials(secret: &str) -> Credentials { + Credentials::new("AKIDEXAMPLE", secret, None, None, "test") +} diff --git a/packages/website/src/content/docs/aether/settings/mcp-servers.mdx b/packages/website/src/content/docs/aether/settings/mcp-servers.mdx index a2f7b0639..263afc224 100644 --- a/packages/website/src/content/docs/aether/settings/mcp-servers.mdx +++ b/packages/website/src/content/docs/aether/settings/mcp-servers.mdx @@ -125,6 +125,7 @@ If a server entry has a `command` and no `type`, it is treated as `"type": "stdi | `url` | `string` | Streamable HTTP endpoint URL. | | `headers` | `object` | Optional headers. Aether currently forwards the `Authorization` header to the transport. | | `oauth` | `object` | Optional custom CIMD or pre-registered public OAuth client settings. See [OAuth clients](#oauth-clients). | + | `sigv4` | `object` | Optional AWS SigV4 request signing. Cannot be combined with `oauth`. See [AWS SigV4](#aws-sigv4). | | `deferTools` | `boolean \| object` | `true` proxies every tool; an object partitions tools with `include` and `exclude` patterns. | @@ -157,7 +158,35 @@ Set `oauth.clientMetadataUrl` to use a custom CIMD, or `oauth.clientId` for a pu | `oauth.clientMetadataUrl` | `string` | HTTPS URL of a custom Client ID Metadata Document. Supports variable expansion. | | `oauth.callbackPort` | `number` | Loopback callback port; defaults to `3118`. The advertised URI is `http://localhost:/` and must exactly match the client metadata or registration. | -When the authorization server does not advertise CIMD support, Aether falls back to deprecated Dynamic Client Registration for compatibility. OAuth discovery, PKCE, scopes, `offline_access`, resource indicators, refresh, and issuer validation are handled by `rmcp`. An explicit bearer `Authorization` header remains authoritative and bypasses OAuth. +When the authorization server does not advertise CIMD support, Aether falls back to deprecated Dynamic Client Registration for compatibility. OAuth discovery, PKCE, scopes, `offline_access`, resource indicators, refresh, and issuer validation are handled by `rmcp`. An `Authorization` header on its own turns off the default OAuth sign-in. A server can set at most one of `oauth`, `sigv4` or an `Authorization` header. + +## AWS SigV4 + +Set `sigv4` to sign every request with AWS Signature Version 4, for example to reach an MCP gateway behind API Gateway IAM auth or a Lambda function URL. A gateway like this can proxy many upstream servers while holding their tokens itself, so the agent only needs AWS credentials: + +```json title=".aether/mcp.json" +{ + "servers": { + "gateway": { + "type": "http", + "url": "https://abc123.execute-api.us-east-1.amazonaws.com/mcp", + "sigv4": { + "service": "execute-api", + "region": "us-east-1", + "profile": "dev" + } + } + } +} +``` + +| Field | Type | Description | +|-------|------|-------------| +| `sigv4.service` | `string` | Required. Signing name of the AWS service in front of the server, e.g. `execute-api` (API Gateway) or `lambda` (function URLs). | +| `sigv4.region` | `string` | Required. Region to sign for. `$AWS_REGION` takes it from the environment, but the config fails to load when that variable is unset. | +| `sigv4.profile` | `string` | Optional named profile. Omit to use the default AWS credential chain (environment, shared config, SSO, web identity, ECS, IMDS), which already honors `AWS_PROFILE`. | + +All fields support variable expansion. A server can set at most one of `oauth`, `sigv4` or an `Authorization` header; `sigv4` sends its signature in the `Authorization` header. ## Built-in servers