diff --git a/Cargo.lock b/Cargo.lock index 88bd0f0..f4ae014 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -565,18 +565,28 @@ checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" name = "contextforge-gateway-rs" version = "0.1.0" dependencies = [ + "axum", "clap", + "contextforge-gateway-rs-apis", "contextforge-gateway-rs-cpex", "contextforge-gateway-rs-lib", "cpex-payload-marker", "cpex-text-prefixer", "cpex-tool-namespace", + "http", + "jsonwebtoken 10.4.0", "num_cpus", + "openport", "opentelemetry", "opentelemetry-otlp", "opentelemetry_sdk", + "redis", + "reqwest", "rmcp", + "rmp-serde", "rustls", + "secrets_detection_rust", + "serde_json", "tikv-jemallocator", "tokio", "tonic", @@ -631,7 +641,7 @@ dependencies = [ "http", "hyper", "hyper-util", - "jsonwebtoken", + "jsonwebtoken 11.0.0", "lru_time_cache", "openport", "opentelemetry", @@ -643,6 +653,7 @@ dependencies = [ "rustls", "rustls-pki-types", "secret-string", + "secrets_detection_rust", "serde", "serde_json", "test-log", @@ -777,6 +788,15 @@ dependencies = [ "cpex", ] +[[package]] +name = "cpex-sdk" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba20e81ab4924231882adc31a437000fa045f890b80d7cce4fc1396848e710f0" +dependencies = [ + "cpex-core", +] + [[package]] name = "cpex-text-prefixer" version = "0.1.0" @@ -1861,6 +1881,30 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonwebtoken" +version = "10.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eba32bfb4ffdeaca3e34431072faf01745c9b26d25504aa7a6cf5684334fc4fc" +dependencies = [ + "base64 0.22.1", + "ed25519-dalek", + "getrandom 0.2.17", + "hmac", + "js-sys", + "p256", + "p384", + "pem", + "rand 0.8.5", + "rsa", + "serde", + "serde_json", + "sha2", + "signature", + "simple_asn1", + "zeroize", +] + [[package]] name = "jsonwebtoken" version = "11.0.0" @@ -3075,6 +3119,18 @@ version = "0.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "068d960589a13890bdb2a9ea5e7ae0974b89249101117f2e9a5e166bf3af44b9" +[[package]] +name = "secrets_detection_rust" +version = "0.0.0" +source = "git+https://github.com/IBM/cpex-plugins?rev=6ff7af74587574fe6115ce87427519b63f6062da#6ff7af74587574fe6115ce87427519b63f6062da" +dependencies = [ + "cpex-core", + "cpex-sdk", + "regex", + "serde", + "serde_json", +] + [[package]] name = "security-framework" version = "3.7.0" diff --git a/Cargo.toml b/Cargo.toml index 58a5ba2..ea9a3f8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -52,8 +52,11 @@ rmp-serde = "1.3.1" async-trait = "0.1.89" reqwest = "0.13" rustls = { version = "0.23", features = ["ring"] } -cpex = "0.2.0" +cpex = "=0.2.2" uuid = { version = "1.23.1", features = ["v4"] } +axum = "0.8" +openport = { version = "0.4.0", features = ["rand"] } +secrets_detection_rust = { git = "https://github.com/IBM/cpex-plugins", rev = "6ff7af74587574fe6115ce87427519b63f6062da", package = "secrets_detection_rust" } [profile.release] codegen-units = 1 diff --git a/crates/contextforge-gateway-rs-lib/Cargo.toml b/crates/contextforge-gateway-rs-lib/Cargo.toml index 411d107..09ce27e 100644 --- a/crates/contextforge-gateway-rs-lib/Cargo.toml +++ b/crates/contextforge-gateway-rs-lib/Cargo.toml @@ -19,7 +19,7 @@ tracing-opentelemetry.workspace = true opentelemetry.workspace = true tokio.workspace = true tokio-util = "0.7" -axum = "0.8" +axum.workspace = true axum-otel-metrics = "0.14" tower-http = { version = "0.7.0", features = ["full"] } tower = "0.5.3" @@ -52,7 +52,8 @@ with_tools = [] [dev-dependencies] opentelemetry_sdk.workspace = true cpex.workspace = true -openport = { version = "0.4.0", features = ["rand"] } +openport.workspace = true +secrets_detection_rust.workspace = true test-log = "0.2.20" axum-server = { version = "0.8.0", features = ["tls-rustls"] } diff --git a/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs b/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs index e0bf611..8273937 100644 --- a/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs +++ b/crates/contextforge-gateway-rs-lib/tests/gateway_plugins.rs @@ -4,6 +4,7 @@ use std::sync::{Arc, Mutex as StdMutex}; use contextforge_gateway_rs_cpex::CpexRuntimeRegistry; use cpex::cpex_core::cmf::Role; +use cpex::cpex_core::config::CpexConfig; use cpex::cpex_core::hooks::types::cmf_hook_names; use rmcp::{ ClientHandler, @@ -13,7 +14,7 @@ use rmcp::{ }, service::{NotificationContext, PeerRequestOptions, RequestHandle, RoleClient, RunningService}, }; -use serde_json::Value; +use serde_json::{Map, Value, json}; use support::{ POST_DENY_ERROR_CODE, PRE_DENY_ERROR_CODE, REWRITTEN_SUM_A, REWRITTEN_SUM_B, RunningGateway, TEST_USER_ID, @@ -124,6 +125,43 @@ fn raw_tool_call(tool_name: &str, request_id: i64, progress_token: &str) -> Valu }) } +fn fake_aws_access_key(suffix: &str) -> String { + ["AKIA", suffix].concat() +} + +fn sum_request_with_secret(secret_field: &str, secret: String) -> CallToolRequestParams { + let mut request = sum_request("sum", 1, 2); + request + .arguments + .as_mut() + .expect("sum request has arguments") + .insert(secret_field.to_owned(), Value::String(secret)); + request +} + +fn reflect_text_request(text: String) -> CallToolRequestParams { + CallToolRequestParams::new("reflect_text") + .with_arguments(Map::from_iter([("text".to_owned(), Value::String(text))])) +} + +async fn runtime_with_secrets_detection(hooks: Vec<&'static str>, plugin_config: Value) -> Arc { + let mut runtime = CpexRuntimeRegistry::default(); + runtime + .register_factory(secrets_detection_rust::KIND, Box::new(secrets_detection_rust::SecretsDetectionFactory)) + .expect("secrets detection factory registers"); + let config: CpexConfig = serde_json::from_value(json!({ + "plugins": [{ + "name": "secrets-detection", + "kind": secrets_detection_rust::KIND, + "hooks": hooks, + "config": plugin_config, + }] + })) + .expect("secrets detection CPEX config parses"); + runtime.apply_config(Some(config)).await.expect("secrets detection runtime applies"); + Arc::new(runtime) +} + fn sse_data_values(body: &str) -> Vec { let values = body .lines() @@ -333,6 +371,139 @@ async fn disabled_runtime_does_not_invoke_registered_plugin() { assert_eq!(0, post_observations.lock().expect("observations lock poisoned").post_calls); } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn secrets_detection_pre_hook_redacts_tool_arguments_before_backend_call() { + let runtime = runtime_with_secrets_detection( + vec![cmf_hook_names::TOOL_PRE_INVOKE], + json!({ + "redact": true, + "redaction_text": "[redacted]", + "block_on_detection": false, + }), + ) + .await; + let gateway = start_gateway("admin@example.com", true, runtime).await; + let service = gateway.connect("admin@example.com").await; + + let result = service + .call_tool(sum_request_with_secret("credential", fake_aws_access_key("1111111111111111"))) + .await + .expect("secret argument is redacted and call succeeds"); + + assert_eq!("3", text(&result)); + let backend_calls = gateway.backend_state.calls.lock().expect("backend calls lock poisoned"); + assert_eq!(1, backend_calls.len()); + assert_eq!( + Some(&Value::from("[redacted]")), + backend_calls[0].args.as_ref().and_then(|args| args.get("credential")) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn secrets_detection_clean_tool_payload_passes_through_unchanged() { + let runtime = runtime_with_secrets_detection( + vec![cmf_hook_names::TOOL_PRE_INVOKE, cmf_hook_names::TOOL_POST_INVOKE], + json!({ + "redact": true, + "redaction_text": "[redacted]", + "block_on_detection": true, + }), + ) + .await; + let gateway = start_gateway("admin@example.com", true, runtime).await; + let service = gateway.connect("admin@example.com").await; + + let result = + service.call_tool(sum_request("sum", 1, 2)).await.expect("clean argument payload passes through unchanged"); + + let result_text = text(&result); + assert_eq!("3", result_text.as_str()); + assert!(!result_text.contains("[redacted]")); + let backend_calls = gateway.backend_state.calls.lock().expect("backend calls lock poisoned"); + assert_eq!(1, backend_calls.len()); + assert_eq!("sum", backend_calls[0].tool_name); + let args = backend_calls[0].args.as_ref().expect("backend call has args"); + assert_eq!(2, args.len()); + assert_eq!(Some(&Value::from(1)), args.get("a")); + assert_eq!(Some(&Value::from(2)), args.get("b")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn secrets_detection_pre_hook_blocks_tool_arguments_before_backend_call() { + let runtime = runtime_with_secrets_detection( + vec![cmf_hook_names::TOOL_PRE_INVOKE], + json!({ + "redact": false, + "block_on_detection": true, + }), + ) + .await; + let gateway = start_gateway("admin@example.com", true, runtime).await; + let service = gateway.connect("admin@example.com").await; + + let error = service + .call_tool(sum_request_with_secret("credential", fake_aws_access_key("2222222222222222"))) + .await + .expect_err("secret argument blocks the call"); + + assert_eq!(ErrorCode::INVALID_REQUEST, error_code(error)); + assert!(gateway.backend_state.calls.lock().expect("backend calls lock poisoned").is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn secrets_detection_post_hook_redacts_tool_result_before_client_response() { + let runtime = runtime_with_secrets_detection( + vec![cmf_hook_names::TOOL_POST_INVOKE], + json!({ + "redact": true, + "redaction_text": "[redacted]", + "block_on_detection": false, + }), + ) + .await; + let gateway = start_gateway("admin@example.com", true, runtime).await; + let service = gateway.connect("admin@example.com").await; + + let result = service + .call_tool(reflect_text_request(fake_aws_access_key("3333333333333333"))) + .await + .expect("secret result is redacted and call succeeds"); + + assert_eq!("[redacted]", text(&result)); + assert_eq!(1, gateway.backend_state.calls.lock().expect("backend calls lock poisoned").len()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn secrets_detection_pre_hook_respects_field_allowlist() { + let runtime = runtime_with_secrets_detection( + vec![cmf_hook_names::TOOL_PRE_INVOKE], + json!({ + "redact": true, + "redaction_text": "[redacted]", + "block_on_detection": false, + "field_allowlist": ["credential"], + }), + ) + .await; + let gateway = start_gateway("admin@example.com", true, runtime).await; + let service = gateway.connect("admin@example.com").await; + let ignored_secret = fake_aws_access_key("4444444444444444"); + let mut request = sum_request_with_secret("credential", fake_aws_access_key("5555555555555555")); + request + .arguments + .as_mut() + .expect("sum request has arguments") + .insert("ignored".to_owned(), Value::String(ignored_secret.clone())); + + let result = service.call_tool(request).await.expect("allowed field is redacted"); + + assert_eq!("3", text(&result)); + let backend_calls = gateway.backend_state.calls.lock().expect("backend calls lock poisoned"); + let args = backend_calls[0].args.as_ref().expect("backend call has args"); + assert_eq!(Some(&Value::from("[redacted]")), args.get("credential")); + assert_eq!(Some(&Value::from(ignored_secret)), args.get("ignored")); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn pre_hook_modifies_backend_arguments_without_rerouting_tool() { let plugin = Arc::new(TestPlugin::new("pre", vec![cmf_hook_names::TOOL_PRE_INVOKE]).with_pre_rewrite()); diff --git a/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs b/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs index e7680c3..edf75a5 100644 --- a/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs +++ b/crates/contextforge-gateway-rs-lib/tests/support/plugin_gateway.rs @@ -127,6 +127,15 @@ impl ServerHandler for TestBackend { } Ok(CallToolResult::success(vec![ContentBlock::text("completed 4 packages")])) }, + "reflect_text" => { + let text = request + .arguments + .as_ref() + .and_then(|args| args.get("text")) + .and_then(Value::as_str) + .ok_or_else(|| ErrorData::invalid_params("reflect_text requires text", None))?; + Ok(CallToolResult::success(vec![ContentBlock::text(text.to_owned())])) + }, "wait_for_cancellation" => { cx.ct.cancelled().await; self.state diff --git a/crates/contextforge-gateway-rs/Cargo.toml b/crates/contextforge-gateway-rs/Cargo.toml index 4991618..07918a4 100644 --- a/crates/contextforge-gateway-rs/Cargo.toml +++ b/crates/contextforge-gateway-rs/Cargo.toml @@ -13,6 +13,7 @@ contextforge-gateway-rs-lib = { path = "../contextforge-gateway-rs-lib" } cpex-payload-marker = { git = "https://github.com/contextforge-gateway-rs/cpex-plugins-rs", rev = "ab47801daccfbba44ea07b033034a347e7b5afdd", optional = true } cpex-text-prefixer = { git = "https://github.com/contextforge-gateway-rs/cpex-plugins-rs", rev = "ab47801daccfbba44ea07b033034a347e7b5afdd", optional = true } cpex-tool-namespace = { git = "https://github.com/contextforge-gateway-rs/cpex-plugins-rs", rev = "ab47801daccfbba44ea07b033034a347e7b5afdd", optional = true } +secrets_detection_rust = { workspace = true, optional = true } clap.workspace = true tracing.workspace = true tracing-appender = "0.2.3" @@ -29,7 +30,19 @@ tikv-jemallocator = "0.7.0" rustls.workspace = true [features] +secrets-detection-plugin = ["dep:secrets_detection_rust"] test-plugins = ["dep:cpex-payload-marker", "dep:cpex-text-prefixer", "dep:cpex-tool-namespace"] +[dev-dependencies] +axum.workspace = true +contextforge-gateway-rs-apis.workspace = true +http.workspace = true +jsonwebtoken = { version = "10.3.0", features = ["rust_crypto"] } +openport.workspace = true +redis.workspace = true +reqwest.workspace = true +rmp-serde.workspace = true +serde_json.workspace = true + [lints] workspace = true diff --git a/crates/contextforge-gateway-rs/src/main.rs b/crates/contextforge-gateway-rs/src/main.rs index 4219d64..30a5734 100644 --- a/crates/contextforge-gateway-rs/src/main.rs +++ b/crates/contextforge-gateway-rs/src/main.rs @@ -58,11 +58,25 @@ fn plugin_runtime_from_config( ) -> Result> { let redis_client = RedisClient::try_from(RedisConfig::try_from(config)?)?; let plugin_runtime = CpexRuntimeRegistry::with_redis_config(redis_client); + #[cfg(any(feature = "test-plugins", feature = "secrets-detection-plugin"))] + let plugin_runtime = register_builtin_factories(plugin_runtime)?; + Ok(plugin_runtime) +} + +#[cfg(any(feature = "test-plugins", feature = "secrets-detection-plugin"))] +fn register_builtin_factories( + mut plugin_runtime: CpexRuntimeRegistry, +) -> Result> { #[cfg(feature = "test-plugins")] - let plugin_runtime = { - let mut plugin_runtime = plugin_runtime; + { test_plugins::register(&mut plugin_runtime)?; - plugin_runtime - }; + } + #[cfg(feature = "secrets-detection-plugin")] + { + plugin_runtime.register_factory( + secrets_detection_rust::KIND, + Box::new(secrets_detection_rust::SecretsDetectionFactory), + )?; + } Ok(plugin_runtime) } diff --git a/crates/contextforge-gateway-rs/tests/secrets_detection_e2e.rs b/crates/contextforge-gateway-rs/tests/secrets_detection_e2e.rs new file mode 100644 index 0000000..a9b35ef --- /dev/null +++ b/crates/contextforge-gateway-rs/tests/secrets_detection_e2e.rs @@ -0,0 +1,496 @@ +// Copyright 2026 +// SPDX-License-Identifier: Apache-2.0 + +#![cfg(feature = "secrets-detection-plugin")] + +use std::{ + collections::HashMap, + fs, + net::TcpStream as StdTcpStream, + path::PathBuf, + process::{Child, Command, Stdio}, + sync::{Arc, Mutex}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use contextforge_gateway_rs_apis::{ + User, + runtime_plugin_config::{RUNTIME_PLUGIN_CONFIG_KEY, RUNTIME_PLUGIN_CONFIG_VERSION}, + user_store::{BackendMCPGateway, Transport, UserConfig, VirtualHost}, +}; +use http::{HeaderMap, HeaderValue}; +use jsonwebtoken::{Algorithm, EncodingKey, Header, encode}; +use redis::aio::ConnectionManagerConfig; +use rmcp::{ + ErrorData, RoleClient, RoleServer, ServerHandler, ServiceExt, + model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ClientCapabilities, ContentBlock, ErrorCode, + Implementation, InitializeRequestParams, InitializeResult, ServerCapabilities, + }, + service::{RequestContext, ServiceError}, + transport::{ + StreamableHttpClientTransport, StreamableHttpServerConfig, StreamableHttpService, + streamable_http_client::StreamableHttpClientTransportConfig, + streamable_http_server::session::local::LocalSessionManager, + }, +}; +use serde_json::{Map, Value, json}; +use tokio::net::TcpListener; + +const TEST_USER_ID: &str = "11111111-1111-1111-1111-111111111111"; +const TEST_USER_EMAIL: &str = "admin@example.com"; +const TEST_VIRTUAL_HOST_ID: &str = "vh-secrets-e2e"; +const TEST_TOKEN_TTL_SECS: u64 = 60 * 60; +const REDACTED: &str = "[redacted]"; + +#[derive(Clone, Debug)] +struct BackendObservation { + args: Option>, +} + +#[derive(Clone, Default)] +struct BackendState { + calls: Arc>>, +} + +#[derive(Clone)] +struct TestBackend { + state: BackendState, +} + +impl ServerHandler for TestBackend { + async fn initialize( + &self, + _request: InitializeRequestParams, + _cx: RequestContext, + ) -> Result { + Ok(InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("secrets-e2e-backend", "0.1.0"))) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _cx: RequestContext, + ) -> Result { + self.state + .calls + .lock() + .expect("backend calls lock poisoned") + .push(BackendObservation { args: request.arguments.clone() }); + + let result = match request.name.as_ref() { + "sum" => { + let args = request + .arguments + .as_ref() + .ok_or_else(|| ErrorData::invalid_params("sum requires arguments", None))?; + let a = args + .get("a") + .and_then(Value::as_i64) + .ok_or_else(|| ErrorData::invalid_params("sum requires numeric a", None))?; + let b = args + .get("b") + .and_then(Value::as_i64) + .ok_or_else(|| ErrorData::invalid_params("sum requires numeric b", None))?; + Ok(CallToolResult::success(vec![ContentBlock::text((a + b).to_string())])) + }, + "reflect_text" => { + let text = request + .arguments + .as_ref() + .and_then(|args| args.get("text")) + .and_then(Value::as_str) + .ok_or_else(|| ErrorData::invalid_params("reflect_text requires text", None))?; + Ok(CallToolResult::success(vec![ContentBlock::text(text.to_owned())])) + }, + _ => Err(ErrorData { + code: ErrorCode::METHOD_NOT_FOUND, + message: format!("unknown tool {}", request.name).into(), + data: None, + }), + }?; + + Ok(result.into()) + } +} + +struct ChildProcess { + name: &'static str, + child: Child, + temp_dir: Option, + port: Option, +} + +impl ChildProcess { + fn new(name: &'static str, child: Child) -> Self { + Self { name, child, temp_dir: None, port: None } + } + + fn with_temp_dir(mut self, temp_dir: PathBuf) -> Self { + self.temp_dir = Some(temp_dir); + self + } + + fn with_port(mut self, port: u16) -> Self { + self.port = Some(port); + self + } + + fn port(&self) -> u16 { + self.port.expect("child process records a port") + } +} + +impl Drop for ChildProcess { + fn drop(&mut self) { + if self.child.try_wait().ok().flatten().is_none() { + let _ = self.child.kill(); + } + let _ = self.child.wait(); + if let Some(temp_dir) = &self.temp_dir { + let _ = fs::remove_dir_all(temp_dir); + } + } +} + +struct RunningBackend { + url: String, + state: BackendState, + handle: tokio::task::JoinHandle<()>, +} + +impl Drop for RunningBackend { + fn drop(&mut self) { + self.handle.abort(); + } +} + +struct E2eEnvironment { + gateway_url: String, + backend: RunningBackend, + _redis: ChildProcess, + _gateway: ChildProcess, +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "spawns redis-server and the contextforge-gateway-rs binary"] +async fn binary_e2e_redacts_tool_arguments_and_results() { + let backend = start_backend().await; + let env = start_environment( + backend, + json!({ + "redact": true, + "redaction_text": REDACTED, + "block_on_detection": false, + }), + ) + .await; + let service = connect_client(&env.gateway_url).await; + + let result = service + .call_tool(sum_request_with_secret("credential", fake_aws_access_key("1111111111111111"))) + .await + .expect("secret argument is redacted and call succeeds"); + + assert_eq!("3", tool_text(&result)); + { + let backend_calls = env.backend.state.calls.lock().expect("backend calls lock poisoned"); + assert_eq!(1, backend_calls.len()); + assert_eq!( + Some(&Value::from(REDACTED)), + backend_calls[0].args.as_ref().and_then(|args| args.get("credential")) + ); + } + + let result = service + .call_tool(reflect_text_request(fake_aws_access_key("2222222222222222"))) + .await + .expect("secret result is redacted and call succeeds"); + + assert_eq!(REDACTED, tool_text(&result)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "spawns redis-server and the contextforge-gateway-rs binary"] +async fn binary_e2e_blocks_tool_arguments_before_backend_call() { + let backend = start_backend().await; + let env = start_environment( + backend, + json!({ + "redact": false, + "block_on_detection": true, + }), + ) + .await; + let service = connect_client(&env.gateway_url).await; + + let error = service + .call_tool(sum_request_with_secret("credential", fake_aws_access_key("3333333333333333"))) + .await + .expect_err("secret argument blocks the call"); + + assert_eq!(ErrorCode::INVALID_REQUEST, mcp_error_code(error)); + assert!(env.backend.state.calls.lock().expect("backend calls lock poisoned").is_empty()); +} + +async fn start_environment(backend: RunningBackend, plugin_config: Value) -> E2eEnvironment { + let redis = start_redis().await; + write_redis_config(redis.port(), &backend).await; + write_runtime_plugin_config(redis.port(), plugin_config).await; + + let gateway_port = openport::pick_random_unused_port().expect("gateway port"); + let mut gateway = start_gateway_process(gateway_port, redis.port()); + wait_for_port(gateway_port, &mut gateway).await; + + E2eEnvironment { + gateway_url: format!("http://127.0.0.1:{gateway_port}/contextforge-rs/servers/{TEST_VIRTUAL_HOST_ID}/mcp"), + backend, + _redis: redis, + _gateway: gateway, + } +} + +async fn start_redis() -> ChildProcess { + let port = openport::pick_random_unused_port().expect("redis port"); + let temp_dir = std::env::temp_dir().join(format!("contextforge-data-plane-redis-{}-{port}", std::process::id())); + fs::create_dir_all(&temp_dir).expect("redis temp dir is created"); + + let child = Command::new("redis-server") + .args([ + "--port", + &port.to_string(), + "--save", + "", + "--appendonly", + "no", + "--dir", + temp_dir.to_str().expect("redis temp dir is UTF-8"), + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("redis-server starts; install redis-server to run this ignored E2E test"); + let mut redis = ChildProcess::new("redis-server", child).with_temp_dir(temp_dir).with_port(port); + wait_for_redis(port, &mut redis).await; + redis +} + +async fn start_backend() -> RunningBackend { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("backend binds"); + let port = listener.local_addr().expect("backend address").port(); + let state = BackendState::default(); + let backend_state = state.clone(); + let service = StreamableHttpService::new( + move || Ok(TestBackend { state: backend_state.clone() }), + LocalSessionManager::default().into(), + StreamableHttpServerConfig::default(), + ); + let router = axum::Router::new().route_service("/mcp", service); + let handle = tokio::spawn(async move { + axum::serve(listener, router).await.expect("backend serves"); + }); + RunningBackend { url: format!("http://127.0.0.1:{port}/mcp"), state, handle } +} + +fn start_gateway_process(gateway_port: u16, redis_port: u16) -> ChildProcess { + let binary = env!("CARGO_BIN_EXE_contextforge-gateway-rs"); + let child = Command::new(binary) + .current_dir(env!("CARGO_MANIFEST_DIR")) + .args([ + "--address", + &format!("127.0.0.1:{gateway_port}"), + "--redis-port", + &redis_port.to_string(), + "--redis-address", + "127.0.0.1", + "--token-verification-public-key", + "../../assets/jwt.key.pub", + "--number-of-cpus", + "1", + "--redis-mode", + "plain-text", + "--upstream-connection-mode", + "plain-text-or-tls", + "--runtime-plugins-enabled", + "true", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("contextforge-gateway-rs binary starts"); + ChildProcess::new("contextforge-gateway-rs", child) +} + +async fn wait_for_redis(port: u16, child: &mut ChildProcess) { + let client = redis::Client::open(format!("redis://127.0.0.1:{port}/")).expect("redis client opens"); + let deadline = Instant::now() + Duration::from_secs(5); + loop { + assert_child_running(child); + if let Ok(mut connection) = client.get_connection_manager_with_config(ConnectionManagerConfig::default()).await + && redis::cmd("PING").query_async::(&mut connection).await.is_ok() + { + return; + } + assert!(Instant::now() < deadline, "redis-server did not start on port {port}"); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +async fn wait_for_port(port: u16, child: &mut ChildProcess) { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + assert_child_running(child); + if StdTcpStream::connect(("127.0.0.1", port)).is_ok() { + return; + } + assert!(Instant::now() < deadline, "{} did not start on port {port}", child.name); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +fn assert_child_running(child: &mut ChildProcess) { + match child.child.try_wait() { + Ok(None) => {}, + Ok(Some(status)) => panic!("{} exited before the E2E test completed: {status}", child.name), + Err(error) => panic!("failed to inspect {} process: {error}", child.name), + } +} + +async fn write_redis_config(redis_port: u16, backend: &RunningBackend) { + let client = redis::Client::open(format!("redis://127.0.0.1:{redis_port}/")).expect("redis client opens"); + let mut connection = client + .get_connection_manager_with_config(ConnectionManagerConfig::default()) + .await + .expect("redis connection opens"); + redis::cmd("FLUSHDB").query_async::(&mut connection).await.expect("redis flush succeeds"); + + let key = rmp_serde::encode::to_vec(&User::new(TEST_USER_ID)).expect("user key encodes"); + let config = UserConfig { + virtual_hosts: HashMap::from([( + TEST_VIRTUAL_HOST_ID.to_owned(), + VirtualHost { + backends: HashMap::from([( + "backend".to_owned(), + BackendMCPGateway { + name: "backend".to_owned(), + url: backend.url.parse().expect("backend URL parses"), + transport: Transport::StreamableHttp, + 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(), + }, + )]), + }, + )]), + }; + let encoded = rmp_serde::encode::to_vec(&config).expect("user config encodes"); + redis::cmd("SET") + .arg(key) + .arg(encoded) + .query_async::(&mut connection) + .await + .expect("user config is written"); +} + +async fn write_runtime_plugin_config(redis_port: u16, plugin_config: Value) { + let client = redis::Client::open(format!("redis://127.0.0.1:{redis_port}/")).expect("redis client opens"); + let mut connection = client + .get_connection_manager_with_config(ConnectionManagerConfig::default()) + .await + .expect("redis connection opens"); + let document = json!({ + "version": RUNTIME_PLUGIN_CONFIG_VERSION, + "cpex": { + "plugins": [{ + "name": "secrets-detection", + "kind": "validator/secrets-detection", + "hooks": ["cmf.tool_pre_invoke", "cmf.tool_post_invoke"], + "config": plugin_config, + }] + } + }); + redis::cmd("SET") + .arg(RUNTIME_PLUGIN_CONFIG_KEY) + .arg(serde_json::to_vec(&document).expect("runtime plugin config serializes")) + .query_async::(&mut connection) + .await + .expect("runtime plugin config is written"); +} + +async fn connect_client(gateway_url: &str) -> rmcp::service::RunningService { + let mut headers = HeaderMap::new(); + headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {}", token(TEST_USER_ID))).expect("valid auth header"), + ); + let client = reqwest::Client::builder().default_headers(headers).build().expect("client builds"); + let transport = StreamableHttpClientTransport::with_client( + client, + StreamableHttpClientTransportConfig::with_uri(gateway_url.to_owned()), + ); + InitializeRequestParams::new(ClientCapabilities::default(), Implementation::new("secrets-e2e-client", "0.1.0")) + .serve(transport) + .await + .expect("client connects to gateway") +} + +fn sum_request_with_secret(secret_field: &str, secret: String) -> CallToolRequestParams { + let mut arguments = Map::from_iter([("a".to_owned(), Value::from(1)), ("b".to_owned(), Value::from(2))]); + arguments.insert(secret_field.to_owned(), Value::String(secret)); + CallToolRequestParams::new("sum").with_arguments(arguments) +} + +fn reflect_text_request(text: String) -> CallToolRequestParams { + CallToolRequestParams::new("reflect_text") + .with_arguments(Map::from_iter([("text".to_owned(), Value::String(text))])) +} + +fn fake_aws_access_key(suffix: &str) -> String { + ["AKIA", suffix].concat() +} + +fn tool_text(result: &CallToolResult) -> String { + result.content.iter().filter_map(|content| content.as_text()).map(|text| text.text.as_str()).collect() +} + +fn mcp_error_code(error: ServiceError) -> ErrorCode { + let ServiceError::McpError(error) = error else { + panic!("expected MCP error, got {error:?}"); + }; + error.code +} + +fn token(user_id: &str) -> String { + let key = EncodingKey::from_rsa_pem(&fs::read("../../assets/jwt.key").expect("jwt key")).expect("encoding key"); + let mut header = Header::new(Algorithm::RS256); + header.kid = Some("test".to_owned()); + let now = SystemTime::now().duration_since(UNIX_EPOCH).expect("system clock").as_secs(); + let claims = json!({ + "iss": "mcpgateway", + "sub": user_id, + "aud": "mcpgateway-api", + "exp": now + TEST_TOKEN_TTL_SECS, + "iat": now, + "jti": "test-token", + "token_use": "api", + "teams": ["team_awesome"], + "user": { + "email": TEST_USER_EMAIL, + "full_name": "API Token User", + "is_admin": true, + "auth_provider": "api_token" + }, + "scopes": { + "server_id": "my_id", + "permissions": ["tools.read", "servers.use"], + "ip_restrictions": ["192.169.1.0/24"], + "time_restrictions": null + }, + }); + encode(&header, &claims, &key).expect("jwt token") +}