diff --git a/Cargo.lock b/Cargo.lock index 18e6591..9a295a6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -382,9 +382,9 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-lc-rs" -version = "1.18.0" +version = "1.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce2b2dcc879c3bae0d371e77c99f2238400ef24ec001394befa67b6e543add9e" +checksum = "b281d307588d634de920874890732659e2e7672f72b5e10e81badc1a8a83621e" dependencies = [ "aws-lc-sys", "untrusted 0.7.1", @@ -393,9 +393,9 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.44.0" +version = "0.45.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f09fae7be8bb3174e05c6afdb34199e6dc0c7c04ba9fa237b1967adfbde27483" +checksum = "9bff6c3b54fad79a2e60b8102caf565819711497c1f5f092f49508e2f5c31b27" dependencies = [ "cc", "cmake", @@ -2921,7 +2921,7 @@ dependencies = [ "tokio-rustls", "tokio-util", "tower", - "tower-http 0.7.0", + "tower-http 0.7.1", "tower-service", "tracing", "tracing-subscriber", @@ -3165,7 +3165,7 @@ dependencies = [ "tokio", "tokio-util", "tower", - "tower-http 0.7.0", + "tower-http 0.7.1", "tracing", "tracing-core", "tracing-serde", @@ -3535,29 +3535,26 @@ dependencies = [ [[package]] name = "imcp2" -version = "0.2.0" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5cd98dd4a89084c09983e2bdaa843604b35dd2668a6faa08f89ea105041c4ad3" +checksum = "5741b97d700d370235feca646911491528bb32d4e3e844f66a2bfea27534ccd1" dependencies = [ "anyhow", "axum", "base64 0.22.1", "candid", - "candid_parser", - "crc32fast", "getrandom 0.3.4", "hex", "ic-agent", + "imcp2-core", "prometheus", - "regex", - "reqwest 0.13.4", "rmcp", "serde", "serde_json", "sha2 0.11.0", "tokio", "tokio-util", - "tower-http 0.7.0", + "tower-http 0.7.1", "tracing", "tracing-subscriber", "url", @@ -3565,6 +3562,31 @@ dependencies = [ "uuid", ] +[[package]] +name = "imcp2-core" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "547f93d45d847e5dc7dbf6b131caadb1f7cb87fac58afd3c64866055845515c4" +dependencies = [ + "base64 0.22.1", + "candid", + "candid_parser", + "crc32fast", + "getrandom 0.3.4", + "hex", + "ic-agent", + "regex", + "reqwest 0.13.4", + "rmcp", + "serde", + "serde_json", + "sha2 0.11.0", + "tokio", + "tracing", + "url", + "urlencoding", +] + [[package]] name = "indexmap" version = "1.9.3" @@ -4028,9 +4050,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libredox" -version = "0.1.21" +version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7955dfc218a8afb29dfeffd540e3a6e96baeb94fe7138228dd7cc6937fbbf96" +checksum = "8d8f1ea3f21fd3405dcaf6c9b5c1630af9afc422d9073ea39c5f6d6c772e08ed" dependencies = [ "libc", ] @@ -4504,25 +4526,6 @@ dependencies = [ "libm", ] -[[package]] -name = "oauth2" -version = "5.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" -dependencies = [ - "base64 0.22.1", - "chrono", - "getrandom 0.2.17", - "http 1.5.0", - "rand 0.8.8", - "serde", - "serde_json", - "serde_path_to_error", - "sha2 0.10.9", - "thiserror 1.0.69", - "url", -] - [[package]] name = "objc2" version = "0.6.4" @@ -5680,11 +5683,9 @@ dependencies = [ "http 1.5.0", "http-body", "http-body-util", - "oauth2", "pastey", "pin-project-lite", "rand 0.10.2", - "reqwest 0.13.4", "rmcp-macros", "schemars 1.2.2", "serde", @@ -5696,7 +5697,6 @@ dependencies = [ "tokio-util", "tower-service", "tracing", - "url", "uuid", ] @@ -6418,9 +6418,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.15.2" +version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +checksum = "b9be42f50aa861c555654aa3a37f52f4b1074bacf4e48fe0ef7fa584e80f1f0f" [[package]] name = "smtp-proto" @@ -7145,9 +7145,9 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.7.0" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b11f75e912b0c2be01b63d8cf8057b8c3f97cf34abb3d431a3a4c8675498e233" +checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c" dependencies = [ "async-compression", "bitflags 2.13.1", @@ -7392,7 +7392,6 @@ dependencies = [ "idna", "percent-encoding", "serde", - "serde_derive", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 8a036ee..75d4744 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,7 +47,7 @@ ic-bn-lib = { version = "0.8.1", features = [ ] } ic-http-certification = { version = "3.1.0", optional = true } ic-http-gateway-protocol = { package = "ic-http-gateway-protocol", git = "https://github.com/dfinity/ic-http-gateway-protocol", tag = "v0.6.0" } -imcp2 = { version = "0.2", optional = true } +imcp2 = { version = "0.3", optional = true } isbot = "0.1" itertools = "0.15.0" lazy_static = "1.5.0" diff --git a/src/mcp.rs b/src/mcp.rs index 008e7e4..5541615 100644 --- a/src/mcp.rs +++ b/src/mcp.rs @@ -2,7 +2,7 @@ use std::{ fmt::Display, str::FromStr, sync::Arc, - time::{SystemTime, UNIX_EPOCH}, + time::{Duration, SystemTime, UNIX_EPOCH}, }; use anyhow::{Context, Error, anyhow}; @@ -14,6 +14,7 @@ use axum::{ response::{IntoResponse, Redirect, Response}, }; use candid::Principal; +use http::Uri; use ic_bn_lib::{ http::extract_authority, tasks::{Run, TaskManager}, @@ -89,6 +90,16 @@ impl Run for McpWrapper { } } +struct MetricsWrapper(Metrics); + +#[async_trait] +impl Run for MetricsWrapper { + async fn run(&self, _token: CancellationToken) -> Result<(), Error> { + self.0.refresh().await; + Ok(()) + } +} + pub struct McpState { hostname: String, router: Router, @@ -156,7 +167,11 @@ pub fn setup_mcp( ) .context("unable to create MCP metrics")?; - let mcp_redirect_url = cli.mcp_root_redirect.to_string(); + let mcp_redirect_url = cli + .mcp_root_redirect + .to_string() + .trim_end_matches('/') + .to_string(); let hostname = cli .mcp_public_url .as_ref() @@ -170,10 +185,23 @@ pub fn setup_mcp( .merge(mcp.well_known_router()) .merge(mcp.root_well_known_router()) .merge(auth_callbacks_router(&[&mcp])) - .fallback(|| async move { Redirect::permanent(&mcp_redirect_url) }) - .layer(from_fn_with_state(metrics, write_request_metrics)); + .fallback(|uri: Uri| async move { + // Preserve path & query of the original request + let redirect_url = uri.query().map_or_else( + || format!("{}{}", mcp_redirect_url, uri.path()), + |query| format!("{}{}?{query}", mcp_redirect_url, uri.path()), + ); + + Redirect::permanent(&redirect_url) + }) + .layer(from_fn_with_state(metrics.clone(), write_request_metrics)); tasks.add("mcp", Arc::new(McpWrapper(mcp))); + tasks.add_interval( + "mcp-metrics-refresh", + Arc::new(MetricsWrapper(metrics)), + Duration::from_secs(5), + ); Ok(McpState { hostname, router }) } diff --git a/src/routing/mod.rs b/src/routing/mod.rs index a3e6d46..efcc286 100644 --- a/src/routing/mod.rs +++ b/src/routing/mod.rs @@ -26,11 +26,7 @@ use ic_bn_lib::{ http::{ Client, ClientHttp, cache::{CacheBuilder, KeyExtractorUriRange}, - middleware::{ - rate_limiter::{Bypasser, NeverBypasser, TokenBypasser}, - request_meta, - waf::WafLayer, - }, + middleware::{request_meta, waf::WafLayer}, shed::{ ShardedOptions, ShedResponse, TypeExtractor, sharded::ShardedLittleLoadShedderLayer, @@ -186,23 +182,6 @@ impl TypeExtractor for RequestTypeExtractor { } } -/// Uses either `TokenBypasser` or `NeverBypasser` depending on whether a -/// bypass token is configured. Needed because bypasser is generic and we need a single type. -#[derive(Clone)] -enum RateLimitBypasser { - Token(TokenBypasser), - Never(NeverBypasser), -} - -impl Bypasser for RateLimitBypasser { - fn should_bypass(&self, req: &Request) -> bool { - match self { - Self::Token(b) => b.should_bypass(req), - Self::Never(b) => b.should_bypass(req), - } - } -} - // TODO: make it less horrible by using maybe builder pattern or just a struct #[allow(clippy::too_many_arguments)] #[allow(clippy::cognitive_complexity)] @@ -625,6 +604,25 @@ pub async fn setup_router( #[cfg(all(target_os = "linux", feature = "sev-snp"))] if cli.sev_snp.sev_snp_enable { + use ic_bn_lib::http::middleware::rate_limiter::{Bypasser, NeverBypasser, TokenBypasser}; + + /// Uses either `TokenBypasser` or `NeverBypasser` depending on whether a + /// bypass token is configured. Needed because bypasser is generic and we need a single type. + #[derive(Clone)] + enum RateLimitBypasser { + Token(TokenBypasser), + Never(NeverBypasser), + } + + impl Bypasser for RateLimitBypasser { + fn should_bypass(&self, req: &Request) -> bool { + match self { + Self::Token(b) => b.should_bypass(req), + Self::Never(b) => b.should_bypass(req), + } + } + } + let router_sev_snp = Router::new().route( "/sev-snp/report", post(ic_bn_lib::sev_snp::handler) @@ -885,23 +883,71 @@ mod test { // Any request to the MCP hostname that doesn't match a real MCP/OAuth/well-known // route falls through to a permanent redirect to the configured root-redirect URL — - // regardless of method, path, or query string (the fallback matches on path only). - for (method, path) in [ - (Method::GET, "/"), - (Method::POST, "/"), - (Method::GET, "/?foo=bar"), - (Method::GET, "/some/unknown/path"), - (Method::GET, "/MCP"), // path matching is case-sensitive: doesn't hit the /mcp mount + // regardless of method (the fallback matches on path only). The path and query of + // the original request are appended to the root-redirect URL, so the target site + // sees where the client was actually headed. + for (method, path, location) in [ + (Method::GET, "/", "https://internetcomputer.org/mcp/"), + (Method::POST, "/", "https://internetcomputer.org/mcp/"), + ( + Method::GET, + "/?foo=bar", + "https://internetcomputer.org/mcp/?foo=bar", + ), + ( + Method::GET, + "/some/unknown/path?a=1&b=2", + "https://internetcomputer.org/mcp/some/unknown/path?a=1&b=2", + ), + // Path matching is case-sensitive: doesn't hit the /mcp mount. + (Method::GET, "/MCP", "https://internetcomputer.org/mcp/MCP"), ] { let resp = router.call(request(method, MCP_HOST, path)).await.unwrap(); assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT, "path {path}"); assert_eq!( resp.headers().get(LOCATION).unwrap(), - "https://internetcomputer.org/mcp", + location, "path {path}", ); } + // A trailing slash on the configured root-redirect URL is trimmed before the path + // is appended, so the redirect target never ends up with a doubled slash. + { + let mut tasks = TaskManager::new(); + let (mut router, _domains) = setup_test_router_with_http_client( + &mut tasks, + Arc::new(TestClient(512)), + &[ + "--mcp-ii-instance", + "prod", + "--mcp-public-url", + "https://mcp.example.com", + "--mcp-state-dir", + "/tmp/ic-gateway-test-mcp-state-slash", + "--mcp-root-redirect", + "https://internetcomputer.org/", + ], + ) + .await; + + for (path, location) in [ + ("/", "https://internetcomputer.org/"), + ("/foo?a=1", "https://internetcomputer.org/foo?a=1"), + ] { + let resp = router + .call(request(Method::GET, MCP_HOST, path)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::PERMANENT_REDIRECT, "path {path}"); + assert_eq!( + resp.headers().get(LOCATION).unwrap(), + location, + "path {path}", + ); + } + } + // A request to some other hostname is untouched by MCP entirely: it falls through to // the ordinary gateway logic instead (here, the base-domain-root dashboard redirect). let resp = router