From 80d8ada57623a126d72ea980f4740ed2dfc0e399 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Tue, 11 Aug 2026 11:47:57 +0100 Subject: [PATCH 1/2] Add stateless backend-scoped task handles Signed-off-by: lucarlig --- Cargo.lock | 2 + _context/wiki/config.md | 1 + _context/wiki/routing.md | 7 + _context/wiki/security.md | 5 +- crates/contextforge-data-plane-lib/Cargo.toml | 2 + .../contextforge-data-plane-lib/src/common.rs | 6 + crates/contextforge-data-plane-lib/src/lib.rs | 1 + .../src/task_handle.rs | 390 ++++++++++++++++++ 8 files changed, 413 insertions(+), 1 deletion(-) create mode 100644 crates/contextforge-data-plane-lib/src/task_handle.rs diff --git a/Cargo.lock b/Cargo.lock index 6abd30c3..d9db2ce4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -615,6 +615,7 @@ dependencies = [ "axum", "axum-otel-metrics", "axum-server", + "base64 0.22.1", "chrono", "clap", "contextforge-data-plane-apis", @@ -632,6 +633,7 @@ dependencies = [ "opentelemetry_sdk", "redis", "reqwest", + "ring", "rmcp", "rmp-serde", "rustls", diff --git a/_context/wiki/config.md b/_context/wiki/config.md index b7999a96..7ce2ad23 100644 --- a/_context/wiki/config.md +++ b/_context/wiki/config.md @@ -18,6 +18,7 @@ Plus at least: `--address` or `--tls-address`, `--token-verification-public-key` | `--server-private-key` | `TLS_SERVER_PRIVATE_KEY` | — | With `--tls-address` | | `--token-verification-public-key` | `TOKEN_VERIFICATION_PUBLIC_KEY` | — | RSA (RS256/384/512) | | `--token-verification-secret` | `TOKEN_SECRET` | — | HMAC (HS256/384/512) | +| `--task-handle-key` | `TASK_HANDLE_KEY` | — | Optional; base64url-no-pad 32-byte AES key. Share across replicas. Rotation invalidates handles. | | `--redis-address` | `REDIS_HOSTNAME` | **required** | | | `--redis-port` | `REDIS_PORT` | **required** | | | `--redis-mode` | `REDIS_CONNECTION_MODE` | **required** | `plain-text` \| `tls` \| `mtls` | diff --git a/_context/wiki/routing.md b/_context/wiki/routing.md index 4a681a85..a47ce34a 100644 --- a/_context/wiki/routing.md +++ b/_context/wiki/routing.md @@ -47,6 +47,13 @@ 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 Handles + +- Never expose an upstream task ID directly. +- Encode it as `cfth1.` with AES-256-GCM and a random nonce. +- Bind the payload to JWT `sub`, virtual host, and backend; reject mismatches and removed backends as `invalid task ID`. +- 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 9d0906da..63afa2f0 100644 --- a/_context/wiki/security.md +++ b/_context/wiki/security.md @@ -12,7 +12,7 @@ ## Identity And Authorization -Authentication is bearer-JWT only: +The control plane authenticates users and mints JWTs. The dataplane implements no login or IAM; it only validates bearer JWTs and applies routing scope: - Accepted algorithms: `RS256/RS384/RS512` (public key configured) or `HS256/HS384/HS512` (shared secret configured). Anything else is rejected. - `iss` must be `mcpgateway`, `aud` must be `mcpgateway-api`, and `exp` is validated. @@ -25,6 +25,7 @@ Authentication is bearer-JWT only: | 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. | @@ -51,4 +52,6 @@ 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. +- Task handles are authenticated, encrypted, and scoped to JWT `sub` + virtual host + backend. They do not replace JWT validation. - Never log: tokens, authorization headers, secrets, Redis key/value bytes, full `UserConfig` documents, or backend credentials. +- Treat task handles as opaque; do not log them. diff --git a/crates/contextforge-data-plane-lib/Cargo.toml b/crates/contextforge-data-plane-lib/Cargo.toml index 6505887b..72400621 100644 --- a/crates/contextforge-data-plane-lib/Cargo.toml +++ b/crates/contextforge-data-plane-lib/Cargo.toml @@ -46,6 +46,8 @@ tokio-rustls = "0.26.4" typed-builder = "0.23.2" url = { workspace = true, features = ["serde"] } secret-string = "0.0.2" +base64 = "0.22.1" +ring = "0.17.14" [features] diff --git a/crates/contextforge-data-plane-lib/src/common.rs b/crates/contextforge-data-plane-lib/src/common.rs index b736b35e..de6421e2 100644 --- a/crates/contextforge-data-plane-lib/src/common.rs +++ b/crates/contextforge-data-plane-lib/src/common.rs @@ -16,6 +16,7 @@ use thiserror::Error; use typed_builder::TypedBuilder; use url::Url; +use crate::task_handle::TaskHandleKey; use crate::user_config_store::UserConfigStore; #[derive(Clone)] @@ -156,6 +157,11 @@ pub struct Config { #[arg(long, env = "CONTEXTFORGE_DATA_PLANE_TOKEN_SECRET")] pub token_verification_secret: Option>, + /// Shared AES-256 key used to protect stateless task handles. The value is + /// URL-safe base64 without padding and must decode to exactly 32 bytes. + #[arg(long, env = "CONTEXTFORGE_DATA_PLANE_TASK_HANDLE_KEY")] + pub task_handle_key: Option, + #[arg(long, env = "CONTEXTFORGE_DATA_PLANE_ENABLE_OPEN_TELEMETRY")] pub enable_open_telemetry: Option, diff --git a/crates/contextforge-data-plane-lib/src/lib.rs b/crates/contextforge-data-plane-lib/src/lib.rs index 4a7f9b32..d9d51151 100644 --- a/crates/contextforge-data-plane-lib/src/lib.rs +++ b/crates/contextforge-data-plane-lib/src/lib.rs @@ -14,6 +14,7 @@ mod common; mod const_values; mod gateway; mod layers; +pub mod task_handle; mod telemetry; mod transports; 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 00000000..35471093 --- /dev/null +++ b/crates/contextforge-data-plane-lib/src/task_handle.rs @@ -0,0 +1,390 @@ +//! 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 authenticated +//! subject and virtual-host scope prevents a handle from being replayed through +//! a different caller or route, while the backend lookup ensures that removed +//! backends fail closed. + +use std::{fmt, str::FromStr}; + +use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use contextforge_data_plane_apis::user_store::VirtualHost; +use ring::{ + aead::{AES_256_GCM, Aad, LessSafeKey, Nonce, UnboundKey}, + rand::{SecureRandom, SystemRandom}, +}; +use rmcp::{ErrorData, model::ErrorCode}; +use secret_string::SecretString; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +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 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(SecretString); + +impl TaskHandleKey { + fn bytes(&self) -> Result<[u8; KEY_LEN], TaskHandleKeyError> { + let bytes = URL_SAFE_NO_PAD.decode(self.0.value()).map_err(|_| TaskHandleKeyError)?; + bytes.try_into().map_err(|_| TaskHandleKeyError) + } +} + +impl FromStr for TaskHandleKey { + type Err = TaskHandleKeyError; + + fn from_str(value: &str) -> Result { + let key = Self(SecretString::new(value.to_owned())); + key.bytes()?; + Ok(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; + +/// Authenticated request scope to which a task handle is bound. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TaskHandleScope<'a> { + subject: &'a str, + virtual_host_id: &'a str, +} + +impl<'a> TaskHandleScope<'a> { + /// Creates a scope from the authenticated subject and path virtual-host ID. + pub fn new(subject: &'a str, virtual_host_id: &'a str) -> Self { + Self { subject, virtual_host_id } + } +} + +/// The route recovered from a valid task handle. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TaskHandleRoute { + backend_name: String, + upstream_task_id: String, +} + +impl TaskHandleRoute { + /// Backend map key in the caller's current virtual-host configuration. + pub fn backend_name(&self) -> &str { + &self.backend_name + } + + /// Original task identifier expected by the upstream backend. + pub fn upstream_task_id(&self) -> &str { + &self.upstream_task_id + } +} + +/// 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 { + #[error("failed to create task handle")] + Encode, + #[error("invalid task handle")] + Invalid, + #[error("invalid task handle")] + UnsupportedVersion, + #[error("invalid task handle")] + WrongScope, + #[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(Debug, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct TaskHandlePayload { + subject: String, + virtual_host_id: String, + backend_name: String, + upstream_task_id: String, +} + +/// Encodes and decodes versioned task handles shared across dataplane replicas. +#[derive(Clone)] +pub struct TaskHandleCodec { + key: LessSafeKey, + random: SystemRandom, +} + +impl TaskHandleCodec { + /// Creates a codec from a validated shared key. + pub fn new(key: &TaskHandleKey) -> Result { + let key = UnboundKey::new(&AES_256_GCM, &key.bytes()?).map_err(|_| TaskHandleKeyError)?; + Ok(Self { key: LessSafeKey::new(key), random: SystemRandom::new() }) + } + + /// Creates an opaque handle for one upstream task. + pub fn encode( + &self, + scope: TaskHandleScope<'_>, + backend_name: &str, + upstream_task_id: &str, + ) -> Result { + let payload = TaskHandlePayload { + subject: scope.subject.to_owned(), + virtual_host_id: scope.virtual_host_id.to_owned(), + backend_name: backend_name.to_owned(), + upstream_task_id: upstream_task_id.to_owned(), + }; + let mut ciphertext = serde_json::to_vec(&payload).map_err(|_| TaskHandleError::Encode)?; + + let mut nonce_bytes = [0_u8; NONCE_LEN]; + self.random.fill(&mut nonce_bytes).map_err(|_| TaskHandleError::Encode)?; + let nonce = Nonce::assume_unique_for_key(nonce_bytes); + self.key + .seal_in_place_append_tag(nonce, Aad::from(HANDLE_PREFIX), &mut ciphertext) + .map_err(|_| TaskHandleError::Encode)?; + + let mut protected = Vec::with_capacity(NONCE_LEN + ciphertext.len()); + protected.extend_from_slice(&nonce_bytes); + 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 + /// caller's current virtual host and references a currently configured + /// backend. + pub fn decode( + &self, + handle: &str, + expected_scope: TaskHandleScope<'_>, + virtual_host: &VirtualHost, + ) -> Result { + 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::assume_unique_for_key(nonce_bytes); + let mut ciphertext = ciphertext.to_vec(); + let plaintext = self + .key + .open_in_place(nonce, Aad::from(HANDLE_PREFIX), &mut ciphertext) + .map_err(|_| TaskHandleError::Invalid)?; + let payload: TaskHandlePayload = serde_json::from_slice(plaintext).map_err(|_| TaskHandleError::Invalid)?; + + if payload.subject != expected_scope.subject || payload.virtual_host_id != expected_scope.virtual_host_id { + return Err(TaskHandleError::WrongScope); + } + if !virtual_host.backends.contains_key(&payload.backend_name) { + return Err(TaskHandleError::UnavailableBackend); + } + + Ok(TaskHandleRoute { backend_name: payload.backend_name, upstream_task_id: 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::collections::HashMap; + + use contextforge_data_plane_apis::user_store::BackendMCPGateway; + use url::Url; + + use super::*; + + const KEY: &str = "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8"; + + fn codec() -> TaskHandleCodec { + TaskHandleCodec::new(&KEY.parse().expect("test key is valid")).expect("test key initializes AES-256-GCM") + } + + fn backend(name: &str) -> BackendMCPGateway { + BackendMCPGateway { + name: name.to_owned(), + url: Url::parse(&format!("https://{name}.example.com/mcp")).expect("test URL is valid"), + passthrough_headers: Vec::new(), + add_headers: HashMap::new(), + remove_headers: Vec::new(), + allowed_tool_names: Vec::new(), + tool_name_aliases: HashMap::new(), + allowed_resource_names: Vec::new(), + allowed_prompt_names: Vec::new(), + } + } + + fn virtual_host(names: &[&str]) -> VirtualHost { + VirtualHost { backends: names.iter().map(|name| ((*name).to_owned(), backend(name))).collect() } + } + + fn scope<'a>(subject: &'a str, virtual_host_id: &'a str) -> TaskHandleScope<'a> { + TaskHandleScope::new(subject, virtual_host_id) + } + + #[test] + fn arbitrary_upstream_task_ids_round_trip_without_loss() { + let codec = codec(); + let virtual_host = virtual_host(&["backend-a"]); + let task_ids = ["", "simple", "with/slashes?and=query", "nul\0byte", "emoji-🦀", "line\nbreak"]; + + for task_id in task_ids { + let handle = codec.encode(scope("caller", "host-a"), "backend-a", task_id).expect("handle encodes"); + let route = codec.decode(&handle, scope("caller", "host-a"), &virtual_host).expect("handle decodes"); + + assert_eq!(route.backend_name(), "backend-a"); + assert_eq!(route.upstream_task_id(), task_id); + } + } + + #[test] + fn identical_task_ids_from_different_backends_remain_isolated() { + let codec = codec(); + let virtual_host = virtual_host(&["backend-a", "backend-b"]); + let handle_a = codec.encode(scope("caller", "host-a"), "backend-a", "same-id").expect("handle A encodes"); + let handle_b = codec.encode(scope("caller", "host-a"), "backend-b", "same-id").expect("handle B encodes"); + + let route_a = codec.decode(&handle_a, scope("caller", "host-a"), &virtual_host).expect("handle A decodes"); + let route_b = codec.decode(&handle_b, scope("caller", "host-a"), &virtual_host).expect("handle B decodes"); + + assert_ne!(handle_a, handle_b); + assert_eq!(route_a.backend_name(), "backend-a"); + assert_eq!(route_b.backend_name(), "backend-b"); + } + + #[test] + fn handle_decodes_on_another_codec_with_the_same_key() { + let first_replica = codec(); + let second_replica = codec(); + let virtual_host = virtual_host(&["backend-a"]); + let handle = first_replica.encode(scope("caller", "host-a"), "backend-a", "task-42").expect("handle encodes"); + + let route = second_replica.decode(&handle, scope("caller", "host-a"), &virtual_host).expect("handle decodes"); + + assert_eq!(route.upstream_task_id(), "task-42"); + } + + #[test] + fn malformed_and_tampered_handles_fail_closed() { + let codec = codec(); + let virtual_host = virtual_host(&["backend-a"]); + let handle = codec.encode(scope("caller", "host-a"), "backend-a", "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!( + codec.decode("not-a-handle", scope("caller", "host-a"), &virtual_host), + Err(TaskHandleError::Invalid) + ); + assert_eq!(codec.decode("cfth1.AA", scope("caller", "host-a"), &virtual_host), Err(TaskHandleError::Invalid)); + assert_eq!(codec.decode(&tampered, scope("caller", "host-a"), &virtual_host), Err(TaskHandleError::Invalid)); + } + + #[test] + fn unsupported_handle_versions_are_rejected() { + let codec = codec(); + let virtual_host = virtual_host(&["backend-a"]); + + assert_eq!( + codec.decode("cfth2.AA", scope("caller", "host-a"), &virtual_host), + Err(TaskHandleError::UnsupportedVersion) + ); + assert_eq!(codec.decode("cfthx.AA", scope("caller", "host-a"), &virtual_host), Err(TaskHandleError::Invalid)); + } + + #[test] + fn caller_and_virtual_host_scope_are_enforced() { + let codec = codec(); + let virtual_host = virtual_host(&["backend-a"]); + let handle = codec.encode(scope("caller-a", "host-a"), "backend-a", "task-42").expect("handle encodes"); + + assert_eq!(codec.decode(&handle, scope("caller-b", "host-a"), &virtual_host), Err(TaskHandleError::WrongScope)); + assert_eq!(codec.decode(&handle, scope("caller-a", "host-b"), &virtual_host), Err(TaskHandleError::WrongScope)); + } + + #[test] + fn removed_backends_are_rejected_without_exposing_the_backend_name() { + let codec = codec(); + let original_virtual_host = virtual_host(&["backend-a"]); + let current_virtual_host = virtual_host(&["backend-b"]); + let handle = codec.encode(scope("caller", "host-a"), "backend-a", "task-42").expect("handle encodes"); + + assert!(codec.decode(&handle, scope("caller", "host-a"), &original_virtual_host).is_ok()); + let error = + codec.decode(&handle, scope("caller", "host-a"), ¤t_virtual_host).expect_err("backend was removed"); + + assert_eq!(error, TaskHandleError::UnavailableBackend); + assert_eq!(error.to_string(), "invalid task handle"); + assert!(!error.to_string().contains("backend-a")); + } + + #[test] + fn invalid_keys_are_rejected_and_debug_output_is_redacted() { + assert_eq!("short".parse::(), Err(TaskHandleKeyError)); + let key: TaskHandleKey = KEY.parse().expect("test key is valid"); + + assert_eq!(format!("{key:?}"), "TaskHandleKey([REDACTED])"); + assert!(!format!("{key:?}").contains(KEY)); + } + + #[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"); + } +} From 4653f8f97e81b5edcae23c66f04fc0fe94ec6364 Mon Sep 17 00:00:00 2001 From: lucarlig Date: Tue, 11 Aug 2026 14:45:23 +0100 Subject: [PATCH 2/2] Stabilize gateway test port allocation Signed-off-by: lucarlig --- .../tests/support/list_tools_gateway.rs | 104 ++++++++++++------ 1 file changed, 68 insertions(+), 36 deletions(-) diff --git a/crates/contextforge-data-plane-lib/tests/support/list_tools_gateway.rs b/crates/contextforge-data-plane-lib/tests/support/list_tools_gateway.rs index 074bd9b1..f283bd23 100644 --- a/crates/contextforge-data-plane-lib/tests/support/list_tools_gateway.rs +++ b/crates/contextforge-data-plane-lib/tests/support/list_tools_gateway.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, net::SocketAddr, sync::Arc}; +use std::{collections::HashMap, sync::Arc}; use contextforge_data_plane_apis::{ User, @@ -41,14 +41,19 @@ pub(crate) fn plaintext_config(gateway_port: u16) -> Config { } pub(crate) fn create_ports(ports: usize) -> Vec { - (0..ports).map(|_| openport::pick_random_unused_port().expect("Expecting to find port")).collect() + let mut selected = Vec::with_capacity(ports); + while selected.len() < ports { + let port = openport::pick_random_unused_port().expect("Expecting to find port"); + if !selected.contains(&port) { + selected.push(port); + } + } + selected } pub(crate) async fn create_gateway_with_four_counters(user: &str, config: Config) -> Result { let mocked_user_config_store = MemoryUserConfigStore::default(); - - let gateway_one_ports = create_ports(2); - let gateway_two_ports = create_ports(2); + let gateway_port = config.address.ok_or("Invalid configuration")?.port(); let service = StreamableHttpService::new( || Ok(mock_counter::Counter::new()), @@ -58,6 +63,9 @@ pub(crate) async fn create_gateway_with_four_counters(user: &str, config: Config let router = axum::Router::new().route_service("/mcp", service); + let (gateway_one_ports, servers_one) = create_axum_servers(2, gateway_port, &router).await?; + let (gateway_two_ports, servers_two) = create_axum_servers(2, gateway_port, &router).await?; + assert_ne!(gateway_one_ports, gateway_two_ports); let gateway_one_backends = create_backends(&gateway_one_ports, false); @@ -102,8 +110,6 @@ pub(crate) async fn create_gateway_with_four_counters(user: &str, config: Config if let Some(address) = config.address.as_ref() { let gateway_url = format!("http://{address}/contextforge-rs/servers/{virtual_host_one_id}/mcp"); - let servers_one = create_axum_servers(&gateway_one_ports, &router); - let servers_two = create_axum_servers(&gateway_two_ports, &router); let handle = tokio::spawn(futures::future::join_all(vec![gateway].into_iter().chain(servers_one).chain(servers_two))); @@ -125,9 +131,7 @@ pub(crate) async fn create_tls_gateway_with_four_tls_counters( config: Config, ) -> Result { let mocked_user_config_store = MemoryUserConfigStore::default(); - - let gateway_one_ports = create_ports(2); - let gateway_two_ports = create_ports(2); + let gateway_port = config.tls_address.ok_or("Invalid configuration")?.port(); let service = StreamableHttpService::new( || Ok(mock_counter::Counter::new()), @@ -137,6 +141,9 @@ pub(crate) async fn create_tls_gateway_with_four_tls_counters( let router = axum::Router::new().route_service("/mcp", service); + let (gateway_one_ports, servers_one) = create_axum_tls_servers(2, gateway_port, router.clone()).await?; + let (gateway_two_ports, servers_two) = create_axum_tls_servers(2, gateway_port, router).await?; + assert_ne!(gateway_one_ports, gateway_two_ports); let gateway_one_backends = create_backends(&gateway_one_ports, true); @@ -181,8 +188,6 @@ pub(crate) async fn create_tls_gateway_with_four_tls_counters( if let Some(address) = config.tls_address.as_ref() { let gateway_url = format!("https://{address}/contextforge-rs/servers/{virtual_host_one_id}/mcp"); - let servers_one = create_axum_tls_servers(&gateway_one_ports, router.clone()).await; - let servers_two = create_axum_tls_servers(&gateway_two_ports, router.clone()).await; let handle = tokio::spawn(futures::future::join_all(vec![gateway].into_iter().chain(servers_one).chain(servers_two))); @@ -272,41 +277,68 @@ fn create_resource_template_uris(ports: &[u16]) -> Vec { .collect() } -fn create_axum_servers(ports: &[u16], router: &axum::Router) -> Vec>> { - ports - .iter() - .map(|port| { - let addr = format!("127.0.0.1:{port}"); - let router = router.clone(); - async { - let listener = tokio::net::TcpListener::bind(addr).await.expect("Expect this to work"); - axum::serve(listener, router).await.expect("server runs"); +async fn create_axum_servers( + server_count: usize, + gateway_port: u16, + router: &axum::Router, +) -> Result<(Vec, Vec>>)> { + let mut ports = Vec::with_capacity(server_count); + let mut servers = Vec::with_capacity(server_count); + + while ports.len() < server_count { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let port = listener.local_addr()?.port(); + if port == gateway_port { + continue; + } + + let router = router.clone(); + ports.push(port); + servers.push( + async move { + axum::serve(listener, router).await?; Ok(()) } - .boxed() - }) - .collect() + .boxed(), + ); + } + + Ok((ports, servers)) } -async fn create_axum_tls_servers(ports: &[u16], router: axum::Router) -> Vec>> { +async fn create_axum_tls_servers( + server_count: usize, + gateway_port: u16, + router: axum::Router, +) -> Result<(Vec, Vec>>)> { let config = axum_server::tls_rustls::RustlsConfig::from_pem_file( "../../assets/contextforgeCA/contextforge-server.cert.pem", "../../assets/contextforgeCA/contextforge-server.key.pem", ) .await .expect("Expect this to work"); - - ports - .iter() - .map(|port| { - let router = router.clone(); - let addr: SocketAddr = format!("127.0.0.1:{port}").parse().expect("Expect this to work"); - let config = config.clone(); + let mut ports = Vec::with_capacity(server_count); + let mut servers = Vec::with_capacity(server_count); + + while ports.len() < server_count { + let listener = std::net::TcpListener::bind("127.0.0.1:0")?; + let port = listener.local_addr()?.port(); + if port == gateway_port { + continue; + } + + listener.set_nonblocking(true)?; + let server = axum_server::from_tcp_rustls(listener, config.clone())?; + let router = router.clone(); + ports.push(port); + servers.push( async move { - _ = axum_server::bind_rustls(addr, config).serve(router.into_make_service()).await; + server.serve(router.into_make_service()).await?; Ok(()) } - .boxed() - }) - .collect() + .boxed(), + ); + } + + Ok((ports, servers)) }