diff --git a/Cargo.lock b/Cargo.lock index a1f9d8c..e64ebc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,6 +8,42 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aead" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1973cfbc1a2daf9cf550e74e1f088c28e7f7d8c1e1418fb6c9dc5184b7e84c99" +dependencies = [ + "crypto-common 0.2.2", + "inout", +] + +[[package]] +name = "aes" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8eb277bec05f56a0e0591f155a484cbd0f4f07ff2905051a48c72f004f7ed58" +dependencies = [ + "cipher", + "cpubits", + "cpufeatures 0.3.0", +] + +[[package]] +name = "aes-gcm-siv" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f437e5b075722bda3f54039e95d60d1c142f140e6c45b718e4f9fca3e5a1514" +dependencies = [ + "aead", + "aes", + "cipher", + "ctr", + "polyval", + "subtle", + "zeroize", +] + [[package]] name = "aho-corasick" version = "1.1.5" @@ -364,6 +400,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "block2" version = "0.6.2" @@ -455,6 +500,17 @@ dependencies = [ "windows-link", ] +[[package]] +name = "cipher" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf2a2c93cd704877c0858356ed03480ff301ee950b43f1cbe4573b088bfa6c" +dependencies = [ + "block-buffer 0.12.1", + "crypto-common 0.2.2", + "inout", +] + [[package]] name = "clap" version = "4.6.6" @@ -504,6 +560,12 @@ dependencies = [ "cc", ] +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + [[package]] name = "colorchoice" version = "1.0.5" @@ -611,10 +673,12 @@ dependencies = [ name = "contextforge-data-plane-lib" version = "0.1.0" dependencies = [ + "aes-gcm-siv", "async-trait", "axum", "axum-otel-metrics", "axum-server", + "base64 0.22.1", "chrono", "clap", "contextforge-data-plane-apis", @@ -651,6 +715,7 @@ dependencies = [ "typed-builder", "url", "uuid", + "zeroize", ] [[package]] @@ -799,6 +864,12 @@ dependencies = [ "cpex", ] +[[package]] +name = "cpubits" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -851,6 +922,26 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "getrandom 0.4.3", + "hybrid-array", + "rand_core 0.10.1", +] + +[[package]] +name = "ctr" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baaca1c4b237092596f64d571e9db6ce4109c4ef9742e27590f1709594461f21" +dependencies = [ + "cipher", +] + [[package]] name = "ctrlc" version = "3.5.2" @@ -862,6 +953,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + [[package]] name = "darling" version = "0.24.0" @@ -914,8 +1014,8 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", ] [[package]] @@ -1395,6 +1495,15 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "hybrid-array" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "707114b52a152fa7bdb290cd7cd5912d9467273b6d74e21b8d81aca1f8533f6b" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.11.0" @@ -1613,6 +1722,15 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "hybrid-array", +] + [[package]] name = "ipnet" version = "2.12.1" @@ -2125,6 +2243,17 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "polyval" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0fa31d631f2b2cb2a544d0aa321ce847a94764d701ca2becc411138b93d49cd" +dependencies = [ + "cpubits", + "cpufeatures 0.3.0", + "universal-hash", +] + [[package]] name = "portable-atomic" version = "1.14.0" @@ -3670,6 +3799,16 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "universal-hash" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4987bdc12753382e0bec4a65c50738ffaabc998b9cdd1f952fb5f39b0048a96" +dependencies = [ + "crypto-common 0.2.2", + "ctutils", +] + [[package]] name = "unsafe-libyaml" version = "0.2.11" diff --git a/_context/wiki/mcp-capability-allocation.md b/_context/wiki/mcp-capability-allocation.md index 2cfc061..be40dbf 100644 --- a/_context/wiki/mcp-capability-allocation.md +++ b/_context/wiki/mcp-capability-allocation.md @@ -138,6 +138,25 @@ authorization context. dataplane, and integration tests must define together. The current coarse `sub`-only implementation is not the Phase 3 target. +### Stateless Task Handles + +Modern Tasks lifecycle calls are targeted operations under the same +authorization invariants. The dataplane exposes an encrypted, stateless handle +rather than an upstream task ID or process-local mapping. The handle binds the +upstream ID and route to a trusted authorization-context ID, virtual server, +configuration revision, immutable backend ID, and backend generation. + +Every `tasks/get`, `tasks/update`, and `tasks/cancel` request independently +derives its authorization context from verified claims and the validated route, +loads the current effective configuration, enforces method scope and compiled +policy, and accepts the handle only when its backend ID still resolves to the +same generation. A mismatch returns the same invalid-task error as malformed +input and makes no upstream call. Dataplane replicas share the handle key so +decoding does not require Redis task state or session affinity. + +The codec is a prerequisite only; task creation and lifecycle proxy handlers +remain separate implementation work and are not current routing behavior. + ## MCP Work Allocation | Work | Target owner and behavior | diff --git a/_context/wiki/routing.md b/_context/wiki/routing.md index 1013311..7430d90 100644 --- a/_context/wiki/routing.md +++ b/_context/wiki/routing.md @@ -52,6 +52,19 @@ The gateway wraps per-backend cursors inside its own opaque token (JSON, treated **Known limitation:** if backend set changes between pages, removed backend's cursor is silently dropped. +## Task-Handle Codec (not wired) + +The library contains the codec prerequisite for modern Tasks, but no current +handler emits task handles or proxies `tasks/get`, `tasks/update`, or +`tasks/cancel` yet. + +- Never expose or log an upstream task ID directly; decoded-route debug output redacts it. +- Encode it as `cfth1.` with misuse-resistant AES-256-GCM-SIV and a random nonce. +- Bind the payload to the trusted authorization-context ID, virtual host, configuration revision, backend ID, and backend generation. +- On decode, independently derive the current authorization scope and accept the backend only when the effective configuration resolves the same ID and generation. +- Return every scope, revision, route, malformed-input, and version mismatch as `invalid task ID` without an upstream call. +- Handles are stateless. Replicas must share the key; key rotation invalidates outstanding handles. + ## Session State (local process) Backend RMCP services are stored in `BackendTransports` keyed by: diff --git a/_context/wiki/security.md b/_context/wiki/security.md index bf0ac53..66af57e 100644 --- a/_context/wiki/security.md +++ b/_context/wiki/security.md @@ -45,6 +45,7 @@ dataplane contract. | If this is compromised | Impact | | --- | --- | | JWT signing key or HMAC secret | Attacker mints tokens for any subject and reaches that subject's backends. Rotate the key and restart; no revocation exists. | +| Task-handle key | Attacker decrypts or forges upstream task routes. Rotate the key; outstanding handles become invalid. | | Redis write access | Attacker rewrites routing (arbitrary backend URLs receive caller traffic) and, if runtime plugins are enabled, chooses which registered hooks run on payloads. Protect Redis with TLS/mTLS and control-plane-only write access. | | A backend MCP server | Attacker sees requests routed to that backend and controls its responses; the namespace prefix limits blast radius to that backend's objects. | | The gateway process | Full compromise: it holds the decoding keys in memory and live backend sessions. | @@ -102,4 +103,7 @@ These routes are registered **outside the authentication middleware** — unauth ## Secrets Handling - The HMAC secret is held as a `SecretString`; key and certificate material is read from disk paths at startup. +- The not-yet-wired task-handle codec encrypts upstream task IDs and binds them to a trusted authorization-context ID, virtual host, configuration revision, backend ID, and backend generation. +- Task handles do not replace per-request JWT validation, effective-configuration lookup, scope/RBAC checks, or current backend-generation validation. - Never log: tokens, authorization headers, secrets, Redis key/value bytes, full `UserConfig` documents, or backend credentials. +- Treat task handles and decoded upstream task IDs as opaque secrets; do not log them. Decoded-route debug output must remain redacted. diff --git a/crates/contextforge-data-plane-lib/Cargo.toml b/crates/contextforge-data-plane-lib/Cargo.toml index 6505887..aef4647 100644 --- a/crates/contextforge-data-plane-lib/Cargo.toml +++ b/crates/contextforge-data-plane-lib/Cargo.toml @@ -46,6 +46,9 @@ tokio-rustls = "0.26.4" typed-builder = "0.23.2" url = { workspace = true, features = ["serde"] } secret-string = "0.0.2" +base64 = "0.22.1" +aes-gcm-siv = { version = "0.12.0", features = ["zeroize"] } +zeroize = "1.9.0" [features] diff --git a/crates/contextforge-data-plane-lib/src/gateway/mcp_call_validator.rs b/crates/contextforge-data-plane-lib/src/gateway/mcp_call_validator.rs index 3b88a6e..8f596e5 100644 --- a/crates/contextforge-data-plane-lib/src/gateway/mcp_call_validator.rs +++ b/crates/contextforge-data-plane-lib/src/gateway/mcp_call_validator.rs @@ -27,7 +27,7 @@ impl<'a> AuthorizedCallValidator<'a> { let has_user_config = maybe_user_config.is_some(); let virtual_hosts = maybe_user_config.map_or(0, |user_config| user_config.virtual_hosts.len()); let has_claims = maybe_claims.is_some(); - let virtual_host_id = maybe_virtual_host_id.map_or("", |id| id.value().as_str()); + let virtual_host_id = maybe_virtual_host_id.map_or("", VirtualHostId::as_str); debug!( "AuthorizedCallValidator::validate - mcp call validation call_name = {call_name} has_user_config = {has_user_config} virtual_hosts = {virtual_hosts} has_claims = {has_claims} virtual_host_id = {virtual_host_id}" ); @@ -48,9 +48,9 @@ impl<'a> AuthorizedCallValidator<'a> { }); }; - let Some(virtual_host) = user_config.virtual_hosts.get(virtual_host_id.value()) else { + let Some(virtual_host) = user_config.virtual_hosts.get(virtual_host_id.as_str()) else { let call_name = self.call_name; - let virtual_host_id = virtual_host_id.value(); + let virtual_host_id = virtual_host_id.as_str(); let virtual_hosts = user_config.virtual_hosts.len(); debug!( "AuthorizedCallValidator::validate - mcp virtual host config missing call_name = {call_name} virtual_host_id = {virtual_host_id} virtual_hosts = {virtual_hosts}" @@ -85,7 +85,7 @@ impl<'a> AuthorizedCallValidator<'a> { let virtual_hosts = maybe_user_config.map_or(0, |user_config| user_config.virtual_hosts.len()); let has_session_id = maybe_session_id.is_some(); let has_claims = maybe_claims.is_some(); - let virtual_host_id = maybe_virtual_host_id.map_or("", |id| id.value().as_str()); + let virtual_host_id = maybe_virtual_host_id.map_or("", VirtualHostId::as_str); debug!( "AuthorizedCallValidator::validate - mcp call validation call_name = {call_name} has_user_config = {has_user_config} virtual_hosts = {virtual_hosts} has_session_id = {has_session_id} has_claims = {has_claims} virtual_host_id = {virtual_host_id}" ); @@ -114,9 +114,9 @@ impl<'a> AuthorizedCallValidator<'a> { }); }; - let Some(virtual_host) = user_config.virtual_hosts.get(virtual_host_id.value()) else { + let Some(virtual_host) = user_config.virtual_hosts.get(virtual_host_id.as_str()) else { let call_name = self.call_name; - let virtual_host_id = virtual_host_id.value(); + let virtual_host_id = virtual_host_id.as_str(); let virtual_hosts = user_config.virtual_hosts.len(); debug!( "AuthorizedCallValidator::validate - mcp virtual host config missing call_name = {call_name} virtual_host_id = {virtual_host_id} virtual_hosts = {virtual_hosts}" @@ -160,7 +160,7 @@ impl<'a> InitializeCallValidator<'a> { let virtual_hosts = maybe_user_config.map_or(0, |user_config| user_config.virtual_hosts.len()); let has_session_id = true; let has_claims = maybe_claims.is_some(); - let virtual_host_id = maybe_virtual_host_id.map_or("", |id| id.value().as_str()); + let virtual_host_id = maybe_virtual_host_id.map_or("", VirtualHostId::as_str); debug!( "InitializeCallValidator::validate - mcp call validation call_name = {call_name} has_user_config = {has_user_config} virtual_hosts = {virtual_hosts} has_session_id = {has_session_id} has_claims = {has_claims} virtual_host_id = {virtual_host_id}" ); @@ -181,9 +181,9 @@ impl<'a> InitializeCallValidator<'a> { }); }; - let Some(virtual_host) = user_config.virtual_hosts.get(virtual_host_id.value()) else { + let Some(virtual_host) = user_config.virtual_hosts.get(virtual_host_id.as_str()) else { let call_name = "initialize"; - let virtual_host_id = virtual_host_id.value(); + let virtual_host_id = virtual_host_id.as_str(); let virtual_hosts = user_config.virtual_hosts.len(); debug!( "InitializeCallValidator::validate - mcp virtual host config missing call_name = {call_name} virtual_host_id = {virtual_host_id} virtual_hosts = {virtual_hosts}" diff --git a/crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs b/crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs index 7634e58..c6a18c0 100644 --- a/crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs +++ b/crates/contextforge-data-plane-lib/src/layers/virtual_host_config.rs @@ -14,7 +14,7 @@ pub async fn virtual_host_config_layer(request: http::Request, if let (Some(virtual_host_id), Some(user_config)) = (virtual_host_id, user_config) && !has_virtual_host(user_config, virtual_host_id) { - let virtual_host_id = virtual_host_id.value(); + let virtual_host_id = virtual_host_id.as_str(); let virtual_hosts = user_config.virtual_hosts.len(); debug!( "virtual_host_config_layer - virtual host config missing virtual_host_id = {virtual_host_id} virtual_hosts = {virtual_hosts}" @@ -26,7 +26,7 @@ pub async fn virtual_host_config_layer(request: http::Request, } fn has_virtual_host(user_config: &UserConfig, virtual_host_id: &VirtualHostId) -> bool { - user_config.virtual_hosts.contains_key(virtual_host_id.value()) + user_config.virtual_hosts.contains_key(virtual_host_id.as_str()) } fn server_not_found_response() -> Response { diff --git a/crates/contextforge-data-plane-lib/src/layers/virtual_host_id.rs b/crates/contextforge-data-plane-lib/src/layers/virtual_host_id.rs index c7e1517..5271fc3 100644 --- a/crates/contextforge-data-plane-lib/src/layers/virtual_host_id.rs +++ b/crates/contextforge-data-plane-lib/src/layers/virtual_host_id.rs @@ -2,17 +2,20 @@ use axum::{body::Body, middleware::Next, response::Response}; use http::{StatusCode, header}; use tracing::debug; -#[derive(Clone, Debug, PartialEq, PartialOrd)] +/// Virtual-server identifier extracted from the downstream MCP route. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct VirtualHostId { value: String, } impl VirtualHostId { - pub(crate) fn new(value: String) -> Self { - Self { value } + /// Creates a virtual-server identifier from its canonical route value. + pub fn new(value: impl Into) -> Self { + Self { value: value.into() } } - pub fn value(&self) -> &String { + /// Returns the canonical route value. + pub fn as_str(&self) -> &str { &self.value } } @@ -40,7 +43,7 @@ fn extract_virtual_host_id(path: &str) -> Option { let l1 = "/servers/".len(); let l2 = path.len() - "/mcp".len(); let vh = &path[l1..l2]; - VirtualHostId::new(vh.to_owned()) + VirtualHostId::new(vh) }) } else { None @@ -52,12 +55,12 @@ mod tests { use crate::layers::virtual_host_id::{VirtualHostId, extract_virtual_host_id}; #[test] - fn test_virtual_host_extractor() { + fn extracts_only_virtual_host_from_server_mcp_routes() { assert_eq!(None, extract_virtual_host_id("/mcp/servers")); assert_eq!(None, extract_virtual_host_id("/servers")); assert_eq!(None, extract_virtual_host_id("/servers/12345_abcd-efgh/mcp/dkfjk")); assert_eq!( - Some(VirtualHostId { value: "12345_abcd-efgh".to_owned() }), + Some(VirtualHostId::new("12345_abcd-efgh")), extract_virtual_host_id("/servers/12345_abcd-efgh/mcp") ); assert_eq!(None, extract_virtual_host_id("/12345_abcd-efgh/12345_abcd-efgh/mcp")); diff --git a/crates/contextforge-data-plane-lib/src/lib.rs b/crates/contextforge-data-plane-lib/src/lib.rs index 8804520..2d88d4f 100644 --- a/crates/contextforge-data-plane-lib/src/lib.rs +++ b/crates/contextforge-data-plane-lib/src/lib.rs @@ -15,6 +15,7 @@ mod const_values; mod gateway; mod layers; mod mcp_standard_headers; +pub mod task_handle; mod telemetry; mod transports; @@ -25,6 +26,7 @@ mod user_config_store; pub use common::{RedisClient, RedisConfig, UpstreamConnectionMode}; use gateway::{BackendTransports, McpService}; use layers::session_id::SessionId; +pub use layers::virtual_host_id::VirtualHostId; use tower_http::cors::{Any, CorsLayer}; use tower_http::trace::TraceLayer; use transports::{DownstreamTls, Tcp}; diff --git a/crates/contextforge-data-plane-lib/src/task_handle.rs b/crates/contextforge-data-plane-lib/src/task_handle.rs new file mode 100644 index 0000000..38eead2 --- /dev/null +++ b/crates/contextforge-data-plane-lib/src/task_handle.rs @@ -0,0 +1,543 @@ +//! Stateless routing handles for the MCP Tasks extension. +//! +//! Handles are encrypted and authenticated because their payload contains an +//! upstream task identifier that may be a bearer token. The trusted +//! authorization context, virtual host, configuration revision, and backend +//! generation prevent replay through a different caller, policy snapshot, or +//! upstream route. + +use std::{fmt, str::FromStr}; + +use aes_gcm_siv::{ + Aes256GcmSiv, Key, Nonce, + aead::{Aead, Generate, KeyInit, Payload as AeadPayload}, +}; +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use rmcp::{ErrorData, model::ErrorCode}; +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use zeroize::Zeroizing; + +use crate::VirtualHostId; + +const HANDLE_PREFIX: &str = "cfth1"; +const HANDLE_FAMILY_PREFIX: &str = "cfth"; +const KEY_LEN: usize = 32; +const NONCE_LEN: usize = 12; +const TAG_LEN: usize = 16; + +/// A validated AES-256-GCM-SIV key for task-handle protection. +/// +/// The textual form is URL-safe base64 without padding and must decode to +/// exactly 32 bytes. Its [`Debug`] output is always redacted. +#[derive(Clone, PartialEq, Eq)] +pub struct TaskHandleKey(Zeroizing<[u8; KEY_LEN]>); + +impl FromStr for TaskHandleKey { + type Err = TaskHandleKeyError; + + fn from_str(value: &str) -> Result { + let mut key = Zeroizing::new([0_u8; KEY_LEN]); + let decoded_len = URL_SAFE_NO_PAD.decode_slice(value, key.as_mut()).map_err(|_| TaskHandleKeyError)?; + if decoded_len != KEY_LEN { + return Err(TaskHandleKeyError); + } + Ok(Self(key)) + } +} + +impl fmt::Debug for TaskHandleKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("TaskHandleKey([REDACTED])") + } +} + +/// Returned when task-handle key material is not a URL-safe base64-encoded +/// 256-bit key. +#[derive(Debug, Error, PartialEq, Eq)] +#[error("task handle key must be URL-safe base64 without padding and decode to exactly 32 bytes")] +pub struct TaskHandleKeyError; + +macro_rules! string_identifier { + ($name:ident, $doc:literal) => { + #[doc = $doc] + #[derive(Debug, Clone, PartialEq, Eq, Hash)] + pub struct $name(String); + + impl $name { + /// Creates an identifier from its canonical trusted value. + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + /// Returns the canonical value. + pub fn as_str(&self) -> &str { + &self.0 + } + } + }; +} + +string_identifier!( + AuthorizationContextId, + "Canonical identity for the authenticated tenant, principal, team, and scope set." +); +string_identifier!(ConfigurationRevision, "Revision of the validated effective-configuration snapshot."); +string_identifier!(BackendId, "Stable identity of a backend in effective configuration."); +string_identifier!(BackendGeneration, "Generation of a backend's routing material."); + +/// Authenticated request scope to which a task handle is bound. +/// +/// The authorization-context ID and configuration revision must come from the +/// verified JWT and the validated effective-configuration snapshot. Client +/// metadata and MCP params are not trusted sources for either value. +#[derive(Clone, Copy, PartialEq, Eq)] +pub struct TaskHandleScope<'a> { + authorization_context_id: &'a AuthorizationContextId, + virtual_host_id: &'a VirtualHostId, + configuration_revision: &'a ConfigurationRevision, +} + +impl<'a> TaskHandleScope<'a> { + /// Creates a scope from trusted authorization and routing context. + pub fn new( + authorization_context_id: &'a AuthorizationContextId, + virtual_host_id: &'a VirtualHostId, + configuration_revision: &'a ConfigurationRevision, + ) -> Self { + Self { authorization_context_id, virtual_host_id, configuration_revision } + } +} + +/// Stable backend identity stored in a task handle. +/// +/// `generation` must change when the routing key is reassigned to a different +/// upstream or when routing material changes in a way that invalidates +/// outstanding upstream task IDs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TaskHandleBackend<'a> { + id: &'a BackendId, + generation: &'a BackendGeneration, +} + +impl<'a> TaskHandleBackend<'a> { + /// Creates backend identity from trusted effective configuration. + pub fn new(id: &'a BackendId, generation: &'a BackendGeneration) -> Self { + Self { id, generation } + } + + /// Stable backend identifier used for routing. + pub fn id(self) -> &'a BackendId { + self.id + } + + /// Generation of the backend routing material. + pub fn generation(self) -> &'a BackendGeneration { + self.generation + } +} + +/// The route recovered from a valid task handle. +#[derive(Clone, PartialEq, Eq)] +pub struct TaskHandleRoute { + backend_id: BackendId, + upstream_task_id: Zeroizing, +} + +impl TaskHandleRoute { + /// Stable backend ID in the caller's current effective configuration. + pub fn backend_id(&self) -> &BackendId { + &self.backend_id + } + + /// Original task identifier expected by the upstream backend. + pub fn upstream_task_id(&self) -> &str { + self.upstream_task_id.as_str() + } +} + +impl fmt::Debug for TaskHandleRoute { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TaskHandleRoute") + .field("backend_id", &self.backend_id.as_str()) + .field("upstream_task_id", &"[REDACTED]") + .finish() + } +} + +/// Errors produced while encoding or decoding task handles. +/// +/// Decode errors intentionally have the same display text so a caller cannot +/// distinguish a malformed handle from a valid handle outside its scope. +#[derive(Debug, Error, PartialEq, Eq)] +pub enum TaskHandleError { + /// The protected handle could not be created. + #[error("failed to create task handle")] + Encode, + /// The handle is malformed or fails authentication. + #[error("invalid task handle")] + Invalid, + /// The handle belongs to another codec version. + #[error("invalid task handle")] + UnsupportedVersion, + /// The handle does not belong to the authenticated request scope. + #[error("invalid task handle")] + WrongScope, + /// The referenced backend and generation are not currently routable. + #[error("invalid task handle")] + UnavailableBackend, +} + +impl From for ErrorData { + fn from(error: TaskHandleError) -> Self { + match error { + TaskHandleError::Encode => ErrorData::new(ErrorCode::INTERNAL_ERROR, "failed to create task handle", None), + TaskHandleError::Invalid + | TaskHandleError::UnsupportedVersion + | TaskHandleError::WrongScope + | TaskHandleError::UnavailableBackend => ErrorData::new(ErrorCode::INVALID_PARAMS, "invalid task ID", None), + } + } +} + +#[derive(Serialize)] +struct TaskHandlePayload<'a> { + authorization_context_id: &'a str, + virtual_host_id: &'a str, + configuration_revision: &'a str, + backend_id: &'a str, + backend_generation: &'a str, + upstream_task_id: &'a str, +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct DecodedTaskHandlePayload { + authorization_context_id: String, + virtual_host_id: String, + configuration_revision: String, + backend_id: String, + backend_generation: String, + upstream_task_id: String, +} + +/// Encodes and decodes versioned task handles shared across dataplane replicas. +#[derive(Clone)] +pub struct TaskHandleCodec { + cipher: Aes256GcmSiv, +} + +impl TaskHandleCodec { + /// Creates a codec from a validated shared key. + pub fn new(key: &TaskHandleKey) -> Self { + let key: &Key = (&*key.0).into(); + let cipher = Aes256GcmSiv::new(key); + Self { cipher } + } + + /// Creates an opaque handle for one upstream task. + pub fn encode( + &self, + scope: TaskHandleScope<'_>, + backend: TaskHandleBackend<'_>, + upstream_task_id: &str, + ) -> Result { + let payload = TaskHandlePayload { + authorization_context_id: scope.authorization_context_id.as_str(), + virtual_host_id: scope.virtual_host_id.as_str(), + configuration_revision: scope.configuration_revision.as_str(), + backend_id: backend.id.as_str(), + backend_generation: backend.generation.as_str(), + upstream_task_id, + }; + let plaintext = Zeroizing::new(serde_json::to_vec(&payload).map_err(|_| TaskHandleError::Encode)?); + + let nonce = Nonce::try_generate().map_err(|_| TaskHandleError::Encode)?; + let ciphertext = self + .cipher + .encrypt(&nonce, AeadPayload { msg: plaintext.as_slice(), aad: HANDLE_PREFIX.as_bytes() }) + .map_err(|_| TaskHandleError::Encode)?; + + let mut protected = Vec::with_capacity(NONCE_LEN + ciphertext.len()); + protected.extend_from_slice(&nonce); + protected.extend_from_slice(&ciphertext); + Ok(format!("{HANDLE_PREFIX}.{}", URL_SAFE_NO_PAD.encode(protected))) + } + + /// Decodes a handle and verifies that it belongs to the authenticated + /// authorization context and references the same backend generation in + /// current effective configuration. + pub fn decode( + &self, + handle: &str, + expected_scope: TaskHandleScope<'_>, + backend_is_current: F, + ) -> Result + where + F: FnOnce(TaskHandleBackend<'_>) -> bool, + { + let (prefix, protected) = handle.split_once('.').ok_or(TaskHandleError::Invalid)?; + if prefix != HANDLE_PREFIX { + return if is_other_version(prefix) { + Err(TaskHandleError::UnsupportedVersion) + } else { + Err(TaskHandleError::Invalid) + }; + } + + let protected = URL_SAFE_NO_PAD.decode(protected).map_err(|_| TaskHandleError::Invalid)?; + if protected.len() < NONCE_LEN + TAG_LEN { + return Err(TaskHandleError::Invalid); + } + let (nonce_bytes, ciphertext) = protected.split_at(NONCE_LEN); + let nonce_bytes: [u8; NONCE_LEN] = nonce_bytes.try_into().map_err(|_| TaskHandleError::Invalid)?; + let nonce = Nonce::from(nonce_bytes); + let plaintext = Zeroizing::new( + self.cipher + .decrypt(&nonce, AeadPayload { msg: ciphertext, aad: HANDLE_PREFIX.as_bytes() }) + .map_err(|_| TaskHandleError::Invalid)?, + ); + let payload: DecodedTaskHandlePayload = + serde_json::from_slice(&plaintext).map_err(|_| TaskHandleError::Invalid)?; + + if payload.authorization_context_id != expected_scope.authorization_context_id.as_str() + || payload.virtual_host_id != expected_scope.virtual_host_id.as_str() + || payload.configuration_revision != expected_scope.configuration_revision.as_str() + { + return Err(TaskHandleError::WrongScope); + } + let backend_id = BackendId::new(payload.backend_id); + let backend_generation = BackendGeneration::new(payload.backend_generation); + if !backend_is_current(TaskHandleBackend::new(&backend_id, &backend_generation)) { + return Err(TaskHandleError::UnavailableBackend); + } + + Ok(TaskHandleRoute { backend_id, upstream_task_id: Zeroizing::new(payload.upstream_task_id) }) + } +} + +fn is_other_version(prefix: &str) -> bool { + prefix + .strip_prefix(HANDLE_FAMILY_PREFIX) + .is_some_and(|version| !version.is_empty() && version.bytes().all(|byte| byte.is_ascii_digit())) +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use super::*; + + const KEY: &str = "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8"; // pragma: allowlist secret + const AUTHORIZATION_CONTEXT_VALUE: &str = "tenant-a:principal-a:team-a:scope-set-a"; + const VIRTUAL_HOST_ID_VALUE: &str = "host-a"; + const CONFIGURATION_REVISION_VALUE: &str = "revision-7"; + const BACKEND_ID_VALUE: &str = "backend-a"; + const BACKEND_GENERATION_VALUE: &str = "generation-3"; + + static AUTHORIZATION_CONTEXT: LazyLock = + LazyLock::new(|| AuthorizationContextId::new(AUTHORIZATION_CONTEXT_VALUE)); + static VIRTUAL_HOST_ID: LazyLock = LazyLock::new(|| VirtualHostId::new(VIRTUAL_HOST_ID_VALUE)); + static CONFIGURATION_REVISION: LazyLock = + LazyLock::new(|| ConfigurationRevision::new(CONFIGURATION_REVISION_VALUE)); + static BACKEND_ID: LazyLock = LazyLock::new(|| BackendId::new(BACKEND_ID_VALUE)); + static BACKEND_GENERATION: LazyLock = + LazyLock::new(|| BackendGeneration::new(BACKEND_GENERATION_VALUE)); + + fn codec() -> TaskHandleCodec { + TaskHandleCodec::new(&KEY.parse().expect("test key is valid")) + } + + fn scope<'a>( + authorization_context_id: &'a AuthorizationContextId, + virtual_host_id: &'a VirtualHostId, + configuration_revision: &'a ConfigurationRevision, + ) -> TaskHandleScope<'a> { + TaskHandleScope::new(authorization_context_id, virtual_host_id, configuration_revision) + } + + fn current_scope() -> TaskHandleScope<'static> { + scope(&AUTHORIZATION_CONTEXT, &VIRTUAL_HOST_ID, &CONFIGURATION_REVISION) + } + + fn backend<'a>(id: &'a BackendId, generation: &'a BackendGeneration) -> TaskHandleBackend<'a> { + TaskHandleBackend::new(id, generation) + } + + fn current_backend() -> TaskHandleBackend<'static> { + backend(&BACKEND_ID, &BACKEND_GENERATION) + } + + fn decode_current( + codec: &TaskHandleCodec, + handle: &str, + expected_scope: TaskHandleScope<'_>, + expected_backend: TaskHandleBackend<'_>, + ) -> Result { + codec.decode(handle, expected_scope, |decoded_backend| { + decoded_backend.id() == expected_backend.id() + && decoded_backend.generation() == expected_backend.generation() + }) + } + + #[test] + fn arbitrary_upstream_task_ids_round_trip_without_loss() { + let codec = codec(); + let task_ids = ["", "simple", "with/slashes?and=query", "nul\0byte", "emoji-🦀", "line\nbreak"]; + + for task_id in task_ids { + let handle = codec.encode(current_scope(), current_backend(), task_id).expect("handle encodes"); + let route = decode_current(&codec, &handle, current_scope(), current_backend()).expect("handle decodes"); + + assert_eq!(route.backend_id().as_str(), BACKEND_ID_VALUE); + assert_eq!(route.upstream_task_id(), task_id); + } + } + + #[test] + fn identical_task_ids_from_different_backends_remain_isolated() { + let codec = codec(); + let first_identity = (BackendId::new("backend-a"), BackendGeneration::new("generation-a")); + let second_identity = (BackendId::new("backend-b"), BackendGeneration::new("generation-b")); + let first_backend = backend(&first_identity.0, &first_identity.1); + let second_backend = backend(&second_identity.0, &second_identity.1); + let first_handle = codec.encode(current_scope(), first_backend, "same-id").expect("first handle encodes"); + let second_handle = codec.encode(current_scope(), second_backend, "same-id").expect("second handle encodes"); + + let first_route = + decode_current(&codec, &first_handle, current_scope(), first_backend).expect("first handle decodes"); + let second_route = + decode_current(&codec, &second_handle, current_scope(), second_backend).expect("second handle decodes"); + + assert_ne!(first_handle, second_handle); + assert_eq!(first_route.backend_id().as_str(), "backend-a"); + assert_eq!(second_route.backend_id().as_str(), "backend-b"); + } + + #[test] + fn handle_decodes_on_another_codec_with_the_same_key() { + let first_replica = codec(); + let second_replica = codec(); + let handle = first_replica.encode(current_scope(), current_backend(), "task-42").expect("handle encodes"); + + let route = + decode_current(&second_replica, &handle, current_scope(), current_backend()).expect("handle decodes"); + + assert_eq!(route.upstream_task_id(), "task-42"); + } + + #[test] + fn malformed_and_tampered_handles_fail_closed() { + let codec = codec(); + let handle = codec.encode(current_scope(), current_backend(), "task-42").expect("handle encodes"); + let mut tampered = handle.into_bytes(); + let last = tampered.last_mut().expect("handle is non-empty"); + *last = if *last == b'A' { b'B' } else { b'A' }; + let tampered = String::from_utf8(tampered).expect("tampered handle remains UTF-8"); + + assert_eq!( + decode_current(&codec, "not-a-handle", current_scope(), current_backend()), + Err(TaskHandleError::Invalid) + ); + assert_eq!( + decode_current(&codec, "cfth1.AA", current_scope(), current_backend()), + Err(TaskHandleError::Invalid) + ); + assert_eq!( + decode_current(&codec, &tampered, current_scope(), current_backend()), + Err(TaskHandleError::Invalid) + ); + } + + #[test] + fn unsupported_handle_versions_are_rejected() { + let codec = codec(); + + assert_eq!( + decode_current(&codec, "cfth2.AA", current_scope(), current_backend()), + Err(TaskHandleError::UnsupportedVersion) + ); + assert_eq!( + decode_current(&codec, "cfthx.AA", current_scope(), current_backend()), + Err(TaskHandleError::Invalid) + ); + } + + #[test] + fn authorization_virtual_host_and_revision_scope_are_enforced() { + let codec = codec(); + let handle = codec.encode(current_scope(), current_backend(), "task-42").expect("handle encodes"); + let other_authorization_context = AuthorizationContextId::new("tenant-b:principal-a:team-a:scope-set-a"); + let other_virtual_host_id = VirtualHostId::new("host-b"); + let other_configuration_revision = ConfigurationRevision::new("revision-8"); + + for wrong_scope in [ + scope(&other_authorization_context, &VIRTUAL_HOST_ID, &CONFIGURATION_REVISION), + scope(&AUTHORIZATION_CONTEXT, &other_virtual_host_id, &CONFIGURATION_REVISION), + scope(&AUTHORIZATION_CONTEXT, &VIRTUAL_HOST_ID, &other_configuration_revision), + ] { + assert_eq!( + decode_current(&codec, &handle, wrong_scope, current_backend()), + Err(TaskHandleError::WrongScope) + ); + } + } + + #[test] + fn removed_or_reassigned_backends_are_rejected_without_exposing_identity() { + let codec = codec(); + let handle = codec.encode(current_scope(), current_backend(), "task-42").expect("handle encodes"); + + assert!(decode_current(&codec, &handle, current_scope(), current_backend()).is_ok()); + let removed = codec.decode(&handle, current_scope(), |_| false).expect_err("removed backend is rejected"); + let reassigned_generation = BackendGeneration::new("generation-4"); + let reassigned = decode_current(&codec, &handle, current_scope(), backend(&BACKEND_ID, &reassigned_generation)) + .expect_err("reassigned backend is rejected"); + + for error in [removed, reassigned] { + assert_eq!(error, TaskHandleError::UnavailableBackend); + assert_eq!(error.to_string(), "invalid task handle"); + assert!(!error.to_string().contains(BACKEND_ID_VALUE)); + assert!(!error.to_string().contains(BACKEND_GENERATION_VALUE)); + } + } + + #[test] + fn invalid_keys_and_sensitive_debug_output_are_redacted() { + assert!("short".parse::().is_err()); + assert!(format!("{KEY}=").parse::().is_err()); + assert!(format!("{KEY}AA").parse::().is_err()); + let key: TaskHandleKey = KEY.parse().expect("test key is valid"); + let codec = TaskHandleCodec::new(&key); + let handle = codec.encode(current_scope(), current_backend(), "bearer-task-id").expect("handle encodes"); + let route = decode_current(&codec, &handle, current_scope(), current_backend()).expect("handle decodes"); + + assert_eq!(format!("{key:?}"), "TaskHandleKey([REDACTED])"); + assert!(!format!("{key:?}").contains(KEY)); + assert_eq!( + format!("{route:?}"), + "TaskHandleRoute { backend_id: \"backend-a\", upstream_task_id: \"[REDACTED]\" }" + ); + assert!(!format!("{route:?}").contains("bearer-task-id")); + } + + #[test] + fn decode_errors_map_to_indistinguishable_invalid_params_errors() { + for error in [ + TaskHandleError::Invalid, + TaskHandleError::UnsupportedVersion, + TaskHandleError::WrongScope, + TaskHandleError::UnavailableBackend, + ] { + let protocol_error = ErrorData::from(error); + + assert_eq!(protocol_error.code, ErrorCode::INVALID_PARAMS); + assert_eq!(protocol_error.message, "invalid task ID"); + assert_eq!(protocol_error.data, None); + } + + let protocol_error = ErrorData::from(TaskHandleError::Encode); + assert_eq!(protocol_error.code, ErrorCode::INTERNAL_ERROR); + assert_eq!(protocol_error.message, "failed to create task handle"); + } +}