From 92bedc525ce67e4f754faceba02e97a0331a9b37 Mon Sep 17 00:00:00 2001 From: daniel Date: Fri, 23 Oct 2026 20:42:43 +0100 Subject: [PATCH 1/5] feat(auth): support workload JWT assertions for MCP and inference --- docs/reference/workload-authentication.md | 36 +++ .../user/getting-started-and-configuration.md | 3 + docs/user/mcp.md | 35 +++ docs/user/speakeasy-workload-auth.md | 47 +++ src/lib.rs | 1 + src/provider/adapter.rs | 30 +- src/provider/mod.rs | 1 + src/provider/speakeasy_workload.rs | 168 +++++++++++ src/tools/mcp.rs | 20 +- src/tools/mcp/auth.rs | 17 +- src/tools/mcp/workload.rs | 273 +++++++++++++++++ src/workload_auth.rs | 280 ++++++++++++++++++ src/workload_auth/tests.rs | 140 +++++++++ 13 files changed, 1042 insertions(+), 9 deletions(-) create mode 100644 docs/reference/workload-authentication.md create mode 100644 docs/user/speakeasy-workload-auth.md create mode 100644 src/provider/speakeasy_workload.rs create mode 100644 src/tools/mcp/workload.rs create mode 100644 src/workload_auth.rs create mode 100644 src/workload_auth/tests.rs diff --git a/docs/reference/workload-authentication.md b/docs/reference/workload-authentication.md new file mode 100644 index 00000000..33af5bf0 --- /dev/null +++ b/docs/reference/workload-authentication.md @@ -0,0 +1,36 @@ +# Workload assertion authentication + +Use workload authentication when a deployment supplies a short-lived JWT in a mounted file instead of a user login or API key. The issuer must accept the clientless RFC 7523 JWT bearer grant for the intended resource. Kit does not register a client, send a client secret, use a refresh token, or open a browser in this mode. + +The shared configuration is: + +```json +{ + "tokenFile": "/run/identity/token", + "expectedIssuer": "https://auth.example.com/tenant", + "tokenEndpoint": "https://auth.example.com/tenant/token", + "resource": "https://service.example.com/mcp" +} +``` + +All fields are required. `tokenFile` must be absolute. URLs must use HTTPS without user information or fragments; the token endpoint must share the expected issuer's origin. Kit reads RFC 8414 authorization-server metadata and requires exact matches for the configured issuer and token endpoint before opening the assertion file. Redirects are not followed. For a path-based issuer, metadata is requested at `/.well-known/oauth-authorization-server/`. + +Treat the configuration, its containing directories, and the token mount as deployment-owned inputs. Do not let a model, repository dependency, or untrusted MCP response choose them. Protect the source assertion from unrelated processes. Kit follows symlinks so that atomically rotated projected token volumes work; it checks the opened handle is a regular file and reads at most 8 KiB. It checks compact JWT syntax but does not decode claims or use unverified claims for authorization. + +Every exchange reopens the file and sends exactly `grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer`, `assertion`, and `resource` as a form body. Configure independent resources for MCP and inference; their sessions do not share access-token caches. + +## Expiry, retries, and replay policy + +Access tokens are kept only in process memory. Kit requires a positive `expires_in` and bearer token type, refreshes before the returned expiry, and caps its cache lifetime at one day. Acquisition is serialized per session and bounded to 15 seconds, including waiting for another acquisition. Missing files, malformed assertions, metadata mismatches, and exchange failures fail closed with redacted errors. + +Metadata discovery retries a 503 response once. Token POSTs are not automatically retried because a lost response may already have consumed the assertion's replay identifier. No refresh token from a response is retained or used. + +The issuer's replay policy still applies. A projected assertion can remain unchanged between exchanges or across independently scoped MCP and inference sessions. Such deployments require an explicitly configured reusable-assertion policy on the server, or a source that rotates the assertion before each exchange. Kit does not bypass or weaken the server's default single-use policy. An exchange refusal requires checking issuer admission, audience, resource, expiry, and replay configuration; API keys are not a fallback. + +## Restoring sessions and process snapshots + +Ordinary transcript restoration constructs fresh authentication clients and reacquires workload tokens. No workload assertion or access token is written to Kit's credential store. + +A process-memory snapshot is different: it can retain tokens, HTTP pools, and work in flight. Kit does not expose a CLI, signal, or gateway resume hook for such snapshots. Applications must not resume a snapshotted executable with resident authentication state. Restart Kit after restore and restore its transcript instead. + +An application embedding the Kit library can integrate `kit::workload_auth::invalidate_after_restore()`. Quiesce requests before the snapshot and keep them quiesced through restore; remount the current assertion source, invoke the hook, and only then admit new requests. The hook invalidates all shared workload token caches and fences acquisitions from the old generation. The application remains responsible for discarding or reconnecting snapshot-resident transports and ensuring no pre-restore request is replayed. Calling the hook without that application lifecycle barrier is not a complete snapshot recovery protocol. diff --git a/docs/user/getting-started-and-configuration.md b/docs/user/getting-started-and-configuration.md index 97840857..97f4475d 100644 --- a/docs/user/getting-started-and-configuration.md +++ b/docs/user/getting-started-and-configuration.md @@ -146,6 +146,9 @@ or behavior configuration. ### Speakeasy AI Control Plane +For deployment-provided JWT assertions instead of a user login, see +[Speakeasy workload authentication](speakeasy-workload-auth.md). + Sign in through the Speakeasy dashboard, then use the same persistent credential store for runtime commands: diff --git a/docs/user/mcp.md b/docs/user/mcp.md index 4c87912d..f9bd9961 100644 --- a/docs/user/mcp.md +++ b/docs/user/mcp.md @@ -6,6 +6,41 @@ 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" + } + } + } +} +``` + +The resource must match the HTTPS MCP URL. Do not combine this mode with a static +bearer token, an Authorization header, or a Host override. Each HTTP operation, +including SSE reconnection, acquires or reuses a valid resource-scoped token. +Authentication failure does not open a browser or fall back to saved credentials. +The file is reopened at each token exchange, including symlink-based rotation. + +Keep this configuration and the assertion mount under deployment control. See +[workload authentication](../reference/workload-authentication.md) for issuer +metadata requirements, bounded acquisition, replay-policy prerequisites, and the +required process-memory restore barrier. Native inference uses its own +[Speakeasy workload configuration](speakeasy-workload-auth.md) and resource. + ## 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/docs/user/speakeasy-workload-auth.md b/docs/user/speakeasy-workload-auth.md new file mode 100644 index 00000000..d860583c --- /dev/null +++ b/docs/user/speakeasy-workload-auth.md @@ -0,0 +1,47 @@ +# Speakeasy workload authentication + +Use native Speakeasy inference without a saved API key or browser sign-in. Supply +a short-lived JWT assertion file and an HTTPS authorization server that accepts +the RFC 7523 JWT bearer grant. Ask your administrator for the issuer, token +endpoint, and inference resource. The token endpoint must share the issuer origin. + +Create a deployment-owned file such as `/etc/kit/speakeasy-workload.json`: + +```json +{ + "project": "your-gram-project", + "workload": { + "tokenFile": "/run/identity/assertion.jwt", + "expectedIssuer": "https://identity.example.com", + "tokenEndpoint": "https://identity.example.com/oauth/token", + "resource": "https://ai.speakeasy.com" + } +} +``` + +Set `KIT_SPEAKEASY_WORKLOAD_AUTH=/etc/kit/speakeasy-workload.json` in the process +environment. The configuration path must be absolute and name a regular file of +at most 64 KiB. Select a `speakeasy:` model as usual. This explicitly selects +workload mode: invalid configuration or token acquisition fails inference, without +API-key, browser, or refresh-token fallback. Unset the variable to restore normal +Speakeasy authentication. + +Keep configuration under deployment control, not in model-generated input or MCP +server responses. `tokenFile` must be absolute. The issuer, endpoint, and resource +must be HTTPS URLs. The resource must use the `https://ai.speakeasy.com` origin on port 443; MCP access uses separate token +sessions even if the assertion source is shared. Tokens are not saved to the +credential store. + +Requests send `Authorization: Bearer` and `Gram-Project`, not `Gram-Key`. The bearer +is restricted to `https://ai.speakeasy.com` on port 443. Redirects and custom +inference origins are unsupported. A 401 invalidates the token for the next +request, without automatically replaying inference. + +Cold starts construct fresh token sessions. Embedders restoring process-memory +snapshots must quiesce requests and call +`kit::workload_auth::invalidate_after_restore()` before resuming. Memory restore +without that hook is unsupported for workload authentication. + +See [workload authentication and restore requirements](../reference/workload-authentication.md) +for replay-policy prerequisites, expiry, and the application snapshot barrier. +See [Configuration and authentication](getting-started-and-configuration.md). diff --git a/src/lib.rs b/src/lib.rs index 2fcf2437..5546dab8 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; +pub mod workload_auth; /// Shared internal filesystem and process-lifetime recovery controls. pub mod resilient_fs { pub use crate::filesystem::*; diff --git a/src/provider/adapter.rs b/src/provider/adapter.rs index a94b4fcc..246464d6 100644 --- a/src/provider/adapter.rs +++ b/src/provider/adapter.rs @@ -742,15 +742,22 @@ impl KitAdapter { })) } ProviderKind::Speakeasy => { - let credentials = speakeasy_auth::load(&credential_storage)? - .ok_or_else(|| SPEAKEASY_AUTH_REQUIRED.to_string())?; + let workload = super::speakeasy_workload::Config::from_env()?; + let (api_key, project) = match &workload { + Some(config) => (String::new(), config.project.clone()), + None => { + let credentials = speakeasy_auth::load(&credential_storage)? + .ok_or_else(|| SPEAKEASY_AUTH_REQUIRED.to_string())?; + (credentials.api_key.clone(), credentials.project.clone()) + } + }; let mut config = OpenRouterConfig::new("unused", model).with_base_url(SPEAKEASY_COMPLETIONS_URL); apply_openrouter_reasoning_effort(&mut config, reasoning_effort); let provider = SpeakeasyProvider { openrouter: OpenRouterProvider::from(config), - api_key: credentials.api_key.clone(), - project: credentials.project.clone(), + api_key, + project, chat_id: None, }; let client = reqwest::Client::builder() @@ -758,9 +765,15 @@ impl KitAdapter { .connect_timeout(Duration::from_secs(10)) .build() .map_err(|_| "could not build Speakeasy completions client".to_string())?; + let client = match workload { + Some(config) => agentkit_http::Http::new( + super::speakeasy_workload::Client::new(client, config)?, + ), + None => agentkit_http::Http::new(client), + }; Ok(Self::Speakeasy(Box::new(SpeakeasyKitAdapter { provider, - client: agentkit_http::Http::new(client), + client, }))) } } @@ -1548,7 +1561,12 @@ fn available_catalog_providers( { providers.push(ProviderKind::OpenRouter); } - if speakeasy_auth::load(storage).ok().flatten().is_some() { + let speakeasy_available = match super::speakeasy_workload::Config::from_env() { + Ok(Some(_)) => true, + Ok(None) => speakeasy_auth::load(storage).ok().flatten().is_some(), + Err(_) => false, + }; + if speakeasy_available { providers.push(ProviderKind::Speakeasy); } if super::anthropic_auth::credentials(storage) diff --git a/src/provider/mod.rs b/src/provider/mod.rs index 522f1fd6..7ae071e6 100644 --- a/src/provider/mod.rs +++ b/src/provider/mod.rs @@ -8,6 +8,7 @@ pub mod chatgpt; pub(crate) mod openai_auth; mod openrouter_auth; mod speakeasy_auth; +mod speakeasy_workload; pub(crate) mod typesafe_auth; pub mod usage; diff --git a/src/provider/speakeasy_workload.rs b/src/provider/speakeasy_workload.rs new file mode 100644 index 00000000..f217719f --- /dev/null +++ b/src/provider/speakeasy_workload.rs @@ -0,0 +1,168 @@ +//! Explicit, deployment-owned workload credentials for native inference. +use agentkit_http::{HttpClient, HttpError, HttpRequest, HttpResponse}; +use async_trait::async_trait; +use serde::Deserialize; +use std::io::Read; + +const CONFIG_ENV: &str = "KIT_SPEAKEASY_WORKLOAD_AUTH"; + +#[derive(Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub(super) struct Config { + pub project: String, + pub workload: crate::workload_auth::Config, +} + +impl Config { + pub(super) fn from_env() -> Result, String> { + let Some(path) = std::env::var_os(CONFIG_ENV) else { + return Ok(None); + }; + let path = std::path::PathBuf::from(path); + if !path.is_absolute() { + return Err(format!("{CONFIG_ENV}: configuration path must be absolute")); + } + let mut options = std::fs::OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.custom_flags(libc::O_NONBLOCK); + } + let file = options + .open(path) + .map_err(|_| format!("{CONFIG_ENV}: cannot open workload configuration"))?; + let metadata = file + .metadata() + .map_err(|_| format!("{CONFIG_ENV}: cannot inspect workload configuration"))?; + if !metadata.is_file() || metadata.len() > 65536 { + return Err(format!( + "{CONFIG_ENV}: configuration must be a regular file of at most 64 KiB" + )); + } + let mut bytes = Vec::new(); + file.take(65537) + .read_to_end(&mut bytes) + .map_err(|_| format!("{CONFIG_ENV}: cannot read workload configuration"))?; + if bytes.len() > 65536 { + return Err(format!("{CONFIG_ENV}: configuration exceeds 64 KiB")); + } + Self::parse(&bytes).map(Some) + } + + fn parse(bytes: &[u8]) -> Result { + let config: Self = serde_json::from_slice(bytes) + .map_err(|_| format!("{CONFIG_ENV}: invalid workload configuration"))?; + if config.project.trim().is_empty() + || agentkit_http::HeaderValue::from_str(&config.project).is_err() + { + return Err(format!( + "{CONFIG_ENV}: project must be a nonempty HTTP header value" + )); + } + config.workload.validate()?; + if !inference_origin(&config.workload.resource) { + return Err(format!( + "{CONFIG_ENV}: resource must belong to https://ai.speakeasy.com" + )); + } + Ok(config) + } +} + +pub(super) struct Client { + client: reqwest::Client, + session: crate::workload_auth::Session, +} + +impl Client { + pub(super) fn new(client: reqwest::Client, config: Config) -> Result { + Ok(Self { + client, + session: crate::workload_auth::Session::new(config.workload)?, + }) + } +} + +fn inference_origin(url: &str) -> bool { + let Ok(url) = reqwest::Url::parse(url) else { + return false; + }; + url.scheme() == "https" + && url.host_str() == Some("ai.speakeasy.com") + && url.port_or_known_default() == Some(443) + && url.username().is_empty() + && url.password().is_none() +} + +#[async_trait] +impl HttpClient for Client { + async fn execute(&self, mut request: HttpRequest) -> Result { + // Check before reading or exchanging the assertion. Never follow redirects. + if !inference_origin(&request.url) { + return Err(HttpError::Other( + "workload inference requires https://ai.speakeasy.com".into(), + )); + } + let token = zeroize::Zeroizing::new(self.session.token().await.map_err(HttpError::Other)?); + let mut authorization = + agentkit_http::HeaderValue::from_str(&format!("Bearer {}", token.as_str())) + .map_err(|_| HttpError::Other("invalid workload access token".into()))?; + authorization.set_sensitive(true); + request.headers.remove("Gram-Key"); + request + .headers + .insert(agentkit_http::header::AUTHORIZATION, authorization); + let response = HttpClient::execute(&self.client, request).await?; + if response.status() == agentkit_http::StatusCode::UNAUTHORIZED { + self.session.invalidate().await; + } + Ok(response) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn workload_inference_is_pinned_to_https_origin() { + assert!(inference_origin( + "https://ai.speakeasy.com/chat/completions" + )); + assert!(inference_origin( + "https://ai.speakeasy.com:443/chat/completions" + )); + for url in [ + "http://ai.speakeasy.com/chat/completions", + "https://ai.speakeasy.com:444/chat/completions", + "https://ai.speakeasy.com.evil.example/chat/completions", + "https://evil.example/chat/completions", + "https://user@ai.speakeasy.com/chat/completions", + ] { + assert!(!inference_origin(url), "{url}"); + } + } + + #[test] + fn workload_resource_must_be_inference_origin() { + let mut value = serde_json::json!({ + "project": "test", + "workload": { + "tokenFile": "/run/assertion.jwt", + "expectedIssuer": "https://issuer.example", + "tokenEndpoint": "https://issuer.example/token", + "resource": "https://ai.speakeasy.com" + } + }); + assert!(Config::parse(&serde_json::to_vec(&value).unwrap()).is_ok()); + value["workload"]["resource"] = serde_json::json!("https://mcp.example.com"); + assert!(Config::parse(&serde_json::to_vec(&value).unwrap()).is_err()); + } + + #[test] + fn workload_configuration_is_explicit_and_strict() { + assert!(Config::parse(b"{}").is_err()); + assert!(Config::parse(br#"{"project":"test","workload":{},"apiKey":"secret"}"#).is_err()); + } +} 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..c0a20cb1 --- /dev/null +++ b/src/tools/mcp/workload.rs @@ -0,0 +1,273 @@ +//! Per-request workload authentication, including catalog reads and SSE reconnects. +use std::{collections::HashMap, sync::Arc}; + +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::{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 = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|_| "could not build workload MCP HTTP client")?; + Ok(Self { + destination: destination.to_owned(), + session: Session::new(config)?, + http, + }) + } + + async fn authorization( + &self, + uri: &str, + headers: &HashMap, + ) -> Result, Error> { + 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.token().await.map(Some).map_err(Self::error) + } + + fn error(message: impl Into) -> Error { + StreamableHttpError::Io(std::io::Error::other(message.into())) + } + + async fn finish(&self, 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 rejected { + // Do not replay a possibly side-effecting POST or follow challenge URLs. + self.session.invalidate().await; + } + result + } +} + +#[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 { + let authorization = self.authorization(&uri, &custom_headers).await?; + self.finish( + self.http + .post_message(uri, message, session_id, authorization, custom_headers) + .await, + ) + .await + } + + async fn delete_session( + &self, + uri: Arc, + session_id: Arc, + _auth_header: Option, + custom_headers: HashMap, + ) -> Result<(), Error> { + let authorization = self.authorization(&uri, &custom_headers).await?; + self.finish( + self.http + .delete_session(uri, session_id, authorization, 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 authorization = self.authorization(&uri, &custom_headers).await?; + self.finish( + self.http + .get_stream( + uri, + session_id, + last_event_id, + authorization, + custom_headers, + ) + .await, + ) + .await + } +} + +#[cfg(test)] +#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +mod tests { + use super::*; + + 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 + .unwrap_err(); + 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 + .unwrap_err() + .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 { + unreachable!() + }; + 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..345c475d --- /dev/null +++ b/src/workload_auth.rs @@ -0,0 +1,280 @@ +//! Clientless RFC 7523 workload authentication. Assertions never enter credential storage. +use std::{ + path::PathBuf, + sync::atomic::{AtomicU64, Ordering}, + time::{Duration, SystemTime}, +}; + +use serde::Deserialize; +use tokio::{io::AsyncReadExt, sync::Mutex}; +use zeroize::Zeroizing; + +static GENERATION: AtomicU64 = AtomicU64::new(0); +const MAX_BYTES: usize = 8192; +const TIMEOUT: Duration = Duration::from_secs(15); + +#[cfg(test)] +#[allow(clippy::unwrap_used, clippy::expect_used)] +mod tests; + +/// Embedders restoring a process-memory snapshot must call this while requests are +/// quiesced, before resuming work. Transcript restores create new sessions instead. +/// This invalidates every workload cache, including acquisitions started earlier. +pub fn invalidate_after_restore() { + GENERATION.fetch_add(1, Ordering::AcqRel); +} + +/// 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) +} + +struct Cached { + value: Zeroizing, + acquired: SystemTime, + refresh_at: SystemTime, + generation: u64, +} + +/// 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()) + .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(&self) { + *self.cached.lock().await = None; + } + + pub async fn token(&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 generation = GENERATION.load(Ordering::Acquire); + let now = SystemTime::now(); + if let Some(token) = cached.as_ref() + && token.generation == generation + && now >= token.acquired + && now < token.refresh_at + { + return Ok(token.value.to_string()); + } + *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")?; + if GENERATION.load(Ordering::Acquire) != generation || SystemTime::now() >= refresh_at { + return Err("workload authentication invalidated or expired during exchange".into()); + } + let value = Zeroizing::new(token.access_token); + let result = value.to_string(); + *cached = Some(Cached { + value, + acquired: started, + refresh_at, + generation, + }); + 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())) +} diff --git a/src/workload_auth/tests.rs b/src/workload_auth/tests.rs new file mode 100644 index 00000000..3deb428a --- /dev/null +++ b/src/workload_auth/tests.rs @@ -0,0 +1,140 @@ +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_restore_generation_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: Zeroizing::new("must-not-escape".into()), + acquired: now - Duration::from_secs(60), + refresh_at: if expired { + now + } else { + now + Duration::from_secs(60) + }, + generation: if expired { + GENERATION.load(Ordering::Acquire) + } else { + GENERATION.load(Ordering::Acquire).wrapping_sub(1) + }, + }); + // 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/inference".into(); + let second = Session::new(second_config).unwrap(); + let now = SystemTime::now(); + *first.cached.lock().await = Some(Cached { + value: Zeroizing::new("first-resource-token".into()), + acquired: now, + refresh_at: now + Duration::from_secs(60), + generation: GENERATION.load(Ordering::Acquire), + }); + 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().await; + assert!(first.cached.lock().await.is_none()); +} From 997478db0e49e17250850e985cf1cc184ac3e89a Mon Sep 17 00:00:00 2001 From: daniel Date: Fri, 23 Oct 2026 20:49:08 +0100 Subject: [PATCH 2/5] test(auth): apply test-only lint policy to workload coverage --- docs/user/speakeasy-workload-auth.md | 2 +- src/provider/speakeasy_workload.rs | 7 +++++++ src/tools/mcp/workload.rs | 10 ++++++++-- src/workload_auth.rs | 1 - src/workload_auth/tests.rs | 10 ++++++++++ 5 files changed, 26 insertions(+), 4 deletions(-) diff --git a/docs/user/speakeasy-workload-auth.md b/docs/user/speakeasy-workload-auth.md index d860583c..cc19d9fa 100644 --- a/docs/user/speakeasy-workload-auth.md +++ b/docs/user/speakeasy-workload-auth.md @@ -9,7 +9,7 @@ Create a deployment-owned file such as `/etc/kit/speakeasy-workload.json`: ```json { - "project": "your-gram-project", + "project": "your-project", "workload": { "tokenFile": "/run/identity/assertion.jwt", "expectedIssuer": "https://identity.example.com", diff --git a/src/provider/speakeasy_workload.rs b/src/provider/speakeasy_workload.rs index f217719f..640796c0 100644 --- a/src/provider/speakeasy_workload.rs +++ b/src/provider/speakeasy_workload.rs @@ -122,6 +122,13 @@ impl HttpClient for Client { } #[cfg(test)] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros +)] mod tests { use super::*; diff --git a/src/tools/mcp/workload.rs b/src/tools/mcp/workload.rs index c0a20cb1..180f430e 100644 --- a/src/tools/mcp/workload.rs +++ b/src/tools/mcp/workload.rs @@ -145,7 +145,13 @@ impl McpHttpClient for Client { } #[cfg(test)] -#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +#[allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros +)] mod tests { use super::*; @@ -216,7 +222,7 @@ mod tests { let (binding, _) = super::super::prepare_server("workload", server, Vec::new(), None).unwrap(); let super::super::McpTransportBinding::StreamableHttp(transport) = binding else { - unreachable!() + panic!("expected workload authentication configuration") }; let client = transport.http_client.unwrap(); assert!( diff --git a/src/workload_auth.rs b/src/workload_auth.rs index 345c475d..450ff2e9 100644 --- a/src/workload_auth.rs +++ b/src/workload_auth.rs @@ -14,7 +14,6 @@ const MAX_BYTES: usize = 8192; const TIMEOUT: Duration = Duration::from_secs(15); #[cfg(test)] -#[allow(clippy::unwrap_used, clippy::expect_used)] mod tests; /// Embedders restoring a process-memory snapshot must call this while requests are diff --git a/src/workload_auth/tests.rs b/src/workload_auth/tests.rs index 3deb428a..dce3bb79 100644 --- a/src/workload_auth/tests.rs +++ b/src/workload_auth/tests.rs @@ -1,3 +1,13 @@ +#![cfg_attr( + test, + allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::disallowed_methods, + clippy::disallowed_macros + ) +)] use super::*; fn config() -> Config { From fea9a9941e0a676d33aed24e86ddcba74b4d8178 Mon Sep 17 00:00:00 2001 From: daniel Date: Fri, 23 Oct 2026 20:57:28 +0100 Subject: [PATCH 3/5] refactor(auth): scope workload identity to generic MCP authentication --- Cargo.lock | 1 + Cargo.toml | 1 + docs/reference/workload-authentication.md | 36 -- .../user/getting-started-and-configuration.md | 3 - docs/user/mcp.md | 20 +- docs/user/speakeasy-workload-auth.md | 47 -- fixtures/workload-auth/README.md | 17 + fixtures/workload-auth/ca.pem | 17 + fixtures/workload-auth/server-key.der | Bin 0 -> 1218 bytes fixtures/workload-auth/server.der | Bin 0 -> 767 bytes src/lib.rs | 2 +- src/provider/adapter.rs | 30 +- src/provider/mod.rs | 1 - src/provider/speakeasy_workload.rs | 175 ------- src/workload_auth.rs | 21 +- src/workload_auth/https_tests.rs | 438 ++++++++++++++++++ src/workload_auth/tests.rs | 16 +- 17 files changed, 503 insertions(+), 322 deletions(-) delete mode 100644 docs/reference/workload-authentication.md delete mode 100644 docs/user/speakeasy-workload-auth.md create mode 100644 fixtures/workload-auth/README.md create mode 100644 fixtures/workload-auth/ca.pem create mode 100644 fixtures/workload-auth/server-key.der create mode 100644 fixtures/workload-auth/server.der delete mode 100644 src/provider/speakeasy_workload.rs create mode 100644 src/workload_auth/https_tests.rs 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/reference/workload-authentication.md b/docs/reference/workload-authentication.md deleted file mode 100644 index 33af5bf0..00000000 --- a/docs/reference/workload-authentication.md +++ /dev/null @@ -1,36 +0,0 @@ -# Workload assertion authentication - -Use workload authentication when a deployment supplies a short-lived JWT in a mounted file instead of a user login or API key. The issuer must accept the clientless RFC 7523 JWT bearer grant for the intended resource. Kit does not register a client, send a client secret, use a refresh token, or open a browser in this mode. - -The shared configuration is: - -```json -{ - "tokenFile": "/run/identity/token", - "expectedIssuer": "https://auth.example.com/tenant", - "tokenEndpoint": "https://auth.example.com/tenant/token", - "resource": "https://service.example.com/mcp" -} -``` - -All fields are required. `tokenFile` must be absolute. URLs must use HTTPS without user information or fragments; the token endpoint must share the expected issuer's origin. Kit reads RFC 8414 authorization-server metadata and requires exact matches for the configured issuer and token endpoint before opening the assertion file. Redirects are not followed. For a path-based issuer, metadata is requested at `/.well-known/oauth-authorization-server/`. - -Treat the configuration, its containing directories, and the token mount as deployment-owned inputs. Do not let a model, repository dependency, or untrusted MCP response choose them. Protect the source assertion from unrelated processes. Kit follows symlinks so that atomically rotated projected token volumes work; it checks the opened handle is a regular file and reads at most 8 KiB. It checks compact JWT syntax but does not decode claims or use unverified claims for authorization. - -Every exchange reopens the file and sends exactly `grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer`, `assertion`, and `resource` as a form body. Configure independent resources for MCP and inference; their sessions do not share access-token caches. - -## Expiry, retries, and replay policy - -Access tokens are kept only in process memory. Kit requires a positive `expires_in` and bearer token type, refreshes before the returned expiry, and caps its cache lifetime at one day. Acquisition is serialized per session and bounded to 15 seconds, including waiting for another acquisition. Missing files, malformed assertions, metadata mismatches, and exchange failures fail closed with redacted errors. - -Metadata discovery retries a 503 response once. Token POSTs are not automatically retried because a lost response may already have consumed the assertion's replay identifier. No refresh token from a response is retained or used. - -The issuer's replay policy still applies. A projected assertion can remain unchanged between exchanges or across independently scoped MCP and inference sessions. Such deployments require an explicitly configured reusable-assertion policy on the server, or a source that rotates the assertion before each exchange. Kit does not bypass or weaken the server's default single-use policy. An exchange refusal requires checking issuer admission, audience, resource, expiry, and replay configuration; API keys are not a fallback. - -## Restoring sessions and process snapshots - -Ordinary transcript restoration constructs fresh authentication clients and reacquires workload tokens. No workload assertion or access token is written to Kit's credential store. - -A process-memory snapshot is different: it can retain tokens, HTTP pools, and work in flight. Kit does not expose a CLI, signal, or gateway resume hook for such snapshots. Applications must not resume a snapshotted executable with resident authentication state. Restart Kit after restore and restore its transcript instead. - -An application embedding the Kit library can integrate `kit::workload_auth::invalidate_after_restore()`. Quiesce requests before the snapshot and keep them quiesced through restore; remount the current assertion source, invoke the hook, and only then admit new requests. The hook invalidates all shared workload token caches and fences acquisitions from the old generation. The application remains responsible for discarding or reconnecting snapshot-resident transports and ensuring no pre-restore request is replayed. Calling the hook without that application lifecycle barrier is not a complete snapshot recovery protocol. diff --git a/docs/user/getting-started-and-configuration.md b/docs/user/getting-started-and-configuration.md index 97f4475d..97840857 100644 --- a/docs/user/getting-started-and-configuration.md +++ b/docs/user/getting-started-and-configuration.md @@ -146,9 +146,6 @@ or behavior configuration. ### Speakeasy AI Control Plane -For deployment-provided JWT assertions instead of a user login, see -[Speakeasy workload authentication](speakeasy-workload-auth.md). - Sign in through the Speakeasy dashboard, then use the same persistent credential store for runtime commands: diff --git a/docs/user/mcp.md b/docs/user/mcp.md index f9bd9961..7297fa89 100644 --- a/docs/user/mcp.md +++ b/docs/user/mcp.md @@ -29,17 +29,15 @@ mounted JWT without browser login, client registration, or a client secret: } ``` -The resource must match the HTTPS MCP URL. Do not combine this mode with a static -bearer token, an Authorization header, or a Host override. Each HTTP operation, -including SSE reconnection, acquires or reuses a valid resource-scoped token. -Authentication failure does not open a browser or fall back to saved credentials. -The file is reopened at each token exchange, including symlink-based rotation. - -Keep this configuration and the assertion mount under deployment control. See -[workload authentication](../reference/workload-authentication.md) for issuer -metadata requirements, bounded acquisition, replay-policy prerequisites, and the -required process-memory restore barrier. Native inference uses its own -[Speakeasy workload configuration](speakeasy-workload-auth.md) and resource. +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 diff --git a/docs/user/speakeasy-workload-auth.md b/docs/user/speakeasy-workload-auth.md deleted file mode 100644 index cc19d9fa..00000000 --- a/docs/user/speakeasy-workload-auth.md +++ /dev/null @@ -1,47 +0,0 @@ -# Speakeasy workload authentication - -Use native Speakeasy inference without a saved API key or browser sign-in. Supply -a short-lived JWT assertion file and an HTTPS authorization server that accepts -the RFC 7523 JWT bearer grant. Ask your administrator for the issuer, token -endpoint, and inference resource. The token endpoint must share the issuer origin. - -Create a deployment-owned file such as `/etc/kit/speakeasy-workload.json`: - -```json -{ - "project": "your-project", - "workload": { - "tokenFile": "/run/identity/assertion.jwt", - "expectedIssuer": "https://identity.example.com", - "tokenEndpoint": "https://identity.example.com/oauth/token", - "resource": "https://ai.speakeasy.com" - } -} -``` - -Set `KIT_SPEAKEASY_WORKLOAD_AUTH=/etc/kit/speakeasy-workload.json` in the process -environment. The configuration path must be absolute and name a regular file of -at most 64 KiB. Select a `speakeasy:` model as usual. This explicitly selects -workload mode: invalid configuration or token acquisition fails inference, without -API-key, browser, or refresh-token fallback. Unset the variable to restore normal -Speakeasy authentication. - -Keep configuration under deployment control, not in model-generated input or MCP -server responses. `tokenFile` must be absolute. The issuer, endpoint, and resource -must be HTTPS URLs. The resource must use the `https://ai.speakeasy.com` origin on port 443; MCP access uses separate token -sessions even if the assertion source is shared. Tokens are not saved to the -credential store. - -Requests send `Authorization: Bearer` and `Gram-Project`, not `Gram-Key`. The bearer -is restricted to `https://ai.speakeasy.com` on port 443. Redirects and custom -inference origins are unsupported. A 401 invalidates the token for the next -request, without automatically replaying inference. - -Cold starts construct fresh token sessions. Embedders restoring process-memory -snapshots must quiesce requests and call -`kit::workload_auth::invalidate_after_restore()` before resuming. Memory restore -without that hook is unsupported for workload authentication. - -See [workload authentication and restore requirements](../reference/workload-authentication.md) -for replay-policy prerequisites, expiry, and the application snapshot barrier. -See [Configuration and authentication](getting-started-and-configuration.md). 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 0000000000000000000000000000000000000000..c10b3ca42acb77afe24d9aa1cf2349f136984360 GIT binary patch literal 1218 zcmV;z1U>sOf&{(-0RS)!1_>&LNQUrs4#*Aqyhl|0)hbn0I<>(Bl2|9 zM~JJkWvbKu^S)CmsJ`2g8G{0Gx32i6-M$MFq<`)P-Q<~579`XsB&1d>Z0XF zhTmv6Y{dt9@Vg4174aa47rR~d?p;UJ7M})DXpKjmymN`EAiqqgavxy}Ak99d`Bac- z=F|NW3zwZ{$b2dt^tHp(-9yEFg$&1#%WPEq|IJ6u`=tWd)w*|M(EU1vB$h}WElqJ@Z!^_bxm^2JUuR9s%!DJa^d6{ifMA@{s%nz^$ z7Q$UI=i?%>X>9>dXqL#9m4Xq@vuIH;$dK&LaIgNZUHWM6`W+O*h3-x4`Z3noN`Xuo z8koJAab{NKUvk|>3-IcTE_+5)NR>JDoRX>tcG z-%7WE0DCfHNcJxFf$dHh=c(<`M}q=^fdJ0~H9RDYZ3Z|<@T0WtDFPM5i-LC&tK_O3 zpc0)6NReifMhS#Zd4& zoBI08n9YeOBxq&Y(NbdQ;Tk{b#4QZAByJOkC?r76*1`UEL>!lOoh;08%w;zqa~lGI zfdH7>ph=`aJE}6Tmw0Ka?^K=T4Y1d;6r*}xAW`cXe=?nRaeYZy&09;VKD1Fga%@ub zsgI)zL*&blpm46f=?!hNgGozJZJj!GR8G2wx`ecTb@gfJFn>t}QpA}hp53|DF|NLv zz@yvj?`f6#a#5`r#|eTj1qjtmhqXNL8tnpsfdIHg5y1?Dmz2KCX0bw_4KCL6l7G6u zdlXx}Z9K$aL*}Ir9R?%kTgrcxxfaxD6#~Y?NxBv|!cQe2`sDm{nV;upy zH`@Y%fHJt{<)j-;|El>(Fte1L?|(xV$y`R&oVhZY)(<6ne^_pV{AxoVztx8sLd~Cb zv@X7l8c_Fz=zs?P_~%BS<<`-T(jq literal 0 HcmV?d00001 diff --git a/fixtures/workload-auth/server.der b/fixtures/workload-auth/server.der new file mode 100644 index 0000000000000000000000000000000000000000..cee77e446355001ccf6b317a57dd878bb5b6d6e3 GIT binary patch literal 767 zcmXqLV)|{+#Q1mtGZP~d6CR@3)im-mVS_YS$@W>M}PLl9uf2VD)&_BK#bFoPv<5vtJc+=XZg(WSk005!2Ya$ zef!&g7Jgb>qwrH#Y+wGTSyxwtTmRg4Nta{Cnl*Dbg-Q8DOxPqhZP!1?^D_>F-JaYV z)$#6CuTbYkY57yrLnVJV{hIaM_+z7isM{x5gW0At0#fbc`7e4bT2U0S`PHJQZf*B7 z%yW*iSAE#cHDC0jLc93x*l+J*-LHtvXA8{eb)UPZxO1h#exDUZ@`>CEXYH2$44IJe z?DAhB?&)*WPSj}0e%X5X%3bH9wXHnIC!Ed+`Sbs*``O=1nQmO&Ri1p2iJ6gsak007 zr-2kO>|_O5gbf6n5aHL(QqRD^Xut`Q6J}vDU@+hYG5J|on3-4?7zl$nsw_MPTx=ZL zY>cd|?9Av<3k+^X2AhTx3EO|}?OprkwRO;n)?E%}_KD|DD7lONiTcB{X^)eU=B0Jz zDsx?@8M1C#-@BsNj^k3m?f%KH46gq2oEKhqaSKOh^bEhp#tc);o*B$l7J7L9`UA~* z{oDd?P2D|Bj&V;}^=Pzd|-Jem-Dee?d1(R${XJ6wa7M=&GU%&)V!9prF&}@Z|C { - let workload = super::speakeasy_workload::Config::from_env()?; - let (api_key, project) = match &workload { - Some(config) => (String::new(), config.project.clone()), - None => { - let credentials = speakeasy_auth::load(&credential_storage)? - .ok_or_else(|| SPEAKEASY_AUTH_REQUIRED.to_string())?; - (credentials.api_key.clone(), credentials.project.clone()) - } - }; + let credentials = speakeasy_auth::load(&credential_storage)? + .ok_or_else(|| SPEAKEASY_AUTH_REQUIRED.to_string())?; let mut config = OpenRouterConfig::new("unused", model).with_base_url(SPEAKEASY_COMPLETIONS_URL); apply_openrouter_reasoning_effort(&mut config, reasoning_effort); let provider = SpeakeasyProvider { openrouter: OpenRouterProvider::from(config), - api_key, - project, + api_key: credentials.api_key.clone(), + project: credentials.project.clone(), chat_id: None, }; let client = reqwest::Client::builder() @@ -765,15 +758,9 @@ impl KitAdapter { .connect_timeout(Duration::from_secs(10)) .build() .map_err(|_| "could not build Speakeasy completions client".to_string())?; - let client = match workload { - Some(config) => agentkit_http::Http::new( - super::speakeasy_workload::Client::new(client, config)?, - ), - None => agentkit_http::Http::new(client), - }; Ok(Self::Speakeasy(Box::new(SpeakeasyKitAdapter { provider, - client, + client: agentkit_http::Http::new(client), }))) } } @@ -1561,12 +1548,7 @@ fn available_catalog_providers( { providers.push(ProviderKind::OpenRouter); } - let speakeasy_available = match super::speakeasy_workload::Config::from_env() { - Ok(Some(_)) => true, - Ok(None) => speakeasy_auth::load(storage).ok().flatten().is_some(), - Err(_) => false, - }; - if speakeasy_available { + if speakeasy_auth::load(storage).ok().flatten().is_some() { providers.push(ProviderKind::Speakeasy); } if super::anthropic_auth::credentials(storage) diff --git a/src/provider/mod.rs b/src/provider/mod.rs index 7ae071e6..522f1fd6 100644 --- a/src/provider/mod.rs +++ b/src/provider/mod.rs @@ -8,7 +8,6 @@ pub mod chatgpt; pub(crate) mod openai_auth; mod openrouter_auth; mod speakeasy_auth; -mod speakeasy_workload; pub(crate) mod typesafe_auth; pub mod usage; diff --git a/src/provider/speakeasy_workload.rs b/src/provider/speakeasy_workload.rs deleted file mode 100644 index 640796c0..00000000 --- a/src/provider/speakeasy_workload.rs +++ /dev/null @@ -1,175 +0,0 @@ -//! Explicit, deployment-owned workload credentials for native inference. -use agentkit_http::{HttpClient, HttpError, HttpRequest, HttpResponse}; -use async_trait::async_trait; -use serde::Deserialize; -use std::io::Read; - -const CONFIG_ENV: &str = "KIT_SPEAKEASY_WORKLOAD_AUTH"; - -#[derive(Clone, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub(super) struct Config { - pub project: String, - pub workload: crate::workload_auth::Config, -} - -impl Config { - pub(super) fn from_env() -> Result, String> { - let Some(path) = std::env::var_os(CONFIG_ENV) else { - return Ok(None); - }; - let path = std::path::PathBuf::from(path); - if !path.is_absolute() { - return Err(format!("{CONFIG_ENV}: configuration path must be absolute")); - } - let mut options = std::fs::OpenOptions::new(); - options.read(true); - #[cfg(unix)] - { - use std::os::unix::fs::OpenOptionsExt; - options.custom_flags(libc::O_NONBLOCK); - } - let file = options - .open(path) - .map_err(|_| format!("{CONFIG_ENV}: cannot open workload configuration"))?; - let metadata = file - .metadata() - .map_err(|_| format!("{CONFIG_ENV}: cannot inspect workload configuration"))?; - if !metadata.is_file() || metadata.len() > 65536 { - return Err(format!( - "{CONFIG_ENV}: configuration must be a regular file of at most 64 KiB" - )); - } - let mut bytes = Vec::new(); - file.take(65537) - .read_to_end(&mut bytes) - .map_err(|_| format!("{CONFIG_ENV}: cannot read workload configuration"))?; - if bytes.len() > 65536 { - return Err(format!("{CONFIG_ENV}: configuration exceeds 64 KiB")); - } - Self::parse(&bytes).map(Some) - } - - fn parse(bytes: &[u8]) -> Result { - let config: Self = serde_json::from_slice(bytes) - .map_err(|_| format!("{CONFIG_ENV}: invalid workload configuration"))?; - if config.project.trim().is_empty() - || agentkit_http::HeaderValue::from_str(&config.project).is_err() - { - return Err(format!( - "{CONFIG_ENV}: project must be a nonempty HTTP header value" - )); - } - config.workload.validate()?; - if !inference_origin(&config.workload.resource) { - return Err(format!( - "{CONFIG_ENV}: resource must belong to https://ai.speakeasy.com" - )); - } - Ok(config) - } -} - -pub(super) struct Client { - client: reqwest::Client, - session: crate::workload_auth::Session, -} - -impl Client { - pub(super) fn new(client: reqwest::Client, config: Config) -> Result { - Ok(Self { - client, - session: crate::workload_auth::Session::new(config.workload)?, - }) - } -} - -fn inference_origin(url: &str) -> bool { - let Ok(url) = reqwest::Url::parse(url) else { - return false; - }; - url.scheme() == "https" - && url.host_str() == Some("ai.speakeasy.com") - && url.port_or_known_default() == Some(443) - && url.username().is_empty() - && url.password().is_none() -} - -#[async_trait] -impl HttpClient for Client { - async fn execute(&self, mut request: HttpRequest) -> Result { - // Check before reading or exchanging the assertion. Never follow redirects. - if !inference_origin(&request.url) { - return Err(HttpError::Other( - "workload inference requires https://ai.speakeasy.com".into(), - )); - } - let token = zeroize::Zeroizing::new(self.session.token().await.map_err(HttpError::Other)?); - let mut authorization = - agentkit_http::HeaderValue::from_str(&format!("Bearer {}", token.as_str())) - .map_err(|_| HttpError::Other("invalid workload access token".into()))?; - authorization.set_sensitive(true); - request.headers.remove("Gram-Key"); - request - .headers - .insert(agentkit_http::header::AUTHORIZATION, authorization); - let response = HttpClient::execute(&self.client, request).await?; - if response.status() == agentkit_http::StatusCode::UNAUTHORIZED { - self.session.invalidate().await; - } - Ok(response) - } -} - -#[cfg(test)] -#[allow( - clippy::unwrap_used, - clippy::expect_used, - clippy::panic, - clippy::disallowed_methods, - clippy::disallowed_macros -)] -mod tests { - use super::*; - - #[test] - fn workload_inference_is_pinned_to_https_origin() { - assert!(inference_origin( - "https://ai.speakeasy.com/chat/completions" - )); - assert!(inference_origin( - "https://ai.speakeasy.com:443/chat/completions" - )); - for url in [ - "http://ai.speakeasy.com/chat/completions", - "https://ai.speakeasy.com:444/chat/completions", - "https://ai.speakeasy.com.evil.example/chat/completions", - "https://evil.example/chat/completions", - "https://user@ai.speakeasy.com/chat/completions", - ] { - assert!(!inference_origin(url), "{url}"); - } - } - - #[test] - fn workload_resource_must_be_inference_origin() { - let mut value = serde_json::json!({ - "project": "test", - "workload": { - "tokenFile": "/run/assertion.jwt", - "expectedIssuer": "https://issuer.example", - "tokenEndpoint": "https://issuer.example/token", - "resource": "https://ai.speakeasy.com" - } - }); - assert!(Config::parse(&serde_json::to_vec(&value).unwrap()).is_ok()); - value["workload"]["resource"] = serde_json::json!("https://mcp.example.com"); - assert!(Config::parse(&serde_json::to_vec(&value).unwrap()).is_err()); - } - - #[test] - fn workload_configuration_is_explicit_and_strict() { - assert!(Config::parse(b"{}").is_err()); - assert!(Config::parse(br#"{"project":"test","workload":{},"apiKey":"secret"}"#).is_err()); - } -} diff --git a/src/workload_auth.rs b/src/workload_auth.rs index 450ff2e9..f496c1d8 100644 --- a/src/workload_auth.rs +++ b/src/workload_auth.rs @@ -1,7 +1,6 @@ //! Clientless RFC 7523 workload authentication. Assertions never enter credential storage. use std::{ path::PathBuf, - sync::atomic::{AtomicU64, Ordering}, time::{Duration, SystemTime}, }; @@ -9,20 +8,12 @@ use serde::Deserialize; use tokio::{io::AsyncReadExt, sync::Mutex}; use zeroize::Zeroizing; -static GENERATION: AtomicU64 = AtomicU64::new(0); const MAX_BYTES: usize = 8192; const TIMEOUT: Duration = Duration::from_secs(15); #[cfg(test)] mod tests; -/// Embedders restoring a process-memory snapshot must call this while requests are -/// quiesced, before resuming work. Transcript restores create new sessions instead. -/// This invalidates every workload cache, including acquisitions started earlier. -pub fn invalidate_after_restore() { - GENERATION.fetch_add(1, Ordering::AcqRel); -} - /// Trusted deployment configuration, not input supplied by an MCP server or model. #[derive(Clone, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -65,7 +56,6 @@ struct Cached { value: Zeroizing, acquired: SystemTime, refresh_at: SystemTime, - generation: u64, } /// One resource and one credential source per cache. No refresh tokens or fallback. @@ -105,10 +95,8 @@ impl Session { async fn acquire(&self) -> Result { let mut cached = self.cached.lock().await; - let generation = GENERATION.load(Ordering::Acquire); let now = SystemTime::now(); if let Some(token) = cached.as_ref() - && token.generation == generation && now >= token.acquired && now < token.refresh_at { @@ -153,8 +141,9 @@ impl Session { let refresh_at = started .checked_add(usable) .ok_or("invalid workload token expiry")?; - if GENERATION.load(Ordering::Acquire) != generation || SystemTime::now() >= refresh_at { - return Err("workload authentication invalidated or expired during exchange".into()); + let completed = SystemTime::now(); + if completed < started || completed >= refresh_at { + return Err("workload authentication expired during exchange".into()); } let value = Zeroizing::new(token.access_token); let result = value.to_string(); @@ -162,7 +151,6 @@ impl Session { value, acquired: started, refresh_at, - generation, }); Ok(result) } @@ -277,3 +265,6 @@ async fn read_assertion(path: &std::path::Path) -> Result, Str } 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..aa83ce78 --- /dev/null +++ b/src/workload_auth/https_tests.rs @@ -0,0 +1,438 @@ +//! 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(), + }; + config.validate().unwrap(); + Session { + config, + client: reqwest::Client::builder() + .tls_certs_only([reqwest::Certificate::from_pem(CA).unwrap()]) + .redirect(reqwest::redirect::Policy::none()) + .no_proxy() + .timeout(Duration::from_secs(5)) + .build() + .unwrap(), + cached: Mutex::new(None), + } + } + 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().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().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_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 index dce3bb79..ea12d267 100644 --- a/src/workload_auth/tests.rs +++ b/src/workload_auth/tests.rs @@ -102,22 +102,21 @@ fn response_requires_expiry_and_bearer_fields() { } #[tokio::test] -async fn expiry_and_restore_generation_reacquire_without_stale_fallback() { +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: Zeroizing::new("must-not-escape".into()), - acquired: now - Duration::from_secs(60), - refresh_at: if expired { - now + acquired: if expired { + now - Duration::from_secs(60) } else { now + Duration::from_secs(60) }, - generation: if expired { - GENERATION.load(Ordering::Acquire) + refresh_at: if expired { + now } else { - GENERATION.load(Ordering::Acquire).wrapping_sub(1) + now + Duration::from_secs(60) }, }); // A failed discovery must not expose a stale cache. This invalid URL @@ -132,14 +131,13 @@ async fn expiry_and_restore_generation_reacquire_without_stale_fallback() { 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/inference".into(); + 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: Zeroizing::new("first-resource-token".into()), acquired: now, refresh_at: now + Duration::from_secs(60), - generation: GENERATION.load(Ordering::Acquire), }); assert!(second.cached.lock().await.is_none()); let (a, b) = tokio::join!(first.token(), first.token()); From 0e63f1a278223ac8ab09690b8bb3fdc83a5363cb Mon Sep 17 00:00:00 2001 From: daniel Date: Fri, 23 Oct 2026 21:21:45 +0100 Subject: [PATCH 4/5] fix(mcp): invalidate only rejected workload acquisitions --- src/tools/mcp/workload.rs | 380 +++++++++++++++++++++++++++++-- src/workload_auth.rs | 55 ++++- src/workload_auth/https_tests.rs | 36 ++- src/workload_auth/tests.rs | 8 +- 4 files changed, 443 insertions(+), 36 deletions(-) diff --git a/src/tools/mcp/workload.rs b/src/tools/mcp/workload.rs index 180f430e..35081ff3 100644 --- a/src/tools/mcp/workload.rs +++ b/src/tools/mcp/workload.rs @@ -1,5 +1,11 @@ //! Per-request workload authentication, including catalog reads and SSE reconnects. -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::HashMap, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; use agentkit_mcp::{ClientJsonRpcMessage, McpHttpClient, McpSseStream}; use async_trait::async_trait; @@ -8,7 +14,7 @@ use rmcp::transport::streamable_http_client::{ StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse, }; -use crate::workload_auth::{Config, Session}; +use crate::workload_auth::{Acquisition, Config, Session}; type Error = StreamableHttpError; @@ -36,6 +42,7 @@ impl Client { } let http = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) + .retry(reqwest::retry::never()) .build() .map_err(|_| "could not build workload MCP HTTP client")?; Ok(Self { @@ -45,11 +52,34 @@ impl Client { }) } + async fn post_with_client( + &self, + uri: Arc, + message: ClientJsonRpcMessage, + session_id: Option>, + custom_headers: HashMap, + builder: reqwest::ClientBuilder, + ) -> Result { + let acquisition = self.authorization(&uri, &custom_headers).await?; + let status = PostStatus::default(); + let http = status.client(builder)?; + let result = 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, Error> { + ) -> Result { if uri != self.destination || headers .keys() @@ -59,14 +89,19 @@ impl Client { "workload MCP transport destination or authorization override rejected", )); } - self.session.token().await.map(Some).map_err(Self::error) + 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, result: Result) -> Result { + async fn finish( + &self, + acquisition: &Acquisition, + raw_rejection: bool, + result: Result, + ) -> Result { let rejected = match &result { Err( StreamableHttpError::AuthRequired(_) | StreamableHttpError::InsufficientScope(_), @@ -77,14 +112,46 @@ impl Client { ), _ => false, }; - if rejected { + if raw_rejection || rejected { // Do not replay a possibly side-effecting POST or follow challenge URLs. - self.session.invalidate().await; + self.session.invalidate_acquisition(acquisition).await; } result } } +// rmcp 3.5.1 consumes POST statuses before returning JSON/SSE results. Observe +// the raw status without replacing its private protocol parser. This client and +// flag belong to one POST, so concurrent responses cannot overwrite each other. +#[derive(Default)] +struct PostStatus(Arc); + +impl PostStatus { + fn client(&self, builder: reqwest::ClientBuilder) -> Result { + let rejected = Arc::clone(&self.0); + let policy = reqwest::retry::never().classify_fn(move |response| { + if matches!( + response.status(), + Some(reqwest::StatusCode::UNAUTHORIZED | reqwest::StatusCode::FORBIDDEN) + ) { + 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) + } + + fn rejected(&self) -> bool { + self.0.load(Ordering::Acquire) + } +} + #[async_trait] impl McpHttpClient for Client { async fn post_message( @@ -95,11 +162,12 @@ impl McpHttpClient for Client { _auth_header: Option, custom_headers: HashMap, ) -> Result { - let authorization = self.authorization(&uri, &custom_headers).await?; - self.finish( - self.http - .post_message(uri, message, session_id, authorization, custom_headers) - .await, + self.post_with_client( + uri, + message, + session_id, + custom_headers, + reqwest::Client::builder(), ) .await } @@ -111,10 +179,17 @@ impl McpHttpClient for Client { _auth_header: Option, custom_headers: HashMap, ) -> Result<(), Error> { - let authorization = self.authorization(&uri, &custom_headers).await?; + let acquisition = self.authorization(&uri, &custom_headers).await?; self.finish( + &acquisition, + false, self.http - .delete_session(uri, session_id, authorization, custom_headers) + .delete_session( + uri, + session_id, + Some(acquisition.token().to_owned()), + custom_headers, + ) .await, ) .await @@ -128,14 +203,16 @@ impl McpHttpClient for Client { _auth_header: Option, custom_headers: HashMap, ) -> Result { - let authorization = self.authorization(&uri, &custom_headers).await?; + let acquisition = self.authorization(&uri, &custom_headers).await?; self.finish( + &acquisition, + false, self.http .get_stream( uri, session_id, last_event_id, - authorization, + Some(acquisition.token().to_owned()), custom_headers, ) .await, @@ -155,6 +232,271 @@ impl McpHttpClient for Client { mod tests { use super::*; + async fn tls_script( + replies: impl FnOnce(&str) -> Vec<(u16, String, String)>, + ) -> (String, tokio::task::JoinHandle>) { + 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(); + for (status, content_type, body) in replies { + let response = format!( + "HTTP/1.1 {status} Test\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + let (socket, _) = listener.accept().await.unwrap(); + let mut stream = acceptor.accept(socket).await.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(); + let _ = stream.shutdown().await; + requests.push(String::from_utf8(bytes).unwrap()); + } + requests + }); + (uri, task) + } + + async fn tls_reply( + status: u16, + content_type: &str, + body: &str, + ) -> (String, tokio::task::JoinHandle>) { + 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 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: tls_client().build().unwrap(), + }; + let message = || { + serde_json::from_value( + serde_json::json!({"jsonrpc":"2.0","id":1,"method":"ping"}), + ) + .unwrap() + }; + let first = client + .post_with_client( + uri.clone().into(), + message(), + None, + HashMap::new(), + tls_client(), + ) + .await; + assert_eq!(first.is_ok(), json); + let second = client + .post_with_client(uri.into(), message(), None, HashMap::new(), tls_client()) + .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 = status + .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 = 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] + .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 = status + .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 = 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", @@ -185,7 +527,8 @@ mod tests { let error = client .authorization("https://other.example/mcp", &HashMap::new()) .await - .unwrap_err(); + .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"))]); @@ -193,7 +536,8 @@ mod tests { client .authorization("https://mcp.example/mcp", &headers) .await - .unwrap_err() + .err() + .unwrap() .to_string() .contains("override rejected") ); diff --git a/src/workload_auth.rs b/src/workload_auth.rs index f496c1d8..a58cb702 100644 --- a/src/workload_auth.rs +++ b/src/workload_auth.rs @@ -1,6 +1,7 @@ //! Clientless RFC 7523 workload authentication. Assertions never enter credential storage. use std::{ path::PathBuf, + sync::Arc, time::{Duration, SystemTime}, }; @@ -14,6 +15,28 @@ 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)] @@ -52,8 +75,17 @@ fn https_url(value: &str) -> Result { Ok(url) } +#[derive(Clone)] +pub struct Acquisition(Arc>); + +impl Acquisition { + pub fn token(&self) -> &str { + self.0.as_str() + } +} + struct Cached { - value: Zeroizing, + value: Acquisition, acquired: SystemTime, refresh_at: SystemTime, } @@ -70,6 +102,7 @@ impl Session { 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() @@ -81,11 +114,17 @@ impl Session { }) } - pub async fn invalidate(&self) { - *self.cached.lock().await = 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 token(&self) -> Result { + 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()) @@ -93,14 +132,14 @@ impl Session { .map_err(|_| "workload authentication timed out".to_string())? } - async fn acquire(&self) -> Result { + 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.to_string()); + return Ok(token.value.clone()); } *cached = None; self.verify_issuer().await?; @@ -145,8 +184,8 @@ impl Session { if completed < started || completed >= refresh_at { return Err("workload authentication expired during exchange".into()); } - let value = Zeroizing::new(token.access_token); - let result = value.to_string(); + let value = Acquisition(Arc::new(Zeroizing::new(token.access_token))); + let result = value.clone(); *cached = Some(Cached { value, acquired: started, diff --git a/src/workload_auth/https_tests.rs b/src/workload_auth/https_tests.rs index aa83ce78..fc308611 100644 --- a/src/workload_auth/https_tests.rs +++ b/src/workload_auth/https_tests.rs @@ -165,18 +165,18 @@ impl Server { token_endpoint: format!("{}/token", self.origin), resource: resource.into(), }; - config.validate().unwrap(); - Session { + test_support::session( config, - client: reqwest::Client::builder() + 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(), - cached: Mutex::new(None), - } + ) + .unwrap() } fn drain(&mut self) -> Vec { let mut requests = Vec::new(); @@ -277,7 +277,9 @@ async fn https_expiry_and_invalidation_reopen_rotated_source() { server.exchange("renewed-access", 3600); assert_eq!(session.token().await.unwrap(), "renewed-access"); std::fs::write(&path, "invalidated.payload.signature").unwrap(); - session.invalidate().await; + 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(); @@ -316,7 +318,9 @@ async fn https_projected_symlink_rotation_is_reopened() { assert_eq!(session.token().await.unwrap(), "second-access"); std::os::unix::fs::symlink(&first, &replacement).unwrap(); std::fs::rename(&replacement, &link).unwrap(); - session.invalidate().await; + 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 @@ -376,6 +380,24 @@ async fn https_post_failures_are_not_retried_or_leaked() { } } +#[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] { diff --git a/src/workload_auth/tests.rs b/src/workload_auth/tests.rs index ea12d267..b3e144c3 100644 --- a/src/workload_auth/tests.rs +++ b/src/workload_auth/tests.rs @@ -107,7 +107,7 @@ async fn expiry_and_clock_rollback_reacquire_without_stale_fallback() { let mut session = Session::new(config()).unwrap(); let now = SystemTime::now(); *session.cached.lock().await = Some(Cached { - value: Zeroizing::new("must-not-escape".into()), + value: Acquisition(Arc::new(Zeroizing::new("must-not-escape".into()))), acquired: if expired { now - Duration::from_secs(60) } else { @@ -135,7 +135,7 @@ async fn caches_are_resource_local_and_invalidatable() { let second = Session::new(second_config).unwrap(); let now = SystemTime::now(); *first.cached.lock().await = Some(Cached { - value: Zeroizing::new("first-resource-token".into()), + value: Acquisition(Arc::new(Zeroizing::new("first-resource-token".into()))), acquired: now, refresh_at: now + Duration::from_secs(60), }); @@ -143,6 +143,8 @@ async fn caches_are_resource_local_and_invalidatable() { 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().await; + first + .invalidate_acquisition(&first.acquisition().await.unwrap()) + .await; assert!(first.cached.lock().await.is_none()); } From 0606bd5d0e1e61d2d82b6eeed3d18d70fc78050a Mon Sep 17 00:00:00 2001 From: daniel Date: Fri, 23 Oct 2026 21:41:50 +0100 Subject: [PATCH 5/5] fix(mcp): retain pooling for workload-authenticated requests --- src/tools/mcp/workload.rs | 280 +++++++++++++++++++++++++++++--------- 1 file changed, 213 insertions(+), 67 deletions(-) diff --git a/src/tools/mcp/workload.rs b/src/tools/mcp/workload.rs index 35081ff3..49cbf093 100644 --- a/src/tools/mcp/workload.rs +++ b/src/tools/mcp/workload.rs @@ -40,10 +40,7 @@ impl Client { "workload MCP URL must be HTTPS and match the configured resource exactly".into(), ); } - let http = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .retry(reqwest::retry::never()) - .build() + let http = PostStatus::client(reqwest::Client::builder()) .map_err(|_| "could not build workload MCP HTTP client")?; Ok(Self { destination: destination.to_owned(), @@ -58,19 +55,17 @@ impl Client { message: ClientJsonRpcMessage, session_id: Option>, custom_headers: HashMap, - builder: reqwest::ClientBuilder, ) -> Result { let acquisition = self.authorization(&uri, &custom_headers).await?; let status = PostStatus::default(); - let http = status.client(builder)?; - let result = http - .post_message( + 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 } @@ -120,21 +115,25 @@ impl Client { } } -// rmcp 3.5.1 consumes POST statuses before returning JSON/SSE results. Observe -// the raw status without replacing its private protocol parser. This client and -// flag belong to one POST, so concurrent responses cannot overwrite each other. +// 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(&self, builder: reqwest::ClientBuilder) -> Result { - let rejected = Arc::clone(&self.0); + 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) ) { - rejected.store(true, Ordering::Release); + 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. @@ -147,6 +146,10 @@ impl PostStatus { .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) } @@ -162,14 +165,8 @@ impl McpHttpClient for Client { _auth_header: Option, custom_headers: HashMap, ) -> Result { - self.post_with_client( - uri, - message, - session_id, - custom_headers, - reqwest::Client::builder(), - ) - .await + self.post_with_client(uri, message, session_id, custom_headers) + .await } async fn delete_session( @@ -234,7 +231,14 @@ mod tests { async fn tls_script( replies: impl FnOnce(&str) -> Vec<(u16, String, String)>, - ) -> (String, tokio::task::JoinHandle>) { + ) -> (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(); @@ -259,13 +263,20 @@ mod tests { 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: close\r\n\r\n{body}", + "HTTP/1.1 {status} Test\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: {connection_header}\r\n\r\n{body}", body.len() ); - let (socket, _) = listener.accept().await.unwrap(); - let mut stream = acceptor.accept(socket).await.unwrap(); + 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]; @@ -289,10 +300,13 @@ mod tests { } } stream.write_all(response.as_bytes()).await.unwrap(); - let _ = stream.shutdown().await; + if !keep_alive { + let _ = stream.shutdown().await; + connection = None; + } requests.push(String::from_utf8(bytes).unwrap()); } - requests + (requests, accepted) }); (uri, task) } @@ -301,7 +315,7 @@ mod tests { status: u16, content_type: &str, body: &str, - ) -> (String, tokio::task::JoinHandle>) { + ) -> (String, tokio::task::JoinHandle<(Vec, usize)>) { tls_script(|_| vec![(status, content_type.to_owned(), body.to_owned())]).await } @@ -317,6 +331,134 @@ mod tests { .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] { @@ -350,7 +492,7 @@ mod tests { let client = Client { destination: uri.clone(), session, - http: tls_client().build().unwrap(), + http: PostStatus::client(tls_client()).unwrap(), }; let message = || { serde_json::from_value( @@ -359,20 +501,14 @@ mod tests { .unwrap() }; let first = client - .post_with_client( - uri.clone().into(), - message(), - None, - HashMap::new(), - tls_client(), - ) + .post_message(uri.clone().into(), message(), None, None, HashMap::new()) .await; assert_eq!(first.is_ok(), json); let second = client - .post_with_client(uri.into(), message(), None, HashMap::new(), tls_client()) + .post_message(uri.into(), message(), None, None, HashMap::new()) .await; assert!(matches!(second, Ok(StreamableHttpPostResponse::Json(_, _)))); - let requests = server.await.unwrap(); + let (requests, _) = server.await.unwrap(); assert_eq!(requests.len(), 6); assert!(requests[1].starts_with("POST /token ")); assert!( @@ -410,23 +546,28 @@ mod tests { ) .await; let status = PostStatus::default(); - let http = status - .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 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 = http - .post_message(uri.into(), message, None, Some("T1".into()), HashMap::new()) + 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 { @@ -438,7 +579,7 @@ mod tests { )); } assert!( - server.await.unwrap()[0] + server.await.unwrap().0[0] .to_lowercase() .contains("authorization: bearer t1") ); @@ -466,22 +607,27 @@ mod tests { ) .await; let status = PostStatus::default(); - let http = status - .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 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 = http - .post_message(uri.into(), message, None, Some("T1".into()), HashMap::new()) + let result = status + .observe(http.post_message( + uri.into(), + message, + None, + Some("T1".into()), + HashMap::new(), + )) .await .unwrap(); assert!(!status.rejected());