diff --git a/Cargo.lock b/Cargo.lock index 93a3768c..f4aadb43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3185,6 +3185,7 @@ dependencies = [ "tar", "tempfile", "tokio", + "tokio-rustls", "tokio-tungstenite 0.30.0", "tokio-util", "toml", diff --git a/Cargo.toml b/Cargo.toml index 14ca7859..37655eab 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -127,6 +127,7 @@ windows-sys = { version = "=0.61.2", features = [ ] } [dev-dependencies] +tokio-rustls = { version = "=0.26.4", default-features = false, features = ["aws_lc_rs", "tls12"] } jsonwebtoken = { version = "=11.1.0", default-features = false, features = ["aws_lc_rs", "use_pem"] } tokio = { version = "=1.53.3", features = ["test-util"] } diff --git a/docs/user/mcp.md b/docs/user/mcp.md index 4c87912d..7297fa89 100644 --- a/docs/user/mcp.md +++ b/docs/user/mcp.md @@ -6,6 +6,39 @@ Kit merges servers from Agent Plugins, `mcp_config` in `~/.kit/config.toml`, `.m Use `kit config set mcp_config ~/.kit/mcp.json` to select the global MCP file, `kit config get mcp_config` to inspect the saved path, or `kit config unset mcp_config` to remove the override. These commands edit `~/.kit/config.toml`, not the MCP JSON file or its server entries. See [configuration editing](getting-started-and-configuration.md#edit-configuration-from-the-command-line). +## Use a deployment-provided workload assertion + +Configure a trusted server with `auth.type: "workload_identity"` to exchange a +mounted JWT without browser login, client registration, or a client secret: + +```json +{ + "mcpServers": { + "workload-service": { + "type": "streamable-http", + "url": "https://service.example.com/mcp", + "auth": { + "type": "workload_identity", + "tokenFile": "/run/identity/token", + "expectedIssuer": "https://auth.example.com/tenant", + "tokenEndpoint": "https://auth.example.com/tenant/token", + "resource": "https://service.example.com/mcp" + } + } + } +} +``` + +Keep this configuration and token mount under deployment control. `tokenFile` +must be absolute; symlink rotation is supported. The resource must match the HTTPS +MCP URL, and the token endpoint must share the expected issuer's origin. Kit checks +issuer metadata before sending the assertion, disables redirects, and refreshes +before expiry. No browser or saved-credential fallback occurs. + +Do not add static authorization or Host headers. Reusing an unchanged assertion +requires the issuer's explicit replay-policy support. After a process-memory +restore, restart Kit and resume the transcript rather than reuse resident tokens. + ## Ask the agent to add a server Kit does not provide an MCP-specific slash command. In a `kit tui` session, describe the server and the scope that you want: diff --git a/fixtures/workload-auth/README.md b/fixtures/workload-auth/README.md new file mode 100644 index 00000000..8dcdfce6 --- /dev/null +++ b/fixtures/workload-auth/README.md @@ -0,0 +1,17 @@ +# Local HTTPS test identity + +These fixtures are exclusively for `src/workload_auth/https_tests.rs`. The +private key is a public, disposable test key, not a deployment credential. +Never trust this CA or use this identity outside tests. The issuing CA's +private key was discarded after signing. + +- `ca.pem`: self-signed test CA, trusted only by the test HTTP client. +- `server.der`: leaf certificate for `localhost` and `127.0.0.1`. +- `server-key.der`: unencrypted PKCS#8 RSA key for the leaf certificate. + +Generated with OpenSSL using a 2048-bit RSA CA and leaf key, a 100-year +validity period to avoid routine expiry, and leaf extensions +`subjectAltName=DNS:localhost,IP:127.0.0.1`, `basicConstraints=CA:FALSE`, +`keyUsage=digitalSignature,keyEncipherment`, and `extendedKeyUsage=serverAuth`. +The local server uses an ephemeral port and the test client still validates +the chain and hostname; it does not disable TLS verification. diff --git a/fixtures/workload-auth/ca.pem b/fixtures/workload-auth/ca.pem new file mode 100644 index 00000000..940366c4 --- /dev/null +++ b/fixtures/workload-auth/ca.pem @@ -0,0 +1,17 @@ +-----BEGIN CERTIFICATE----- +MIICvDCCAaQCCQD9JAqIqRQZ5DANBgkqhkiG9w0BAQsFADAfMR0wGwYDVQQDDBRL +aXQgd29ya2xvYWQgdGVzdCBDQTAgFw0yNjEwMjMxOTQ1MjlaGA8yMTI2MDkyOTE5 +NDUyOVowHzEdMBsGA1UEAwwUS2l0IHdvcmtsb2FkIHRlc3QgQ0EwggEiMA0GCSqG +SIb3DQEBAQUAA4IBDwAwggEKAoIBAQDF/ak/Vx1B5nzTI2neCGzt1J1n1lXGJ2aY +dBsruuRiuFxNvkd8rxEVH4tqSRJo3JBC/c1guMR0yaH55NJYmO0heNCqiTUTsgZi +yqiWS90AS+Hk/WWA8E9+ninMwlasEnbkarsK5m2MQilDZjSxzlLBJXmPYBqjeLVB +ucr2lUrFk8X8mkKNPIh40Z2AgfeYbuh6RYLn+/weRxavOFesFMt1GPE9vH62rkxU +HGwf12p+wo18Q8vLFOvcEgZKvAifMxjt2y3TWzJ5AEUFYjj48TpNOol9iXRMkCVC +vgIATVm/hP5S9PJdHRiJv7Ui4OUqimmPk62xKP5p8Gr86Flux6shAgMBAAEwDQYJ +KoZIhvcNAQELBQADggEBAEqpDya+Y7CRaSBzKh1xJeEVlphm7b3FLCPtEX7WY6fP +9IPtKM6UO6NzATfC/fYLKgaBv2jgyXZoR9Sl51la1LmQEbI4VepjnuJU9kbPXUZW +S7G6OSKlNJVCivOFqd7rRj/6THt9SaHIk47S2kRETGbCXhPquWwLt5UWtvcnQkz3 +MWOsvCsP5Y/7IZLkMH7iOXnKqYjlj+gIYDfwsPi+uxXTY7C559aja+AOZpDpYkKs +q1QX+KiZOVabDUtbHOIrTgFoTniEu/eqQu7uIs+yOq/GPbervq+707eRFnBslnty +TNGe+kZUGzvZQnyV+Nrp0ycIIdj+kSWOQMHyHXGYvzw= +-----END CERTIFICATE----- diff --git a/fixtures/workload-auth/server-key.der b/fixtures/workload-auth/server-key.der new file mode 100644 index 00000000..c10b3ca4 Binary files /dev/null and b/fixtures/workload-auth/server-key.der differ diff --git a/fixtures/workload-auth/server.der b/fixtures/workload-auth/server.der new file mode 100644 index 00000000..cee77e44 Binary files /dev/null and b/fixtures/workload-auth/server.der differ diff --git a/src/lib.rs b/src/lib.rs index 2fcf2437..4c7c58db 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -38,6 +38,7 @@ pub mod protocols; pub mod provider; pub mod request_budget; pub mod runtime; +mod workload_auth; /// Shared internal filesystem and process-lifetime recovery controls. pub mod resilient_fs { pub use crate::filesystem::*; diff --git a/src/tools/mcp.rs b/src/tools/mcp.rs index 2d838890..3dd9b5bc 100644 --- a/src/tools/mcp.rs +++ b/src/tools/mcp.rs @@ -12,6 +12,7 @@ mod refresh; )] pub(crate) mod refresh_tests; mod subscriptions; +mod workload; pub use subscriptions::McpSubscriptionTool; mod schema; pub use schema::ToolSchema; @@ -922,6 +923,17 @@ impl McpAuthResponder for AuthRecorder { .check_available() .map_err(McpError::AuthResolution)?; } + if record + .as_ref() + .and_then(|record| record.oauth.as_ref()) + .and_then(auth::Config::workload) + .is_some() + { + return Err(McpError::AuthResolution( + "workload MCP authentication failed; check the provisioned workload identity" + .into(), + )); + } let session = self.oauth_sessions.lock().await.get(&server).cloned(); if let Some(session) = &session { session @@ -1239,6 +1251,12 @@ fn prepare_server( .with_header(name.as_str(), value.as_str()) .map_err(|error| format!("invalid MCP config for {id}: {error}"))?; } + if let Some(config) = server.auth.as_ref().and_then(auth::Config::workload) { + transport = transport.with_http_client(Arc::new(workload::Client::new( + &server.url, + config.clone(), + )?)); + } Ok(( McpTransportBinding::StreamableHttp(transport), ServerRecord { @@ -2752,7 +2770,7 @@ impl McpRuntime { ))); } let oauth = record.oauth.clone().unwrap_or_default(); - if oauth.machine().is_some() { + if oauth.machine().is_some() || oauth.workload().is_some() { return Err(ToolError::InvalidInput("machine MCP authentication is managed by provisioned credential files, not interactive auth".into())); } diff --git a/src/tools/mcp/auth.rs b/src/tools/mcp/auth.rs index fb3907ad..71117546 100644 --- a/src/tools/mcp/auth.rs +++ b/src/tools/mcp/auth.rs @@ -33,6 +33,8 @@ pub enum Config { Oauth(InteractiveConfig), #[serde(rename = "client_credentials")] Machine(super::machine::Config), + #[serde(rename = "workload_identity")] + Workload(crate::workload_auth::Config), } impl Default for Config { @@ -49,13 +51,20 @@ impl Config { pub fn machine(&self) -> Option<&super::machine::Config> { match self { Self::Machine(config) => Some(config), - Self::Oauth(_) => None, + Self::Oauth(_) | Self::Workload(_) => None, } } + pub fn workload(&self) -> Option<&crate::workload_auth::Config> { + match self { + Self::Workload(config) => Some(config), + _ => None, + } + } + fn interactive(&self) -> Result<&InteractiveConfig, String> { match self { Self::Oauth(config) => Ok(config), - Self::Machine(_) => { + Self::Machine(_) | Self::Workload(_) => { Err("machine MCP authentication does not support interactive authorization".into()) } } @@ -143,6 +152,10 @@ pub async fn restore( config: &Config, credential_storage: &CredentialStorage, ) -> Result, String> { + if config.workload().is_some() { + // Workload acquisition belongs to the HTTP client, never the credential store. + return Ok(None); + } if let Some(machine) = config.machine() { return super::machine::authenticate(resource_url, machine) .await diff --git a/src/tools/mcp/workload.rs b/src/tools/mcp/workload.rs new file mode 100644 index 00000000..49cbf093 --- /dev/null +++ b/src/tools/mcp/workload.rs @@ -0,0 +1,769 @@ +//! Per-request workload authentication, including catalog reads and SSE reconnects. +use std::{ + collections::HashMap, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; + +use agentkit_mcp::{ClientJsonRpcMessage, McpHttpClient, McpSseStream}; +use async_trait::async_trait; +use reqwest::header::{HeaderName, HeaderValue}; +use rmcp::transport::streamable_http_client::{ + StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, +}; + +use crate::workload_auth::{Acquisition, Config, Session}; + +type Error = StreamableHttpError; + +pub(super) struct Client { + destination: String, + session: Session, + http: reqwest::Client, +} + +impl Client { + pub(super) fn new(destination: &str, config: Config) -> Result { + config.validate()?; + let url = url::Url::parse(destination).map_err(|_| "invalid workload MCP URL")?; + let resource = + url::Url::parse(&config.resource).map_err(|_| "invalid workload MCP resource")?; + if url.scheme() != "https" + || !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + || url != resource + { + return Err( + "workload MCP URL must be HTTPS and match the configured resource exactly".into(), + ); + } + let http = PostStatus::client(reqwest::Client::builder()) + .map_err(|_| "could not build workload MCP HTTP client")?; + Ok(Self { + destination: destination.to_owned(), + session: Session::new(config)?, + http, + }) + } + + async fn post_with_client( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + custom_headers: HashMap, + ) -> Result { + let acquisition = self.authorization(&uri, &custom_headers).await?; + let status = PostStatus::default(); + let result = status + .observe(self.http.post_message( + uri, + message, + session_id, + Some(acquisition.token().to_owned()), + custom_headers, + )) + .await; + self.finish(&acquisition, status.rejected(), result).await + } + + async fn authorization( + &self, + uri: &str, + headers: &HashMap, + ) -> Result { + if uri != self.destination + || headers + .keys() + .any(|name| name == reqwest::header::AUTHORIZATION || name == reqwest::header::HOST) + { + return Err(Self::error( + "workload MCP transport destination or authorization override rejected", + )); + } + self.session.acquisition().await.map_err(Self::error) + } + + fn error(message: impl Into) -> Error { + StreamableHttpError::Io(std::io::Error::other(message.into())) + } + + async fn finish( + &self, + acquisition: &Acquisition, + raw_rejection: bool, + result: Result, + ) -> Result { + let rejected = match &result { + Err( + StreamableHttpError::AuthRequired(_) | StreamableHttpError::InsufficientScope(_), + ) => true, + Err(StreamableHttpError::Client(error)) => matches!( + error.status(), + Some(reqwest::StatusCode::UNAUTHORIZED | reqwest::StatusCode::FORBIDDEN) + ), + _ => false, + }; + if raw_rejection || rejected { + // Do not replay a possibly side-effecting POST or follow challenge URLs. + self.session.invalidate_acquisition(acquisition).await; + } + result + } +} + +// Pinned reqwest 0.13.5 / tower 0.5.3 call the classifier synchronously in +// ResponseFuture::poll, inside the scoped POST future. Tokio restores this scope +// for each poll, including worker migration; concurrent POSTs never share a flag. +// Revalidate this seam when upgrading reqwest or tower. rmcp retains its parser. +tokio::task_local! { + static POST_REJECTION: Arc; +} + +#[derive(Default)] +struct PostStatus(Arc); + +impl PostStatus { + fn client(builder: reqwest::ClientBuilder) -> Result { + let policy = reqwest::retry::never().classify_fn(move |response| { + if matches!( + response.status(), + Some(reqwest::StatusCode::UNAUTHORIZED | reqwest::StatusCode::FORBIDDEN) + ) { + let _ = POST_REJECTION.try_with(|rejected| rejected.store(true, Ordering::Release)); + } + // Always return success: this observes status, never retries a POST. + // Do not set max_retries_per_request(0): tower then skips classification. + response.success() + }); + builder + .redirect(reqwest::redirect::Policy::none()) + .retry(policy) + .build() + .map_err(StreamableHttpError::Client) + } + + async fn observe(&self, future: F) -> F::Output { + POST_REJECTION.scope(Arc::clone(&self.0), future).await + } + + fn rejected(&self) -> bool { + self.0.load(Ordering::Acquire) + } +} + +#[async_trait] +impl McpHttpClient for Client { + async fn post_message( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + _auth_header: Option, + custom_headers: HashMap, + ) -> Result { + self.post_with_client(uri, message, session_id, custom_headers) + .await + } + + async fn delete_session( + &self, + uri: Arc, + session_id: Arc, + _auth_header: Option, + custom_headers: HashMap, + ) -> Result<(), Error> { + let acquisition = self.authorization(&uri, &custom_headers).await?; + self.finish( + &acquisition, + false, + self.http + .delete_session( + uri, + session_id, + Some(acquisition.token().to_owned()), + custom_headers, + ) + .await, + ) + .await + } + + async fn get_stream( + &self, + uri: Arc, + session_id: Option>, + last_event_id: Option, + _auth_header: Option, + custom_headers: HashMap, + ) -> Result { + let acquisition = self.authorization(&uri, &custom_headers).await?; + self.finish( + &acquisition, + false, + self.http + .get_stream( + uri, + session_id, + last_event_id, + Some(acquisition.token().to_owned()), + custom_headers, + ) + .await, + ) + .await + } +} + +#[cfg(test)] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros +)] +mod tests { + use super::*; + + async fn tls_script( + replies: impl FnOnce(&str) -> Vec<(u16, String, String)>, + ) -> (String, tokio::task::JoinHandle<(Vec, usize)>) { + tls_connections(replies, false).await + } + + async fn tls_connections( + replies: impl FnOnce(&str) -> Vec<(u16, String, String)>, + keep_alive: bool, + ) -> (String, tokio::task::JoinHandle<(Vec, usize)>) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio_rustls::{TlsAcceptor, rustls}; + let provider = rustls::crypto::aws_lc_rs::default_provider(); + let config = rustls::ServerConfig::builder_with_provider(Arc::new(provider)) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert( + vec![rustls::pki_types::CertificateDer::from( + include_bytes!("../../../fixtures/workload-auth/server.der").to_vec(), + )], + rustls::pki_types::PrivatePkcs8KeyDer::from( + include_bytes!("../../../fixtures/workload-auth/server-key.der").to_vec(), + ) + .into(), + ) + .unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let uri = format!("https://{}/mcp", listener.local_addr().unwrap()); + let origin = format!("https://{}", listener.local_addr().unwrap()); + let replies = replies(&origin); + let task = tokio::spawn(async move { + let acceptor = TlsAcceptor::from(Arc::new(config)); + let mut requests = Vec::new(); + let mut connection = None; + let mut accepted = 0; + for (status, content_type, body) in replies { + let connection_header = if keep_alive { "keep-alive" } else { "close" }; + let response = format!( + "HTTP/1.1 {status} Test\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: {connection_header}\r\n\r\n{body}", + body.len() + ); + if connection.is_none() { + let (socket, _) = listener.accept().await.unwrap(); + connection = Some(acceptor.accept(socket).await.unwrap()); + accepted += 1; + } + let stream = connection.as_mut().unwrap(); + let mut bytes = Vec::new(); + loop { + let mut chunk = [0; 4096]; + let read = stream.read(&mut chunk).await.unwrap(); + assert!(read > 0); + bytes.extend_from_slice(&chunk[..read]); + assert!(bytes.len() < 32768); + if let Some(end) = bytes.windows(4).position(|part| part == b"\r\n\r\n") { + let headers = std::str::from_utf8(&bytes[..end]).unwrap(); + let length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap_or(0); + if bytes.len() >= end + 4 + length { + break; + } + } + } + stream.write_all(response.as_bytes()).await.unwrap(); + if !keep_alive { + let _ = stream.shutdown().await; + connection = None; + } + requests.push(String::from_utf8(bytes).unwrap()); + } + (requests, accepted) + }); + (uri, task) + } + + async fn tls_reply( + status: u16, + content_type: &str, + body: &str, + ) -> (String, tokio::task::JoinHandle<(Vec, usize)>) { + tls_script(|_| vec![(status, content_type.to_owned(), body.to_owned())]).await + } + + fn tls_client() -> reqwest::ClientBuilder { + reqwest::Client::builder() + .tls_certs_only([reqwest::Certificate::from_pem(include_bytes!( + "../../../fixtures/workload-auth/ca.pem" + )) + .unwrap()]) + .no_proxy() + .timeout(std::time::Duration::from_secs(5)) + .redirect(reqwest::redirect::Policy::none()) + .retry(reqwest::retry::never()) + } + + #[tokio::test] + async fn sequential_posts_reuse_one_tls_connection() { + let reply = r#"{"jsonrpc":"2.0","id":1,"result":{}}"#; + let (uri, mcp) = tls_connections( + |_| { + vec![ + (200, "application/json".into(), reply.into()), + (200, "application/json".into(), reply.into()), + ] + }, + true, + ) + .await; + let (issuer, auth) = tls_script(|origin| { + vec![ + ( + 200, + "application/json".into(), + serde_json::json!({"issuer":origin,"token_endpoint":format!("{origin}/token")}) + .to_string(), + ), + ( + 200, + "application/json".into(), + serde_json::json!({"access_token":"T1","token_type":"Bearer","expires_in":600}) + .to_string(), + ), + ] + }) + .await; + let directory = tempfile::tempdir().unwrap(); + let token_file = directory.path().join("assertion.jwt"); + std::fs::write(&token_file, "header.payload.signature").unwrap(); + let origin = issuer.strip_suffix("/mcp").unwrap(); + let session = crate::workload_auth::test_support::session( + Config { + token_file, + expected_issuer: origin.into(), + token_endpoint: format!("{origin}/token"), + resource: uri.clone(), + }, + tls_client().build().unwrap(), + ) + .unwrap(); + let client = Client { + destination: uri.clone(), + session, + http: PostStatus::client(tls_client()).unwrap(), + }; + for _ in 0..2 { + let message = + serde_json::from_value(serde_json::json!({"jsonrpc":"2.0","id":1,"method":"ping"})) + .unwrap(); + let result = client + .post_message(uri.clone().into(), message, None, None, HashMap::new()) + .await + .unwrap(); + assert!(matches!(result, StreamableHttpPostResponse::Json(_, _))); + } + let (requests, accepted) = mcp.await.unwrap(); + assert_eq!( + accepted, 1, + "both POSTs must use the same accepted TLS connection" + ); + assert_eq!(requests.len(), 2); + assert!( + requests + .iter() + .all(|request| request.starts_with("POST /mcp ") + && request.to_lowercase().contains("authorization: bearer t1")) + ); + auth.await.unwrap(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn concurrent_post_status_scopes_are_isolated() { + let body = r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"rejected"}}"#; + let (rejected_uri, rejected_server) = tls_reply(401, "application/json", body).await; + let (successful_uri, successful_server) = tls_reply(200, "application/json", body).await; + let http = PostStatus::client(tls_client()).unwrap(); + let rejected = PostStatus::default(); + let successful = PostStatus::default(); + let message = || { + serde_json::from_value(serde_json::json!({"jsonrpc":"2.0","id":1,"method":"ping"})) + .unwrap() + }; + // Poll both scoped futures in one task: a thread-local/unscoped flag would + // let the rejection contaminate its concurrently polled neighbour. + let (first, second) = tokio::join!( + rejected.observe(http.post_message( + rejected_uri.into(), + message(), + None, + Some("T1".into()), + HashMap::new() + )), + successful.observe(http.post_message( + successful_uri.into(), + message(), + None, + Some("T2".into()), + HashMap::new() + )), + ); + assert!(matches!(first, Ok(StreamableHttpPostResponse::Json(_, _)))); + assert!(matches!(second, Ok(StreamableHttpPostResponse::Json(_, _)))); + assert!(rejected.rejected()); + assert!(!successful.rejected()); + assert!(POST_REJECTION.try_with(|_| ()).is_err()); + rejected_server.await.unwrap(); + successful_server.await.unwrap(); + + // Calls without a POST scope (including GET/DELETE) have no rejection + // slot to update, even when the pooled client's classifier runs. + let (uri, server) = tls_reply(403, "application/json", body).await; + let _ = http + .post_message( + uri.into(), + message(), + None, + Some("T3".into()), + HashMap::new(), + ) + .await; + assert!(!successful.rejected()); + server.await.unwrap(); + } + + #[tokio::test] + async fn bare_post_rejection_reacquires_on_next_operation() { + for code in [401, 403] { + for json in [false, true] { + let (uri, server) = tls_script(|origin| { + let discovery = serde_json::json!({"issuer": origin, "token_endpoint": format!("{origin}/token")}).to_string(); + let token = |value| serde_json::json!({"access_token": value, "token_type": "Bearer", "expires_in": 600}).to_string(); + vec![ + (200, "application/json".into(), discovery.clone()), + (200, "application/json".into(), token("T1")), + (code, if json { "application/json" } else { "text/plain" }.into(), if json { r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"rejected"}}"# } else { "rejected" }.into()), + (200, "application/json".into(), discovery), + (200, "application/json".into(), token("T2")), + (200, "application/json".into(), r#"{"jsonrpc":"2.0","id":1,"result":{}}"#.into()), + ] + }).await; + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("assertion.jwt"); + std::fs::write(&path, "header.payload.signature").unwrap(); + let origin = uri.strip_suffix("/mcp").unwrap(); + let session = crate::workload_auth::test_support::session( + Config { + token_file: path, + expected_issuer: origin.into(), + token_endpoint: format!("{origin}/token"), + resource: uri.clone(), + }, + tls_client().build().unwrap(), + ) + .unwrap(); + let client = Client { + destination: uri.clone(), + session, + http: PostStatus::client(tls_client()).unwrap(), + }; + let message = || { + serde_json::from_value( + serde_json::json!({"jsonrpc":"2.0","id":1,"method":"ping"}), + ) + .unwrap() + }; + let first = client + .post_message(uri.clone().into(), message(), None, None, HashMap::new()) + .await; + assert_eq!(first.is_ok(), json); + let second = client + .post_message(uri.into(), message(), None, None, HashMap::new()) + .await; + assert!(matches!(second, Ok(StreamableHttpPostResponse::Json(_, _)))); + let (requests, _) = server.await.unwrap(); + assert_eq!(requests.len(), 6); + assert!(requests[1].starts_with("POST /token ")); + assert!( + requests[2] + .to_lowercase() + .contains("authorization: bearer t1") + ); + assert!(requests[4].starts_with("POST /token ")); + assert!( + requests[5] + .to_lowercase() + .contains("authorization: bearer t2") + ); + } + } + } + + #[tokio::test] + async fn bare_post_rejections_preserve_raw_status_without_replaying() { + for code in [401, 403] { + for json in [false, true] { + let body = if json { + r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"rejected"}}"# + } else { + "rejected" + }; + let (uri, server) = tls_reply( + code, + if json { + "application/json" + } else { + "text/plain" + }, + body, + ) + .await; + let status = PostStatus::default(); + let http = PostStatus::client( + reqwest::Client::builder() + .tls_certs_only([reqwest::Certificate::from_pem(include_bytes!( + "../../../fixtures/workload-auth/ca.pem" + )) + .unwrap()]) + .no_proxy() + .timeout(std::time::Duration::from_secs(5)), + ) + .unwrap(); + let message = serde_json::from_value( + serde_json::json!({"jsonrpc":"2.0","id":1,"method":"ping"}), + ) + .unwrap(); + let result = status + .observe(http.post_message( + uri.into(), + message, + None, + Some("T1".into()), + HashMap::new(), + )) + .await; + assert!(status.rejected(), "bare HTTP {code} was lost"); + if json { + assert!(matches!(result, Ok(StreamableHttpPostResponse::Json(_, _)))); + } else { + assert!(matches!( + result, + Err(StreamableHttpError::UnexpectedServerResponse(_)) + )); + } + assert!( + server.await.unwrap().0[0] + .to_lowercase() + .contains("authorization: bearer t1") + ); + } + } + } + + #[tokio::test] + async fn post_status_observation_preserves_json_and_sse_success() { + use futures_util::StreamExt; + for sse in [false, true] { + let body = if sse { + "data: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}\n\n" + } else { + r#"{"jsonrpc":"2.0","id":1,"result":{}}"# + }; + let (uri, server) = tls_reply( + 200, + if sse { + "text/event-stream" + } else { + "application/json" + }, + body, + ) + .await; + let status = PostStatus::default(); + let http = PostStatus::client( + reqwest::Client::builder() + .tls_certs_only([reqwest::Certificate::from_pem(include_bytes!( + "../../../fixtures/workload-auth/ca.pem" + )) + .unwrap()]) + .no_proxy() + .timeout(std::time::Duration::from_secs(5)), + ) + .unwrap(); + let message = + serde_json::from_value(serde_json::json!({"jsonrpc":"2.0","id":1,"method":"ping"})) + .unwrap(); + let result = status + .observe(http.post_message( + uri.into(), + message, + None, + Some("T1".into()), + HashMap::new(), + )) + .await + .unwrap(); + assert!(!status.rejected()); + if sse { + let StreamableHttpPostResponse::Sse(mut events, _) = result else { + panic!("expected SSE") + }; + assert!(events.next().await.unwrap().is_ok()); + } else { + assert!(matches!(result, StreamableHttpPostResponse::Json(_, _))); + } + server.await.unwrap(); + } + } + + fn config() -> Config { + serde_json::from_value(serde_json::json!({ + "tokenFile": "/nonexistent/kit-workload-test-token", + "expectedIssuer": "https://issuer.example", + "tokenEndpoint": "https://issuer.example/token", + "resource": "https://mcp.example/mcp" + })) + .unwrap() + } + + #[test] + fn workload_destination_is_bound_to_resource() { + for destination in [ + "http://mcp.example/mcp", + "https://other.example/mcp", + "https://mcp.example/other", + "https://user@mcp.example/mcp", + "https://mcp.example/mcp#fragment", + ] { + assert!(Client::new(destination, config()).is_err()); + } + assert!(Client::new("https://mcp.example/mcp", config()).is_ok()); + } + + #[tokio::test] + async fn workload_transport_rejects_overrides_before_acquisition() { + let client = Client::new("https://mcp.example/mcp", config()).unwrap(); + let error = client + .authorization("https://other.example/mcp", &HashMap::new()) + .await + .err() + .unwrap(); + assert!(error.to_string().contains("override rejected")); + for name in [reqwest::header::AUTHORIZATION, reqwest::header::HOST] { + let headers = HashMap::from([(name, HeaderValue::from_static("override"))]); + assert!( + client + .authorization("https://mcp.example/mcp", &headers) + .await + .err() + .unwrap() + .to_string() + .contains("override rejected") + ); + } + assert!( + client + .authorization("https://mcp.example/mcp", &HashMap::new()) + .await + .is_err() + ); + } + + #[tokio::test] + async fn workload_prepared_transport_requires_credentials_before_http() { + let value = serde_json::json!({ + "url": "https://mcp.example/mcp", + "auth": { + "type": "workload_identity", + "tokenFile": "/nonexistent/kit-workload-test-token", + "expectedIssuer": "https://issuer.example", + "tokenEndpoint": "https://issuer.example/token", + "resource": "https://mcp.example/mcp" + } + }); + let server = serde_json::from_value(value.clone()).unwrap(); + let (binding, _) = + super::super::prepare_server("workload", server, Vec::new(), None).unwrap(); + let super::super::McpTransportBinding::StreamableHttp(transport) = binding else { + panic!("expected workload authentication configuration") + }; + let client = transport.http_client.unwrap(); + assert!( + client + .get_stream( + Arc::from("https://mcp.example/mcp"), + None, + None, + None, + HashMap::new() + ) + .await + .is_err() + ); + for field in ["bearerToken", "headers"] { + let mut conflict = value.clone(); + conflict[field] = if field == "headers" { + serde_json::json!({"Authorization": "Bearer static"}) + } else { + serde_json::json!("static") + }; + let server = serde_json::from_value(conflict).unwrap(); + assert!(super::super::prepare_server("workload", server, Vec::new(), None).is_err()); + } + } + + #[tokio::test] + async fn workload_configuration_never_uses_stored_oauth() { + let mut value = serde_json::json!({ + "type": "workload_identity", + "tokenFile": "/nonexistent/kit-workload-test-token", + "expectedIssuer": "https://issuer.example", + "tokenEndpoint": "https://issuer.example/token", + "resource": "https://mcp.example/mcp" + }); + let config: super::super::auth::Config = serde_json::from_value(value.clone()).unwrap(); + assert!(config.workload().is_some()); + assert!(config.machine().is_none()); + let storage = super::super::CredentialStorage::Memory; + assert!( + super::super::auth::restore("https://mcp.example/mcp", &config, &storage) + .await + .unwrap() + .is_none() + ); + assert!( + super::super::auth::begin("https://mcp.example/mcp", &config, &storage, None, None) + .await + .is_err() + ); + value["clientId"] = serde_json::json!("not-supported"); + assert!(serde_json::from_value::(value).is_err()); + } +} diff --git a/src/workload_auth.rs b/src/workload_auth.rs new file mode 100644 index 00000000..a58cb702 --- /dev/null +++ b/src/workload_auth.rs @@ -0,0 +1,309 @@ +//! Clientless RFC 7523 workload authentication. Assertions never enter credential storage. +use std::{ + path::PathBuf, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use serde::Deserialize; +use tokio::{io::AsyncReadExt, sync::Mutex}; +use zeroize::Zeroizing; + +const MAX_BYTES: usize = 8192; +const TIMEOUT: Duration = Duration::from_secs(15); + +#[cfg(test)] +mod tests; + +#[cfg(test)] +pub(crate) mod test_support { + pub(crate) fn session( + config: super::Config, + client: reqwest::Client, + ) -> Result { + config.validate()?; + Ok(super::Session { + config, + client, + cached: tokio::sync::Mutex::new(None), + }) + } + impl super::Session { + pub async fn token(&self) -> Result { + self.acquisition() + .await + .map(|value| value.token().to_owned()) + } + } +} + +/// Trusted deployment configuration, not input supplied by an MCP server or model. +#[derive(Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Config { + pub token_file: PathBuf, + pub expected_issuer: String, + pub token_endpoint: String, + pub resource: String, +} + +impl Config { + pub fn validate(&self) -> Result<(), String> { + if !self.token_file.is_absolute() { + return Err("workload tokenFile must be an absolute path".into()); + } + let issuer = https_url(&self.expected_issuer)?; + let endpoint = https_url(&self.token_endpoint)?; + https_url(&self.resource)?; + if issuer.query().is_some() || issuer.origin() != endpoint.origin() { + return Err("workload token endpoint must share the expected issuer origin".into()); + } + Ok(()) + } +} + +fn https_url(value: &str) -> Result { + let url = url::Url::parse(value).map_err(|_| "invalid workload HTTPS URL")?; + if url.scheme() != "https" + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + { + return Err("workload URLs must be HTTPS without credentials or fragments".into()); + } + Ok(url) +} + +#[derive(Clone)] +pub struct Acquisition(Arc>); + +impl Acquisition { + pub fn token(&self) -> &str { + self.0.as_str() + } +} + +struct Cached { + value: Acquisition, + acquired: SystemTime, + refresh_at: SystemTime, +} + +/// One resource and one credential source per cache. No refresh tokens or fallback. +pub struct Session { + config: Config, + client: reqwest::Client, + cached: Mutex>, +} + +impl Session { + pub fn new(config: Config) -> Result { + config.validate()?; + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .retry(reqwest::retry::never()) + .timeout(TIMEOUT) + .connect_timeout(Duration::from_secs(5)) + .build() + .map_err(|_| "could not create workload HTTP client")?; + Ok(Self { + config, + client, + cached: Mutex::new(None), + }) + } + + pub async fn invalidate_acquisition(&self, rejected: &Acquisition) { + let mut cached = self.cached.lock().await; + if cached + .as_ref() + .is_some_and(|current| Arc::ptr_eq(¤t.value.0, &rejected.0)) + { + *cached = None; + } + } + + pub async fn acquisition(&self) -> Result { + // Serialize acquisition, not merely publication, to avoid consuming a + // single-use assertion concurrently. Waiting and all I/O are bounded. + tokio::time::timeout(TIMEOUT, self.acquire()) + .await + .map_err(|_| "workload authentication timed out".to_string())? + } + + async fn acquire(&self) -> Result { + let mut cached = self.cached.lock().await; + let now = SystemTime::now(); + if let Some(token) = cached.as_ref() + && now >= token.acquired + && now < token.refresh_at + { + return Ok(token.value.clone()); + } + *cached = None; + self.verify_issuer().await?; + // Reopen on every exchange: projected files rotate through symlinks. + let assertion = read_assertion(&self.config.token_file).await?; + let started = SystemTime::now(); + // Never automatically retry a POST: the server may already have consumed + // this assertion's replay identifier, even if the response was lost. + let response = self + .client + .post(&self.config.token_endpoint) + .form(&[ + ("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"), + ("assertion", assertion.as_str()), + ("resource", self.config.resource.as_str()), + ]) + .send() + .await + .map_err(|_| "workload token exchange failed")?; + if !response.status().is_success() { + return Err("workload token exchange refused".into()); + } + let bytes = bounded_body(response).await?; + let token: Token = + serde_json::from_slice(&bytes).map_err(|_| "invalid workload token response")?; + if !token.token_type.eq_ignore_ascii_case("bearer") + || token.expires_in == 0 + || token.access_token.is_empty() + || token.access_token.len() > MAX_BYTES + || !token.access_token.bytes().all(|c| c.is_ascii_graphic()) + { + return Err("invalid workload token response".into()); + } + // Cap long issuer lifetimes; subtract a proportional margin for short ones. + let lifetime = token.expires_in.min(86400); + let usable = Duration::from_secs(lifetime) + .saturating_sub(Duration::from_millis((lifetime * 100).min(30000))); + let refresh_at = started + .checked_add(usable) + .ok_or("invalid workload token expiry")?; + let completed = SystemTime::now(); + if completed < started || completed >= refresh_at { + return Err("workload authentication expired during exchange".into()); + } + let value = Acquisition(Arc::new(Zeroizing::new(token.access_token))); + let result = value.clone(); + *cached = Some(Cached { + value, + acquired: started, + refresh_at, + }); + Ok(result) + } + + async fn verify_issuer(&self) -> Result<(), String> { + let issuer = https_url(&self.config.expected_issuer)?; + let path = issuer.path().trim_end_matches('/'); + let metadata_url = format!( + "{}/.well-known/oauth-authorization-server{}", + issuer.origin().ascii_serialization(), + path + ); + // Discovery is safe to retry once; no assertion has been read or sent. + let mut response = self + .client + .get(&metadata_url) + .send() + .await + .map_err(|_| "workload issuer discovery failed")?; + if response.status() == reqwest::StatusCode::SERVICE_UNAVAILABLE { + tokio::time::sleep(Duration::from_millis(100)).await; + response = self + .client + .get(&metadata_url) + .send() + .await + .map_err(|_| "workload issuer discovery failed")?; + } + if !response.status().is_success() { + return Err("workload issuer discovery refused".into()); + } + let bytes = bounded_body(response).await?; + let metadata: Metadata = + serde_json::from_slice(&bytes).map_err(|_| "invalid workload issuer metadata")?; + if metadata.issuer != self.config.expected_issuer + || metadata.token_endpoint != self.config.token_endpoint + { + return Err( + "workload issuer or token endpoint does not match trusted configuration".into(), + ); + } + Ok(()) + } +} + +#[derive(Deserialize)] +struct Metadata { + issuer: String, + token_endpoint: String, +} +#[derive(Deserialize)] +struct Token { + access_token: String, + token_type: String, + expires_in: u64, +} + +async fn bounded_body(mut response: reqwest::Response) -> Result>, String> { + let mut bytes = Zeroizing::new(Vec::new()); + while let Some(chunk) = response + .chunk() + .await + .map_err(|_| "could not read workload response")? + { + if bytes.len().saturating_add(chunk.len()) > 65536 { + return Err("workload response exceeds size limit".into()); + } + bytes.extend_from_slice(&chunk); + } + Ok(bytes) +} + +async fn read_assertion(path: &std::path::Path) -> Result, String> { + let mut options = tokio::fs::OpenOptions::new(); + options.read(true); + #[cfg(unix)] + options.custom_flags(libc::O_NONBLOCK); + let file = options + .open(path) + .await + .map_err(|_| "could not open workload token file")?; + if !file + .metadata() + .await + .map_err(|_| "could not inspect workload token file")? + .is_file() + { + return Err("workload token must be a regular file".into()); + } + let mut bytes = Zeroizing::new(Vec::new()); + file.take((MAX_BYTES + 1) as u64) + .read_to_end(&mut bytes) + .await + .map_err(|_| "could not read workload token file")?; + if bytes.len() > MAX_BYTES { + return Err("workload assertion exceeds size limit".into()); + } + let value = std::str::from_utf8(&bytes) + .map_err(|_| "invalid workload assertion")? + .trim_end_matches(['\r', '\n']); + // Syntax only. Claims are never decoded or treated as authorization. + let parts: Vec<_> = value.split('.').collect(); + if parts.len() != 3 + || parts.iter().any(|p| { + p.is_empty() + || !p + .bytes() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, b'-' | b'_')) + }) + { + return Err("invalid workload assertion".into()); + } + Ok(Zeroizing::new(value.to_string())) +} + +#[cfg(test)] +mod https_tests; diff --git a/src/workload_auth/https_tests.rs b/src/workload_auth/https_tests.rs new file mode 100644 index 00000000..fc308611 --- /dev/null +++ b/src/workload_auth/https_tests.rs @@ -0,0 +1,460 @@ +//! Real TLS exchanges against a scripted localhost authorization server. +#![allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros +)] + +use super::*; +use std::{collections::BTreeMap, sync::Arc}; +use tokio::{io::AsyncWriteExt, net::TcpListener, sync::mpsc, task::JoinHandle}; +use tokio_rustls::{TlsAcceptor, rustls}; + +const CA: &[u8] = include_bytes!("../../fixtures/workload-auth/ca.pem"); +const CERT: &[u8] = include_bytes!("../../fixtures/workload-auth/server.der"); +const KEY: &[u8] = include_bytes!("../../fixtures/workload-auth/server-key.der"); +const ASSERTION: &str = "private-header.private-payload.private-signature"; + +struct Reply { + status: u16, + body: String, +} +impl Reply { + fn json(body: serde_json::Value) -> Self { + Self { + status: 200, + body: body.to_string(), + } + } + fn token(value: &str, expiry: u64) -> Self { + Self::json(serde_json::json!({ + "access_token": value, "token_type": "Bearer", "expires_in": expiry, + "refresh_token": "ignored-refresh-secret" + })) + } + fn status(status: u16) -> Self { + Self { + status, + body: format!("server echoed {ASSERTION} access-secret"), + } + } +} + +#[derive(Debug)] +struct Request { + method: String, + path: String, + headers: String, + body: String, +} +impl Request { + fn form(&self) -> BTreeMap { + url::form_urlencoded::parse(self.body.as_bytes()) + .into_owned() + .collect() + } +} + +struct Server { + origin: String, + replies: mpsc::UnboundedSender, + requests: mpsc::UnboundedReceiver, + task: JoinHandle<()>, +} +impl Drop for Server { + fn drop(&mut self) { + self.task.abort(); + } +} +impl Server { + async fn start() -> Self { + let provider = rustls::crypto::aws_lc_rs::default_provider(); + let config = rustls::ServerConfig::builder_with_provider(Arc::new(provider)) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert( + vec![rustls::pki_types::CertificateDer::from(CERT.to_vec())], + rustls::pki_types::PrivatePkcs8KeyDer::from(KEY.to_vec()).into(), + ) + .unwrap(); + let acceptor = TlsAcceptor::from(Arc::new(config)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin = format!("https://{}", listener.local_addr().unwrap()); + let (replies, mut responses) = mpsc::unbounded_channel::(); + let (records, requests) = mpsc::unbounded_channel(); + let task = tokio::spawn(async move { + loop { + let (socket, _) = listener.accept().await.unwrap(); + let mut stream = acceptor.accept(socket).await.unwrap(); + let mut data = Vec::new(); + let (header_end, length) = loop { + let mut chunk = [0; 4096]; + let n = stream.read(&mut chunk).await.unwrap(); + assert!(n > 0, "request ended before headers"); + data.extend_from_slice(&chunk[..n]); + assert!(data.len() < 32768); + if let Some(end) = data.windows(4).position(|w| w == b"\r\n\r\n") { + let headers = std::str::from_utf8(&data[..end]).unwrap(); + let length = headers + .lines() + .find_map(|line| { + let (key, value) = line.split_once(':')?; + key.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap_or(0); + break (end + 4, length); + } + }; + while data.len() < header_end + length { + let mut chunk = [0; 4096]; + let n = stream.read(&mut chunk).await.unwrap(); + assert!(n > 0); + data.extend_from_slice(&chunk[..n]); + } + let headers = String::from_utf8(data[..header_end].to_vec()).unwrap(); + let mut line = headers.lines().next().unwrap().split_whitespace(); + let request = Request { + method: line.next().unwrap().into(), + path: line.next().unwrap().into(), + body: String::from_utf8(data[header_end..header_end + length].to_vec()) + .unwrap(), + headers, + }; + records.send(request).unwrap(); + // An unexpected retry receives a refusal, so a broken client fails + // promptly rather than hanging while waiting for a scripted reply. + let reply = responses.try_recv().unwrap_or_else(|_| Reply::status(500)); + let response = format!( + "HTTP/1.1 {} Test\r\nContent-Type: application/json\r\nLocation: /redirect-target\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + reply.status, + reply.body.len(), + reply.body + ); + stream.write_all(response.as_bytes()).await.unwrap(); + let _ = stream.shutdown().await; + } + }); + Self { + origin, + replies, + requests, + task, + } + } + fn enqueue(&self, reply: Reply) { + self.replies.send(reply).unwrap(); + } + fn discovery(&self) { + self.enqueue(Reply::json(serde_json::json!({ + "issuer": format!("{}/tenant", self.origin), + "token_endpoint": format!("{}/token", self.origin) + }))); + } + fn exchange(&self, token: &str, expiry: u64) { + self.discovery(); + self.enqueue(Reply::token(token, expiry)); + } + fn session(&self, path: PathBuf, resource: &str) -> Session { + let config = Config { + token_file: path, + expected_issuer: format!("{}/tenant", self.origin), + token_endpoint: format!("{}/token", self.origin), + resource: resource.into(), + }; + test_support::session( + config, + reqwest::Client::builder() + .tls_certs_only([reqwest::Certificate::from_pem(CA).unwrap()]) + .redirect(reqwest::redirect::Policy::none()) + .retry(reqwest::retry::never()) + .no_proxy() + .timeout(Duration::from_secs(5)) + .build() + .unwrap(), + ) + .unwrap() + } + fn drain(&mut self) -> Vec { + let mut requests = Vec::new(); + while let Ok(request) = self.requests.try_recv() { + requests.push(request); + } + requests + } +} + +fn source() -> (tempfile::TempDir, PathBuf) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("assertion"); + std::fs::write(&path, format!("{ASSERTION}\n")).unwrap(); + (dir, path) +} + +#[tokio::test] +async fn https_fixture_chain_and_hostname_are_verified() { + let server = Server::start().await; + let (_dir, path) = source(); + let session = server.session(path, "https://api.example/mcp"); + server.enqueue(Reply::status(200)); + // Keep transport diagnostics in test code; production errors stay redacted. + let response = session.client.get(&server.origin).send().await.unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::OK); +} + +#[tokio::test] +async fn https_grant_single_flight_and_resource_isolation() { + let mut server = Server::start().await; + let (_dir, path) = source(); + let session = Arc::new(server.session(path.clone(), "https://api.example/one")); + server.exchange("first-access", 3600); + let mut jobs = tokio::task::JoinSet::new(); + let barrier = Arc::new(tokio::sync::Barrier::new(16)); + for _ in 0..16 { + let session = session.clone(); + let barrier = barrier.clone(); + jobs.spawn(async move { + barrier.wait().await; + session.token().await + }); + } + while let Some(result) = jobs.join_next().await { + assert_eq!(result.unwrap().unwrap(), "first-access"); + } + assert_eq!(session.token().await.unwrap(), "first-access"); + let requests = server.drain(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].method, "GET"); + assert_eq!( + requests[0].path, + "/.well-known/oauth-authorization-server/tenant" + ); + assert!(requests[0].body.is_empty()); + assert!(!requests[0].headers.contains(ASSERTION)); + assert_eq!(requests[1].method, "POST"); + assert_eq!(requests[1].path, "/token"); + assert!( + !requests[1] + .headers + .to_ascii_lowercase() + .contains("authorization:") + ); + assert_eq!( + requests[1].form(), + BTreeMap::from([ + ( + "grant_type".into(), + "urn:ietf:params:oauth:grant-type:jwt-bearer".into() + ), + ("assertion".into(), ASSERTION.into()), + ("resource".into(), "https://api.example/one".into()), + ]) + ); + let other = server.session(path, "https://api.example/two"); + server.exchange("second-access", 3600); + assert_eq!(other.token().await.unwrap(), "second-access"); + assert_eq!( + server.drain()[1].form()["resource"], + "https://api.example/two" + ); + assert_eq!(session.token().await.unwrap(), "first-access"); + assert!(server.drain().is_empty()); +} + +#[tokio::test] +async fn https_expiry_and_invalidation_reopen_rotated_source() { + let mut server = Server::start().await; + let (_dir, path) = source(); + let session = server.session(path.clone(), "https://api.example/mcp"); + server.exchange("short-access", 1); + assert_eq!(session.token().await.unwrap(), "short-access"); + std::fs::write(&path, "rotated.payload.signature").unwrap(); + assert_eq!(session.token().await.unwrap(), "short-access"); + tokio::time::sleep(Duration::from_millis(1100)).await; + server.exchange("renewed-access", 3600); + assert_eq!(session.token().await.unwrap(), "renewed-access"); + std::fs::write(&path, "invalidated.payload.signature").unwrap(); + session + .invalidate_acquisition(&session.acquisition().await.unwrap()) + .await; + server.exchange("invalidated-access", 3600); + assert_eq!(session.token().await.unwrap(), "invalidated-access"); + let requests = server.drain(); + let assertions: Vec<_> = requests + .iter() + .filter(|r| r.method == "POST") + .map(|r| r.form()["assertion"].clone()) + .collect(); + assert_eq!( + assertions, + [ + ASSERTION, + "rotated.payload.signature", + "invalidated.payload.signature" + ] + ); +} + +#[cfg(unix)] +#[tokio::test] +async fn https_projected_symlink_rotation_is_reopened() { + let mut server = Server::start().await; + let (dir, first) = source(); + let second = dir.path().join("second"); + std::fs::write(&second, "second.payload.signature").unwrap(); + let link = dir.path().join("projected"); + std::os::unix::fs::symlink(&first, &link).unwrap(); + let session = server.session(link.clone(), "https://api.example/mcp"); + server.exchange("first-access", 1); + assert_eq!(session.token().await.unwrap(), "first-access"); + let replacement = dir.path().join("replacement"); + std::os::unix::fs::symlink(&second, &replacement).unwrap(); + std::fs::rename(&replacement, &link).unwrap(); + tokio::time::sleep(Duration::from_millis(1100)).await; + server.exchange("second-access", 3600); + assert_eq!(session.token().await.unwrap(), "second-access"); + std::os::unix::fs::symlink(&first, &replacement).unwrap(); + std::fs::rename(&replacement, &link).unwrap(); + session + .invalidate_acquisition(&session.acquisition().await.unwrap()) + .await; + server.exchange("third-access", 3600); + assert_eq!(session.token().await.unwrap(), "third-access"); + let assertions: Vec<_> = server + .drain() + .iter() + .filter(|r| r.method == "POST") + .map(|r| r.form()["assertion"].clone()) + .collect(); + assert_eq!( + assertions, + [ASSERTION, "second.payload.signature", ASSERTION] + ); +} + +#[tokio::test] +async fn https_discovery_redirect_and_mismatch_never_send_assertion() { + for failure in ["redirect", "issuer", "endpoint"] { + let mut server = Server::start().await; + let (_dir, path) = source(); + let session = server.session(path, "https://api.example/mcp"); + if failure == "redirect" { + server.enqueue(Reply::status(302)); + } else { + server.enqueue(Reply::json(serde_json::json!({ + "issuer": if failure == "issuer" { "https://wrong.example".into() } else { format!("{}/tenant", server.origin) }, + "token_endpoint": if failure == "endpoint" { "https://wrong.example/token".into() } else { format!("{}/token", server.origin) } + }))); + } + let error = session.token().await.unwrap_err(); + assert!(!error.contains(ASSERTION)); + let requests = server.drain(); + assert_eq!(requests.len(), 1, "{failure}"); + assert_eq!(requests[0].method, "GET"); + assert!(requests[0].body.is_empty()); + assert!(!requests[0].headers.contains(ASSERTION)); + } +} + +#[tokio::test] +async fn https_post_failures_are_not_retried_or_leaked() { + for status in [302, 307, 400, 401, 503] { + let mut server = Server::start().await; + let (_dir, path) = source(); + let session = server.session(path, "https://api.example/mcp"); + server.discovery(); + server.enqueue(Reply::status(status)); + assert_eq!( + session.token().await.unwrap_err(), + "workload token exchange refused" + ); + let requests = server.drain(); + assert_eq!(requests.len(), 2, "status {status}"); + assert_eq!(requests[1].method, "POST"); + // A later explicit acquisition can succeed: failures are not cached. + server.exchange("recovered-access", 3600); + assert_eq!(session.token().await.unwrap(), "recovered-access"); + } +} + +#[tokio::test] +async fn https_delayed_rejection_preserves_newer_acquisition() { + let mut server = Server::start().await; + let (_dir, path) = source(); + let session = server.session(path, "https://api.example/mcp"); + server.exchange("same-access-value", 3600); + let first = session.acquisition().await.unwrap(); + let delayed = first.clone(); + session.invalidate_acquisition(&first).await; + // An issuer may return the same token bytes for a distinct exchange. + server.exchange("same-access-value", 3600); + let replacement = session.acquisition().await.unwrap(); + session.invalidate_acquisition(&delayed).await; + let retained = session.acquisition().await.unwrap(); + assert!(Arc::ptr_eq(&replacement.0, &retained.0)); + assert_eq!(server.drain().len(), 4); +} + +#[tokio::test] +async fn https_discovery_503_retries_once_only() { + for recover in [false, true] { + let mut server = Server::start().await; + let (_dir, path) = source(); + let session = server.session(path, "https://api.example/mcp"); + server.enqueue(Reply::status(503)); + if recover { + server.exchange("access", 3600); + } else { + server.enqueue(Reply::status(503)); + } + let result = session.token().await; + if recover { + assert_eq!(result.unwrap(), "access"); + } else { + assert_eq!(result.unwrap_err(), "workload issuer discovery refused"); + } + let requests = server.drain(); + assert_eq!(requests.iter().filter(|r| r.method == "GET").count(), 2); + assert_eq!( + requests.iter().filter(|r| r.method == "POST").count(), + usize::from(recover) + ); + } +} + +#[tokio::test] +async fn https_invalid_responses_are_rejected_without_caching_or_leaks() { + let valid = serde_json::json!({"access_token":"access-secret", "token_type":"Bearer", "expires_in":3600}); + let mut bodies = vec![format!("not-json {ASSERTION}"), "{}".into()]; + for (field, value) in [ + ("token_type", serde_json::json!("Basic")), + ("expires_in", serde_json::json!(0)), + ("expires_in", serde_json::json!(-1)), + ("expires_in", serde_json::json!("3600")), + ("access_token", serde_json::json!("")), + ("access_token", serde_json::json!("access-secret\n")), + ("access_token", serde_json::json!("a".repeat(8193))), + ] { + let mut body = valid.clone(); + body[field] = value; + bodies.push(body.to_string()); + } + for body in bodies { + let mut server = Server::start().await; + let (_dir, path) = source(); + let session = server.session(path, "https://api.example/mcp"); + server.discovery(); + server.enqueue(Reply { status: 200, body }); + assert_eq!( + session.token().await.unwrap_err(), + "invalid workload token response" + ); + assert_eq!(server.drain().len(), 2); + server.exchange("valid-access", 3600); + assert_eq!(session.token().await.unwrap(), "valid-access"); + assert_eq!(server.drain().len(), 2); + } +} diff --git a/src/workload_auth/tests.rs b/src/workload_auth/tests.rs new file mode 100644 index 00000000..b3e144c3 --- /dev/null +++ b/src/workload_auth/tests.rs @@ -0,0 +1,150 @@ +#![cfg_attr( + test, + allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros + ) +)] +use super::*; + +fn config() -> Config { + Config { + token_file: PathBuf::from("/run/identity/token"), + expected_issuer: "https://issuer.example/tenant".into(), + token_endpoint: "https://issuer.example/tenant/token".into(), + resource: "https://resource.example/mcp".into(), + } +} + +#[test] +fn rejects_unsafe_configuration_before_credentials() { + assert!(config().validate().is_ok()); + for endpoint in [ + "http://issuer.example/token", + "https://other.example/token", + "https://secret@issuer.example/token", + "https://issuer.example/token#fragment", + ] { + let mut candidate = config(); + candidate.token_endpoint = endpoint.into(); + assert!(Session::new(candidate).is_err()); + } + let mut candidate = config(); + candidate.token_file = PathBuf::from("relative/token"); + assert!(Session::new(candidate).is_err()); +} + +#[tokio::test] +async fn file_reopens_and_malformed_tokens_are_redacted() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("token"); + assert!(read_assertion(&path).await.is_err()); + for value in ["aaa.bbb.ccc", "ddd.eee.fff"] { + std::fs::write(&path, format!("{value}\n")).unwrap(); + assert_eq!(read_assertion(&path).await.unwrap().as_str(), value); + } + for value in ["", "secret", "a..b", "a.b.c.d", "a. b.c", "a.b.c\nsecret"] { + std::fs::write(&path, value).unwrap(); + assert_eq!( + read_assertion(&path).await.err().unwrap(), + "invalid workload assertion" + ); + } + std::fs::write(&path, vec![b'a'; MAX_BYTES + 1]).unwrap(); + assert_eq!( + read_assertion(&path).await.err().unwrap(), + "workload assertion exceeds size limit" + ); + assert!(read_assertion(dir.path()).await.is_err()); +} + +#[cfg(unix)] +#[tokio::test] +async fn projected_symlink_rotation_is_supported() { + use std::os::unix::fs::symlink; + let dir = tempfile::tempdir().unwrap(); + for (version, value) in [("v1", "aaa.bbb.ccc"), ("v2", "ddd.eee.fff")] { + std::fs::create_dir(dir.path().join(version)).unwrap(); + std::fs::write(dir.path().join(version).join("token"), value).unwrap(); + } + symlink("v1", dir.path().join("..data")).unwrap(); + symlink("..data/token", dir.path().join("token")).unwrap(); + assert_eq!( + read_assertion(&dir.path().join("token")) + .await + .unwrap() + .as_str(), + "aaa.bbb.ccc" + ); + symlink("v2", dir.path().join("..data-next")).unwrap(); + std::fs::rename(dir.path().join("..data-next"), dir.path().join("..data")).unwrap(); + assert_eq!( + read_assertion(&dir.path().join("token")) + .await + .unwrap() + .as_str(), + "ddd.eee.fff" + ); +} + +#[test] +fn response_requires_expiry_and_bearer_fields() { + for body in [ + r#"{"access_token":"secret","token_type":"Bearer"}"#, + r#"{"access_token":"secret","expires_in":60}"#, + r#"{"access_token":"secret","token_type":"Bearer","expires_in":-1}"#, + ] { + assert!(serde_json::from_str::(body).is_err()); + } +} + +#[tokio::test] +async fn expiry_and_clock_rollback_reacquire_without_stale_fallback() { + for expired in [false, true] { + let mut session = Session::new(config()).unwrap(); + let now = SystemTime::now(); + *session.cached.lock().await = Some(Cached { + value: Acquisition(Arc::new(Zeroizing::new("must-not-escape".into()))), + acquired: if expired { + now - Duration::from_secs(60) + } else { + now + Duration::from_secs(60) + }, + refresh_at: if expired { + now + } else { + now + Duration::from_secs(60) + }, + }); + // A failed discovery must not expose a stale cache. This invalid URL + // exercises the real pre-network validation boundary without an issuer. + session.config.expected_issuer = "http://issuer.example".into(); + assert!(session.token().await.is_err()); + assert!(session.cached.lock().await.is_none()); + } +} + +#[tokio::test] +async fn caches_are_resource_local_and_invalidatable() { + let first = Session::new(config()).unwrap(); + let mut second_config = config(); + second_config.resource = "https://resource.example/other-mcp".into(); + let second = Session::new(second_config).unwrap(); + let now = SystemTime::now(); + *first.cached.lock().await = Some(Cached { + value: Acquisition(Arc::new(Zeroizing::new("first-resource-token".into()))), + acquired: now, + refresh_at: now + Duration::from_secs(60), + }); + assert!(second.cached.lock().await.is_none()); + let (a, b) = tokio::join!(first.token(), first.token()); + assert_eq!(a.unwrap(), "first-resource-token"); + assert_eq!(b.unwrap(), "first-resource-token"); + first + .invalidate_acquisition(&first.acquisition().await.unwrap()) + .await; + assert!(first.cached.lock().await.is_none()); +}