From 490983fe604e836399fc195a354737a0f7a25ba6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 25 Aug 2026 14:14:57 +0000 Subject: [PATCH 01/10] Stage 2: imcp2-local, the stdio binary with built-in II browser login MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The local half of the design merged in #148: a single-user binary an AI tool spawns and talks to over stdio, serving the same tool surface as the hosted server against the same mainnet and production Internet Identity, with none of the hosted server's OAuth machinery. imcp2-local (new crate, publish = false; distributed as release binaries per the design's Stage 3): - login driver: `authenticate` mints a FRESH session id + registration key, binds a transient loopback listener on 127.0.0.1:0, and returns the id.ai connect link in-band (plus a best-effort `open::that` browser launch), never blocking on the callback. The listener serves exactly three routes — the pinned fragment-reading callback page (shared with hosted, same CSP + hardening headers), the slim redeem, and the #4091 allow-list (CORS-readable, no-store) — and is torn down on redeem, on the 10-minute handshake timeout, or when a fresh flow replaces it. The redeem mirrors the hosted order (validate state, parse before claiming, single-flight claim, one mcp_register_v2 call) minus the OAuth tail: success answers {"done": true} and the page renders "Signed in — you can close this tab" in place. - wrapper handler: the two local-only login tools live on LocalServer's own #[tool_router]; tools/list is the merged surface (26 core + 2), tools/call dispatches to the login router first and forwards everything else (calls, resources, get_tool) to the inner IcTools handler verbatim. get_info identifies as imcp2-local and appends the sign-in guidance to core's instructions. - main: mainnet agent (IMCP2_IC_URL override; IMCP2_FETCH_ROOT_KEY only honoured for loopback endpoints — the local-replica test configuration), production II (II_URL_PROD/II_CANISTER_ID_PROD overrides), management identity derived at the hosted server's origin by default (IMCP2_MANAGEMENT_ORIGIN) so the same anchor keeps the same controller principal across deployments, IMCP2_NO_OPEN, all logging to stderr (stdout is the JSON-RPC channel). imcp2-core (the two seams Stage 1 left for the login driver): - SessionSource::Singleton now holds a SessionSlot — a shared, REPLACEABLE holder, not a fixed id: the II contract requires a fresh session key after Unauthorized, so each successful login writes a fresh session id and reauthentication is a same-process step. - the pinned callback page's script gains the {"done": true} success arm (hosted keeps answering `redirect`), with a `.done` terminal state in connect.css that drops the spinner without the error line. - regression guard: the core router must never carry the login tools (it IS the hosted tools/list). Tests: 218 pass workspace-wide (9 new in imcp2-local, 3 new in core), including a real MCP round-trip over an in-process duplex pipe (merged tools/list, login lifecycle through tools/call, forwarding into core tools) and browser-played listener tests (pinned page + CSP, allow-list verbatim-match + CORS, redeem refusal paths staying retryable). The built binary answers initialize/tools/list over real stdio. No behavior change to the hosted server. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- Cargo.lock | 64 +- Cargo.toml | 19 +- crates/imcp2-core/src/assets/connect.css | 4 + crates/imcp2-core/src/iiconnect.rs | 40 +- crates/imcp2-core/src/lib.rs | 2 +- crates/imcp2-core/src/tools.rs | 78 ++- crates/imcp2-local/Cargo.toml | 45 ++ crates/imcp2-local/src/login.rs | 778 +++++++++++++++++++++++ crates/imcp2-local/src/main.rs | 161 +++++ crates/imcp2-local/src/server.rs | 372 +++++++++++ 10 files changed, 1535 insertions(+), 28 deletions(-) create mode 100644 crates/imcp2-local/Cargo.toml create mode 100644 crates/imcp2-local/src/login.rs create mode 100644 crates/imcp2-local/src/main.rs create mode 100644 crates/imcp2-local/src/server.rs diff --git a/Cargo.lock b/Cargo.lock index 1e9e2ba..b0a889b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -996,7 +996,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -1823,6 +1823,25 @@ dependencies = [ "urlencoding", ] +[[package]] +name = "imcp2-local" +version = "0.2.0" +dependencies = [ + "anyhow", + "axum", + "imcp2-core", + "open", + "reqwest 0.13.4", + "rmcp", + "serde", + "serde_json", + "tokio", + "tracing", + "tracing-subscriber", + "url", + "uuid", +] + [[package]] name = "indexmap" version = "2.14.0" @@ -1850,6 +1869,25 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "is-docker" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "928bae27f42bc99b60d9ac7334e3a21d10ad8f1835a4e12ec3ec0464765ed1b3" +dependencies = [ + "once_cell", +] + +[[package]] +name = "is-wsl" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "173609498df190136aa7dea1a91db051746d339e18476eed5ca40521f02d7aa5" +dependencies = [ + "is-docker", + "once_cell", +] + [[package]] name = "itertools" version = "0.11.0" @@ -2176,7 +2214,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -2250,6 +2288,16 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" +[[package]] +name = "open" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ade3be4664bc1ef537ce133015f04c176b737815c2ba9fd60edf212d6e90dd55" +dependencies = [ + "is-wsl", + "libc", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -2999,7 +3047,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3057,7 +3105,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3457,7 +3505,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -3499,7 +3547,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -3613,7 +3661,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4326,7 +4374,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 321c1b3..3b850f7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,13 +34,18 @@ exclude = [ ] [workspace] -members = [".", "crates/imcp2-core"] +members = [".", "crates/imcp2-core", "crates/imcp2-local"] -# Single source of truth for the version pins both crates share, so the agent, -# codec, and candid lines can never drift between imcp2-core and imcp2. +# Single source of truth for the version pins the crates share, so the agent, +# codec, and candid lines can never drift between imcp2-core, imcp2, and +# imcp2-local. [workspace.dependencies] imcp2-core = { path = "crates/imcp2-core", version = "0.2.0" } rmcp = "1.7" +axum = "0.8" +anyhow = "1" +uuid = { version = "1", features = ["v4"] } +tracing-subscriber = { version = "0.3", features = ["env-filter"] } # The Agent is INJECTED by the embedding application (see McpServer::new), so a # host shares one boundary-node client and one ic-agent across the process. # Deliberately a direct dependency for now: sourcing ic-agent through ic-bn-lib @@ -83,14 +88,14 @@ rmcp = { workspace = true, features = [ ic-agent = { workspace = true } candid = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread", "signal", "time"] } -axum = "0.8" -anyhow = "1" +axum = { workspace = true } +anyhow = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } tracing = { workspace = true } -tracing-subscriber = { version = "0.3", features = ["env-filter"] } +tracing-subscriber = { workspace = true } tokio-util = "0.7" -uuid = { version = "1", features = ["v4"] } +uuid = { workspace = true } tower-http = { version = "0.7", features = ["cors"] } # Metrics exposition at /metrics. Pinned to the version dfinity/ic uses: these # series are destined for the same Prometheus / Victoria Metrics clusters that diff --git a/crates/imcp2-core/src/assets/connect.css b/crates/imcp2-core/src/assets/connect.css index b20f671..e1c2c45 100644 --- a/crates/imcp2-core/src/assets/connect.css +++ b/crates/imcp2-core/src/assets/connect.css @@ -132,6 +132,10 @@ body { /* Terminal error: the spinner tile is dropped and the message carries the state. The pinned page's script adds `.error` to `.screen`; the error page ships it. */ .screen.error .spinner-tile { display: none; } +/* Terminal success ("Signed in — you can close this tab", the local redeem's + `done` answer): the spinner is likewise dropped, but the contact line stays + hidden — nothing went wrong. */ +.screen.done .spinner-tile { display: none; } @keyframes connect-spin { to { transform: rotate(360deg); } } @media (prefers-reduced-motion: reduce) { .spinner { animation: none; } } @media (prefers-color-scheme: dark) { diff --git a/crates/imcp2-core/src/iiconnect.rs b/crates/imcp2-core/src/iiconnect.rs index 4c41758..7466965 100644 --- a/crates/imcp2-core/src/iiconnect.rs +++ b/crates/imcp2-core/src/iiconnect.rs @@ -99,11 +99,14 @@ const PINNED_PAGE_HTML: &str = include_str!("assets/connect-callback.html"); /// one dynamic value, the redeem URL, is spliced in by replacing /// `__REDEEM_URL__`, which sits inside a quoted JS string literal below. const PINNED_PAGE_JS: &str = r#"(function () { - function show(t, err) { + // Swap the status line to `t` and move the screen into a terminal state: + // 'error' drops the spinner and reveals the contact line, 'done' just drops + // the spinner (see connect.css). + function show(t, cls) { document.getElementById('m').textContent = t; - if (err) { + if (cls) { var c = document.querySelector('.screen'); - if (c) { c.classList.add('error'); } + if (c) { c.classList.add(cls); } } } // II delivers #delegation=&state=: the two-hop chain plus @@ -129,14 +132,19 @@ const PINNED_PAGE_JS: &str = r#"(function () { }) .then(function (r) { return r.json().catch(function () { return {}; }); }) .then(function (d) { + // Two success shapes, per deployment: the hosted redeem answers with the + // `redirect` that continues its OAuth flow; a local redeem has nothing to + // navigate to, so it answers `done` and this page IS the terminal state. if (d && d.redirect) { location.replace(d.redirect); + } else if (d && d.done) { + show("Signed in — you can close this tab.", 'done'); } else { - show((d && d.error) || "We couldn't finish the connection. Restart from your client.", true); + show((d && d.error) || "We couldn't finish the connection. Restart from your client.", 'error'); } }) .catch(function () { - show("We couldn't reach the server. Restart from your client.", true); + show("We couldn't reach the server. Restart from your client.", 'error'); }); })();"#; @@ -488,6 +496,28 @@ mod tests { assert!(url.contains("&state=sess-1&ttl=3600®istration_key=AQID"), "{url}"); } + // The page's script resolves the redeem answer three ways: `redirect` + // (hosted — navigate into the OAuth continuation), `done` (local — render + // the terminal "Signed in" state in place; there is nothing to navigate + // to), and everything else as the error screen. The `done` state drops the + // spinner without revealing the contact line (see connect.css). Guards the + // contract the local redeem answers with `{"done": true}` against. + #[test] + fn pinned_page_script_handles_redirect_done_and_error() { + let page = super::pinned_callback_page("/redeem", "mcp@dfinity.org"); + assert!(page.html.contains("d.redirect"), "hosted success arm"); + assert!(page.html.contains("d.done"), "local success arm"); + assert!( + page.html.contains("Signed in \u{2014} you can close this tab."), + "the done arm renders the terminal signed-in state" + ); + assert!(page.html.contains("'done'") && page.html.contains("'error'")); + assert!( + super::CONNECT_PAGE_CSS.contains(".screen.done .spinner-tile"), + "the done state must drop the spinner (connect.css)" + ); + } + #[test] fn pinned_page_binds_one_nonce_into_html_and_csp() { let page = super::pinned_callback_page("/oauth/connect/redeem", "mcp@dfinity.org"); diff --git a/crates/imcp2-core/src/lib.rs b/crates/imcp2-core/src/lib.rs index 29e9b4e..85407ff 100644 --- a/crates/imcp2-core/src/lib.rs +++ b/crates/imcp2-core/src/lib.rs @@ -36,7 +36,7 @@ mod discover; mod management; pub use identities::{IiInstance, SessionGauges}; -pub use tools::{AuthedSession, IcTools, SessionSource}; +pub use tools::{AuthedSession, IcTools, SessionSlot, SessionSource}; /// The IC [`Agent`] type the components are built around, re-exported so /// callers construct the injected agent from the exact `ic-agent` version this /// crate links. diff --git a/crates/imcp2-core/src/tools.rs b/crates/imcp2-core/src/tools.rs index c7b4b75..6b2f20a 100644 --- a/crates/imcp2-core/src/tools.rs +++ b/crates/imcp2-core/src/tools.rs @@ -18,6 +18,7 @@ use rmcp::{ }; use crate::{calls, discover, identities, identities::Identities, management, skills}; +use std::sync::Arc; /// Cap on the per-canister Candid probes open_app / discover_app_canisters run to /// fill in OQL / api-doc capability flags (#3). Discovery output is already bounded, @@ -76,16 +77,16 @@ impl IcTools { /// The session id this tool call acts as, per the deployment's /// [`SessionSource`]: looked up from the request's [`AuthedSession`] - /// extension under `Bearer` (multi-user, hosted), or the fixed - /// login-established session under `Singleton` (single-user, local). + /// extension under `Bearer` (multi-user, hosted), or read from the shared + /// login-filled [`SessionSlot`] under `Singleton` (single-user, local). /// `None` means the call is unauthenticated ("needs an authenticated /// session" errors at the call sites are unchanged). fn current_session_id(&self, ctx: &RequestContext) -> Option { match &self.session { SessionSource::Bearer => authed_session(ctx).map(|s| s.session_id), - SessionSource::Singleton(id) => { + SessionSource::Singleton(slot) => { let _ = ctx; - Some(id.clone()) + slot.get() } } } @@ -1521,9 +1522,40 @@ pub enum SessionSource { /// Multi-user (the hosted server): resolve the session per request from /// the [`AuthedSession`] extension its bearer-token gate injected. Bearer, - /// Single-user (a local server): every call acts as this one session, - /// established by the binary's own login flow. - Singleton(String), + /// Single-user (a local server): every call acts as the one session in + /// the shared slot, established — and replaced on each re-login — by the + /// binary's own login flow. + Singleton(SessionSlot), +} + +/// The single-user session holder behind [`SessionSource::Singleton`]: a +/// shared, **replaceable** slot rather than a fixed id. The Internet Identity +/// contract requires a FRESH session key after an `Unauthorized` (a reused id +/// would reuse its keys), so each successful login writes a fresh session id +/// here — same-process reauthentication after grant expiry or revocation, with +/// no client restart. Clones share the slot: the login flow holds a writing +/// handle, [`IcTools`] reads it per call. Starts empty (signed out). +#[derive(Clone, Debug, Default)] +pub struct SessionSlot(Arc>>); + +impl SessionSlot { + /// A fresh, empty slot (no session — every session-needing tool reports + /// "needs an authenticated session" until the login flow fills it). + pub fn new() -> Self { + Self::default() + } + + /// The current session id, if a login has completed. The critical section + /// is a clone — never held across an await. + pub fn get(&self) -> Option { + self.0.read().expect("session slot poisoned").clone() + } + + /// Replace the session id after a successful login (a fresh id per the II + /// fresh-session-key contract). + pub fn set(&self, session_id: String) { + *self.0.write().expect("session slot poisoned") = Some(session_id); + } } /// The verified session id of an authenticated MCP session, injected into the @@ -2454,6 +2486,38 @@ mod tests { assert_eq!(cc.destructive_hint, Some(true)); } + // The local binary's login tools (`authenticate`/`auth_status`) live on its + // OWN wrapper handler, never in this router: this router IS what the hosted + // server advertises on `tools/list`, so a login tool landing here would ship + // to every hosted client (which logs in via OAuth instead). Regression guard + // for that boundary. + #[test] + fn the_core_router_carries_no_local_login_tools() { + let tools = super::IcTools::tool_router().list_all(); + for name in ["authenticate", "auth_status"] { + assert!( + tools.iter().all(|t| &*t.name != name), + "{name} must not be in the core router (it would ship on the hosted tools/list)" + ); + } + } + + // The Singleton session slot is shared and REPLACEABLE: it starts empty + // (signed out), a login fills it, and a re-login (fresh session id, per the + // II fresh-session-key contract) replaces it — through every clone, since + // clones share the holder. This is what makes reauthentication after grant + // expiry a same-process step for the local binary. + #[test] + fn the_session_slot_is_shared_and_replaceable() { + let slot = super::SessionSlot::new(); + let reader = slot.clone(); + assert_eq!(reader.get(), None, "a fresh slot is signed out"); + slot.set("sess-1".into()); + assert_eq!(reader.get(), Some("sess-1".into()), "clones share the slot"); + slot.set("sess-2".into()); + assert_eq!(reader.get(), Some("sess-2".into()), "a re-login replaces the id"); + } + // EVERY tool must declare an outputSchema so a model knows the shape of its // reply — and MCP requires that schema to be object-rooted. This guards the // whole surface: a new tool added without an output schema fails here. diff --git a/crates/imcp2-local/Cargo.toml b/crates/imcp2-local/Cargo.toml new file mode 100644 index 0000000..ae0e975 --- /dev/null +++ b/crates/imcp2-local/Cargo.toml @@ -0,0 +1,45 @@ +[package] +name = "imcp2-local" +version = "0.2.0" +edition = "2021" +description = "Local MCP server for the Internet Computer: the imcp2 tool surface over stdio, with a built-in Internet Identity browser login." +license = "Apache-2.0" +repository = "https://github.com/dfinity/imcp2" +# Distributed as signed release binaries (the design's end-user setup and +# signing components), not through crates.io — flip deliberately if that ever +# changes. This also keeps the tag-lockstep publish workflow scoped to +# imcp2-core + imcp2. +publish = false + +[dependencies] +# The whole server is composed from imcp2-core: the IcTools surface (served +# over stdio here), the II session engine, and the connect-handshake +# primitives the login listener wraps. The IC agent is constructed from +# core's `ic_agent` re-export, so the process links exactly one ic-agent. +imcp2-core = { workspace = true } +rmcp = { workspace = true, features = ["server", "macros", "transport-io"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "net", "time"] } +# The transient loopback login listener (component 5 of the design): three +# routes for one handshake, then torn down. The MCP tool surface never rides +# HTTP here. +axum = { workspace = true } +# Best-effort browser launch for the sign-in link, as in the ICP CLI's +# `icp identity link web`. The flow never depends on it succeeding. +open = "5" +anyhow = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tracing = { workspace = true } +# stdout is the JSON-RPC channel, so ALL diagnostics go to stderr. +tracing-subscriber = { workspace = true } +uuid = { workspace = true } +url = { workspace = true } + +[dev-dependencies] +# The login-flow tests play the browser: fetch the callback page, the +# allow-list, and POST redeems against the loopback listener. +reqwest = { workspace = true, features = ["json"] } +# The wrapper tests run a REAL MCP round-trip: the server served over an +# in-process duplex pipe (the same shape as stdio), driven by rmcp's client. +rmcp = { workspace = true, features = ["client"] } +tokio = { workspace = true, features = ["io-util"] } diff --git a/crates/imcp2-local/src/login.rs b/crates/imcp2-local/src/login.rs new file mode 100644 index 0000000..8851ad5 --- /dev/null +++ b/crates/imcp2-local/src/login.rs @@ -0,0 +1,778 @@ +//! The **browser Internet Identity login** for the local binary: mint a fresh +//! session, stand up a transient loopback listener, hand the user an id.ai +//! sign-in link, and redeem the delegation II navigates back with. On success +//! the shared [`SessionSlot`] is filled and every tool call acts as the user; +//! the listener is torn down (redeem or timeout) and the grant lives in memory +//! only. +//! +//! The listener is the unavoidable minimum, not a design slip: II delivers the +//! delegation by *navigating the browser* to the callback (a URL fragment only +//! a served page can read and POST back), and II's #4091 check fetches +//! `/.well-known/ii-auth-callbacks` from the callback's origin before honoring +//! it — both require a real HTTP origin. This is the standard native-app +//! loopback redirect (RFC 8252 §7.3), the same shape as the ICP CLI's +//! `icp identity link web`. It serves exactly three routes for one handshake +//! and never serves the MCP tool surface, which rides stdio. + +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use axum::extract::State; +use axum::response::{Html, IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use imcp2_core::identities::Identities; +use imcp2_core::iiconnect::{self, RedeemBody, AUTH_CALLBACKS_WELL_KNOWN}; +use imcp2_core::SessionSlot; +use serde_json::json; +use tokio::sync::{Mutex, Notify}; + +/// How long one sign-in handshake may take, from the `authenticate` call to +/// the browser's redeem, before the pending flow and its listener are torn +/// down. Matches the hosted server's pending-connect TTL. +const HANDSHAKE_TTL: Duration = Duration::from_secs(600); + +/// `ttl` (seconds) requested for the II grant in the connect link, matching +/// the hosted server. II clamps it to [600, 2592000], and the effective +/// session duration is whatever the user picks on II's consent screen. +const GRANT_TTL_SECS: u64 = 3600; + +/// Where the pinned callback page's error state points users ("contact us to +/// report it") — DFINITY ships this binary, so failures report to the same +/// address as the hosted server's. +const CONTACT: &str = "mcp@dfinity.org"; + +fn now_ns() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() as u64 +} + +/// The login flow's shared handle: the `authenticate`/`auth_status` tools call +/// [`begin`](Self::begin)/[`status`](Self::status), and the loopback redeem +/// route completes the flow through a clone. Single-user by construction — one +/// pending handshake at a time, serialized by the state lock. +#[derive(Clone)] +pub struct LoginDriver { + inner: Arc, +} + +struct Inner { + identities: Identities, + /// The shared, replaceable session slot `IcTools` reads under + /// `SessionSource::Singleton`: each successful redeem writes the fresh + /// session id here (the II contract requires a fresh session key after + /// `Unauthorized`, and a fresh id gets fresh keys). + slot: SessionSlot, + /// Best-effort `open::that` on the sign-in link. The flow never depends on + /// it succeeding — the link is always returned in-band by `authenticate`. + auto_open: bool, + state: Mutex, +} + +#[derive(Default)] +struct LoginState { + pending: Option, + /// The most recent successful login, kept (even past expiry) so + /// `auth_status` can say *whose* session expired. + grant: Option, +} + +struct Pending { + /// Doubles as the connect `state` in the II link — II echoes it to the + /// callback, and the redeem accepts only this value. + session_id: String, + url: String, + /// This flow's loopback callback (`http://127.0.0.1:/callback`) — + /// the origin the listener serves on. + callback_url: String, + started: Instant, + /// Single-flight marker: `true` while a redemption attempt (the + /// `mcp_register_v2` network call) is mid-flight, so a page double-submit + /// can't fire two. Cleared on failure so a genuine retry can proceed. + redeeming: bool, + /// Tears the transient listener down (graceful): notified on redeem + /// success, on handshake timeout, and when a fresh flow replaces this one. + shutdown: Arc, +} + +impl Pending { + fn expired(&self) -> bool { + self.started.elapsed() >= HANDSHAKE_TTL + } +} + +/// A completed login, as `auth_status` reports it. +#[derive(Clone)] +pub struct Grant { + pub session_id: String, + /// The session principal (`self_authenticating(pub(S))`) — attribution, + /// not the user's per-app principal. + pub principal: Option, + /// II's recorded access level: `"queries"` (read-only) or `"all"`. + pub permissions: &'static str, + /// Grant expiration, ns since the Unix epoch (the user's consent-time + /// session choice). + pub expiration_ns: u64, +} + +impl Grant { + pub fn expired(&self) -> bool { + self.expiration_ns <= now_ns() + } + + /// Whole minutes until expiry (0 when expired). + pub fn minutes_left(&self) -> u64 { + self.expiration_ns.saturating_sub(now_ns()) / 60_000_000_000 + } +} + +/// What [`LoginDriver::status`] reports — the states `auth_status` renders. +pub enum LoginStatus { + SignedIn(Grant), + /// A handshake is waiting for the browser. + Pending { + url: String, + }, + /// A previous grant ran out (or was revoked and then expired here); a new + /// `authenticate` gets a fresh session. + Expired(Grant), + SignedOut, +} + +/// What [`LoginDriver::begin`] returns to the `authenticate` tool. +pub enum BeginOutcome { + /// A live grant exists and `refresh` was not requested. + AlreadySignedIn(Grant), + /// A sign-in link is ready. `fresh` distinguishes a newly started flow + /// from an already-pending one being returned again (single-flight). + Pending { url: String, fresh: bool }, +} + +impl LoginDriver { + pub fn new(identities: Identities, slot: SessionSlot, auto_open: bool) -> Self { + Self { + inner: Arc::new(Inner { + identities, + slot, + auto_open, + state: Mutex::new(LoginState::default()), + }), + } + } + + /// Start (or return the already-pending) sign-in: mint a **fresh** session + /// id + registration key, bind the loopback listener, and build the id.ai + /// link — returned immediately, never blocking on the browser. With a live + /// grant and `refresh: false`, reports the grant instead of starting over. + pub async fn begin(&self, refresh: bool) -> Result { + // The lock is held across the (local, network-free) setup, so + // concurrent `authenticate` calls serialize into one flow. + let mut state = self.inner.state.lock().await; + + if let Some(p) = &state.pending { + if !p.expired() { + return Ok(BeginOutcome::Pending { + url: p.url.clone(), + fresh: false, + }); + } + // Expired while nobody was looking: tear it down, then start fresh. + if let Some(p) = state.pending.take() { + p.shutdown.notify_one(); + } + } + if !refresh { + if let Some(g) = &state.grant { + if !g.expired() { + return Ok(BeginOutcome::AlreadySignedIn(g.clone())); + } + } + } + + // A FRESH session id per flow: fresh id ⇒ fresh session + registration + // keys in `Identities`, per the II fresh-session-key contract. + let session_id = uuid::Uuid::new_v4().to_string(); + let reg_pubkey = self + .inner + .identities + .registration_pubkey_b64(&session_id) + .await?; + + // Port 0: the OS picks a free port; both the II link's `callback` and + // the #4091 allow-list entry derive from the ONE resulting origin, so + // they cannot drift (II matches by exact string equality). + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .map_err(|e| format!("could not bind the login callback listener: {e}"))?; + let port = listener + .local_addr() + .map_err(|e| format!("could not read the listener address: {e}"))? + .port(); + let callback_url = format!("http://127.0.0.1:{port}/callback"); + + let url = iiconnect::ii_mcp_url( + &self.inner.identities.instance().ii_url, + &callback_url, + &session_id, + GRANT_TTL_SECS, + ®_pubkey, + ); + + let shutdown = Arc::new(Notify::new()); + let app = login_router(self.clone(), callback_url.clone()); + let sd = shutdown.clone(); + tokio::spawn(async move { + if let Err(e) = axum::serve(listener, app) + .with_graceful_shutdown(async move { sd.notified().await }) + .await + { + tracing::warn!("login listener error: {e}"); + } + }); + + // Watchdog: when the handshake window closes, drop the pending flow + // and the listener (if a redeem didn't already). + let driver = self.clone(); + let sid = session_id.clone(); + tokio::spawn(async move { + tokio::time::sleep(HANDSHAKE_TTL).await; + let mut state = driver.inner.state.lock().await; + let ours = matches!(&state.pending, Some(p) if p.session_id == sid); + if ours { + if let Some(p) = state.pending.take() { + p.shutdown.notify_one(); + tracing::info!( + callback = %p.callback_url, + "sign-in window expired; login listener closed" + ); + } + } + }); + + state.pending = Some(Pending { + session_id, + url: url.clone(), + callback_url, + started: Instant::now(), + redeeming: false, + shutdown, + }); + drop(state); + + if self.inner.auto_open { + // `open::that` can block on the launcher, so it runs off the + // async workers — and only best-effort: the link is in the tool + // result either way. + let link = url.clone(); + tokio::task::spawn_blocking(move || { + if let Err(e) = open::that(&link) { + tracing::debug!("could not auto-open the browser: {e}"); + } + }); + } + + Ok(BeginOutcome::Pending { url, fresh: true }) + } + + /// The pending flow's `(session_id, callback_url)`, for the tests that + /// play the browser against the live listener. + #[cfg(test)] + pub(crate) async fn pending_handshake(&self) -> Option<(String, String)> { + let state = self.inner.state.lock().await; + state + .pending + .as_ref() + .map(|p| (p.session_id.clone(), p.callback_url.clone())) + } + + pub async fn status(&self) -> LoginStatus { + let state = self.inner.state.lock().await; + if let Some(g) = &state.grant { + if !g.expired() { + return LoginStatus::SignedIn(g.clone()); + } + } + if let Some(p) = &state.pending { + if !p.expired() { + return LoginStatus::Pending { url: p.url.clone() }; + } + } + match &state.grant { + Some(g) => LoginStatus::Expired(g.clone()), + None => LoginStatus::SignedOut, + } + } +} + +/// The three transient loopback routes — the login handshake's whole HTTP +/// surface. `callback_url` is this listener's own callback, echoed verbatim in +/// the allow-list document. +fn login_router(driver: LoginDriver, callback_url: String) -> Router { + Router::new() + .route("/callback", get(callback_page)) + .route("/redeem", post(redeem)) + .route(AUTH_CALLBACKS_WELL_KNOWN, get(auth_callbacks)) + .with_state(RouteCtx { + driver, + callback_url, + }) +} + +#[derive(Clone)] +struct RouteCtx { + driver: LoginDriver, + callback_url: String, +} + +/// GET /callback — the pinned fragment-reading page (shared with the hosted +/// server), pointed at this listener's `/redeem`, with the same non-CSP +/// hardening headers the hosted server adds. +async fn callback_page() -> Response { + let page = iiconnect::pinned_callback_page("/redeem", CONTACT); + let mut resp = Html(page.html).into_response(); + let h = resp.headers_mut(); + h.insert( + axum::http::header::CONTENT_SECURITY_POLICY, + axum::http::HeaderValue::from_str(&page.csp).expect("valid CSP"), + ); + h.insert( + axum::http::header::REFERRER_POLICY, + axum::http::HeaderValue::from_static("no-referrer"), + ); + h.insert( + axum::http::header::X_CONTENT_TYPE_OPTIONS, + axum::http::HeaderValue::from_static("nosniff"), + ); + h.insert( + axum::http::header::X_FRAME_OPTIONS, + axum::http::HeaderValue::from_static("DENY"), + ); + resp +} + +/// GET /.well-known/ii-auth-callbacks — the #4091 allow-list: exactly this +/// listener's callback. II fetches it cross-origin (hence the one CORS +/// header) and fail-closed before honoring the callback; `no-store` keeps an +/// intermediary from serving a stale document for a past listener's port. +async fn auth_callbacks(State(ctx): State) -> Response { + let mut resp = Json(json!({ "callbacks": [ctx.callback_url] })).into_response(); + let h = resp.headers_mut(); + h.insert( + axum::http::header::CACHE_CONTROL, + axum::http::HeaderValue::from_static("no-store"), + ); + h.insert( + axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN, + axum::http::HeaderValue::from_static("*"), + ); + resp +} + +fn redeem_err(msg: &str) -> Response { + ( + axum::http::StatusCode::BAD_REQUEST, + Json(json!({ "error": msg })), + ) + .into_response() +} + +/// The local success answer: no OAuth continuation to redirect into, so the +/// pinned page renders its terminal "Signed in — you can close this tab" +/// state in place (the page script's `done` arm). +fn done() -> Response { + Json(json!({ "done": true })).into_response() +} + +/// POST /redeem — the pinned page POSTs the fragment here. The hosted redeem +/// minus its OAuth tail: no initiator cookie (there is no public initiate +/// endpoint to confuse a deputy through — the binary itself mints `X` and this +/// listener is loopback-only), no PKCE code, and no redirect — success answers +/// `{"done": true}` and the page renders its terminal state. Parse before +/// claiming, single-flight claim, one `mcp_register_v2` call, then fill the +/// session slot and tear the listener down. +async fn redeem(State(ctx): State, Json(body): Json) -> Response { + let driver = &ctx.driver; + + // Validate the delivery against the pending flow first (NOT claiming yet), + // so an unknown/expired `state` gets its accurate refusal without parsing. + { + let state = driver.inner.state.lock().await; + if state + .grant + .as_ref() + .is_some_and(|g| g.session_id == body.state) + { + // A retry of an already-successful redeem (II's delegation redeems + // repeatedly within its lifetime): idempotent success. + return done(); + } + match &state.pending { + Some(p) if p.session_id == body.state => { + if p.expired() { + return redeem_err( + "This sign-in expired. Ask the agent to authenticate again.", + ); + } + } + _ => { + return redeem_err( + "This sign-in is unknown or already used. Ask the agent to authenticate again.", + ) + } + } + } + + // Decode the fragment delegation (pure, size-bounded) BEFORE claiming, so + // a malformed delivery never occupies the single-flight slot. + let (user_key, chain) = match iiconnect::parse_registration_delegation(&body.delegation) { + Ok(v) => v, + Err(e) => { + return redeem_err(&format!( + "We couldn't read the sign-in response. Ask the agent to authenticate again. ({e})" + )) + } + }; + + // Claim atomically (re-checking: the flow may have moved while parsing). + { + let mut state = driver.inner.state.lock().await; + if state + .grant + .as_ref() + .is_some_and(|g| g.session_id == body.state) + { + return done(); + } + match state.pending.as_mut() { + Some(p) if p.session_id == body.state => { + if p.expired() { + return redeem_err( + "This sign-in expired. Ask the agent to authenticate again.", + ); + } + if p.redeeming { + return redeem_err( + "This sign-in is already being processed. Wait a moment; \ + if nothing happens, ask the agent to authenticate again.", + ); + } + p.redeeming = true; + } + _ => { + return redeem_err( + "This sign-in is unknown or already used. Ask the agent to authenticate again.", + ) + } + } + } + + // The one network call, outside the lock. + match driver + .inner + .identities + .redeem_registration_delegation(&body.state, user_key, chain) + .await + { + Err(e) => { + let mut state = driver.inner.state.lock().await; + if let Some(p) = state.pending.as_mut() { + if p.session_id == body.state { + p.redeeming = false; // free the claim for a genuine retry + } + } + redeem_err(&e) + } + Ok(outcome) => { + let principal = driver.inner.identities.session_principal(&body.state).await; + let mut state = driver.inner.state.lock().await; + // Fill the shared slot LAST-WRITER-WINS: II bound the grant to this + // session's key, so even if a newer flow replaced the pending entry + // mid-call, this successful login is real and becomes current. + driver.inner.slot.set(body.state.clone()); + state.grant = Some(Grant { + session_id: body.state.clone(), + principal, + permissions: outcome.permissions, + expiration_ns: outcome.expiration_ns, + }); + let ours = matches!(&state.pending, Some(p) if p.session_id == body.state); + if ours { + if let Some(p) = state.pending.take() { + p.shutdown.notify_one(); + } + } + drop(state); + tracing::info!( + permissions = outcome.permissions, + expiration_ns = outcome.expiration_ns, + "signed in: registration delegation redeemed; login listener closing" + ); + done() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use imcp2_core::IiInstance; + + /// A driver against the real mainnet/prod-II config — nothing here touches + /// the network (agent construction, key minting, and the loopback listener + /// are all local), which is exactly what these tests rely on. + fn test_driver() -> (LoginDriver, SessionSlot) { + let agent = imcp2_core::Agent::builder() + .with_url(imcp2_core::IC_URL) + .build() + .expect("agent"); + let identities = Identities::new( + IiInstance::prod().expect("prod II instance"), + "https://mcp.internetcomputer.org".into(), + agent, + ); + let slot = SessionSlot::new(); + ( + LoginDriver::new(identities, slot.clone(), /* auto_open */ false), + slot, + ) + } + + /// A shape-valid two-hop chain (the same wire shape as iiconnect's parser + /// tests) whose signatures are garbage: it must get PAST parsing and be + /// rejected by the redeem path itself — offline, before any network call. + fn fake_chain_json() -> String { + serde_json::json!({ + "delegations": [ + { + "delegation": { "pubkey": "070707", "expiration": "66", "targets": ["aaaaa-aa"] }, + "signature": "040506", + }, + { + "delegation": { "pubkey": "09080706", "expiration": "66" }, + "signature": "010909", + }, + ], + "publicKey": "010203", + }) + .to_string() + } + + // One `authenticate` = one flow: a fresh id.ai link over a fresh loopback + // callback; a second call while it is pending returns the SAME link + // (single-flight) rather than racing a second listener. The session slot + // stays empty until a redeem succeeds — starting a login signs nobody in. + #[tokio::test] + async fn begin_mints_one_flow_and_repeats_it_while_pending() { + let (driver, slot) = test_driver(); + let BeginOutcome::Pending { url, fresh } = driver.begin(false).await.expect("begin") else { + panic!("a fresh driver must start a flow, not report a session") + }; + assert!(fresh); + assert!(url.starts_with("https://id.ai/mcp#callback="), "{url}"); + assert!( + url.contains("http%3A%2F%2F127.0.0.1%3A"), + "loopback callback in the fragment: {url}" + ); + assert!(url.contains("&ttl=3600®istration_key="), "{url}"); + + let BeginOutcome::Pending { url: again, fresh } = driver.begin(false).await.expect("begin") + else { + panic!("a pending flow must be returned, not replaced") + }; + assert!(!fresh, "the second call must join the pending flow"); + assert_eq!(again, url, "same link while the handshake is pending"); + + assert_eq!(slot.get(), None, "no session until a redeem succeeds"); + assert!(matches!(driver.status().await, LoginStatus::Pending { .. })); + } + + // The transient listener's whole surface, driven like the browser: the + // pinned page (strict CSP + the hosted server's hardening headers, script + // pointed at /redeem), and the #4091 allow-list (this callback verbatim, + // CORS-readable, never cached). + #[tokio::test] + async fn the_listener_serves_the_pinned_page_and_the_allow_list() { + let (driver, _slot) = test_driver(); + driver.begin(false).await.expect("begin"); + let (_, callback_url) = driver.pending_handshake().await.expect("pending"); + let origin = callback_url.strip_suffix("/callback").unwrap().to_string(); + let http = reqwest::Client::new(); + + let page = http.get(&callback_url).send().await.expect("GET /callback"); + assert_eq!(page.status(), 200); + let csp = page + .headers() + .get("content-security-policy") + .expect("pinned page ships a CSP") + .to_str() + .unwrap() + .to_string(); + assert!(csp.contains("default-src 'none'"), "{csp}"); + for (name, want) in [ + ("referrer-policy", "no-referrer"), + ("x-content-type-options", "nosniff"), + ("x-frame-options", "DENY"), + ] { + assert_eq!( + page.headers().get(name).and_then(|v| v.to_str().ok()), + Some(want) + ); + } + let html = page.text().await.unwrap(); + assert!( + html.contains("/redeem"), + "the script must POST to this listener's redeem" + ); + assert!( + html.contains("d.done"), + "the local success arm must be in the shipped page" + ); + + let wk = http + .get(format!("{origin}{AUTH_CALLBACKS_WELL_KNOWN}")) + .send() + .await + .expect("GET allow-list"); + assert_eq!(wk.status(), 200); + assert_eq!( + wk.headers() + .get("access-control-allow-origin") + .and_then(|v| v.to_str().ok()), + Some("*"), + "II fetches the allow-list cross-origin" + ); + assert_eq!( + wk.headers() + .get("cache-control") + .and_then(|v| v.to_str().ok()), + Some("no-store"), + "fail-closed infrastructure must never be served stale" + ); + let doc: serde_json::Value = wk.json().await.unwrap(); + assert_eq!( + doc, + serde_json::json!({ "callbacks": [callback_url] }), + "the declared callback must equal the link's callback VERBATIM (exact-match)" + ); + } + + // The redeem's refusal paths, in order: an unknown `state` (or one from a + // replaced flow), a delegation the parser rejects, and a shape-valid chain + // with garbage signatures that the redeem itself rejects — all 400 with an + // {"error"} body the pinned page renders, all offline, and none of them + // consumes the pending flow (a genuine retry stays possible). + #[tokio::test] + async fn redeem_refuses_bad_deliveries_and_keeps_the_flow_retryable() { + let (driver, slot) = test_driver(); + let BeginOutcome::Pending { url, .. } = driver.begin(false).await.expect("begin") else { + panic!("fresh flow") + }; + let (state, callback_url) = driver.pending_handshake().await.expect("pending"); + let origin = callback_url.strip_suffix("/callback").unwrap().to_string(); + let redeem_url = format!("{origin}/redeem"); + let http = reqwest::Client::new(); + + let post = |body: serde_json::Value| { + let http = http.clone(); + let redeem_url = redeem_url.clone(); + async move { + http.post(&redeem_url) + .json(&body) + .send() + .await + .expect("POST /redeem") + } + }; + + let r = post(serde_json::json!({ "state": "someone-else", "delegation": "" })).await; + assert_eq!(r.status(), 400); + let e: serde_json::Value = r.json().await.unwrap(); + assert!( + e["error"] + .as_str() + .unwrap() + .contains("unknown or already used"), + "{e}" + ); + + let r = post(serde_json::json!({ "state": state, "delegation": "not json" })).await; + assert_eq!(r.status(), 400); + let e: serde_json::Value = r.json().await.unwrap(); + assert!( + e["error"] + .as_str() + .unwrap() + .contains("couldn't read the sign-in response"), + "{e}" + ); + + let r = post(serde_json::json!({ "state": state, "delegation": fake_chain_json() })).await; + assert_eq!( + r.status(), + 400, + "garbage signatures must be rejected by the redeem" + ); + let e: serde_json::Value = r.json().await.unwrap(); + assert!(e["error"].is_string(), "{e}"); + + assert_eq!(slot.get(), None, "no failure may fill the session slot"); + let BeginOutcome::Pending { url: still, fresh } = driver.begin(false).await.expect("begin") + else { + panic!("the flow must survive failed redeems") + }; + assert!(!fresh, "failed redeems must not consume the pending flow"); + assert_eq!(still, url); + } + + // The status lifecycle around a grant (injected directly — a real one + // needs live II): live grant → SignedIn with the wallclock math; past + // expiration → Expired, still naming the principal; a fresh driver is + // simply SignedOut. + #[tokio::test] + async fn status_reports_the_grant_lifecycle() { + let (driver, _slot) = test_driver(); + assert!(matches!(driver.status().await, LoginStatus::SignedOut)); + + let grant = |expiration_ns: u64| Grant { + session_id: "sess".into(), + principal: Some("aaaaa-aa".into()), + permissions: "queries", + expiration_ns, + }; + driver.inner.state.lock().await.grant = Some(grant(now_ns() + 30 * 60_000_000_000)); + match driver.status().await { + LoginStatus::SignedIn(g) => { + assert!( + (25..=30).contains(&g.minutes_left()), + "{}", + g.minutes_left() + ); + } + _ => panic!("a live grant must report SignedIn"), + } + + driver.inner.state.lock().await.grant = Some(grant(now_ns() - 1)); + match driver.status().await { + LoginStatus::Expired(g) => assert_eq!(g.principal.as_deref(), Some("aaaaa-aa")), + _ => panic!("a past-expiry grant must report Expired"), + } + } + + // The handshake window is exactly HANDSHAKE_TTL: a pending flow older than + // that is expired (the watchdog and the redeem gate both read this). + #[test] + fn pending_expires_after_the_handshake_ttl() { + let pending = |age: Duration| Pending { + session_id: "s".into(), + url: "u".into(), + callback_url: "c".into(), + started: Instant::now() - age, + redeeming: false, + shutdown: Arc::new(Notify::new()), + }; + assert!(!pending(Duration::from_secs(1)).expired()); + assert!(pending(HANDSHAKE_TTL + Duration::from_secs(1)).expired()); + } +} diff --git a/crates/imcp2-local/src/main.rs b/crates/imcp2-local/src/main.rs new file mode 100644 index 0000000..9bb98a8 --- /dev/null +++ b/crates/imcp2-local/src/main.rs @@ -0,0 +1,161 @@ +//! `imcp2-local` — the Internet Computer MCP server for **local deployments**: +//! a single-user binary an AI tool (Claude Desktop/Code, Codex, Cursor, …) +//! spawns and talks to over **stdio**. It serves the same tool surface as the +//! hosted `imcp2` server against the same **mainnet** and **production +//! Internet Identity**, and drops the hosted server's entire OAuth 2.1 +//! authorization server: a single-user process reached over a pipe needs no +//! bearer tokens — the OS process boundary is the auth gate. +//! +//! Signing in is a built-in browser handshake (see `login`): the +//! `authenticate` tool returns an id.ai link, a transient loopback listener +//! receives II's callback, and the redeemed session lives in memory only. +//! +//! stdout is the JSON-RPC channel; ALL diagnostics go to stderr. +//! +//! Environment: +//! * `IMCP2_IC_URL` — IC API endpoint (default `https://icp-api.io`). +//! * `IMCP2_FETCH_ROOT_KEY` — truthy opts in to trusting the endpoint's +//! fetched root key, ONLY honoured when `IMCP2_IC_URL` targets loopback +//! (a local replica / PocketIC for integration tests); refused otherwise, +//! so a mis-set environment can never make a binary trust a fetched key +//! against mainnet. +//! * `II_URL_PROD` / `II_CANISTER_ID_PROD` — override production Internet +//! Identity (e.g. to point a test build at beta II or a PocketIC-deployed +//! II canister). +//! * `IMCP2_MANAGEMENT_ORIGIN` — the derivation origin of the canister- +//! management identity (default `https://mcp.internetcomputer.org`, the +//! hosted server's origin, so the same anchor gets the SAME controller/ +//! funder principal locally and hosted — canisters created through one +//! stay manageable through the other). +//! * `IMCP2_NO_OPEN` — truthy disables the best-effort browser auto-open on +//! sign-in (the link is always returned in-band regardless). + +mod login; +mod server; + +use imcp2_core::{identities::Identities, skills, IcTools, IiInstance, SessionSlot, SessionSource}; +use rmcp::ServiceExt; +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; + +fn truthy(var: &str) -> bool { + std::env::var(var) + .map(|v| { + matches!( + v.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) + }) + .unwrap_or(false) +} + +fn ic_url() -> String { + std::env::var("IMCP2_IC_URL").unwrap_or_else(|_| imcp2_core::IC_URL.to_string()) +} + +/// See the crate docs: defaults to the hosted server's origin so management +/// principals stay continuous across the two deployments. +fn management_origin() -> String { + std::env::var("IMCP2_MANAGEMENT_ORIGIN") + .unwrap_or_else(|_| "https://mcp.internetcomputer.org".to_string()) +} + +/// Whether `url` targets a loopback host — the only endpoints +/// `IMCP2_FETCH_ROOT_KEY` may be honoured for. +fn is_loopback_url(url: &str) -> bool { + match url::Url::parse(url) { + Ok(u) => match u.host() { + Some(url::Host::Ipv4(ip)) => ip.is_loopback(), + Some(url::Host::Ipv6(ip)) => ip.is_loopback(), + Some(url::Host::Domain(d)) => d.eq_ignore_ascii_case("localhost"), + None => false, + }, + Err(_) => false, + } +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + // stderr only: stdout carries the MCP JSON-RPC stream, and every client + // routes a stdio server's stderr to a log file/panel, never the chat. + tracing_subscriber::registry() + .with( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "info".to_string().into()), + ) + .with(tracing_subscriber::fmt::layer().with_writer(std::io::stderr)) + .init(); + + let ic_url = ic_url(); + let agent = imcp2_core::Agent::builder().with_url(&ic_url).build()?; + if truthy("IMCP2_FETCH_ROOT_KEY") { + // The guard, not a convenience: fetching a root key means TRUSTING the + // endpoint, which is only ever right for a local test replica. + anyhow::ensure!( + is_loopback_url(&ic_url), + "IMCP2_FETCH_ROOT_KEY is only honoured when IMCP2_IC_URL targets loopback \ + (a local replica); refusing to trust a fetched root key for {ic_url}" + ); + agent.fetch_root_key().await?; + tracing::info!( + "trusting the fetched root key of {ic_url} (local-replica test configuration)" + ); + } + tracing::info!("built ic-agent against {ic_url}"); + + let instance = IiInstance::prod().map_err(anyhow::Error::msg)?; + tracing::info!(ii = %instance.ii_url, "using Internet Identity instance \"{}\"", instance.name); + let identities = Identities::new(instance, management_origin(), agent.clone()); + + // The single-user session seam: the login flow fills (and on re-login + // replaces) the slot; every tool call reads it. + let slot = SessionSlot::new(); + let tools = IcTools::new( + agent, + identities.clone(), + skills::SkillsCatalog::new(), + SessionSource::Singleton(slot.clone()), + ); + let auto_open = !truthy("IMCP2_NO_OPEN"); + let login = login::LoginDriver::new(identities, slot, auto_open); + + let service = server::LocalServer::new(tools, login, auto_open) + .serve(rmcp::transport::stdio()) + .await?; + tracing::info!("imcp2-local serving MCP over stdio"); + // Runs until the client closes the pipe (or the service errors). + service.waiting().await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + // The root-key guard must accept every loopback shape a test harness + // reasonably uses and refuse everything else — mainnet above all. A truthy + // IMCP2_FETCH_ROOT_KEY against a non-loopback endpoint aborts startup + // rather than degrading to "just don't fetch". + #[test] + fn fetch_root_key_guard_accepts_only_loopback() { + for ok in [ + "http://127.0.0.1:4943", + "http://localhost:4943", + "http://[::1]:4943", + "https://127.0.0.1:8080/api", + ] { + assert!(super::is_loopback_url(ok), "{ok} should count as loopback"); + } + for bad in [ + "https://icp-api.io", + "https://ic0.app", + // RFC 5737 TEST-NET-1: a routable-looking address that is NOT + // loopback — private/LAN replicas must be refused too. + "http://192.0.2.1:4943", + "http://replica.internal:4943", + "not a url", + ] { + assert!( + !super::is_loopback_url(bad), + "{bad} must not count as loopback" + ); + } + } +} diff --git a/crates/imcp2-local/src/server.rs b/crates/imcp2-local/src/server.rs new file mode 100644 index 0000000..46eb7ed --- /dev/null +++ b/crates/imcp2-local/src/server.rs @@ -0,0 +1,372 @@ +//! The local MCP server: [`IcTools`] wrapped in a thin handler that adds the +//! two **local-only login tools**. They cannot live in `imcp2-core` — they +//! drive this binary's loopback listener and browser opener, which core cannot +//! call back into — so this wrapper carries them in its own `#[tool_router]`, +//! answers `tools/list` with the merged list, dispatches `tools/call` to its +//! own router first, and forwards everything else to the inner [`IcTools`] +//! handler verbatim. The hosted server contains none of this code, so its tool +//! surface stays login-free by construction (a core-side test guards the +//! boundary too). + +use rmcp::handler::server::tool::ToolCallContext; +use rmcp::handler::server::wrapper::Parameters; +use rmcp::model::{ + CallToolRequestParams, CallToolResult, Content, ListResourcesResult, ListToolsResult, + PaginatedRequestParams, ReadResourceRequestParams, ReadResourceResult, Tool, +}; +use rmcp::service::RequestContext; +use rmcp::{schemars, tool, tool_router, ErrorData as McpError, RoleServer, ServerHandler}; + +use crate::login::{BeginOutcome, Grant, LoginDriver, LoginStatus}; +use imcp2_core::IcTools; + +pub struct LocalServer { + tools: IcTools, + login: LoginDriver, + /// Whether the login flow will also try to open the browser itself — + /// only phrasing: the link is always returned in-band. + auto_open: bool, +} + +/// Arguments of the `authenticate` tool. +#[derive(Debug, serde::Deserialize, schemars::JsonSchema)] +pub struct AuthenticateArgs { + /// Start a fresh sign-in even when a session is already active — to switch + /// accounts, raise the access level, or extend the grant. Default: an + /// active session is reported instead of starting over. + #[serde(default)] + pub refresh: bool, +} + +fn signed_in_line(g: &Grant) -> String { + let who = match &g.principal { + Some(p) => format!(" as {p}"), + None => String::new(), + }; + let access = match g.permissions { + "queries" => "read-only", + other => other, + }; + format!( + "Signed in{who} (access: {access}; the grant expires in about {}).", + human_minutes(g.minutes_left()) + ) +} + +/// "45 min" / "3 hours" / "12 days" — the II consent screen offers session +/// lengths from minutes to a month, so raw minutes would read like a counter. +fn human_minutes(mins: u64) -> String { + if mins >= 48 * 60 { + format!("{} days", mins / (24 * 60)) + } else if mins >= 120 { + format!("{} hours", mins / 60) + } else { + format!("{mins} min") + } +} + +#[tool_router] +impl LocalServer { + pub fn new(tools: IcTools, login: LoginDriver, auto_open: bool) -> Self { + Self { + tools, + login, + auto_open, + } + } + + #[tool( + description = "Sign in with Internet Identity. Returns an id.ai sign-in link for the USER to open in their browser (the server also tries to open it for them), and returns immediately — it never waits for the browser. After the user finishes, auth_status (or simply retrying the tool that needed a session) confirms the session. Call this when a tool answers that it needs an authenticated session. Repeat calls while a sign-in is pending return the same link.", + annotations( + title = "Sign in with Internet Identity", + read_only_hint = false, + destructive_hint = false, + open_world_hint = true + ) + )] + async fn authenticate( + &self, + Parameters(AuthenticateArgs { refresh }): Parameters, + ) -> Result { + let text = match self.login.begin(refresh).await { + Err(e) => { + return Ok(CallToolResult::error(vec![Content::text(format!( + "Could not start the sign-in: {e}" + ))])) + } + Ok(BeginOutcome::AlreadySignedIn(g)) => format!( + "{} Pass refresh=true to sign in again (switch accounts or extend the session).", + signed_in_line(&g) + ), + Ok(BeginOutcome::Pending { url, fresh }) => { + let opener = if self.auto_open { + "I also asked the system to open it in your browser. " + } else { + "" + }; + let lead = if fresh { + "To sign in, open this Internet Identity link in your browser" + } else { + "A sign-in is already waiting for you — open this Internet Identity link in your browser" + }; + format!( + "{lead}:\n\n{url}\n\n{opener}The link is valid for 10 minutes. \ + After you finish in the browser, call auth_status — or just retry \ + the tool you wanted — to confirm the session." + ) + } + }; + Ok(CallToolResult::success(vec![Content::text(text)])) + } + + #[tool( + description = "Report the Internet Identity sign-in state of this local server: signed in (with the session principal, access level, and time to expiry), sign-in pending (with the link to finish it), expired, or signed out.", + annotations( + title = "Check sign-in status", + read_only_hint = true, + destructive_hint = false, + open_world_hint = false + ) + )] + async fn auth_status(&self) -> Result { + let text = match self.login.status().await { + LoginStatus::SignedIn(g) => signed_in_line(&g), + LoginStatus::Pending { url } => format!( + "A sign-in is pending — waiting for the browser handshake to finish at:\n\n{url}" + ), + LoginStatus::Expired(g) => { + let who = match &g.principal { + Some(p) => format!(" (was {p})"), + None => String::new(), + }; + format!("The session expired{who}. Call authenticate to sign in again.") + } + LoginStatus::SignedOut => { + "Not signed in. Call authenticate to get a sign-in link.".to_string() + } + }; + Ok(CallToolResult::success(vec![Content::text(text)])) + } +} + +impl ServerHandler for LocalServer { + fn get_info(&self) -> rmcp::model::ServerInfo { + let mut info = self.tools.get_info(); + info.server_info.name = "imcp2-local".into(); + info.server_info.version = env!("CARGO_PKG_VERSION").into(); + // The core instructions apply verbatim; add how signing in works here + // (the hosted server logs in via OAuth outside the tool surface, so + // core's text says nothing about it). + let core = info.instructions.take().unwrap_or_default(); + info.instructions = Some(format!( + "{core}\n\n\ + SIGNING IN (this local server). Tool calls act as the USER's Internet Identity. \ + When a tool answers that it needs an authenticated session, call `authenticate`: \ + it returns an id.ai sign-in link (and best-effort opens the browser) and never \ + blocks — after the user finishes in the browser, `auth_status` or simply retrying \ + the original tool confirms. Sessions are in-memory: the user signs in again after \ + a restart, when the grant expires, or if they revoke it at id.ai." + )); + info + } + + async fn list_tools( + &self, + request: Option, + context: RequestContext, + ) -> Result { + let mut result = self.tools.list_tools(request, context).await?; + result.tools.extend(Self::tool_router().list_all()); + Ok(result) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + if Self::tool_router().has_route(&request.name) { + let tcc = ToolCallContext::new(self, request, context); + Self::tool_router().call(tcc).await + } else { + self.tools.call_tool(request, context).await + } + } + + fn get_tool(&self, name: &str) -> Option { + Self::tool_router() + .get(name) + .cloned() + .or_else(|| self.tools.get_tool(name)) + } + + async fn list_resources( + &self, + request: Option, + context: RequestContext, + ) -> Result { + self.tools.list_resources(request, context).await + } + + async fn read_resource( + &self, + request: ReadResourceRequestParams, + context: RequestContext, + ) -> Result { + self.tools.read_resource(request, context).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use imcp2_core::{identities::Identities, skills, IiInstance, SessionSlot, SessionSource}; + use rmcp::ServiceExt; + + fn test_server() -> LocalServer { + let agent = imcp2_core::Agent::builder() + .with_url(imcp2_core::IC_URL) + .build() + .expect("agent"); + let identities = Identities::new( + IiInstance::prod().expect("prod II"), + "https://mcp.internetcomputer.org".into(), + agent.clone(), + ); + let slot = SessionSlot::new(); + let tools = IcTools::new( + agent, + identities.clone(), + skills::SkillsCatalog::new(), + SessionSource::Singleton(slot.clone()), + ); + let login = crate::login::LoginDriver::new(identities, slot, /* auto_open */ false); + LocalServer::new(tools, login, /* auto_open */ false) + } + + // The wrapper's own router carries EXACTLY the two login tools — the rest + // of the surface comes from IcTools by forwarding, never by duplication — + // and both declare honest annotations (auth_status is the read-only one). + #[test] + fn the_wrapper_router_carries_exactly_the_login_tools() { + let tools = LocalServer::tool_router().list_all(); + let mut names: Vec<&str> = tools.iter().map(|t| &*t.name).collect(); + names.sort_unstable(); + assert_eq!(names, ["auth_status", "authenticate"]); + let ann = |n: &str| { + tools + .iter() + .find(|t| &*t.name == n) + .and_then(|t| t.annotations.clone()) + .unwrap_or_else(|| panic!("{n} must carry annotations")) + }; + assert_eq!(ann("authenticate").read_only_hint, Some(false)); + assert_eq!(ann("authenticate").destructive_hint, Some(false)); + assert_eq!(ann("auth_status").read_only_hint, Some(true)); + let auth = tools.iter().find(|t| &*t.name == "authenticate").unwrap(); + let schema = serde_json::to_value(&auth.input_schema).unwrap(); + assert!( + schema["properties"]["refresh"].is_object(), + "authenticate must declare its refresh flag: {schema}" + ); + } + + // Expiry phrasing across the II consent screen's range (minutes → a month): + // never a raw five-digit minute counter. + #[test] + fn expiry_is_reported_in_human_units() { + assert_eq!(super::human_minutes(45), "45 min"); + assert_eq!(super::human_minutes(180), "3 hours"); + assert_eq!(super::human_minutes(30 * 24 * 60), "30 days"); + } + + // A REAL MCP round-trip over an in-process duplex pipe (the same + // byte-stream shape stdio serves): initialize; tools/list is the merged + // surface (all 26 core tools + the 2 login tools); tools/call dispatches + // login tools to the wrapper and everything else to IcTools; the login + // lifecycle (signed out → link → pending) runs through the MCP layer — + // all with no network and no browser. + #[tokio::test] + async fn a_real_mcp_client_sees_the_merged_surface_and_the_login_flow() { + let (client_io, server_io) = tokio::io::duplex(64 * 1024); + let server = test_server().serve(server_io); + let client = ().serve(client_io); + let (server, client) = tokio::join!(server, client); + let server = server.expect("server up"); + let client = client.expect("client up"); + + // The server identifies as imcp2-local and teaches the login flow in + // its instructions, on top of core's own text. + let info = client.peer_info().expect("initialized"); + assert_eq!(info.server_info.name, "imcp2-local"); + let instructions = info.instructions.as_deref().unwrap_or_default(); + assert!( + instructions.contains("SIGNING IN"), + "login guidance must be taught" + ); + assert!( + instructions.contains("textual Candid"), + "core guidance must survive the merge" + ); + + let tools = client.list_all_tools().await.expect("tools/list"); + assert_eq!( + tools.len(), + 28, + "26 core tools + authenticate + auth_status" + ); + for expected in [ + "get_canister_candid", + "canister_query", + "authenticate", + "auth_status", + ] { + assert!( + tools.iter().any(|t| &*t.name == expected), + "missing {expected}" + ); + } + + let call = |name: &'static str, args: Option| { + let client = &client; + async move { + let mut params = rmcp::model::CallToolRequestParams::new(name); + params.arguments = args.map(|v| v.as_object().expect("object args").clone()); + let result = client.call_tool(params).await.expect("tools/call"); + let text = result + .content + .iter() + .flat_map(|c| c.as_text().map(|t| t.text.clone())) + .collect::>() + .join("\n"); + (result.is_error.unwrap_or(false), text) + } + }; + + // Login lifecycle through the MCP layer. + let (is_error, text) = call("auth_status", None).await; + assert!(!is_error, "{text}"); + assert!(text.contains("Not signed in"), "{text}"); + + let (is_error, text) = call("authenticate", Some(serde_json::json!({}))).await; + assert!(!is_error, "{text}"); + assert!(text.contains("https://id.ai/mcp#callback="), "{text}"); + + let (is_error, text) = call("auth_status", None).await; + assert!(!is_error, "{text}"); + assert!(text.contains("sign-in is pending"), "{text}"); + + // Forwarding: a core tool answers through the wrapper (its own error + // path, so no network is touched). + let (is_error, text) = call( + "get_canister_candid", + Some(serde_json::json!({ "canister_id": "not-a-canister" })), + ) + .await; + assert!(is_error, "an invalid canister id is a tool error"); + assert!(text.contains("invalid canister id"), "{text}"); + + client.cancel().await.expect("client shutdown"); + server.cancel().await.expect("server shutdown"); + } +} From f2128b4a927afa021e7d20ee1f23d31f764ac561 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 25 Aug 2026 14:43:07 +0000 Subject: [PATCH 02/10] Address review: structured login-tool outputs, status precedence, stale-redeem guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - authenticate/auth_status now declare the object-rooted outputSchema and attach matching structuredContent (one AuthOutput shape for the whole sign-in lifecycle), preserving the all-tools MCP contract the core surface enforces. - auth_status reports a live pending handshake BEFORE a live grant, matching begin's precedence: after authenticate(refresh=true) the truthful answer is "a replacement sign-in is waiting", not "signed in". - The redeem's commit is no longer last-writer-wins: each flow carries an epoch, and a stale redeem (claimed, then outlived by the watchdog while its mcp_register_v2 call ran) is refused when a NEWER pending or completed sign-in exists — while a late success with nothing newer around still lands. Tests: 220 pass workspace-wide (2 new: pending-outranks-grant status, stale-redeem supersession) with structured-content assertions added to the MCP round-trip and the outputSchema check extended to the merged surface. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- crates/imcp2-local/src/login.rs | 172 ++++++++++++++++++++------ crates/imcp2-local/src/server.rs | 204 ++++++++++++++++++++++++++----- 2 files changed, 310 insertions(+), 66 deletions(-) diff --git a/crates/imcp2-local/src/login.rs b/crates/imcp2-local/src/login.rs index 8851ad5..6bc0104 100644 --- a/crates/imcp2-local/src/login.rs +++ b/crates/imcp2-local/src/login.rs @@ -77,6 +77,19 @@ struct LoginState { /// The most recent successful login, kept (even past expiry) so /// `auth_status` can say *whose* session expired. grant: Option, + /// Flow counter: each `begin`-started flow gets the next number, so a + /// redeem that raced a replacement can tell whether something NEWER + /// exists (see [`superseded_by_newer`]). + epoch: u64, +} + +/// Whether a flow newer than `epoch` exists — a pending handshake or a +/// completed grant started after the flow that is asking. A stale redeem +/// (claimed, then outlived by the watchdog while its `mcp_register_v2` call +/// was in flight) must not overwrite a newer sign-in. +fn superseded_by_newer(state: &LoginState, epoch: u64) -> bool { + state.pending.as_ref().is_some_and(|p| p.epoch > epoch) + || state.grant.as_ref().is_some_and(|g| g.epoch > epoch) } struct Pending { @@ -95,6 +108,8 @@ struct Pending { /// Tears the transient listener down (graceful): notified on redeem /// success, on handshake timeout, and when a fresh flow replaces this one. shutdown: Arc, + /// This flow's number (see [`LoginState::epoch`]). + epoch: u64, } impl Pending { @@ -115,6 +130,9 @@ pub struct Grant { /// Grant expiration, ns since the Unix epoch (the user's consent-time /// session choice). pub expiration_ns: u64, + /// The number of the flow that produced this grant (see + /// [`LoginState::epoch`]). + epoch: u64, } impl Grant { @@ -251,6 +269,7 @@ impl LoginDriver { } }); + state.epoch += 1; state.pending = Some(Pending { session_id, url: url.clone(), @@ -258,6 +277,7 @@ impl LoginDriver { started: Instant::now(), redeeming: false, shutdown, + epoch: state.epoch, }); drop(state); @@ -289,16 +309,21 @@ impl LoginDriver { pub async fn status(&self) -> LoginStatus { let state = self.inner.state.lock().await; - if let Some(g) = &state.grant { - if !g.expired() { - return LoginStatus::SignedIn(g.clone()); - } - } + // A live pending handshake outranks a live grant, matching `begin`'s + // precedence: after `authenticate(refresh: true)` the old session may + // still be valid, but the truthful answer is "a replacement sign-in is + // waiting for the browser", not "signed in" — otherwise the client + // reads the refresh as already complete. if let Some(p) = &state.pending { if !p.expired() { return LoginStatus::Pending { url: p.url.clone() }; } } + if let Some(g) = &state.grant { + if !g.expired() { + return LoginStatus::SignedIn(g.clone()); + } + } match &state.grant { Some(g) => LoginStatus::Expired(g.clone()), None => LoginStatus::SignedOut, @@ -435,38 +460,39 @@ async fn redeem(State(ctx): State, Json(body): Json) -> Re } }; - // Claim atomically (re-checking: the flow may have moved while parsing). - { - let mut state = driver.inner.state.lock().await; - if state - .grant - .as_ref() - .is_some_and(|g| g.session_id == body.state) + // Claim atomically (re-checking: the flow may have moved while parsing), + // remembering the claimed flow's number for the commit-time check below. + let epoch = { - return done(); - } - match state.pending.as_mut() { - Some(p) if p.session_id == body.state => { - if p.expired() { - return redeem_err( - "This sign-in expired. Ask the agent to authenticate again.", - ); - } - if p.redeeming { - return redeem_err( - "This sign-in is already being processed. Wait a moment; \ + let mut state = driver.inner.state.lock().await; + if state + .grant + .as_ref() + .is_some_and(|g| g.session_id == body.state) + { + return done(); + } + match state.pending.as_mut() { + Some(p) if p.session_id == body.state => { + if p.expired() { + return redeem_err( + "This sign-in expired. Ask the agent to authenticate again.", + ); + } + if p.redeeming { + return redeem_err( + "This sign-in is already being processed. Wait a moment; \ if nothing happens, ask the agent to authenticate again.", - ); + ); + } + p.redeeming = true; + p.epoch } - p.redeeming = true; - } - _ => { - return redeem_err( + _ => return redeem_err( "This sign-in is unknown or already used. Ask the agent to authenticate again.", - ) + ), } - } - } + }; // The one network call, outside the lock. match driver @@ -487,15 +513,26 @@ async fn redeem(State(ctx): State, Json(body): Json) -> Re Ok(outcome) => { let principal = driver.inner.identities.session_principal(&body.state).await; let mut state = driver.inner.state.lock().await; - // Fill the shared slot LAST-WRITER-WINS: II bound the grant to this - // session's key, so even if a newer flow replaced the pending entry - // mid-call, this successful login is real and becomes current. + // Commit-time check: this flow was claimed before the network + // call, but the watchdog may have expired it and a NEWER sign-in + // (pending or already completed) may exist by now — a stale + // success must not overwrite that. A late success with nothing + // newer around still lands (II did bind the grant; the watchdog + // alone is no reason to make the user redo it). + if superseded_by_newer(&state, epoch) { + drop(state); + return redeem_err( + "This sign-in was superseded by a newer one. Use the most recent \ + sign-in tab, or ask the agent to authenticate again.", + ); + } driver.inner.slot.set(body.state.clone()); state.grant = Some(Grant { session_id: body.state.clone(), principal, permissions: outcome.permissions, expiration_ns: outcome.expiration_ns, + epoch, }); let ours = matches!(&state.pending, Some(p) if p.session_id == body.state); if ours { @@ -740,6 +777,7 @@ mod tests { principal: Some("aaaaa-aa".into()), permissions: "queries", expiration_ns, + epoch: 1, }; driver.inner.state.lock().await.grant = Some(grant(now_ns() + 30 * 60_000_000_000)); match driver.status().await { @@ -760,6 +798,69 @@ mod tests { } } + // `authenticate(refresh: true)` while a live grant exists: the truthful + // status is the PENDING replacement handshake, not the old "signed in" — + // otherwise the client reads the refresh as already complete (the same + // precedence `begin` applies). + #[tokio::test] + async fn a_pending_refresh_outranks_the_live_grant_in_status() { + let (driver, _slot) = test_driver(); + driver.inner.state.lock().await.grant = Some(Grant { + session_id: "old".into(), + principal: Some("aaaaa-aa".into()), + permissions: "all", + expiration_ns: now_ns() + 30 * 60_000_000_000, + epoch: 1, + }); + assert!(matches!(driver.status().await, LoginStatus::SignedIn(_))); + + let BeginOutcome::Pending { fresh, .. } = driver.begin(true).await.expect("refresh") else { + panic!("refresh=true must start a replacement flow past a live grant") + }; + assert!(fresh); + assert!( + matches!(driver.status().await, LoginStatus::Pending { .. }), + "a live replacement handshake must outrank the old grant" + ); + } + + // The stale-redeem guard: a flow that was claimed, then outlived by the + // watchdog while its network call ran, must not overwrite a NEWER sign-in + // (pending or completed) at commit time — but a late success with nothing + // newer around still lands. + #[test] + fn a_stale_redeem_is_superseded_by_any_newer_flow() { + let pending = |epoch: u64| Pending { + session_id: "s".into(), + url: "u".into(), + callback_url: "c".into(), + started: Instant::now(), + redeeming: false, + shutdown: Arc::new(Notify::new()), + epoch, + }; + let grant = |epoch: u64| Grant { + session_id: "g".into(), + principal: None, + permissions: "queries", + expiration_ns: u64::MAX, + epoch, + }; + let state = |p: Option, g: Option| LoginState { + pending: p, + grant: g, + epoch: 9, + }; + // Nothing newer: the late success lands (watchdog expiry alone is no + // reason to make the user redo a login II already bound). + assert!(!superseded_by_newer(&state(None, None), 3)); + assert!(!superseded_by_newer(&state(Some(pending(3)), None), 3)); + assert!(!superseded_by_newer(&state(None, Some(grant(2))), 3)); + // A newer pending handshake or completed grant wins over the stale one. + assert!(superseded_by_newer(&state(Some(pending(4)), None), 3)); + assert!(superseded_by_newer(&state(None, Some(grant(4))), 3)); + } + // The handshake window is exactly HANDSHAKE_TTL: a pending flow older than // that is expired (the watchdog and the redeem gate both read this). #[test] @@ -771,6 +872,7 @@ mod tests { started: Instant::now() - age, redeeming: false, shutdown: Arc::new(Notify::new()), + epoch: 1, }; assert!(!pending(Duration::from_secs(1)).expired()); assert!(pending(HANDSHAKE_TTL + Duration::from_secs(1)).expired()); diff --git a/crates/imcp2-local/src/server.rs b/crates/imcp2-local/src/server.rs index 46eb7ed..c242870 100644 --- a/crates/imcp2-local/src/server.rs +++ b/crates/imcp2-local/src/server.rs @@ -38,17 +38,114 @@ pub struct AuthenticateArgs { pub refresh: bool, } +/// Structured output of `authenticate` and `auth_status` — one shape for the +/// whole sign-in lifecycle, so both tools honour the same all-tools +/// `outputSchema` contract as the core surface (every core tool declares an +/// object-rooted schema and attaches matching `structuredContent`). +#[derive(Debug, serde::Serialize, schemars::JsonSchema)] +pub struct AuthOutput { + /// The sign-in state: `signed_in`, `pending` (a browser handshake is + /// waiting to be finished), `expired`, or `signed_out`. + pub status: String, + /// The id.ai sign-in link to open, present while a handshake is pending. + #[serde(skip_serializing_if = "Option::is_none")] + pub url: Option, + /// The session principal, when signed in (or the principal whose session + /// expired). + #[serde(skip_serializing_if = "Option::is_none")] + pub principal: Option, + /// The grant's access level, when signed in: `read-only` or `all`. + #[serde(skip_serializing_if = "Option::is_none")] + pub access: Option, + /// Whole minutes until the grant expires, when signed in. + #[serde(skip_serializing_if = "Option::is_none")] + pub expires_in_minutes: Option, +} + +impl AuthOutput { + fn signed_in(g: &Grant) -> Self { + Self { + status: "signed_in".into(), + url: None, + principal: g.principal.clone(), + access: Some(access_word(g.permissions).into()), + expires_in_minutes: Some(g.minutes_left()), + } + } + + fn pending(url: String) -> Self { + Self { + status: "pending".into(), + url: Some(url), + principal: None, + access: None, + expires_in_minutes: None, + } + } + + fn expired(g: &Grant) -> Self { + Self { + status: "expired".into(), + url: None, + principal: g.principal.clone(), + access: None, + expires_in_minutes: None, + } + } + + fn signed_out() -> Self { + Self { + status: "signed_out".into(), + url: None, + principal: None, + access: None, + expires_in_minutes: None, + } + } +} + +/// The grant vocabulary (`"queries"`) rendered for humans and the structured +/// `access` field alike. +fn access_word(permissions: &str) -> &str { + match permissions { + "queries" => "read-only", + other => other, + } +} + +/// Same contract as the core surface's result builder: the human-readable +/// text plus `structuredContent` conforming to the declared `outputSchema` +/// (attached only when it serializes to a JSON object, which [`AuthOutput`] +/// always does). +fn ok_structured(text: String, value: &T) -> CallToolResult { + let mut result = CallToolResult::success(vec![Content::text(text)]); + result.structured_content = match serde_json::to_value(value) { + Ok(v @ serde_json::Value::Object(_)) => Some(v), + _ => None, + }; + result +} + +/// Same panic-at-router-construction posture as the core surface: a +/// non-object output schema is a programming error. +fn schema_for_output( +) -> std::sync::Arc { + rmcp::handler::server::tool::schema_for_output::().unwrap_or_else(|e| { + panic!( + "output schema for `{}` must be object-rooted: {e}", + std::any::type_name::() + ) + }) +} + fn signed_in_line(g: &Grant) -> String { let who = match &g.principal { Some(p) => format!(" as {p}"), None => String::new(), }; - let access = match g.permissions { - "queries" => "read-only", - other => other, - }; format!( - "Signed in{who} (access: {access}; the grant expires in about {}).", + "Signed in{who} (access: {}; the grant expires in about {}).", + access_word(g.permissions), human_minutes(g.minutes_left()) ) } @@ -82,21 +179,25 @@ impl LocalServer { read_only_hint = false, destructive_hint = false, open_world_hint = true - ) + ), + output_schema = schema_for_output::(), )] async fn authenticate( &self, Parameters(AuthenticateArgs { refresh }): Parameters, ) -> Result { - let text = match self.login.begin(refresh).await { + let (text, output) = match self.login.begin(refresh).await { Err(e) => { return Ok(CallToolResult::error(vec![Content::text(format!( "Could not start the sign-in: {e}" ))])) } - Ok(BeginOutcome::AlreadySignedIn(g)) => format!( - "{} Pass refresh=true to sign in again (switch accounts or extend the session).", - signed_in_line(&g) + Ok(BeginOutcome::AlreadySignedIn(g)) => ( + format!( + "{} Pass refresh=true to sign in again (switch accounts or extend the session).", + signed_in_line(&g) + ), + AuthOutput::signed_in(&g), ), Ok(BeginOutcome::Pending { url, fresh }) => { let opener = if self.auto_open { @@ -109,14 +210,17 @@ impl LocalServer { } else { "A sign-in is already waiting for you — open this Internet Identity link in your browser" }; - format!( - "{lead}:\n\n{url}\n\n{opener}The link is valid for 10 minutes. \ - After you finish in the browser, call auth_status — or just retry \ - the tool you wanted — to confirm the session." + ( + format!( + "{lead}:\n\n{url}\n\n{opener}The link is valid for 10 minutes. \ + After you finish in the browser, call auth_status — or just retry \ + the tool you wanted — to confirm the session." + ), + AuthOutput::pending(url), ) } }; - Ok(CallToolResult::success(vec![Content::text(text)])) + Ok(ok_structured(text, &output)) } #[tool( @@ -126,26 +230,34 @@ impl LocalServer { read_only_hint = true, destructive_hint = false, open_world_hint = false - ) + ), + output_schema = schema_for_output::(), )] async fn auth_status(&self) -> Result { - let text = match self.login.status().await { - LoginStatus::SignedIn(g) => signed_in_line(&g), - LoginStatus::Pending { url } => format!( - "A sign-in is pending — waiting for the browser handshake to finish at:\n\n{url}" + let (text, output) = match self.login.status().await { + LoginStatus::SignedIn(g) => (signed_in_line(&g), AuthOutput::signed_in(&g)), + LoginStatus::Pending { url } => ( + format!( + "A sign-in is pending — waiting for the browser handshake to finish at:\n\n{url}" + ), + AuthOutput::pending(url), ), LoginStatus::Expired(g) => { let who = match &g.principal { Some(p) => format!(" (was {p})"), None => String::new(), }; - format!("The session expired{who}. Call authenticate to sign in again.") - } - LoginStatus::SignedOut => { - "Not signed in. Call authenticate to get a sign-in link.".to_string() + ( + format!("The session expired{who}. Call authenticate to sign in again."), + AuthOutput::expired(&g), + ) } + LoginStatus::SignedOut => ( + "Not signed in. Call authenticate to get a sign-in link.".to_string(), + AuthOutput::signed_out(), + ), }; - Ok(CallToolResult::success(vec![Content::text(text)])) + Ok(ok_structured(text, &output)) } } @@ -269,6 +381,21 @@ mod tests { schema["properties"]["refresh"].is_object(), "authenticate must declare its refresh flag: {schema}" ); + // The merged surface keeps the core contract: every tool declares an + // object-rooted outputSchema (the core suite enforces it for the 26; + // this enforces it for the two added here). + for t in &tools { + let schema = t + .output_schema + .as_ref() + .unwrap_or_else(|| panic!("{} must declare an outputSchema", t.name)); + assert_eq!( + serde_json::to_value(schema).unwrap()["type"], + "object", + "{}'s outputSchema must be object-rooted", + t.name + ); + } } // Expiry phrasing across the II consent screen's range (minutes → a month): @@ -339,26 +466,41 @@ mod tests { .flat_map(|c| c.as_text().map(|t| t.text.clone())) .collect::>() .join("\n"); - (result.is_error.unwrap_or(false), text) + let structured = result.structured_content.clone(); + (result.is_error.unwrap_or(false), text, structured) } }; - // Login lifecycle through the MCP layer. - let (is_error, text) = call("auth_status", None).await; + // Login lifecycle through the MCP layer — text for humans, structured + // content (per the declared outputSchema) for typed consumers. + let (is_error, text, structured) = call("auth_status", None).await; assert!(!is_error, "{text}"); assert!(text.contains("Not signed in"), "{text}"); + assert_eq!( + structured.expect("structuredContent")["status"], + "signed_out" + ); - let (is_error, text) = call("authenticate", Some(serde_json::json!({}))).await; + let (is_error, text, structured) = call("authenticate", Some(serde_json::json!({}))).await; assert!(!is_error, "{text}"); assert!(text.contains("https://id.ai/mcp#callback="), "{text}"); + let structured = structured.expect("structuredContent"); + assert_eq!(structured["status"], "pending"); + assert!( + structured["url"] + .as_str() + .is_some_and(|u| u.starts_with("https://id.ai/mcp#")), + "{structured}" + ); - let (is_error, text) = call("auth_status", None).await; + let (is_error, text, structured) = call("auth_status", None).await; assert!(!is_error, "{text}"); assert!(text.contains("sign-in is pending"), "{text}"); + assert_eq!(structured.expect("structuredContent")["status"], "pending"); // Forwarding: a core tool answers through the wrapper (its own error // path, so no network is touched). - let (is_error, text) = call( + let (is_error, text, _) = call( "get_canister_candid", Some(serde_json::json!({ "canister_id": "not-a-canister" })), ) From 1b410bedd19f68f80dff8a0aae083d01b24b8e83 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 08:25:20 +0000 Subject: [PATCH 03/10] Address review: serialize registrations instead of guarding the commit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The epoch guard fixed only the local half of the stale-redeem race: II's mcp_register_v2 REPLACES the anchor's previous grant, so a timed-out flow's slow registration landing after a newer flow's left II honoring the old session while the guarded slot kept naming the new one — every subsequent tool call would come back Unauthorized. A local commit-time check cannot undo the remote side effect. Replace the guard with what the race actually needs: a registration mutex that serializes each redeem's mcp_register_v2 AND its local commit, across flows. The last registration to settle at II is the live grant, and under the lock the commit order equals the settle order, so the slot always names the session II actually honors — in both interleavings of an old in-flight registration and its replacement. The per-flow single-flight claim stays (it stops page double-submits); the epoch machinery goes. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- crates/imcp2-local/src/login.rs | 152 ++++++++++---------------------- 1 file changed, 49 insertions(+), 103 deletions(-) diff --git a/crates/imcp2-local/src/login.rs b/crates/imcp2-local/src/login.rs index 6bc0104..e1de004 100644 --- a/crates/imcp2-local/src/login.rs +++ b/crates/imcp2-local/src/login.rs @@ -69,6 +69,14 @@ struct Inner { /// it succeeding — the link is always returned in-band by `authenticate`. auto_open: bool, state: Mutex, + /// Serializes each redeem's `mcp_register_v2` **and** its local commit, + /// across flows. II's register REPLACES the anchor's previous grant (the + /// last registration to settle at II is the live one), so the local slot + /// must be committed in the same order the registrations settle: without + /// this, a timed-out flow's slow register could land at II *after* a + /// newer flow's — II would then honor the old session while the slot + /// named the new one, and every tool call would come back `Unauthorized`. + registration: Mutex<()>, } #[derive(Default)] @@ -77,19 +85,6 @@ struct LoginState { /// The most recent successful login, kept (even past expiry) so /// `auth_status` can say *whose* session expired. grant: Option, - /// Flow counter: each `begin`-started flow gets the next number, so a - /// redeem that raced a replacement can tell whether something NEWER - /// exists (see [`superseded_by_newer`]). - epoch: u64, -} - -/// Whether a flow newer than `epoch` exists — a pending handshake or a -/// completed grant started after the flow that is asking. A stale redeem -/// (claimed, then outlived by the watchdog while its `mcp_register_v2` call -/// was in flight) must not overwrite a newer sign-in. -fn superseded_by_newer(state: &LoginState, epoch: u64) -> bool { - state.pending.as_ref().is_some_and(|p| p.epoch > epoch) - || state.grant.as_ref().is_some_and(|g| g.epoch > epoch) } struct Pending { @@ -108,8 +103,6 @@ struct Pending { /// Tears the transient listener down (graceful): notified on redeem /// success, on handshake timeout, and when a fresh flow replaces this one. shutdown: Arc, - /// This flow's number (see [`LoginState::epoch`]). - epoch: u64, } impl Pending { @@ -130,9 +123,6 @@ pub struct Grant { /// Grant expiration, ns since the Unix epoch (the user's consent-time /// session choice). pub expiration_ns: u64, - /// The number of the flow that produced this grant (see - /// [`LoginState::epoch`]). - epoch: u64, } impl Grant { @@ -176,6 +166,7 @@ impl LoginDriver { slot, auto_open, state: Mutex::new(LoginState::default()), + registration: Mutex::new(()), }), } } @@ -269,7 +260,6 @@ impl LoginDriver { } }); - state.epoch += 1; state.pending = Some(Pending { session_id, url: url.clone(), @@ -277,7 +267,6 @@ impl LoginDriver { started: Instant::now(), redeeming: false, shutdown, - epoch: state.epoch, }); drop(state); @@ -460,41 +449,52 @@ async fn redeem(State(ctx): State, Json(body): Json) -> Re } }; - // Claim atomically (re-checking: the flow may have moved while parsing), - // remembering the claimed flow's number for the commit-time check below. - let epoch = + // Claim atomically (re-checking: the flow may have moved while parsing). + { + let mut state = driver.inner.state.lock().await; + if state + .grant + .as_ref() + .is_some_and(|g| g.session_id == body.state) { - let mut state = driver.inner.state.lock().await; - if state - .grant - .as_ref() - .is_some_and(|g| g.session_id == body.state) - { - return done(); - } - match state.pending.as_mut() { - Some(p) if p.session_id == body.state => { - if p.expired() { - return redeem_err( - "This sign-in expired. Ask the agent to authenticate again.", - ); - } - if p.redeeming { - return redeem_err( - "This sign-in is already being processed. Wait a moment; \ + return done(); + } + match state.pending.as_mut() { + Some(p) if p.session_id == body.state => { + if p.expired() { + return redeem_err( + "This sign-in expired. Ask the agent to authenticate again.", + ); + } + if p.redeeming { + return redeem_err( + "This sign-in is already being processed. Wait a moment; \ if nothing happens, ask the agent to authenticate again.", - ); - } - p.redeeming = true; - p.epoch + ); } - _ => return redeem_err( + p.redeeming = true; + } + _ => { + return redeem_err( "This sign-in is unknown or already used. Ask the agent to authenticate again.", - ), + ) } - }; + } + } - // The one network call, outside the lock. + // The network call + local commit, SERIALIZED across flows (see + // `Inner::registration`): II's `mcp_register_v2` replaces the anchor's + // previous grant, so whichever registration settles at II last is the + // live one — and because every register-then-commit runs alone inside + // this lock, the commit order equals the settle order, and the local + // slot always names the session II actually honors. In particular a + // flow that timed out locally while its register was in flight either + // settles BEFORE a replacement (which then overwrites it, at II and + // here) or AFTER it (and then the late registration is the live grant, + // and committing it here is exactly right). The per-flow single-flight + // claim above stays: it stops a page double-submit; this lock orders + // DISTINCT flows. + let _registration = driver.inner.registration.lock().await; match driver .inner .identities @@ -513,26 +513,12 @@ async fn redeem(State(ctx): State, Json(body): Json) -> Re Ok(outcome) => { let principal = driver.inner.identities.session_principal(&body.state).await; let mut state = driver.inner.state.lock().await; - // Commit-time check: this flow was claimed before the network - // call, but the watchdog may have expired it and a NEWER sign-in - // (pending or already completed) may exist by now — a stale - // success must not overwrite that. A late success with nothing - // newer around still lands (II did bind the grant; the watchdog - // alone is no reason to make the user redo it). - if superseded_by_newer(&state, epoch) { - drop(state); - return redeem_err( - "This sign-in was superseded by a newer one. Use the most recent \ - sign-in tab, or ask the agent to authenticate again.", - ); - } driver.inner.slot.set(body.state.clone()); state.grant = Some(Grant { session_id: body.state.clone(), principal, permissions: outcome.permissions, expiration_ns: outcome.expiration_ns, - epoch, }); let ours = matches!(&state.pending, Some(p) if p.session_id == body.state); if ours { @@ -777,7 +763,6 @@ mod tests { principal: Some("aaaaa-aa".into()), permissions: "queries", expiration_ns, - epoch: 1, }; driver.inner.state.lock().await.grant = Some(grant(now_ns() + 30 * 60_000_000_000)); match driver.status().await { @@ -810,7 +795,6 @@ mod tests { principal: Some("aaaaa-aa".into()), permissions: "all", expiration_ns: now_ns() + 30 * 60_000_000_000, - epoch: 1, }); assert!(matches!(driver.status().await, LoginStatus::SignedIn(_))); @@ -824,43 +808,6 @@ mod tests { ); } - // The stale-redeem guard: a flow that was claimed, then outlived by the - // watchdog while its network call ran, must not overwrite a NEWER sign-in - // (pending or completed) at commit time — but a late success with nothing - // newer around still lands. - #[test] - fn a_stale_redeem_is_superseded_by_any_newer_flow() { - let pending = |epoch: u64| Pending { - session_id: "s".into(), - url: "u".into(), - callback_url: "c".into(), - started: Instant::now(), - redeeming: false, - shutdown: Arc::new(Notify::new()), - epoch, - }; - let grant = |epoch: u64| Grant { - session_id: "g".into(), - principal: None, - permissions: "queries", - expiration_ns: u64::MAX, - epoch, - }; - let state = |p: Option, g: Option| LoginState { - pending: p, - grant: g, - epoch: 9, - }; - // Nothing newer: the late success lands (watchdog expiry alone is no - // reason to make the user redo a login II already bound). - assert!(!superseded_by_newer(&state(None, None), 3)); - assert!(!superseded_by_newer(&state(Some(pending(3)), None), 3)); - assert!(!superseded_by_newer(&state(None, Some(grant(2))), 3)); - // A newer pending handshake or completed grant wins over the stale one. - assert!(superseded_by_newer(&state(Some(pending(4)), None), 3)); - assert!(superseded_by_newer(&state(None, Some(grant(4))), 3)); - } - // The handshake window is exactly HANDSHAKE_TTL: a pending flow older than // that is expired (the watchdog and the redeem gate both read this). #[test] @@ -872,7 +819,6 @@ mod tests { started: Instant::now() - age, redeeming: false, shutdown: Arc::new(Notify::new()), - epoch: 1, }; assert!(!pending(Duration::from_secs(1)).expired()); assert!(pending(HANDSHAKE_TTL + Duration::from_secs(1)).expired()); From a5064b13beb82e6d57763967f7a29947fa01991c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 08:31:47 +0000 Subject: [PATCH 04/10] Address review: Host-validate the login listener (anti-DNS-rebinding) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The design's loopback hardening requires it (bind 127.0.0.1 explicitly AND validate the Host header, the local analogue of the hosted server's allowed-hosts list), and the listener accepted any Host: a page on an attacker's hostname that rebinds to 127.0.0.1 could reach the three routes with `Host: attacker.example`. Every URL this flow hands out carries exactly one authority — this listener's 127.0.0.1: — so a middleware now rejects anything else (exact string match, like II's own allow-list check) before routing. Even without it the routes are largely inert to a rebinder (/redeem requires the connect state and a chain targeting the in-process X); this closes the door outright. Tested with raw-socket requests: a rebound hostname and even localhost: bounce with 403; the one advertised authority passes. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- crates/imcp2-local/src/login.rs | 99 +++++++++++++++++++++++++++++++-- 1 file changed, 95 insertions(+), 4 deletions(-) diff --git a/crates/imcp2-local/src/login.rs b/crates/imcp2-local/src/login.rs index e1de004..cde7000 100644 --- a/crates/imcp2-local/src/login.rs +++ b/crates/imcp2-local/src/login.rs @@ -219,7 +219,8 @@ impl LoginDriver { .local_addr() .map_err(|e| format!("could not read the listener address: {e}"))? .port(); - let callback_url = format!("http://127.0.0.1:{port}/callback"); + let authority = format!("127.0.0.1:{port}"); + let callback_url = format!("http://{authority}/callback"); let url = iiconnect::ii_mcp_url( &self.inner.identities.instance().ii_url, @@ -230,7 +231,7 @@ impl LoginDriver { ); let shutdown = Arc::new(Notify::new()); - let app = login_router(self.clone(), callback_url.clone()); + let app = login_router(self.clone(), callback_url.clone(), authority); let sd = shutdown.clone(); tokio::spawn(async move { if let Err(e) = axum::serve(listener, app) @@ -322,18 +323,53 @@ impl LoginDriver { /// The three transient loopback routes — the login handshake's whole HTTP /// surface. `callback_url` is this listener's own callback, echoed verbatim in -/// the allow-list document. -fn login_router(driver: LoginDriver, callback_url: String) -> Router { +/// the allow-list document; `authority` is its `127.0.0.1:` host, the +/// only `Host` header the listener accepts. +fn login_router(driver: LoginDriver, callback_url: String, authority: String) -> Router { Router::new() .route("/callback", get(callback_page)) .route("/redeem", post(redeem)) .route(AUTH_CALLBACKS_WELL_KNOWN, get(auth_callbacks)) + // Anti-DNS-rebinding, per the design's loopback hardening (the local + // analogue of the hosted server's Host allow-list): a page on an + // attacker's hostname that rebinds to 127.0.0.1 reaches this port + // with `Host: attacker.example` — every URL this flow hands out + // carries exactly one authority, so anything else is rejected before + // routing. (Even without this the routes are largely inert to a + // rebinder — `/redeem` requires the connect `state` and a chain + // targeting the in-process `X` — this closes the door outright.) + .layer(axum::middleware::from_fn_with_state( + authority, + require_own_host, + )) .with_state(RouteCtx { driver, callback_url, }) } +/// Reject any request whose `Host` header is not this listener's own +/// `127.0.0.1:` authority (exact string match, like II's allow-list). +async fn require_own_host( + State(expected): State, + req: axum::extract::Request, + next: axum::middleware::Next, +) -> Response { + let ok = req + .headers() + .get(axum::http::header::HOST) + .and_then(|v| v.to_str().ok()) + .is_some_and(|host| host == expected); + if !ok { + return ( + axum::http::StatusCode::FORBIDDEN, + "wrong Host for this login listener", + ) + .into_response(); + } + next.run(req).await +} + #[derive(Clone)] struct RouteCtx { driver: LoginDriver, @@ -808,6 +844,61 @@ mod tests { ); } + // Anti-DNS-rebinding (the design's loopback hardening): a request whose + // `Host` names anything but this listener's own `127.0.0.1:` is + // rejected before routing — a rebinding page fetches this port with the + // attacker's hostname in `Host`, which is exactly what must bounce. The + // true authority (what every handed-out URL carries) passes untouched. + #[tokio::test] + async fn the_listener_rejects_foreign_host_headers() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let (driver, _slot) = test_driver(); + driver.begin(false).await.expect("begin"); + let (_, callback_url) = driver.pending_handshake().await.expect("pending"); + let authority = callback_url + .strip_prefix("http://") + .and_then(|s| s.strip_suffix("/callback")) + .expect("authority"); + let port: u16 = authority.split(':').nth(1).unwrap().parse().unwrap(); + + let request = |host: String| async move { + let mut stream = tokio::net::TcpStream::connect(("127.0.0.1", port)) + .await + .expect("connect"); + stream + .write_all( + format!("GET /callback HTTP/1.1\r\nHost: {host}\r\nConnection: close\r\n\r\n") + .as_bytes(), + ) + .await + .expect("write"); + let mut response = String::new(); + stream.read_to_string(&mut response).await.expect("read"); + response + .split_whitespace() + .nth(1) + .expect("status code") + .to_string() + }; + + assert_eq!( + request("attacker.example".into()).await, + "403", + "rebound hostname" + ); + assert_eq!( + request(format!("localhost:{port}")).await, + "403", + "even localhost: only the one advertised authority is served" + ); + assert_eq!( + request(authority.to_string()).await, + "200", + "the true authority passes" + ); + } + // The handshake window is exactly HANDSHAKE_TTL: a pending flow older than // that is expired (the watchdog and the redeem gate both read this). #[test] From f50ad08e1e48ad1dd20266d636848b4374518f6c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 09:40:20 +0000 Subject: [PATCH 05/10] Address review: readable variable names in the callback page script MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pinned page's inline script used single-character names (t, cls, c, r, d); source stays readable — minification, if ever wanted, is a build transformation, not a source style. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- crates/imcp2-core/src/iiconnect.rs | 34 +++++++++++++++--------------- crates/imcp2-local/src/login.rs | 2 +- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/crates/imcp2-core/src/iiconnect.rs b/crates/imcp2-core/src/iiconnect.rs index 7466965..ca6f528 100644 --- a/crates/imcp2-core/src/iiconnect.rs +++ b/crates/imcp2-core/src/iiconnect.rs @@ -99,14 +99,14 @@ const PINNED_PAGE_HTML: &str = include_str!("assets/connect-callback.html"); /// one dynamic value, the redeem URL, is spliced in by replacing /// `__REDEEM_URL__`, which sits inside a quoted JS string literal below. const PINNED_PAGE_JS: &str = r#"(function () { - // Swap the status line to `t` and move the screen into a terminal state: - // 'error' drops the spinner and reveals the contact line, 'done' just drops - // the spinner (see connect.css). - function show(t, cls) { - document.getElementById('m').textContent = t; - if (cls) { - var c = document.querySelector('.screen'); - if (c) { c.classList.add(cls); } + // Swap the status line to `message` and move the screen into a terminal + // state: 'error' drops the spinner and reveals the contact line, 'done' + // just drops the spinner (see connect.css). + function show(message, stateClass) { + document.getElementById('m').textContent = message; + if (stateClass) { + var screenElement = document.querySelector('.screen'); + if (screenElement) { screenElement.classList.add(stateClass); } } } // II delivers #delegation=&state=: the two-hop chain plus @@ -123,24 +123,24 @@ const PINNED_PAGE_JS: &str = r#"(function () { // Scrub the delegation from the address bar, keeping the path and any query // string the declared callback carries. Best-effort: the POST below works // even if a browser refuses the history call. - try { history.replaceState(null, '', location.pathname + location.search); } catch (e) {} + try { history.replaceState(null, '', location.pathname + location.search); } catch (ignored) {} fetch("__REDEEM_URL__", { method: 'POST', headers: { 'content-type': 'application/json' }, credentials: 'same-origin', body: body }) - .then(function (r) { return r.json().catch(function () { return {}; }); }) - .then(function (d) { + .then(function (response) { return response.json().catch(function () { return {}; }); }) + .then(function (data) { // Two success shapes, per deployment: the hosted redeem answers with the // `redirect` that continues its OAuth flow; a local redeem has nothing to // navigate to, so it answers `done` and this page IS the terminal state. - if (d && d.redirect) { - location.replace(d.redirect); - } else if (d && d.done) { + if (data && data.redirect) { + location.replace(data.redirect); + } else if (data && data.done) { show("Signed in — you can close this tab.", 'done'); } else { - show((d && d.error) || "We couldn't finish the connection. Restart from your client.", 'error'); + show((data && data.error) || "We couldn't finish the connection. Restart from your client.", 'error'); } }) .catch(function () { @@ -505,8 +505,8 @@ mod tests { #[test] fn pinned_page_script_handles_redirect_done_and_error() { let page = super::pinned_callback_page("/redeem", "mcp@dfinity.org"); - assert!(page.html.contains("d.redirect"), "hosted success arm"); - assert!(page.html.contains("d.done"), "local success arm"); + assert!(page.html.contains("data.redirect"), "hosted success arm"); + assert!(page.html.contains("data.done"), "local success arm"); assert!( page.html.contains("Signed in \u{2014} you can close this tab."), "the done arm renders the terminal signed-in state" diff --git a/crates/imcp2-local/src/login.rs b/crates/imcp2-local/src/login.rs index cde7000..73296e1 100644 --- a/crates/imcp2-local/src/login.rs +++ b/crates/imcp2-local/src/login.rs @@ -685,7 +685,7 @@ mod tests { "the script must POST to this listener's redeem" ); assert!( - html.contains("d.done"), + html.contains("data.done"), "the local success arm must be in the shipped page" ); From 1f8c1b76e5ceb70a4bd06e467150d458d0a4403b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 09:54:37 +0000 Subject: [PATCH 06/10] Address review: authentication moves out of core behind a SessionResolver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The SessionSource enum put deployment knowledge (bearer vs singleton) inside core's tool layer. Replace it with an injected resolver: core only asks "which already-validated session does this call act as?", and each binary owns its authentication end to end, in its own package — the hosted crate's require_token middleware validates the bearer and its resolver reads the stashed AuthedSession back (both now defined in imcp2's auth.rs again); the local binary's login flow fills its SessionSlot (now defined in imcp2-local's login.rs) and hands core the slot-reading resolver. Core drops its `http` dependency along the way — the request-extension read was the only use. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- Cargo.lock | 1 - crates/imcp2-core/Cargo.toml | 3 - crates/imcp2-core/src/lib.rs | 13 ++-- crates/imcp2-core/src/tools.rs | 116 +++++-------------------------- crates/imcp2-local/src/login.rs | 65 ++++++++++++++++- crates/imcp2-local/src/main.rs | 10 +-- crates/imcp2-local/src/server.rs | 5 +- src/auth.rs | 23 +++++- src/lib.rs | 7 +- 9 files changed, 123 insertions(+), 120 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b0a889b..265051f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1809,7 +1809,6 @@ dependencies = [ "crc32fast", "getrandom 0.3.4", "hex", - "http", "ic-agent", "regex", "reqwest 0.13.4", diff --git a/crates/imcp2-core/Cargo.toml b/crates/imcp2-core/Cargo.toml index e8f4296..b72e568 100644 --- a/crates/imcp2-core/Cargo.toml +++ b/crates/imcp2-core/Cargo.toml @@ -31,6 +31,3 @@ reqwest = { workspace = true } getrandom = { workspace = true } urlencoding = { workspace = true } url = { workspace = true } -# The request-extension type the Bearer session lookup reads (axum re-exports -# this same type; depending on `http` directly keeps axum out of this crate). -http = { workspace = true } diff --git a/crates/imcp2-core/src/lib.rs b/crates/imcp2-core/src/lib.rs index 85407ff..2d8de6c 100644 --- a/crates/imcp2-core/src/lib.rs +++ b/crates/imcp2-core/src/lib.rs @@ -15,11 +15,12 @@ //! plain values so each binary wraps them in its own HTTP handlers. //! //! How a tool call finds the II session it acts as is the one seam between -//! deployments: the hosted server resolves it per request from the -//! [`AuthedSession`] extension its bearer-token gate injects -//! ([`SessionSource::Bearer`]), while a single-user local server pins one -//! session established at login ([`SessionSource::Singleton`]). The tool -//! implementations are identical in both. +//! deployments, and **authentication itself stays in the embedding binary**: +//! each binary injects a [`SessionResolver`] reporting the already-validated +//! session for a call — the hosted server's reads what its bearer-token +//! middleware stashed on the request, a single-user local server's returns +//! the one session its login flow established. The tool implementations are +//! identical in both. //! //! The IC [`Agent`] is **inherited from the embedding application**, not built //! here: the binary passes in its own agent, so anonymous canister calls go @@ -36,7 +37,7 @@ mod discover; mod management; pub use identities::{IiInstance, SessionGauges}; -pub use tools::{AuthedSession, IcTools, SessionSlot, SessionSource}; +pub use tools::{IcTools, SessionResolver}; /// The IC [`Agent`] type the components are built around, re-exported so /// callers construct the injected agent from the exact `ic-agent` version this /// crate links. diff --git a/crates/imcp2-core/src/tools.rs b/crates/imcp2-core/src/tools.rs index 5cc3b1c..366d263 100644 --- a/crates/imcp2-core/src/tools.rs +++ b/crates/imcp2-core/src/tools.rs @@ -54,18 +54,30 @@ pub struct IcTools { agent: Agent, identities: Identities, skills: skills::SkillsCatalog, - /// Where a tool call finds the Internet Identity session it acts as — the - /// one per-deployment seam (see [`SessionSource`]). - session: SessionSource, + /// The authentication seam (see [`SessionResolver`]): asks the embedding + /// binary which already-validated session this call acts as. + session: SessionResolver, } +/// The one seam between deployments: how a tool call finds the Internet +/// Identity session it acts as. **Authentication itself lives in the embedding +/// binary, not here** — the resolver only reports the outcome of a validation +/// an earlier layer already performed. The hosted server's bearer middleware +/// validates the token per request and stashes the session id where its +/// resolver reads it back; the local binary's login flow keeps the one +/// signed-in session in a slot its resolver returns. `None` means this call is +/// unauthenticated (the "needs an authenticated session" errors at the tool +/// call sites are the caller-facing surface of that). +pub type SessionResolver = + Arc) -> Option + Send + Sync>; + #[tool_router] impl IcTools { pub fn new( agent: Agent, identities: Identities, skills: skills::SkillsCatalog, - session: SessionSource, + session: SessionResolver, ) -> Self { Self { agent, @@ -75,20 +87,10 @@ impl IcTools { } } - /// The session id this tool call acts as, per the deployment's - /// [`SessionSource`]: looked up from the request's [`AuthedSession`] - /// extension under `Bearer` (multi-user, hosted), or read from the shared - /// login-filled [`SessionSlot`] under `Singleton` (single-user, local). - /// `None` means the call is unauthenticated ("needs an authenticated - /// session" errors at the call sites are unchanged). + /// The already-validated session id this tool call acts as, per the + /// embedding binary's [`SessionResolver`]. fn current_session_id(&self, ctx: &RequestContext) -> Option { - match &self.session { - SessionSource::Bearer => authed_session(ctx).map(|s| s.session_id), - SessionSource::Singleton(slot) => { - let _ = ctx; - slot.get() - } - } + (self.session)(ctx) } #[tool( @@ -1507,70 +1509,6 @@ impl IcTools { } } -/// Where [`IcTools`] finds the Internet Identity session a tool call acts as — -/// the one seam between deployments. The tool implementations are identical -/// under both variants. -#[derive(Clone, Debug)] -pub enum SessionSource { - /// Multi-user (the hosted server): resolve the session per request from - /// the [`AuthedSession`] extension its bearer-token gate injected. - Bearer, - /// Single-user (a local server): every call acts as the one session in - /// the shared slot, established — and replaced on each re-login — by the - /// binary's own login flow. - Singleton(SessionSlot), -} - -/// The single-user session holder behind [`SessionSource::Singleton`]: a -/// shared, **replaceable** slot rather than a fixed id. The Internet Identity -/// contract requires a FRESH session key after an `Unauthorized` (a reused id -/// would reuse its keys), so each successful login writes a fresh session id -/// here — same-process reauthentication after grant expiry or revocation, with -/// no client restart. Clones share the slot: the login flow holds a writing -/// handle, [`IcTools`] reads it per call. Starts empty (signed out). -#[derive(Clone, Debug, Default)] -pub struct SessionSlot(Arc>>); - -impl SessionSlot { - /// A fresh, empty slot (no session — every session-needing tool reports - /// "needs an authenticated session" until the login flow fills it). - pub fn new() -> Self { - Self::default() - } - - /// The current session id, if a login has completed. The critical section - /// is a clone — never held across an await. - pub fn get(&self) -> Option { - self.0.read().expect("session slot poisoned").clone() - } - - /// Replace the session id after a successful login (a fresh id per the II - /// fresh-session-key contract). - pub fn set(&self, session_id: String) { - *self.0.write().expect("session slot poisoned") = Some(session_id); - } -} - -/// The verified session id of an authenticated MCP session, injected into the -/// request extensions by the embedding server's auth gate (the hosted binary's -/// bearer-token middleware) and read back under [`SessionSource::Bearer`]. -#[derive(Clone, Debug)] -pub struct AuthedSession { - pub session_id: String, -} - -/// The authenticated MCP session of the calling request, if the embedding -/// server's auth gate injected one. The transport surfaces the HTTP request's -/// `Parts` in the tool context's extensions; `http::request::Parts` is the -/// same type axum re-exports, so this reads what an axum middleware inserted -/// without this crate depending on axum. -fn authed_session(ctx: &RequestContext) -> Option { - ctx.extensions - .get::() - .and_then(|parts| parts.extensions.get::()) - .cloned() -} - /// The rejection for an OQL read (`get_canister_oql_schema`, or `canister_query`'s /// `oql` path) attempted with NO derivation origin. OQL reads per-app data gated by /// the caller's principal, so an anonymous read is signed as 2vxsx-fae and returns @@ -2500,22 +2438,6 @@ mod tests { } } - // The Singleton session slot is shared and REPLACEABLE: it starts empty - // (signed out), a login fills it, and a re-login (fresh session id, per the - // II fresh-session-key contract) replaces it — through every clone, since - // clones share the holder. This is what makes reauthentication after grant - // expiry a same-process step for the local binary. - #[test] - fn the_session_slot_is_shared_and_replaceable() { - let slot = super::SessionSlot::new(); - let reader = slot.clone(); - assert_eq!(reader.get(), None, "a fresh slot is signed out"); - slot.set("sess-1".into()); - assert_eq!(reader.get(), Some("sess-1".into()), "clones share the slot"); - slot.set("sess-2".into()); - assert_eq!(reader.get(), Some("sess-2".into()), "a re-login replaces the id"); - } - // icp_top_up_canister is a compliance-sensitive surface: its description must // present the tool as instructions-only (provided for completeness, executes // nothing on the user's behalf), so directory reviewers and calling models diff --git a/crates/imcp2-local/src/login.rs b/crates/imcp2-local/src/login.rs index 73296e1..e29384b 100644 --- a/crates/imcp2-local/src/login.rs +++ b/crates/imcp2-local/src/login.rs @@ -23,10 +23,48 @@ use axum::routing::{get, post}; use axum::{Json, Router}; use imcp2_core::identities::Identities; use imcp2_core::iiconnect::{self, RedeemBody, AUTH_CALLBACKS_WELL_KNOWN}; -use imcp2_core::SessionSlot; use serde_json::json; use tokio::sync::{Mutex, Notify}; +/// The single-user session holder — a shared, **replaceable** slot rather +/// than a fixed id. The Internet Identity contract requires a FRESH session +/// key after an `Unauthorized` (a reused id would reuse its keys), so each +/// successful login writes a fresh session id here — same-process +/// reauthentication after grant expiry or revocation, with no client restart. +/// Clones share the slot: the login flow holds a writing handle, and the slot +/// answers `IcTools`' session lookups through [`SessionSlot::resolver`] — the +/// local binary's whole authentication layer, living here rather than in +/// core, which only asks the injected resolver. Starts empty (signed out). +#[derive(Clone, Debug, Default)] +pub struct SessionSlot(Arc>>); + +impl SessionSlot { + /// A fresh, empty slot (no session — every session-needing tool reports + /// "needs an authenticated session" until the login flow fills it). + pub fn new() -> Self { + Self::default() + } + + /// The current session id, if a login has completed. The critical section + /// is a clone — never held across an await. + pub fn get(&self) -> Option { + self.0.read().expect("session slot poisoned").clone() + } + + /// Replace the session id after a successful login (a fresh id per the II + /// fresh-session-key contract). + pub fn set(&self, session_id: String) { + *self.0.write().expect("session slot poisoned") = Some(session_id); + } + + /// This slot as the [`imcp2_core::SessionResolver`] the local binary + /// injects into `IcTools`: every call acts as the one signed-in session. + pub fn resolver(&self) -> imcp2_core::SessionResolver { + let slot = self.clone(); + Arc::new(move |_ctx| slot.get()) + } +} + /// How long one sign-in handshake may take, from the `authenticate` call to /// the browser's redeem, before the pending flow and its listener are torn /// down. Matches the hosted server's pending-connect TTL. @@ -60,8 +98,8 @@ pub struct LoginDriver { struct Inner { identities: Identities, - /// The shared, replaceable session slot `IcTools` reads under - /// `SessionSource::Singleton`: each successful redeem writes the fresh + /// The shared, replaceable session slot `IcTools` reads through + /// [`SessionSlot::resolver`]: each successful redeem writes the fresh /// session id here (the II contract requires a fresh session key after /// `Unauthorized`, and a fresh id gets fresh keys). slot: SessionSlot, @@ -899,6 +937,27 @@ mod tests { ); } + // The session slot is shared and REPLACEABLE: it starts empty (signed + // out), a login fills it, and a re-login (fresh session id, per the II + // fresh-session-key contract) replaces it — through every clone and + // through the resolver handed to IcTools, since clones share the holder. + // This is what makes reauthentication after grant expiry a same-process + // step. + #[test] + fn the_session_slot_is_shared_and_replaceable() { + let slot = SessionSlot::new(); + let reader = slot.clone(); + assert_eq!(reader.get(), None, "a fresh slot is signed out"); + slot.set("sess-1".into()); + assert_eq!(reader.get(), Some("sess-1".into()), "clones share the slot"); + slot.set("sess-2".into()); + assert_eq!( + reader.get(), + Some("sess-2".into()), + "a re-login replaces the id" + ); + } + // The handshake window is exactly HANDSHAKE_TTL: a pending flow older than // that is expired (the watchdog and the redeem gate both read this). #[test] diff --git a/crates/imcp2-local/src/main.rs b/crates/imcp2-local/src/main.rs index 9bb98a8..53a8a78 100644 --- a/crates/imcp2-local/src/main.rs +++ b/crates/imcp2-local/src/main.rs @@ -33,7 +33,8 @@ mod login; mod server; -use imcp2_core::{identities::Identities, skills, IcTools, IiInstance, SessionSlot, SessionSource}; +use imcp2_core::{identities::Identities, skills, IcTools, IiInstance}; +use login::SessionSlot; use rmcp::ServiceExt; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; @@ -106,14 +107,15 @@ async fn main() -> anyhow::Result<()> { tracing::info!(ii = %instance.ii_url, "using Internet Identity instance \"{}\"", instance.name); let identities = Identities::new(instance, management_origin(), agent.clone()); - // The single-user session seam: the login flow fills (and on re-login - // replaces) the slot; every tool call reads it. + // The single-user authentication layer, owned by THIS binary: the login + // flow fills (and on re-login replaces) the slot, and the resolver it + // hands core answers every tool call's "which session?" from it. let slot = SessionSlot::new(); let tools = IcTools::new( agent, identities.clone(), skills::SkillsCatalog::new(), - SessionSource::Singleton(slot.clone()), + slot.resolver(), ); let auto_open = !truthy("IMCP2_NO_OPEN"); let login = login::LoginDriver::new(identities, slot, auto_open); diff --git a/crates/imcp2-local/src/server.rs b/crates/imcp2-local/src/server.rs index c242870..11946ab 100644 --- a/crates/imcp2-local/src/server.rs +++ b/crates/imcp2-local/src/server.rs @@ -332,7 +332,8 @@ impl ServerHandler for LocalServer { #[cfg(test)] mod tests { use super::*; - use imcp2_core::{identities::Identities, skills, IiInstance, SessionSlot, SessionSource}; + use crate::login::SessionSlot; + use imcp2_core::{identities::Identities, skills, IiInstance}; use rmcp::ServiceExt; fn test_server() -> LocalServer { @@ -350,7 +351,7 @@ mod tests { agent, identities.clone(), skills::SkillsCatalog::new(), - SessionSource::Singleton(slot.clone()), + slot.resolver(), ); let login = crate::login::LoginDriver::new(identities, slot, /* auto_open */ false); LocalServer::new(tools, login, /* auto_open */ false) diff --git a/src/auth.rs b/src/auth.rs index 1dd8181..733d15d 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -114,7 +114,28 @@ use uuid::Uuid; use imcp2_core::identities::Identities; use imcp2_core::iiconnect::{self, RedeemBody}; pub use imcp2_core::iiconnect::AUTH_CALLBACKS_WELL_KNOWN; -pub use imcp2_core::AuthedSession; + +/// The verified session id of an authenticated MCP session: [`require_token`] +/// validates the bearer token and stashes this on the request; the +/// [`bearer_session_resolver`] hands it to the tool layer. The whole +/// authentication step lives HERE, in the hosted binary — `imcp2-core` only +/// asks its injected [`imcp2_core::SessionResolver`] for the outcome. +#[derive(Clone, Debug)] +pub struct AuthedSession { + pub session_id: String, +} + +/// The hosted binary's [`imcp2_core::SessionResolver`]: read back the +/// [`AuthedSession`] that [`require_token`] injected into the request +/// extensions (rmcp surfaces the HTTP request's `Parts` in the tool context). +pub fn bearer_session_resolver() -> imcp2_core::SessionResolver { + std::sync::Arc::new(|ctx| { + ctx.extensions + .get::() + .and_then(|parts| parts.extensions.get::()) + .map(|session| session.session_id.clone()) + }) +} /// How long an authorization request and its pending II handshake stay valid /// before the user must restart. diff --git a/src/lib.rs b/src/lib.rs index 723e5e2..d96597d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -258,9 +258,10 @@ impl McpServer { agent.clone(), identities.clone(), skills.clone(), - // Multi-user: each request's session comes from the - // bearer-token gate's AuthedSession extension. - imcp2_core::SessionSource::Bearer, + // Multi-user: each request's session is whatever the + // bearer-token gate already validated (auth stays in + // THIS crate; core only asks the injected resolver). + auth::bearer_session_resolver(), )) }, LocalSessionManager::default().into(), From 87d56bd0811f63ce23df95e35d1e371293df0f2f Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 10:03:26 +0000 Subject: [PATCH 07/10] Address review: split the login module, name the contact const, detach the opener - login.rs (~1000 lines) becomes the login/ module: mod.rs (the driver and its state machine), slot.rs (the SessionSlot), routes.rs (the loopback listener's three routes + Host validation), tests.rs. - CONTACT -> CONTACT_EMAIL, saying what kind of string it is. - The browser opener uses open::that_detached: open 5.x already nulls the launcher's stdio either way (so nothing could interleave with the JSON-RPC stream), but detached additionally double-forks the launcher away off macOS, so it can never hold a thread or a handle on us. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- crates/imcp2-local/src/login.rs | 976 ------------------------- crates/imcp2-local/src/login/mod.rs | 325 ++++++++ crates/imcp2-local/src/login/routes.rs | 265 +++++++ crates/imcp2-local/src/login/slot.rs | 44 ++ crates/imcp2-local/src/login/tests.rs | 365 +++++++++ 5 files changed, 999 insertions(+), 976 deletions(-) delete mode 100644 crates/imcp2-local/src/login.rs create mode 100644 crates/imcp2-local/src/login/mod.rs create mode 100644 crates/imcp2-local/src/login/routes.rs create mode 100644 crates/imcp2-local/src/login/slot.rs create mode 100644 crates/imcp2-local/src/login/tests.rs diff --git a/crates/imcp2-local/src/login.rs b/crates/imcp2-local/src/login.rs deleted file mode 100644 index e29384b..0000000 --- a/crates/imcp2-local/src/login.rs +++ /dev/null @@ -1,976 +0,0 @@ -//! The **browser Internet Identity login** for the local binary: mint a fresh -//! session, stand up a transient loopback listener, hand the user an id.ai -//! sign-in link, and redeem the delegation II navigates back with. On success -//! the shared [`SessionSlot`] is filled and every tool call acts as the user; -//! the listener is torn down (redeem or timeout) and the grant lives in memory -//! only. -//! -//! The listener is the unavoidable minimum, not a design slip: II delivers the -//! delegation by *navigating the browser* to the callback (a URL fragment only -//! a served page can read and POST back), and II's #4091 check fetches -//! `/.well-known/ii-auth-callbacks` from the callback's origin before honoring -//! it — both require a real HTTP origin. This is the standard native-app -//! loopback redirect (RFC 8252 §7.3), the same shape as the ICP CLI's -//! `icp identity link web`. It serves exactly three routes for one handshake -//! and never serves the MCP tool surface, which rides stdio. - -use std::sync::Arc; -use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; - -use axum::extract::State; -use axum::response::{Html, IntoResponse, Response}; -use axum::routing::{get, post}; -use axum::{Json, Router}; -use imcp2_core::identities::Identities; -use imcp2_core::iiconnect::{self, RedeemBody, AUTH_CALLBACKS_WELL_KNOWN}; -use serde_json::json; -use tokio::sync::{Mutex, Notify}; - -/// The single-user session holder — a shared, **replaceable** slot rather -/// than a fixed id. The Internet Identity contract requires a FRESH session -/// key after an `Unauthorized` (a reused id would reuse its keys), so each -/// successful login writes a fresh session id here — same-process -/// reauthentication after grant expiry or revocation, with no client restart. -/// Clones share the slot: the login flow holds a writing handle, and the slot -/// answers `IcTools`' session lookups through [`SessionSlot::resolver`] — the -/// local binary's whole authentication layer, living here rather than in -/// core, which only asks the injected resolver. Starts empty (signed out). -#[derive(Clone, Debug, Default)] -pub struct SessionSlot(Arc>>); - -impl SessionSlot { - /// A fresh, empty slot (no session — every session-needing tool reports - /// "needs an authenticated session" until the login flow fills it). - pub fn new() -> Self { - Self::default() - } - - /// The current session id, if a login has completed. The critical section - /// is a clone — never held across an await. - pub fn get(&self) -> Option { - self.0.read().expect("session slot poisoned").clone() - } - - /// Replace the session id after a successful login (a fresh id per the II - /// fresh-session-key contract). - pub fn set(&self, session_id: String) { - *self.0.write().expect("session slot poisoned") = Some(session_id); - } - - /// This slot as the [`imcp2_core::SessionResolver`] the local binary - /// injects into `IcTools`: every call acts as the one signed-in session. - pub fn resolver(&self) -> imcp2_core::SessionResolver { - let slot = self.clone(); - Arc::new(move |_ctx| slot.get()) - } -} - -/// How long one sign-in handshake may take, from the `authenticate` call to -/// the browser's redeem, before the pending flow and its listener are torn -/// down. Matches the hosted server's pending-connect TTL. -const HANDSHAKE_TTL: Duration = Duration::from_secs(600); - -/// `ttl` (seconds) requested for the II grant in the connect link, matching -/// the hosted server. II clamps it to [600, 2592000], and the effective -/// session duration is whatever the user picks on II's consent screen. -const GRANT_TTL_SECS: u64 = 3600; - -/// Where the pinned callback page's error state points users ("contact us to -/// report it") — DFINITY ships this binary, so failures report to the same -/// address as the hosted server's. -const CONTACT: &str = "mcp@dfinity.org"; - -fn now_ns() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_nanos() as u64 -} - -/// The login flow's shared handle: the `authenticate`/`auth_status` tools call -/// [`begin`](Self::begin)/[`status`](Self::status), and the loopback redeem -/// route completes the flow through a clone. Single-user by construction — one -/// pending handshake at a time, serialized by the state lock. -#[derive(Clone)] -pub struct LoginDriver { - inner: Arc, -} - -struct Inner { - identities: Identities, - /// The shared, replaceable session slot `IcTools` reads through - /// [`SessionSlot::resolver`]: each successful redeem writes the fresh - /// session id here (the II contract requires a fresh session key after - /// `Unauthorized`, and a fresh id gets fresh keys). - slot: SessionSlot, - /// Best-effort `open::that` on the sign-in link. The flow never depends on - /// it succeeding — the link is always returned in-band by `authenticate`. - auto_open: bool, - state: Mutex, - /// Serializes each redeem's `mcp_register_v2` **and** its local commit, - /// across flows. II's register REPLACES the anchor's previous grant (the - /// last registration to settle at II is the live one), so the local slot - /// must be committed in the same order the registrations settle: without - /// this, a timed-out flow's slow register could land at II *after* a - /// newer flow's — II would then honor the old session while the slot - /// named the new one, and every tool call would come back `Unauthorized`. - registration: Mutex<()>, -} - -#[derive(Default)] -struct LoginState { - pending: Option, - /// The most recent successful login, kept (even past expiry) so - /// `auth_status` can say *whose* session expired. - grant: Option, -} - -struct Pending { - /// Doubles as the connect `state` in the II link — II echoes it to the - /// callback, and the redeem accepts only this value. - session_id: String, - url: String, - /// This flow's loopback callback (`http://127.0.0.1:/callback`) — - /// the origin the listener serves on. - callback_url: String, - started: Instant, - /// Single-flight marker: `true` while a redemption attempt (the - /// `mcp_register_v2` network call) is mid-flight, so a page double-submit - /// can't fire two. Cleared on failure so a genuine retry can proceed. - redeeming: bool, - /// Tears the transient listener down (graceful): notified on redeem - /// success, on handshake timeout, and when a fresh flow replaces this one. - shutdown: Arc, -} - -impl Pending { - fn expired(&self) -> bool { - self.started.elapsed() >= HANDSHAKE_TTL - } -} - -/// A completed login, as `auth_status` reports it. -#[derive(Clone)] -pub struct Grant { - pub session_id: String, - /// The session principal (`self_authenticating(pub(S))`) — attribution, - /// not the user's per-app principal. - pub principal: Option, - /// II's recorded access level: `"queries"` (read-only) or `"all"`. - pub permissions: &'static str, - /// Grant expiration, ns since the Unix epoch (the user's consent-time - /// session choice). - pub expiration_ns: u64, -} - -impl Grant { - pub fn expired(&self) -> bool { - self.expiration_ns <= now_ns() - } - - /// Whole minutes until expiry (0 when expired). - pub fn minutes_left(&self) -> u64 { - self.expiration_ns.saturating_sub(now_ns()) / 60_000_000_000 - } -} - -/// What [`LoginDriver::status`] reports — the states `auth_status` renders. -pub enum LoginStatus { - SignedIn(Grant), - /// A handshake is waiting for the browser. - Pending { - url: String, - }, - /// A previous grant ran out (or was revoked and then expired here); a new - /// `authenticate` gets a fresh session. - Expired(Grant), - SignedOut, -} - -/// What [`LoginDriver::begin`] returns to the `authenticate` tool. -pub enum BeginOutcome { - /// A live grant exists and `refresh` was not requested. - AlreadySignedIn(Grant), - /// A sign-in link is ready. `fresh` distinguishes a newly started flow - /// from an already-pending one being returned again (single-flight). - Pending { url: String, fresh: bool }, -} - -impl LoginDriver { - pub fn new(identities: Identities, slot: SessionSlot, auto_open: bool) -> Self { - Self { - inner: Arc::new(Inner { - identities, - slot, - auto_open, - state: Mutex::new(LoginState::default()), - registration: Mutex::new(()), - }), - } - } - - /// Start (or return the already-pending) sign-in: mint a **fresh** session - /// id + registration key, bind the loopback listener, and build the id.ai - /// link — returned immediately, never blocking on the browser. With a live - /// grant and `refresh: false`, reports the grant instead of starting over. - pub async fn begin(&self, refresh: bool) -> Result { - // The lock is held across the (local, network-free) setup, so - // concurrent `authenticate` calls serialize into one flow. - let mut state = self.inner.state.lock().await; - - if let Some(p) = &state.pending { - if !p.expired() { - return Ok(BeginOutcome::Pending { - url: p.url.clone(), - fresh: false, - }); - } - // Expired while nobody was looking: tear it down, then start fresh. - if let Some(p) = state.pending.take() { - p.shutdown.notify_one(); - } - } - if !refresh { - if let Some(g) = &state.grant { - if !g.expired() { - return Ok(BeginOutcome::AlreadySignedIn(g.clone())); - } - } - } - - // A FRESH session id per flow: fresh id ⇒ fresh session + registration - // keys in `Identities`, per the II fresh-session-key contract. - let session_id = uuid::Uuid::new_v4().to_string(); - let reg_pubkey = self - .inner - .identities - .registration_pubkey_b64(&session_id) - .await?; - - // Port 0: the OS picks a free port; both the II link's `callback` and - // the #4091 allow-list entry derive from the ONE resulting origin, so - // they cannot drift (II matches by exact string equality). - let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) - .await - .map_err(|e| format!("could not bind the login callback listener: {e}"))?; - let port = listener - .local_addr() - .map_err(|e| format!("could not read the listener address: {e}"))? - .port(); - let authority = format!("127.0.0.1:{port}"); - let callback_url = format!("http://{authority}/callback"); - - let url = iiconnect::ii_mcp_url( - &self.inner.identities.instance().ii_url, - &callback_url, - &session_id, - GRANT_TTL_SECS, - ®_pubkey, - ); - - let shutdown = Arc::new(Notify::new()); - let app = login_router(self.clone(), callback_url.clone(), authority); - let sd = shutdown.clone(); - tokio::spawn(async move { - if let Err(e) = axum::serve(listener, app) - .with_graceful_shutdown(async move { sd.notified().await }) - .await - { - tracing::warn!("login listener error: {e}"); - } - }); - - // Watchdog: when the handshake window closes, drop the pending flow - // and the listener (if a redeem didn't already). - let driver = self.clone(); - let sid = session_id.clone(); - tokio::spawn(async move { - tokio::time::sleep(HANDSHAKE_TTL).await; - let mut state = driver.inner.state.lock().await; - let ours = matches!(&state.pending, Some(p) if p.session_id == sid); - if ours { - if let Some(p) = state.pending.take() { - p.shutdown.notify_one(); - tracing::info!( - callback = %p.callback_url, - "sign-in window expired; login listener closed" - ); - } - } - }); - - state.pending = Some(Pending { - session_id, - url: url.clone(), - callback_url, - started: Instant::now(), - redeeming: false, - shutdown, - }); - drop(state); - - if self.inner.auto_open { - // `open::that` can block on the launcher, so it runs off the - // async workers — and only best-effort: the link is in the tool - // result either way. - let link = url.clone(); - tokio::task::spawn_blocking(move || { - if let Err(e) = open::that(&link) { - tracing::debug!("could not auto-open the browser: {e}"); - } - }); - } - - Ok(BeginOutcome::Pending { url, fresh: true }) - } - - /// The pending flow's `(session_id, callback_url)`, for the tests that - /// play the browser against the live listener. - #[cfg(test)] - pub(crate) async fn pending_handshake(&self) -> Option<(String, String)> { - let state = self.inner.state.lock().await; - state - .pending - .as_ref() - .map(|p| (p.session_id.clone(), p.callback_url.clone())) - } - - pub async fn status(&self) -> LoginStatus { - let state = self.inner.state.lock().await; - // A live pending handshake outranks a live grant, matching `begin`'s - // precedence: after `authenticate(refresh: true)` the old session may - // still be valid, but the truthful answer is "a replacement sign-in is - // waiting for the browser", not "signed in" — otherwise the client - // reads the refresh as already complete. - if let Some(p) = &state.pending { - if !p.expired() { - return LoginStatus::Pending { url: p.url.clone() }; - } - } - if let Some(g) = &state.grant { - if !g.expired() { - return LoginStatus::SignedIn(g.clone()); - } - } - match &state.grant { - Some(g) => LoginStatus::Expired(g.clone()), - None => LoginStatus::SignedOut, - } - } -} - -/// The three transient loopback routes — the login handshake's whole HTTP -/// surface. `callback_url` is this listener's own callback, echoed verbatim in -/// the allow-list document; `authority` is its `127.0.0.1:` host, the -/// only `Host` header the listener accepts. -fn login_router(driver: LoginDriver, callback_url: String, authority: String) -> Router { - Router::new() - .route("/callback", get(callback_page)) - .route("/redeem", post(redeem)) - .route(AUTH_CALLBACKS_WELL_KNOWN, get(auth_callbacks)) - // Anti-DNS-rebinding, per the design's loopback hardening (the local - // analogue of the hosted server's Host allow-list): a page on an - // attacker's hostname that rebinds to 127.0.0.1 reaches this port - // with `Host: attacker.example` — every URL this flow hands out - // carries exactly one authority, so anything else is rejected before - // routing. (Even without this the routes are largely inert to a - // rebinder — `/redeem` requires the connect `state` and a chain - // targeting the in-process `X` — this closes the door outright.) - .layer(axum::middleware::from_fn_with_state( - authority, - require_own_host, - )) - .with_state(RouteCtx { - driver, - callback_url, - }) -} - -/// Reject any request whose `Host` header is not this listener's own -/// `127.0.0.1:` authority (exact string match, like II's allow-list). -async fn require_own_host( - State(expected): State, - req: axum::extract::Request, - next: axum::middleware::Next, -) -> Response { - let ok = req - .headers() - .get(axum::http::header::HOST) - .and_then(|v| v.to_str().ok()) - .is_some_and(|host| host == expected); - if !ok { - return ( - axum::http::StatusCode::FORBIDDEN, - "wrong Host for this login listener", - ) - .into_response(); - } - next.run(req).await -} - -#[derive(Clone)] -struct RouteCtx { - driver: LoginDriver, - callback_url: String, -} - -/// GET /callback — the pinned fragment-reading page (shared with the hosted -/// server), pointed at this listener's `/redeem`, with the same non-CSP -/// hardening headers the hosted server adds. -async fn callback_page() -> Response { - let page = iiconnect::pinned_callback_page("/redeem", CONTACT); - let mut resp = Html(page.html).into_response(); - let h = resp.headers_mut(); - h.insert( - axum::http::header::CONTENT_SECURITY_POLICY, - axum::http::HeaderValue::from_str(&page.csp).expect("valid CSP"), - ); - h.insert( - axum::http::header::REFERRER_POLICY, - axum::http::HeaderValue::from_static("no-referrer"), - ); - h.insert( - axum::http::header::X_CONTENT_TYPE_OPTIONS, - axum::http::HeaderValue::from_static("nosniff"), - ); - h.insert( - axum::http::header::X_FRAME_OPTIONS, - axum::http::HeaderValue::from_static("DENY"), - ); - resp -} - -/// GET /.well-known/ii-auth-callbacks — the #4091 allow-list: exactly this -/// listener's callback. II fetches it cross-origin (hence the one CORS -/// header) and fail-closed before honoring the callback; `no-store` keeps an -/// intermediary from serving a stale document for a past listener's port. -async fn auth_callbacks(State(ctx): State) -> Response { - let mut resp = Json(json!({ "callbacks": [ctx.callback_url] })).into_response(); - let h = resp.headers_mut(); - h.insert( - axum::http::header::CACHE_CONTROL, - axum::http::HeaderValue::from_static("no-store"), - ); - h.insert( - axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN, - axum::http::HeaderValue::from_static("*"), - ); - resp -} - -fn redeem_err(msg: &str) -> Response { - ( - axum::http::StatusCode::BAD_REQUEST, - Json(json!({ "error": msg })), - ) - .into_response() -} - -/// The local success answer: no OAuth continuation to redirect into, so the -/// pinned page renders its terminal "Signed in — you can close this tab" -/// state in place (the page script's `done` arm). -fn done() -> Response { - Json(json!({ "done": true })).into_response() -} - -/// POST /redeem — the pinned page POSTs the fragment here. The hosted redeem -/// minus its OAuth tail: no initiator cookie (there is no public initiate -/// endpoint to confuse a deputy through — the binary itself mints `X` and this -/// listener is loopback-only), no PKCE code, and no redirect — success answers -/// `{"done": true}` and the page renders its terminal state. Parse before -/// claiming, single-flight claim, one `mcp_register_v2` call, then fill the -/// session slot and tear the listener down. -async fn redeem(State(ctx): State, Json(body): Json) -> Response { - let driver = &ctx.driver; - - // Validate the delivery against the pending flow first (NOT claiming yet), - // so an unknown/expired `state` gets its accurate refusal without parsing. - { - let state = driver.inner.state.lock().await; - if state - .grant - .as_ref() - .is_some_and(|g| g.session_id == body.state) - { - // A retry of an already-successful redeem (II's delegation redeems - // repeatedly within its lifetime): idempotent success. - return done(); - } - match &state.pending { - Some(p) if p.session_id == body.state => { - if p.expired() { - return redeem_err( - "This sign-in expired. Ask the agent to authenticate again.", - ); - } - } - _ => { - return redeem_err( - "This sign-in is unknown or already used. Ask the agent to authenticate again.", - ) - } - } - } - - // Decode the fragment delegation (pure, size-bounded) BEFORE claiming, so - // a malformed delivery never occupies the single-flight slot. - let (user_key, chain) = match iiconnect::parse_registration_delegation(&body.delegation) { - Ok(v) => v, - Err(e) => { - return redeem_err(&format!( - "We couldn't read the sign-in response. Ask the agent to authenticate again. ({e})" - )) - } - }; - - // Claim atomically (re-checking: the flow may have moved while parsing). - { - let mut state = driver.inner.state.lock().await; - if state - .grant - .as_ref() - .is_some_and(|g| g.session_id == body.state) - { - return done(); - } - match state.pending.as_mut() { - Some(p) if p.session_id == body.state => { - if p.expired() { - return redeem_err( - "This sign-in expired. Ask the agent to authenticate again.", - ); - } - if p.redeeming { - return redeem_err( - "This sign-in is already being processed. Wait a moment; \ - if nothing happens, ask the agent to authenticate again.", - ); - } - p.redeeming = true; - } - _ => { - return redeem_err( - "This sign-in is unknown or already used. Ask the agent to authenticate again.", - ) - } - } - } - - // The network call + local commit, SERIALIZED across flows (see - // `Inner::registration`): II's `mcp_register_v2` replaces the anchor's - // previous grant, so whichever registration settles at II last is the - // live one — and because every register-then-commit runs alone inside - // this lock, the commit order equals the settle order, and the local - // slot always names the session II actually honors. In particular a - // flow that timed out locally while its register was in flight either - // settles BEFORE a replacement (which then overwrites it, at II and - // here) or AFTER it (and then the late registration is the live grant, - // and committing it here is exactly right). The per-flow single-flight - // claim above stays: it stops a page double-submit; this lock orders - // DISTINCT flows. - let _registration = driver.inner.registration.lock().await; - match driver - .inner - .identities - .redeem_registration_delegation(&body.state, user_key, chain) - .await - { - Err(e) => { - let mut state = driver.inner.state.lock().await; - if let Some(p) = state.pending.as_mut() { - if p.session_id == body.state { - p.redeeming = false; // free the claim for a genuine retry - } - } - redeem_err(&e) - } - Ok(outcome) => { - let principal = driver.inner.identities.session_principal(&body.state).await; - let mut state = driver.inner.state.lock().await; - driver.inner.slot.set(body.state.clone()); - state.grant = Some(Grant { - session_id: body.state.clone(), - principal, - permissions: outcome.permissions, - expiration_ns: outcome.expiration_ns, - }); - let ours = matches!(&state.pending, Some(p) if p.session_id == body.state); - if ours { - if let Some(p) = state.pending.take() { - p.shutdown.notify_one(); - } - } - drop(state); - tracing::info!( - permissions = outcome.permissions, - expiration_ns = outcome.expiration_ns, - "signed in: registration delegation redeemed; login listener closing" - ); - done() - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use imcp2_core::IiInstance; - - /// A driver against the real mainnet/prod-II config — nothing here touches - /// the network (agent construction, key minting, and the loopback listener - /// are all local), which is exactly what these tests rely on. - fn test_driver() -> (LoginDriver, SessionSlot) { - let agent = imcp2_core::Agent::builder() - .with_url(imcp2_core::IC_URL) - .build() - .expect("agent"); - let identities = Identities::new( - IiInstance::prod().expect("prod II instance"), - "https://mcp.internetcomputer.org".into(), - agent, - ); - let slot = SessionSlot::new(); - ( - LoginDriver::new(identities, slot.clone(), /* auto_open */ false), - slot, - ) - } - - /// A shape-valid two-hop chain (the same wire shape as iiconnect's parser - /// tests) whose signatures are garbage: it must get PAST parsing and be - /// rejected by the redeem path itself — offline, before any network call. - fn fake_chain_json() -> String { - serde_json::json!({ - "delegations": [ - { - "delegation": { "pubkey": "070707", "expiration": "66", "targets": ["aaaaa-aa"] }, - "signature": "040506", - }, - { - "delegation": { "pubkey": "09080706", "expiration": "66" }, - "signature": "010909", - }, - ], - "publicKey": "010203", - }) - .to_string() - } - - // One `authenticate` = one flow: a fresh id.ai link over a fresh loopback - // callback; a second call while it is pending returns the SAME link - // (single-flight) rather than racing a second listener. The session slot - // stays empty until a redeem succeeds — starting a login signs nobody in. - #[tokio::test] - async fn begin_mints_one_flow_and_repeats_it_while_pending() { - let (driver, slot) = test_driver(); - let BeginOutcome::Pending { url, fresh } = driver.begin(false).await.expect("begin") else { - panic!("a fresh driver must start a flow, not report a session") - }; - assert!(fresh); - assert!(url.starts_with("https://id.ai/mcp#callback="), "{url}"); - assert!( - url.contains("http%3A%2F%2F127.0.0.1%3A"), - "loopback callback in the fragment: {url}" - ); - assert!(url.contains("&ttl=3600®istration_key="), "{url}"); - - let BeginOutcome::Pending { url: again, fresh } = driver.begin(false).await.expect("begin") - else { - panic!("a pending flow must be returned, not replaced") - }; - assert!(!fresh, "the second call must join the pending flow"); - assert_eq!(again, url, "same link while the handshake is pending"); - - assert_eq!(slot.get(), None, "no session until a redeem succeeds"); - assert!(matches!(driver.status().await, LoginStatus::Pending { .. })); - } - - // The transient listener's whole surface, driven like the browser: the - // pinned page (strict CSP + the hosted server's hardening headers, script - // pointed at /redeem), and the #4091 allow-list (this callback verbatim, - // CORS-readable, never cached). - #[tokio::test] - async fn the_listener_serves_the_pinned_page_and_the_allow_list() { - let (driver, _slot) = test_driver(); - driver.begin(false).await.expect("begin"); - let (_, callback_url) = driver.pending_handshake().await.expect("pending"); - let origin = callback_url.strip_suffix("/callback").unwrap().to_string(); - let http = reqwest::Client::new(); - - let page = http.get(&callback_url).send().await.expect("GET /callback"); - assert_eq!(page.status(), 200); - let csp = page - .headers() - .get("content-security-policy") - .expect("pinned page ships a CSP") - .to_str() - .unwrap() - .to_string(); - assert!(csp.contains("default-src 'none'"), "{csp}"); - for (name, want) in [ - ("referrer-policy", "no-referrer"), - ("x-content-type-options", "nosniff"), - ("x-frame-options", "DENY"), - ] { - assert_eq!( - page.headers().get(name).and_then(|v| v.to_str().ok()), - Some(want) - ); - } - let html = page.text().await.unwrap(); - assert!( - html.contains("/redeem"), - "the script must POST to this listener's redeem" - ); - assert!( - html.contains("data.done"), - "the local success arm must be in the shipped page" - ); - - let wk = http - .get(format!("{origin}{AUTH_CALLBACKS_WELL_KNOWN}")) - .send() - .await - .expect("GET allow-list"); - assert_eq!(wk.status(), 200); - assert_eq!( - wk.headers() - .get("access-control-allow-origin") - .and_then(|v| v.to_str().ok()), - Some("*"), - "II fetches the allow-list cross-origin" - ); - assert_eq!( - wk.headers() - .get("cache-control") - .and_then(|v| v.to_str().ok()), - Some("no-store"), - "fail-closed infrastructure must never be served stale" - ); - let doc: serde_json::Value = wk.json().await.unwrap(); - assert_eq!( - doc, - serde_json::json!({ "callbacks": [callback_url] }), - "the declared callback must equal the link's callback VERBATIM (exact-match)" - ); - } - - // The redeem's refusal paths, in order: an unknown `state` (or one from a - // replaced flow), a delegation the parser rejects, and a shape-valid chain - // with garbage signatures that the redeem itself rejects — all 400 with an - // {"error"} body the pinned page renders, all offline, and none of them - // consumes the pending flow (a genuine retry stays possible). - #[tokio::test] - async fn redeem_refuses_bad_deliveries_and_keeps_the_flow_retryable() { - let (driver, slot) = test_driver(); - let BeginOutcome::Pending { url, .. } = driver.begin(false).await.expect("begin") else { - panic!("fresh flow") - }; - let (state, callback_url) = driver.pending_handshake().await.expect("pending"); - let origin = callback_url.strip_suffix("/callback").unwrap().to_string(); - let redeem_url = format!("{origin}/redeem"); - let http = reqwest::Client::new(); - - let post = |body: serde_json::Value| { - let http = http.clone(); - let redeem_url = redeem_url.clone(); - async move { - http.post(&redeem_url) - .json(&body) - .send() - .await - .expect("POST /redeem") - } - }; - - let r = post(serde_json::json!({ "state": "someone-else", "delegation": "" })).await; - assert_eq!(r.status(), 400); - let e: serde_json::Value = r.json().await.unwrap(); - assert!( - e["error"] - .as_str() - .unwrap() - .contains("unknown or already used"), - "{e}" - ); - - let r = post(serde_json::json!({ "state": state, "delegation": "not json" })).await; - assert_eq!(r.status(), 400); - let e: serde_json::Value = r.json().await.unwrap(); - assert!( - e["error"] - .as_str() - .unwrap() - .contains("couldn't read the sign-in response"), - "{e}" - ); - - let r = post(serde_json::json!({ "state": state, "delegation": fake_chain_json() })).await; - assert_eq!( - r.status(), - 400, - "garbage signatures must be rejected by the redeem" - ); - let e: serde_json::Value = r.json().await.unwrap(); - assert!(e["error"].is_string(), "{e}"); - - assert_eq!(slot.get(), None, "no failure may fill the session slot"); - let BeginOutcome::Pending { url: still, fresh } = driver.begin(false).await.expect("begin") - else { - panic!("the flow must survive failed redeems") - }; - assert!(!fresh, "failed redeems must not consume the pending flow"); - assert_eq!(still, url); - } - - // The status lifecycle around a grant (injected directly — a real one - // needs live II): live grant → SignedIn with the wallclock math; past - // expiration → Expired, still naming the principal; a fresh driver is - // simply SignedOut. - #[tokio::test] - async fn status_reports_the_grant_lifecycle() { - let (driver, _slot) = test_driver(); - assert!(matches!(driver.status().await, LoginStatus::SignedOut)); - - let grant = |expiration_ns: u64| Grant { - session_id: "sess".into(), - principal: Some("aaaaa-aa".into()), - permissions: "queries", - expiration_ns, - }; - driver.inner.state.lock().await.grant = Some(grant(now_ns() + 30 * 60_000_000_000)); - match driver.status().await { - LoginStatus::SignedIn(g) => { - assert!( - (25..=30).contains(&g.minutes_left()), - "{}", - g.minutes_left() - ); - } - _ => panic!("a live grant must report SignedIn"), - } - - driver.inner.state.lock().await.grant = Some(grant(now_ns() - 1)); - match driver.status().await { - LoginStatus::Expired(g) => assert_eq!(g.principal.as_deref(), Some("aaaaa-aa")), - _ => panic!("a past-expiry grant must report Expired"), - } - } - - // `authenticate(refresh: true)` while a live grant exists: the truthful - // status is the PENDING replacement handshake, not the old "signed in" — - // otherwise the client reads the refresh as already complete (the same - // precedence `begin` applies). - #[tokio::test] - async fn a_pending_refresh_outranks_the_live_grant_in_status() { - let (driver, _slot) = test_driver(); - driver.inner.state.lock().await.grant = Some(Grant { - session_id: "old".into(), - principal: Some("aaaaa-aa".into()), - permissions: "all", - expiration_ns: now_ns() + 30 * 60_000_000_000, - }); - assert!(matches!(driver.status().await, LoginStatus::SignedIn(_))); - - let BeginOutcome::Pending { fresh, .. } = driver.begin(true).await.expect("refresh") else { - panic!("refresh=true must start a replacement flow past a live grant") - }; - assert!(fresh); - assert!( - matches!(driver.status().await, LoginStatus::Pending { .. }), - "a live replacement handshake must outrank the old grant" - ); - } - - // Anti-DNS-rebinding (the design's loopback hardening): a request whose - // `Host` names anything but this listener's own `127.0.0.1:` is - // rejected before routing — a rebinding page fetches this port with the - // attacker's hostname in `Host`, which is exactly what must bounce. The - // true authority (what every handed-out URL carries) passes untouched. - #[tokio::test] - async fn the_listener_rejects_foreign_host_headers() { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let (driver, _slot) = test_driver(); - driver.begin(false).await.expect("begin"); - let (_, callback_url) = driver.pending_handshake().await.expect("pending"); - let authority = callback_url - .strip_prefix("http://") - .and_then(|s| s.strip_suffix("/callback")) - .expect("authority"); - let port: u16 = authority.split(':').nth(1).unwrap().parse().unwrap(); - - let request = |host: String| async move { - let mut stream = tokio::net::TcpStream::connect(("127.0.0.1", port)) - .await - .expect("connect"); - stream - .write_all( - format!("GET /callback HTTP/1.1\r\nHost: {host}\r\nConnection: close\r\n\r\n") - .as_bytes(), - ) - .await - .expect("write"); - let mut response = String::new(); - stream.read_to_string(&mut response).await.expect("read"); - response - .split_whitespace() - .nth(1) - .expect("status code") - .to_string() - }; - - assert_eq!( - request("attacker.example".into()).await, - "403", - "rebound hostname" - ); - assert_eq!( - request(format!("localhost:{port}")).await, - "403", - "even localhost: only the one advertised authority is served" - ); - assert_eq!( - request(authority.to_string()).await, - "200", - "the true authority passes" - ); - } - - // The session slot is shared and REPLACEABLE: it starts empty (signed - // out), a login fills it, and a re-login (fresh session id, per the II - // fresh-session-key contract) replaces it — through every clone and - // through the resolver handed to IcTools, since clones share the holder. - // This is what makes reauthentication after grant expiry a same-process - // step. - #[test] - fn the_session_slot_is_shared_and_replaceable() { - let slot = SessionSlot::new(); - let reader = slot.clone(); - assert_eq!(reader.get(), None, "a fresh slot is signed out"); - slot.set("sess-1".into()); - assert_eq!(reader.get(), Some("sess-1".into()), "clones share the slot"); - slot.set("sess-2".into()); - assert_eq!( - reader.get(), - Some("sess-2".into()), - "a re-login replaces the id" - ); - } - - // The handshake window is exactly HANDSHAKE_TTL: a pending flow older than - // that is expired (the watchdog and the redeem gate both read this). - #[test] - fn pending_expires_after_the_handshake_ttl() { - let pending = |age: Duration| Pending { - session_id: "s".into(), - url: "u".into(), - callback_url: "c".into(), - started: Instant::now() - age, - redeeming: false, - shutdown: Arc::new(Notify::new()), - }; - assert!(!pending(Duration::from_secs(1)).expired()); - assert!(pending(HANDSHAKE_TTL + Duration::from_secs(1)).expired()); - } -} diff --git a/crates/imcp2-local/src/login/mod.rs b/crates/imcp2-local/src/login/mod.rs new file mode 100644 index 0000000..bc93480 --- /dev/null +++ b/crates/imcp2-local/src/login/mod.rs @@ -0,0 +1,325 @@ +//! The **browser Internet Identity login** for the local binary: mint a fresh +//! session, stand up a transient loopback listener, hand the user an id.ai +//! sign-in link, and redeem the delegation II navigates back with. On success +//! the shared [`SessionSlot`] is filled and every tool call acts as the user; +//! the listener is torn down (redeem or timeout) and the grant lives in memory +//! only. +//! +//! The listener is the unavoidable minimum, not a design slip: II delivers the +//! delegation by *navigating the browser* to the callback (a URL fragment only +//! a served page can read and POST back), and II's #4091 check fetches +//! `/.well-known/ii-auth-callbacks` from the callback's origin before honoring +//! it — both require a real HTTP origin. This is the standard native-app +//! loopback redirect (RFC 8252 §7.3), the same shape as the ICP CLI's +//! `icp identity link web`. It serves exactly three routes for one handshake +//! and never serves the MCP tool surface, which rides stdio. +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use imcp2_core::identities::Identities; +use imcp2_core::iiconnect; +use tokio::sync::{Mutex, Notify}; + +mod routes; +mod slot; +#[cfg(test)] +mod tests; + +pub use slot::SessionSlot; + +/// How long one sign-in handshake may take, from the `authenticate` call to +/// the browser's redeem, before the pending flow and its listener are torn +/// down. Matches the hosted server's pending-connect TTL. +const HANDSHAKE_TTL: Duration = Duration::from_secs(600); + +/// `ttl` (seconds) requested for the II grant in the connect link, matching +/// the hosted server. II clamps it to [600, 2592000], and the effective +/// session duration is whatever the user picks on II's consent screen. +const GRANT_TTL_SECS: u64 = 3600; + +/// Where the pinned callback page's error state points users ("contact us to +/// report it") — DFINITY ships this binary, so failures report to the same +/// address as the hosted server's. +const CONTACT_EMAIL: &str = "mcp@dfinity.org"; + +fn now_ns() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() as u64 +} + +/// The login flow's shared handle: the `authenticate`/`auth_status` tools call +/// [`begin`](Self::begin)/[`status`](Self::status), and the loopback redeem +/// route completes the flow through a clone. Single-user by construction — one +/// pending handshake at a time, serialized by the state lock. +#[derive(Clone)] +pub struct LoginDriver { + inner: Arc, +} + +struct Inner { + identities: Identities, + /// The shared, replaceable session slot `IcTools` reads through + /// [`SessionSlot::resolver`]: each successful redeem writes the fresh + /// session id here (the II contract requires a fresh session key after + /// `Unauthorized`, and a fresh id gets fresh keys). + slot: SessionSlot, + /// Best-effort `open::that` on the sign-in link. The flow never depends on + /// it succeeding — the link is always returned in-band by `authenticate`. + auto_open: bool, + state: Mutex, + /// Serializes each redeem's `mcp_register_v2` **and** its local commit, + /// across flows. II's register REPLACES the anchor's previous grant (the + /// last registration to settle at II is the live one), so the local slot + /// must be committed in the same order the registrations settle: without + /// this, a timed-out flow's slow register could land at II *after* a + /// newer flow's — II would then honor the old session while the slot + /// named the new one, and every tool call would come back `Unauthorized`. + registration: Mutex<()>, +} + +#[derive(Default)] +struct LoginState { + pending: Option, + /// The most recent successful login, kept (even past expiry) so + /// `auth_status` can say *whose* session expired. + grant: Option, +} + +struct Pending { + /// Doubles as the connect `state` in the II link — II echoes it to the + /// callback, and the redeem accepts only this value. + session_id: String, + url: String, + /// This flow's loopback callback (`http://127.0.0.1:/callback`) — + /// the origin the listener serves on. + callback_url: String, + started: Instant, + /// Single-flight marker: `true` while a redemption attempt (the + /// `mcp_register_v2` network call) is mid-flight, so a page double-submit + /// can't fire two. Cleared on failure so a genuine retry can proceed. + redeeming: bool, + /// Tears the transient listener down (graceful): notified on redeem + /// success, on handshake timeout, and when a fresh flow replaces this one. + shutdown: Arc, +} + +impl Pending { + fn expired(&self) -> bool { + self.started.elapsed() >= HANDSHAKE_TTL + } +} + +/// A completed login, as `auth_status` reports it. +#[derive(Clone)] +pub struct Grant { + pub session_id: String, + /// The session principal (`self_authenticating(pub(S))`) — attribution, + /// not the user's per-app principal. + pub principal: Option, + /// II's recorded access level: `"queries"` (read-only) or `"all"`. + pub permissions: &'static str, + /// Grant expiration, ns since the Unix epoch (the user's consent-time + /// session choice). + pub expiration_ns: u64, +} + +impl Grant { + pub fn expired(&self) -> bool { + self.expiration_ns <= now_ns() + } + + /// Whole minutes until expiry (0 when expired). + pub fn minutes_left(&self) -> u64 { + self.expiration_ns.saturating_sub(now_ns()) / 60_000_000_000 + } +} + +/// What [`LoginDriver::status`] reports — the states `auth_status` renders. +pub enum LoginStatus { + SignedIn(Grant), + /// A handshake is waiting for the browser. + Pending { + url: String, + }, + /// A previous grant ran out (or was revoked and then expired here); a new + /// `authenticate` gets a fresh session. + Expired(Grant), + SignedOut, +} + +/// What [`LoginDriver::begin`] returns to the `authenticate` tool. +pub enum BeginOutcome { + /// A live grant exists and `refresh` was not requested. + AlreadySignedIn(Grant), + /// A sign-in link is ready. `fresh` distinguishes a newly started flow + /// from an already-pending one being returned again (single-flight). + Pending { url: String, fresh: bool }, +} + +impl LoginDriver { + pub fn new(identities: Identities, slot: SessionSlot, auto_open: bool) -> Self { + Self { + inner: Arc::new(Inner { + identities, + slot, + auto_open, + state: Mutex::new(LoginState::default()), + registration: Mutex::new(()), + }), + } + } + + /// Start (or return the already-pending) sign-in: mint a **fresh** session + /// id + registration key, bind the loopback listener, and build the id.ai + /// link — returned immediately, never blocking on the browser. With a live + /// grant and `refresh: false`, reports the grant instead of starting over. + pub async fn begin(&self, refresh: bool) -> Result { + // The lock is held across the (local, network-free) setup, so + // concurrent `authenticate` calls serialize into one flow. + let mut state = self.inner.state.lock().await; + + if let Some(p) = &state.pending { + if !p.expired() { + return Ok(BeginOutcome::Pending { + url: p.url.clone(), + fresh: false, + }); + } + // Expired while nobody was looking: tear it down, then start fresh. + if let Some(p) = state.pending.take() { + p.shutdown.notify_one(); + } + } + if !refresh { + if let Some(g) = &state.grant { + if !g.expired() { + return Ok(BeginOutcome::AlreadySignedIn(g.clone())); + } + } + } + + // A FRESH session id per flow: fresh id ⇒ fresh session + registration + // keys in `Identities`, per the II fresh-session-key contract. + let session_id = uuid::Uuid::new_v4().to_string(); + let reg_pubkey = self + .inner + .identities + .registration_pubkey_b64(&session_id) + .await?; + + // Port 0: the OS picks a free port; both the II link's `callback` and + // the #4091 allow-list entry derive from the ONE resulting origin, so + // they cannot drift (II matches by exact string equality). + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .map_err(|e| format!("could not bind the login callback listener: {e}"))?; + let port = listener + .local_addr() + .map_err(|e| format!("could not read the listener address: {e}"))? + .port(); + let authority = format!("127.0.0.1:{port}"); + let callback_url = format!("http://{authority}/callback"); + + let url = iiconnect::ii_mcp_url( + &self.inner.identities.instance().ii_url, + &callback_url, + &session_id, + GRANT_TTL_SECS, + ®_pubkey, + ); + + let shutdown = Arc::new(Notify::new()); + let app = routes::login_router(self.clone(), callback_url.clone(), authority); + let sd = shutdown.clone(); + tokio::spawn(async move { + if let Err(e) = axum::serve(listener, app) + .with_graceful_shutdown(async move { sd.notified().await }) + .await + { + tracing::warn!("login listener error: {e}"); + } + }); + + // Watchdog: when the handshake window closes, drop the pending flow + // and the listener (if a redeem didn't already). + let driver = self.clone(); + let sid = session_id.clone(); + tokio::spawn(async move { + tokio::time::sleep(HANDSHAKE_TTL).await; + let mut state = driver.inner.state.lock().await; + let ours = matches!(&state.pending, Some(p) if p.session_id == sid); + if ours { + if let Some(p) = state.pending.take() { + p.shutdown.notify_one(); + tracing::info!( + callback = %p.callback_url, + "sign-in window expired; login listener closed" + ); + } + } + }); + + state.pending = Some(Pending { + session_id, + url: url.clone(), + callback_url, + started: Instant::now(), + redeeming: false, + shutdown, + }); + drop(state); + + if self.inner.auto_open { + // Detached: the launcher runs with stdin/stdout/stderr nulled and + // (off macOS) double-forked away, so it can never touch this + // process's stdio — stdout is the MCP JSON-RPC channel — nor hold + // a thread. Best-effort either way: the link is in the tool + // result. (A rare launcher can still block briefly at spawn, so + // it stays off the async workers.) + let link = url.clone(); + tokio::task::spawn_blocking(move || { + if let Err(e) = open::that_detached(&link) { + tracing::debug!("could not auto-open the browser: {e}"); + } + }); + } + + Ok(BeginOutcome::Pending { url, fresh: true }) + } + + /// The pending flow's `(session_id, callback_url)`, for the tests that + /// play the browser against the live listener. + #[cfg(test)] + pub(crate) async fn pending_handshake(&self) -> Option<(String, String)> { + let state = self.inner.state.lock().await; + state + .pending + .as_ref() + .map(|p| (p.session_id.clone(), p.callback_url.clone())) + } + + pub async fn status(&self) -> LoginStatus { + let state = self.inner.state.lock().await; + // A live pending handshake outranks a live grant, matching `begin`'s + // precedence: after `authenticate(refresh: true)` the old session may + // still be valid, but the truthful answer is "a replacement sign-in is + // waiting for the browser", not "signed in" — otherwise the client + // reads the refresh as already complete. + if let Some(p) = &state.pending { + if !p.expired() { + return LoginStatus::Pending { url: p.url.clone() }; + } + } + if let Some(g) = &state.grant { + if !g.expired() { + return LoginStatus::SignedIn(g.clone()); + } + } + match &state.grant { + Some(g) => LoginStatus::Expired(g.clone()), + None => LoginStatus::SignedOut, + } + } +} diff --git a/crates/imcp2-local/src/login/routes.rs b/crates/imcp2-local/src/login/routes.rs new file mode 100644 index 0000000..91be4b7 --- /dev/null +++ b/crates/imcp2-local/src/login/routes.rs @@ -0,0 +1,265 @@ +//! The transient loopback listener's HTTP surface — the three login-handshake +//! routes (pinned callback page, slim redeem, the #4091 allow-list) plus the +//! Host-validation middleware. Up only for the duration of one handshake; the +//! MCP tool surface never rides HTTP. + +use axum::extract::State; +use axum::response::{Html, IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use imcp2_core::iiconnect::{self, RedeemBody, AUTH_CALLBACKS_WELL_KNOWN}; +use serde_json::json; + +use super::{Grant, LoginDriver, CONTACT_EMAIL}; + +/// The three transient loopback routes — the login handshake's whole HTTP +/// surface. `callback_url` is this listener's own callback, echoed verbatim in +/// the allow-list document; `authority` is its `127.0.0.1:` host, the +/// only `Host` header the listener accepts. +pub(super) fn login_router(driver: LoginDriver, callback_url: String, authority: String) -> Router { + Router::new() + .route("/callback", get(callback_page)) + .route("/redeem", post(redeem)) + .route(AUTH_CALLBACKS_WELL_KNOWN, get(auth_callbacks)) + // Anti-DNS-rebinding, per the design's loopback hardening (the local + // analogue of the hosted server's Host allow-list): a page on an + // attacker's hostname that rebinds to 127.0.0.1 reaches this port + // with `Host: attacker.example` — every URL this flow hands out + // carries exactly one authority, so anything else is rejected before + // routing. (Even without this the routes are largely inert to a + // rebinder — `/redeem` requires the connect `state` and a chain + // targeting the in-process `X` — this closes the door outright.) + .layer(axum::middleware::from_fn_with_state( + authority, + require_own_host, + )) + .with_state(RouteCtx { + driver, + callback_url, + }) +} + +/// Reject any request whose `Host` header is not this listener's own +/// `127.0.0.1:` authority (exact string match, like II's allow-list). +async fn require_own_host( + State(expected): State, + req: axum::extract::Request, + next: axum::middleware::Next, +) -> Response { + let ok = req + .headers() + .get(axum::http::header::HOST) + .and_then(|v| v.to_str().ok()) + .is_some_and(|host| host == expected); + if !ok { + return ( + axum::http::StatusCode::FORBIDDEN, + "wrong Host for this login listener", + ) + .into_response(); + } + next.run(req).await +} + +#[derive(Clone)] +struct RouteCtx { + driver: LoginDriver, + callback_url: String, +} + +/// GET /callback — the pinned fragment-reading page (shared with the hosted +/// server), pointed at this listener's `/redeem`, with the same non-CSP +/// hardening headers the hosted server adds. +async fn callback_page() -> Response { + let page = iiconnect::pinned_callback_page("/redeem", CONTACT_EMAIL); + let mut resp = Html(page.html).into_response(); + let h = resp.headers_mut(); + h.insert( + axum::http::header::CONTENT_SECURITY_POLICY, + axum::http::HeaderValue::from_str(&page.csp).expect("valid CSP"), + ); + h.insert( + axum::http::header::REFERRER_POLICY, + axum::http::HeaderValue::from_static("no-referrer"), + ); + h.insert( + axum::http::header::X_CONTENT_TYPE_OPTIONS, + axum::http::HeaderValue::from_static("nosniff"), + ); + h.insert( + axum::http::header::X_FRAME_OPTIONS, + axum::http::HeaderValue::from_static("DENY"), + ); + resp +} + +/// GET /.well-known/ii-auth-callbacks — the #4091 allow-list: exactly this +/// listener's callback. II fetches it cross-origin (hence the one CORS +/// header) and fail-closed before honoring the callback; `no-store` keeps an +/// intermediary from serving a stale document for a past listener's port. +async fn auth_callbacks(State(ctx): State) -> Response { + let mut resp = Json(json!({ "callbacks": [ctx.callback_url] })).into_response(); + let h = resp.headers_mut(); + h.insert( + axum::http::header::CACHE_CONTROL, + axum::http::HeaderValue::from_static("no-store"), + ); + h.insert( + axum::http::header::ACCESS_CONTROL_ALLOW_ORIGIN, + axum::http::HeaderValue::from_static("*"), + ); + resp +} + +fn redeem_err(msg: &str) -> Response { + ( + axum::http::StatusCode::BAD_REQUEST, + Json(json!({ "error": msg })), + ) + .into_response() +} + +/// The local success answer: no OAuth continuation to redirect into, so the +/// pinned page renders its terminal "Signed in — you can close this tab" +/// state in place (the page script's `done` arm). +fn done() -> Response { + Json(json!({ "done": true })).into_response() +} + +/// POST /redeem — the pinned page POSTs the fragment here. The hosted redeem +/// minus its OAuth tail: no initiator cookie (there is no public initiate +/// endpoint to confuse a deputy through — the binary itself mints `X` and this +/// listener is loopback-only), no PKCE code, and no redirect — success answers +/// `{"done": true}` and the page renders its terminal state. Parse before +/// claiming, single-flight claim, one `mcp_register_v2` call, then fill the +/// session slot and tear the listener down. +async fn redeem(State(ctx): State, Json(body): Json) -> Response { + let driver = &ctx.driver; + + // Validate the delivery against the pending flow first (NOT claiming yet), + // so an unknown/expired `state` gets its accurate refusal without parsing. + { + let state = driver.inner.state.lock().await; + if state + .grant + .as_ref() + .is_some_and(|g| g.session_id == body.state) + { + // A retry of an already-successful redeem (II's delegation redeems + // repeatedly within its lifetime): idempotent success. + return done(); + } + match &state.pending { + Some(p) if p.session_id == body.state => { + if p.expired() { + return redeem_err( + "This sign-in expired. Ask the agent to authenticate again.", + ); + } + } + _ => { + return redeem_err( + "This sign-in is unknown or already used. Ask the agent to authenticate again.", + ) + } + } + } + + // Decode the fragment delegation (pure, size-bounded) BEFORE claiming, so + // a malformed delivery never occupies the single-flight slot. + let (user_key, chain) = match iiconnect::parse_registration_delegation(&body.delegation) { + Ok(v) => v, + Err(e) => { + return redeem_err(&format!( + "We couldn't read the sign-in response. Ask the agent to authenticate again. ({e})" + )) + } + }; + + // Claim atomically (re-checking: the flow may have moved while parsing). + { + let mut state = driver.inner.state.lock().await; + if state + .grant + .as_ref() + .is_some_and(|g| g.session_id == body.state) + { + return done(); + } + match state.pending.as_mut() { + Some(p) if p.session_id == body.state => { + if p.expired() { + return redeem_err( + "This sign-in expired. Ask the agent to authenticate again.", + ); + } + if p.redeeming { + return redeem_err( + "This sign-in is already being processed. Wait a moment; \ + if nothing happens, ask the agent to authenticate again.", + ); + } + p.redeeming = true; + } + _ => { + return redeem_err( + "This sign-in is unknown or already used. Ask the agent to authenticate again.", + ) + } + } + } + + // The network call + local commit, SERIALIZED across flows (see + // `Inner::registration`): II's `mcp_register_v2` replaces the anchor's + // previous grant, so whichever registration settles at II last is the + // live one — and because every register-then-commit runs alone inside + // this lock, the commit order equals the settle order, and the local + // slot always names the session II actually honors. In particular a + // flow that timed out locally while its register was in flight either + // settles BEFORE a replacement (which then overwrites it, at II and + // here) or AFTER it (and then the late registration is the live grant, + // and committing it here is exactly right). The per-flow single-flight + // claim above stays: it stops a page double-submit; this lock orders + // DISTINCT flows. + let _registration = driver.inner.registration.lock().await; + match driver + .inner + .identities + .redeem_registration_delegation(&body.state, user_key, chain) + .await + { + Err(e) => { + let mut state = driver.inner.state.lock().await; + if let Some(p) = state.pending.as_mut() { + if p.session_id == body.state { + p.redeeming = false; // free the claim for a genuine retry + } + } + redeem_err(&e) + } + Ok(outcome) => { + let principal = driver.inner.identities.session_principal(&body.state).await; + let mut state = driver.inner.state.lock().await; + driver.inner.slot.set(body.state.clone()); + state.grant = Some(Grant { + session_id: body.state.clone(), + principal, + permissions: outcome.permissions, + expiration_ns: outcome.expiration_ns, + }); + let ours = matches!(&state.pending, Some(p) if p.session_id == body.state); + if ours { + if let Some(p) = state.pending.take() { + p.shutdown.notify_one(); + } + } + drop(state); + tracing::info!( + permissions = outcome.permissions, + expiration_ns = outcome.expiration_ns, + "signed in: registration delegation redeemed; login listener closing" + ); + done() + } + } +} diff --git a/crates/imcp2-local/src/login/slot.rs b/crates/imcp2-local/src/login/slot.rs new file mode 100644 index 0000000..ddcfb16 --- /dev/null +++ b/crates/imcp2-local/src/login/slot.rs @@ -0,0 +1,44 @@ +//! The [`SessionSlot`] — the local binary's whole authentication state: one +//! shared, replaceable holder for the signed-in session id, read by core +//! through the injected resolver and written by the login flow. + +use std::sync::Arc; + +/// The single-user session holder — a shared, **replaceable** slot rather +/// than a fixed id. The Internet Identity contract requires a FRESH session +/// key after an `Unauthorized` (a reused id would reuse its keys), so each +/// successful login writes a fresh session id here — same-process +/// reauthentication after grant expiry or revocation, with no client restart. +/// Clones share the slot: the login flow holds a writing handle, and the slot +/// answers `IcTools`' session lookups through [`SessionSlot::resolver`] — the +/// local binary's whole authentication layer, living here rather than in +/// core, which only asks the injected resolver. Starts empty (signed out). +#[derive(Clone, Debug, Default)] +pub struct SessionSlot(Arc>>); + +impl SessionSlot { + /// A fresh, empty slot (no session — every session-needing tool reports + /// "needs an authenticated session" until the login flow fills it). + pub fn new() -> Self { + Self::default() + } + + /// The current session id, if a login has completed. The critical section + /// is a clone — never held across an await. + pub fn get(&self) -> Option { + self.0.read().expect("session slot poisoned").clone() + } + + /// Replace the session id after a successful login (a fresh id per the II + /// fresh-session-key contract). + pub fn set(&self, session_id: String) { + *self.0.write().expect("session slot poisoned") = Some(session_id); + } + + /// This slot as the [`imcp2_core::SessionResolver`] the local binary + /// injects into `IcTools`: every call acts as the one signed-in session. + pub fn resolver(&self) -> imcp2_core::SessionResolver { + let slot = self.clone(); + Arc::new(move |_ctx| slot.get()) + } +} diff --git a/crates/imcp2-local/src/login/tests.rs b/crates/imcp2-local/src/login/tests.rs new file mode 100644 index 0000000..0e9eff7 --- /dev/null +++ b/crates/imcp2-local/src/login/tests.rs @@ -0,0 +1,365 @@ +//! Offline tests of the login flow: the driver's single-flight and status +//! lifecycle, and the listener driven like the browser (and like an +//! attacker) over real HTTP. + +use super::*; +use imcp2_core::iiconnect::AUTH_CALLBACKS_WELL_KNOWN; +use imcp2_core::IiInstance; + +/// A driver against the real mainnet/prod-II config — nothing here touches +/// the network (agent construction, key minting, and the loopback listener +/// are all local), which is exactly what these tests rely on. +fn test_driver() -> (LoginDriver, SessionSlot) { + let agent = imcp2_core::Agent::builder() + .with_url(imcp2_core::IC_URL) + .build() + .expect("agent"); + let identities = Identities::new( + IiInstance::prod().expect("prod II instance"), + "https://mcp.internetcomputer.org".into(), + agent, + ); + let slot = SessionSlot::new(); + ( + LoginDriver::new(identities, slot.clone(), /* auto_open */ false), + slot, + ) +} + +/// A shape-valid two-hop chain (the same wire shape as iiconnect's parser +/// tests) whose signatures are garbage: it must get PAST parsing and be +/// rejected by the redeem path itself — offline, before any network call. +fn fake_chain_json() -> String { + serde_json::json!({ + "delegations": [ + { + "delegation": { "pubkey": "070707", "expiration": "66", "targets": ["aaaaa-aa"] }, + "signature": "040506", + }, + { + "delegation": { "pubkey": "09080706", "expiration": "66" }, + "signature": "010909", + }, + ], + "publicKey": "010203", + }) + .to_string() +} + +// One `authenticate` = one flow: a fresh id.ai link over a fresh loopback +// callback; a second call while it is pending returns the SAME link +// (single-flight) rather than racing a second listener. The session slot +// stays empty until a redeem succeeds — starting a login signs nobody in. +#[tokio::test] +async fn begin_mints_one_flow_and_repeats_it_while_pending() { + let (driver, slot) = test_driver(); + let BeginOutcome::Pending { url, fresh } = driver.begin(false).await.expect("begin") else { + panic!("a fresh driver must start a flow, not report a session") + }; + assert!(fresh); + assert!(url.starts_with("https://id.ai/mcp#callback="), "{url}"); + assert!( + url.contains("http%3A%2F%2F127.0.0.1%3A"), + "loopback callback in the fragment: {url}" + ); + assert!(url.contains("&ttl=3600®istration_key="), "{url}"); + + let BeginOutcome::Pending { url: again, fresh } = driver.begin(false).await.expect("begin") + else { + panic!("a pending flow must be returned, not replaced") + }; + assert!(!fresh, "the second call must join the pending flow"); + assert_eq!(again, url, "same link while the handshake is pending"); + + assert_eq!(slot.get(), None, "no session until a redeem succeeds"); + assert!(matches!(driver.status().await, LoginStatus::Pending { .. })); +} + +// The transient listener's whole surface, driven like the browser: the +// pinned page (strict CSP + the hosted server's hardening headers, script +// pointed at /redeem), and the #4091 allow-list (this callback verbatim, +// CORS-readable, never cached). +#[tokio::test] +async fn the_listener_serves_the_pinned_page_and_the_allow_list() { + let (driver, _slot) = test_driver(); + driver.begin(false).await.expect("begin"); + let (_, callback_url) = driver.pending_handshake().await.expect("pending"); + let origin = callback_url.strip_suffix("/callback").unwrap().to_string(); + let http = reqwest::Client::new(); + + let page = http.get(&callback_url).send().await.expect("GET /callback"); + assert_eq!(page.status(), 200); + let csp = page + .headers() + .get("content-security-policy") + .expect("pinned page ships a CSP") + .to_str() + .unwrap() + .to_string(); + assert!(csp.contains("default-src 'none'"), "{csp}"); + for (name, want) in [ + ("referrer-policy", "no-referrer"), + ("x-content-type-options", "nosniff"), + ("x-frame-options", "DENY"), + ] { + assert_eq!( + page.headers().get(name).and_then(|v| v.to_str().ok()), + Some(want) + ); + } + let html = page.text().await.unwrap(); + assert!( + html.contains("/redeem"), + "the script must POST to this listener's redeem" + ); + assert!( + html.contains("data.done"), + "the local success arm must be in the shipped page" + ); + + let wk = http + .get(format!("{origin}{AUTH_CALLBACKS_WELL_KNOWN}")) + .send() + .await + .expect("GET allow-list"); + assert_eq!(wk.status(), 200); + assert_eq!( + wk.headers() + .get("access-control-allow-origin") + .and_then(|v| v.to_str().ok()), + Some("*"), + "II fetches the allow-list cross-origin" + ); + assert_eq!( + wk.headers() + .get("cache-control") + .and_then(|v| v.to_str().ok()), + Some("no-store"), + "fail-closed infrastructure must never be served stale" + ); + let doc: serde_json::Value = wk.json().await.unwrap(); + assert_eq!( + doc, + serde_json::json!({ "callbacks": [callback_url] }), + "the declared callback must equal the link's callback VERBATIM (exact-match)" + ); +} + +// The redeem's refusal paths, in order: an unknown `state` (or one from a +// replaced flow), a delegation the parser rejects, and a shape-valid chain +// with garbage signatures that the redeem itself rejects — all 400 with an +// {"error"} body the pinned page renders, all offline, and none of them +// consumes the pending flow (a genuine retry stays possible). +#[tokio::test] +async fn redeem_refuses_bad_deliveries_and_keeps_the_flow_retryable() { + let (driver, slot) = test_driver(); + let BeginOutcome::Pending { url, .. } = driver.begin(false).await.expect("begin") else { + panic!("fresh flow") + }; + let (state, callback_url) = driver.pending_handshake().await.expect("pending"); + let origin = callback_url.strip_suffix("/callback").unwrap().to_string(); + let redeem_url = format!("{origin}/redeem"); + let http = reqwest::Client::new(); + + let post = |body: serde_json::Value| { + let http = http.clone(); + let redeem_url = redeem_url.clone(); + async move { + http.post(&redeem_url) + .json(&body) + .send() + .await + .expect("POST /redeem") + } + }; + + let r = post(serde_json::json!({ "state": "someone-else", "delegation": "" })).await; + assert_eq!(r.status(), 400); + let e: serde_json::Value = r.json().await.unwrap(); + assert!( + e["error"] + .as_str() + .unwrap() + .contains("unknown or already used"), + "{e}" + ); + + let r = post(serde_json::json!({ "state": state, "delegation": "not json" })).await; + assert_eq!(r.status(), 400); + let e: serde_json::Value = r.json().await.unwrap(); + assert!( + e["error"] + .as_str() + .unwrap() + .contains("couldn't read the sign-in response"), + "{e}" + ); + + let r = post(serde_json::json!({ "state": state, "delegation": fake_chain_json() })).await; + assert_eq!( + r.status(), + 400, + "garbage signatures must be rejected by the redeem" + ); + let e: serde_json::Value = r.json().await.unwrap(); + assert!(e["error"].is_string(), "{e}"); + + assert_eq!(slot.get(), None, "no failure may fill the session slot"); + let BeginOutcome::Pending { url: still, fresh } = driver.begin(false).await.expect("begin") + else { + panic!("the flow must survive failed redeems") + }; + assert!(!fresh, "failed redeems must not consume the pending flow"); + assert_eq!(still, url); +} + +// The status lifecycle around a grant (injected directly — a real one +// needs live II): live grant → SignedIn with the wallclock math; past +// expiration → Expired, still naming the principal; a fresh driver is +// simply SignedOut. +#[tokio::test] +async fn status_reports_the_grant_lifecycle() { + let (driver, _slot) = test_driver(); + assert!(matches!(driver.status().await, LoginStatus::SignedOut)); + + let grant = |expiration_ns: u64| Grant { + session_id: "sess".into(), + principal: Some("aaaaa-aa".into()), + permissions: "queries", + expiration_ns, + }; + driver.inner.state.lock().await.grant = Some(grant(now_ns() + 30 * 60_000_000_000)); + match driver.status().await { + LoginStatus::SignedIn(g) => { + assert!( + (25..=30).contains(&g.minutes_left()), + "{}", + g.minutes_left() + ); + } + _ => panic!("a live grant must report SignedIn"), + } + + driver.inner.state.lock().await.grant = Some(grant(now_ns() - 1)); + match driver.status().await { + LoginStatus::Expired(g) => assert_eq!(g.principal.as_deref(), Some("aaaaa-aa")), + _ => panic!("a past-expiry grant must report Expired"), + } +} + +// `authenticate(refresh: true)` while a live grant exists: the truthful +// status is the PENDING replacement handshake, not the old "signed in" — +// otherwise the client reads the refresh as already complete (the same +// precedence `begin` applies). +#[tokio::test] +async fn a_pending_refresh_outranks_the_live_grant_in_status() { + let (driver, _slot) = test_driver(); + driver.inner.state.lock().await.grant = Some(Grant { + session_id: "old".into(), + principal: Some("aaaaa-aa".into()), + permissions: "all", + expiration_ns: now_ns() + 30 * 60_000_000_000, + }); + assert!(matches!(driver.status().await, LoginStatus::SignedIn(_))); + + let BeginOutcome::Pending { fresh, .. } = driver.begin(true).await.expect("refresh") else { + panic!("refresh=true must start a replacement flow past a live grant") + }; + assert!(fresh); + assert!( + matches!(driver.status().await, LoginStatus::Pending { .. }), + "a live replacement handshake must outrank the old grant" + ); +} + +// Anti-DNS-rebinding (the design's loopback hardening): a request whose +// `Host` names anything but this listener's own `127.0.0.1:` is +// rejected before routing — a rebinding page fetches this port with the +// attacker's hostname in `Host`, which is exactly what must bounce. The +// true authority (what every handed-out URL carries) passes untouched. +#[tokio::test] +async fn the_listener_rejects_foreign_host_headers() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let (driver, _slot) = test_driver(); + driver.begin(false).await.expect("begin"); + let (_, callback_url) = driver.pending_handshake().await.expect("pending"); + let authority = callback_url + .strip_prefix("http://") + .and_then(|s| s.strip_suffix("/callback")) + .expect("authority"); + let port: u16 = authority.split(':').nth(1).unwrap().parse().unwrap(); + + let request = |host: String| async move { + let mut stream = tokio::net::TcpStream::connect(("127.0.0.1", port)) + .await + .expect("connect"); + stream + .write_all( + format!("GET /callback HTTP/1.1\r\nHost: {host}\r\nConnection: close\r\n\r\n") + .as_bytes(), + ) + .await + .expect("write"); + let mut response = String::new(); + stream.read_to_string(&mut response).await.expect("read"); + response + .split_whitespace() + .nth(1) + .expect("status code") + .to_string() + }; + + assert_eq!( + request("attacker.example".into()).await, + "403", + "rebound hostname" + ); + assert_eq!( + request(format!("localhost:{port}")).await, + "403", + "even localhost: only the one advertised authority is served" + ); + assert_eq!( + request(authority.to_string()).await, + "200", + "the true authority passes" + ); +} + +// The session slot is shared and REPLACEABLE: it starts empty (signed +// out), a login fills it, and a re-login (fresh session id, per the II +// fresh-session-key contract) replaces it — through every clone and +// through the resolver handed to IcTools, since clones share the holder. +// This is what makes reauthentication after grant expiry a same-process +// step. +#[test] +fn the_session_slot_is_shared_and_replaceable() { + let slot = SessionSlot::new(); + let reader = slot.clone(); + assert_eq!(reader.get(), None, "a fresh slot is signed out"); + slot.set("sess-1".into()); + assert_eq!(reader.get(), Some("sess-1".into()), "clones share the slot"); + slot.set("sess-2".into()); + assert_eq!( + reader.get(), + Some("sess-2".into()), + "a re-login replaces the id" + ); +} + +// The handshake window is exactly HANDSHAKE_TTL: a pending flow older than +// that is expired (the watchdog and the redeem gate both read this). +#[test] +fn pending_expires_after_the_handshake_ttl() { + let pending = |age: Duration| Pending { + session_id: "s".into(), + url: "u".into(), + callback_url: "c".into(), + started: Instant::now() - age, + redeeming: false, + shutdown: Arc::new(Notify::new()), + }; + assert!(!pending(Duration::from_secs(1)).expired()); + assert!(pending(HANDSHAKE_TTL + Duration::from_secs(1)).expired()); +} From f80948b455c8900fc2a306952db9702667f5d4fd Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 10:05:06 +0000 Subject: [PATCH 08/10] Refine the design's tool/session seam: authentication injected, core blind to it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Component 7 now says what the reviewed implementation does: IcTools asks one injected SessionResolver "which already-validated session is this call?" and knows nothing else — not the transport, not how callers are validated, not whether the deployment is single- or multi-user. Each binary owns its authentication end to end in its own package (the hosted bearer middleware + AuthedSession in imcp2's auth.rs; the local login flow's SessionSlot in imcp2-local), and the doc spells out why the seam is a value-returning callback rather than a middleware slot in core: real authentication middleware needs its transport, so the middlewares live with their transports and core receives only their outcome. Core's dependency profile loses `http` accordingly, and the concise doc's component 7 matches. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- docs/local-deployment.md | 11 ++++-- docs/scoping-local-deployment.md | 64 ++++++++++++++++++-------------- 2 files changed, 43 insertions(+), 32 deletions(-) diff --git a/docs/local-deployment.md b/docs/local-deployment.md index c9b4738..9ac8151 100644 --- a/docs/local-deployment.md +++ b/docs/local-deployment.md @@ -74,10 +74,13 @@ machinery is needed. 6. **AI tool clients** — Claude Desktop, Claude Code, Codex, Cursor, Antigravity, and the Perplexity macOS app run `imcp2-local` directly; cloud surfaces keep using hosted. 7. **Session handling** — every tool call acts as the user's Internet Identity session, so - each server must answer "which session?" The hosted server serves many users and resolves - it from each request's OAuth bearer token; the local server has exactly one user, so its - tools read the single in-memory session established at login. The tool implementations - are identical in both. + the tool layer needs one answer per call: "which already-validated session is this?" + Authentication itself stays in each binary, injected into the shared tools as a resolver + for exactly that question: the hosted server's bearer middleware validates each request + and its resolver reports the result; the local server's login flow keeps the one + signed-in session, and its resolver reads it. The shared core knows neither transports + nor whether a deployment is single- or multi-user; the tool implementations are + identical in both. 8. **Security model** — whoever drives the binary acts as the user's real II accounts on mainnet: treat the binary and its client config like a wallet. 9. **End-user setup** — no config editing anywhere: double-click the plugin bundle for diff --git a/docs/scoping-local-deployment.md b/docs/scoping-local-deployment.md index af2d93a..3e6cc5f 100644 --- a/docs/scoping-local-deployment.md +++ b/docs/scoping-local-deployment.md @@ -110,8 +110,8 @@ minimal. `#[tool_handler] impl ServerHandler for IcTools` macros expand into one impl on `IcTools` and **must stay co-located**; both binaries construct `IcTools`, differing in transport + session source (the local binary additionally wraps it in its login - handler, component 6). The `SessionSource` seam and the plain `AuthedSession` struct - (component 7) live here. + handler, component 6). The injected `SessionResolver` seam (component 7) lives here — + authentication itself stays in the binaries. - **new `iiconnect` module** — the II connect-handshake primitives lifted out of `auth.rs` (component 4), re-parameterised to plain values so they carry no `AuthStore`/OAuth state. @@ -119,13 +119,14 @@ minimal. references) and the connect-flow assets under `src/assets/` (page, CSS, logo). - The move is not visibility-neutral: workspace packages share no crate-private access, so the items the binaries compose against — `IcTools`/`IcTools::new` with - `SessionSource`, the `Identities` registration/redeem entry points above, the + its `SessionResolver`, the `Identities` registration/redeem entry points above, the `iiconnect` builders — graduate from `pub(crate)` behind one binary to core's deliberate **public API** (`imcp2` re-exports it below, so embedders see no change). - Deps: rmcp `["server", "macros"]` (no transport), `ic-agent`/`candid`/`candid_parser`, tokio, serde/serde_json, reqwest/url/regex (discovery), the II-connect and management utilities (`base64`, `hex`, `sha2`, `crc32fast`, `getrandom`, `uuid`, `urlencoding`), - `http` (the request-extension type the Bearer session arm reads), tracing. + tracing — and nothing HTTP- or transport-shaped: authentication is injected + (component 7). - **`imcp2`** (existing crate: lib + the hosted `imcp2` binary) — depends on `imcp2-core` and adds the hosted composition: the OAuth 2.1 AS (the OAuth half of `auth.rs`), `McpServer`/`McpConfig`/the routers (`lib.rs`, including `state_dir`/`clients`/ @@ -189,8 +190,7 @@ dependency. - canister management: `sha2` (Wasm hash + ledger AccountIdentifier), `crc32fast`, `base64`, `hex` (`management.rs`). - II connect/delegation: `getrandom` (key seeds + CSP nonce), `urlencoding` (II link), - `uuid` (connect `state`), `base64`/`hex` (delegation chain), `http` (the request-extension - type the session lookup reads). + `uuid` (connect `state`), `base64`/`hex` (delegation chain). - foundation: `tokio`, `serde`, `serde_json`, `tracing`, `anyhow`. **Added by `imcp2` (the hosted stack):** rmcp's `transport-streamable-http-server`; `axum`; @@ -434,27 +434,35 @@ local binary. ### 7. Tool / session seam -Today a tool gets its `session_id` via bearer → `require_token` → `AuthedSession` injected in -the request extensions → `authed_session(ctx)` (`tools.rs:1493`). Under stdio there is one -user and one connection, so this collapses to a **singleton** session id, set at login and -replaced on each re-login. - -Minimal, verified seam (no tool-signature changes, no `identities.rs` changes): -- add `session: SessionSource { Bearer, Singleton(..) }` to `IcTools` (`tools.rs:52`). The - `Singleton` arm holds a **shared, replaceable slot** (an `Arc>>`- - style holder whose exact shape lands with the Stage 2 login driver), not a fixed id: the - II contract requires a **fresh session key** after `Unauthorized`, so each successful - `authenticate` writes a fresh session id into the slot — same-process reauthentication - after grant expiry or revocation, with no client restart; -- one `current_session_id(&self, &ctx) -> Option` method replacing the free fn; -- rewrite the **13** lookup call-sites (`tools.rs:351,809,865,1298,1321,1344,1365,1386,1407, - 1428,1446,1464,1482`) to call it; the `.ok_or(…)` handling is unchanged; -- move the plain `AuthedSession` struct into `imcp2-core` beside the seam; the hosted - `require_token` middleware (in `imcp2`) keeps inserting it into the request extensions. - -The seam involves no conditional compilation: the Bearer arm reads `http::request::Parts` -(the `http` crate — axum merely re-exports it), which core depends on directly. Hosted -constructs `IcTools` with `Bearer`; local with `Singleton`, whose slot the login flow fills. +Every authenticated tool call acts as one Internet Identity session, so the tool layer +needs an answer to exactly one question per call: *which already-validated session is +this?* **Everything else about authentication — the transport it rides, how a caller is +validated, whether the deployment serves many users or one — stays out of `imcp2-core`.** +`IcTools` takes an injected `SessionResolver` (a `Fn(&RequestContext) -> Option`) +at construction and asks it through one internal `current_session_id` helper at the 13 +session-using call sites; `None` means unauthenticated, and the existing "needs an +authenticated session" errors are the caller-facing surface of that. Core carries no HTTP +types and no transport features for this — it cannot tell bearer from singleton, or +streamable-HTTP from stdio. + +Each binary then owns its authentication end to end, in its own package: + +- **Hosted (`imcp2`)** — unchanged mechanics: the `require_token` middleware validates the + bearer token per request and stashes the plain `AuthedSession { session_id }` on the + request; the resolver it injects reads that back out of the tool context (rmcp surfaces + the HTTP request's parts there). Both the struct and the resolver live in `imcp2`'s + `auth.rs`, beside the middleware that gives them meaning. +- **Local (`imcp2-local`)** — the login flow (component 5) owns a `SessionSlot`: a shared, + **replaceable** holder for the one signed-in session id, filled on each successful + redeem. Replaceable, not fixed: the II contract requires a **fresh session key** after + `Unauthorized`, so each re-login writes a fresh session id — same-process + reauthentication after grant expiry or revocation, with no client restart. The resolver + it injects simply reads the slot. + +The seam is deliberately just a value-returning callback rather than a middleware slot in +core: a real authentication middleware needs its transport (the hosted validation is HTTP +middleware; a stdio pipe has no per-request middleware at all), so the middlewares live +with their transports and core only receives their outcome. Tools that are already session-free work unchanged locally: `get_canister_candid`, `get_canister_api_doc`, `open_app`, `resolve_app`, `discover_app_canisters`, the skills/lookup @@ -573,7 +581,7 @@ building and signing these artifacts is a Stage 3 deliverable. `identities/calls/discover/management/skills/tools` + the new `iiconnect` (extracted from `auth.rs`) + `static/` + the connect assets into the new core crate; `imcp2` keeps the OAuth AS, `McpServer`/`McpConfig`/`main.rs`, `metrics`, and the `e2e` harness, depends on core, and -re-exports its public items so embedders keep compiling. Apply the `SessionSource` seam +re-exports its public items so embedders keep compiling. Apply the `SessionResolver` seam (component 7). *Exit:* the `imcp2` binary builds and its tests pass unchanged; `imcp2-core` compiles standalone with no `axum`, no `prometheus`, and no OAuth code in its graph — `tower-http` remains only as rmcp/reqwest's transitive internal, never as the direct CORS From 84f7a27efe16ed202e00d173afebca246481ee8a Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 10:12:25 +0000 Subject: [PATCH 09/10] Address review: session reaper + the PocketIC login e2e moves into this PR - The local binary now runs the same 60-second session reaper the hosted server does: expired grants, abandoned sign-in keys, and their cached delegations leave memory promptly instead of accumulating over a long-lived process (below MAX_PENDING_CONNECTS nothing else scans the map). - The register -> slot/grant commit -> listener-shutdown success path now has an in-repository test ON THIS PR: the PocketIC end-to-end login (a live II canister issues a genuine canister-signed registration delegation; the test plays the browser over real HTTP against the real listener) moves here from the follow-up branch, gated exactly like the hosted crate's harness (cargo feature `e2e` + II_WASM/POCKET_IC_BIN at runtime). Run here and passing. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- Cargo.lock | 4 + crates/imcp2-local/Cargo.toml | 17 + crates/imcp2-local/src/e2e_local_login.rs | 518 ++++++++++++++++++++++ crates/imcp2-local/src/main.rs | 21 + 4 files changed, 560 insertions(+) create mode 100644 crates/imcp2-local/src/e2e_local_login.rs diff --git a/Cargo.lock b/Cargo.lock index 265051f..2b2d537 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1828,8 +1828,12 @@ version = "0.2.0" dependencies = [ "anyhow", "axum", + "base64 0.22.1", + "candid", + "hex", "imcp2-core", "open", + "pocket-ic", "reqwest 0.13.4", "rmcp", "serde", diff --git a/crates/imcp2-local/Cargo.toml b/crates/imcp2-local/Cargo.toml index ae0e975..d64e33e 100644 --- a/crates/imcp2-local/Cargo.toml +++ b/crates/imcp2-local/Cargo.toml @@ -34,6 +34,18 @@ tracing = { workspace = true } tracing-subscriber = { workspace = true } uuid = { workspace = true } url = { workspace = true } +# Only for the PocketIC-backed end-to-end login test (src/e2e_local_login.rs), +# behind the `e2e` feature so the default build never compiles it — the same +# gating as the hosted crate's harness. (An optional dependency, not a +# dev-dependency: cargo has no optional dev-dependencies.) +pocket-ic = { version = "15", optional = true } + +[features] +# Compiles + enables the real II ↔ local-login e2e against a live Internet +# Identity canister in PocketIC (needs II_WASM + POCKET_IC_BIN at runtime; see +# src/e2e_local_login.rs). Off by default: the standard `cargo test` neither +# builds `pocket-ic` nor touches the network. +e2e = ["dep:pocket-ic"] [dev-dependencies] # The login-flow tests play the browser: fetch the callback page, the @@ -43,3 +55,8 @@ reqwest = { workspace = true, features = ["json"] } # in-process duplex pipe (the same shape as stdio), driven by rmcp's client. rmcp = { workspace = true, features = ["client"] } tokio = { workspace = true, features = ["io-util"] } +# The e2e login test speaks the II canister's candid interface directly +# (playing II's frontend) and (de)codes the delegation-chain wire fields. +candid = { workspace = true } +hex = { workspace = true } +base64 = { workspace = true } diff --git a/crates/imcp2-local/src/e2e_local_login.rs b/crates/imcp2-local/src/e2e_local_login.rs new file mode 100644 index 0000000..c7dc6a2 --- /dev/null +++ b/crates/imcp2-local/src/e2e_local_login.rs @@ -0,0 +1,518 @@ +//! Real Internet-Identity ↔ **local login** test, end to end against a live +//! Internet Identity canister running in PocketIC — the local-replica test +//! configuration of the design's component 3, extending the hosted crate's +//! e2e harness (`imcp2/src/e2e_handshake.rs`) to this binary's browser +//! handshake. Nothing is mocked: II runs its production wasm and issues a +//! genuine canister-signed registration delegation; the login driver binds +//! its real loopback listener; this test plays the browser over **real HTTP** +//! (the pinned callback page, II's #4091 allow-list fetch, the redeem POST); +//! and the driver redeems via `mcp_register_v2` on the live II, verifying the +//! canister signature against PocketIC's fetched root key. +//! +//! ## What it exercises for real +//! `authenticate` (a fresh session + the id.ai-shaped link + the transient +//! listener) → the pinned page and the allow-list served over the wire → the +//! II registration-delegation ceremony (`prepare_mcp_registration_delegation` +//! + `get_mcp_registration_delegation`, playing II's frontend) → the **real** +//! delegation POSTed to the listener's `/redeem` → `{"done": true}` → the +//! session slot filled, `auth_status` signed in, and the listener torn down. +//! Only the human's browser clicks are stood in for. +//! +//! ## Gating (default `cargo test` is untouched and offline) +//! Behind the `e2e` cargo feature (so `pocket-ic` is neither compiled nor +//! downloaded by the default build) AND a runtime guard that skips unless +//! both artifacts — which cargo does not fetch — are provided: +//! +//! ```text +//! II_WASM=/abs/internet_identity_backend.wasm.gz \ # release asset, gz as-is +//! POCKET_IC_BIN=/abs/pocket-ic \ # pocket-ic v15 server +//! cargo test -p imcp2-local --features e2e -- --nocapture +//! ``` +//! +//! The candid types below are hand-rolled from `internet_identity.did` (the +//! same subset, at the same tag, as the hosted harness) so the crate stays +//! self-contained. + +// The hand-rolled candid types mirror II's `.did` in full for wire fidelity, +// so some variants/fields are declared but never constructed on our send path. +#![allow(dead_code)] + +use candid::{CandidType, Decode, Encode, Principal}; +use serde::Deserialize; +use std::time::SystemTime; + +use crate::login::{BeginOutcome, LoginDriver, LoginStatus, SessionSlot}; +use imcp2_core::identities::Identities; +use imcp2_core::iiconnect::AUTH_CALLBACKS_WELL_KNOWN; +use imcp2_core::IiInstance; + +// ---- II candid interface (subset), as in the hosted harness ---------------- + +/// `type Permissions = variant { queries; all }`. +#[derive(CandidType, Deserialize, Clone, Copy)] +enum Permissions { + #[serde(rename = "queries")] + Queries, + #[serde(rename = "all")] + All, +} + +/// `type McpConfig = record { enabled : bool; url : opt text }` — II's record +/// of the user's trusted MCP server, a precondition for `prepare_…`. +#[derive(CandidType)] +struct McpConfig { + enabled: bool, + url: Option, +} + +/// `type Delegation = record { pubkey; expiration; targets : opt vec principal; +/// permissions : opt text }` (the registration delegation returns +/// `permissions` absent). +#[derive(CandidType, Deserialize)] +struct Delegation { + pubkey: Vec, + expiration: u64, + targets: Option>, + permissions: Option, +} + +/// `type SignedDelegation = record { delegation; signature }`. +#[derive(CandidType, Deserialize)] +struct SignedDelegation { + delegation: Delegation, + signature: Vec, +} + +/// `Ok` payload of `prepare_mcp_registration_delegation`. +#[derive(CandidType, Deserialize)] +struct PrepareMcpRegistrationDelegation { + user_key: Vec, + expiration: u64, +} + +/// `type Challenge = record { png_base64 : text; challenge_key : text }`. +#[derive(CandidType, Deserialize)] +struct Challenge { + png_base64: String, + challenge_key: String, +} + +/// `type ChallengeResult = record { key : text; chars : text }`. +#[derive(CandidType)] +struct ChallengeResult { + key: String, + chars: String, +} + +#[derive(CandidType, Deserialize)] +enum Purpose { + #[serde(rename = "recovery")] + Recovery, + #[serde(rename = "authentication")] + Authentication, +} + +#[derive(CandidType, Deserialize)] +enum KeyType { + #[serde(rename = "unknown")] + Unknown, + #[serde(rename = "platform")] + Platform, + #[serde(rename = "cross_platform")] + CrossPlatform, + #[serde(rename = "seed_phrase")] + SeedPhrase, + #[serde(rename = "browser_storage_key")] + BrowserStorageKey, +} + +#[derive(CandidType, Deserialize)] +enum DeviceProtection { + #[serde(rename = "protected")] + Protected, + #[serde(rename = "unprotected")] + Unprotected, +} + +/// `type MetadataMap` — recursive; only needed so `DeviceData.metadata` has a +/// type (we send `null`). +#[derive(CandidType, Deserialize)] +enum MetadataVal { + #[serde(rename = "map")] + Map(Vec<(String, MetadataVal)>), + #[serde(rename = "string")] + String(String), + #[serde(rename = "bytes")] + Bytes(Vec), +} + +/// `type DeviceData = record { ... }`. +#[derive(CandidType)] +struct DeviceData { + pubkey: Vec, + alias: String, + credential_id: Option>, + aaguid: Option>, + purpose: Purpose, + key_type: KeyType, + protection: DeviceProtection, + origin: Option, + metadata: Option>, +} + +/// `type RegisterResponse = variant { registered : record { user_number }; +/// canister_full; bad_challenge }`. +#[derive(CandidType, Deserialize)] +enum RegisterResponse { + #[serde(rename = "registered")] + Registered { user_number: u64 }, + #[serde(rename = "canister_full")] + CanisterFull, + #[serde(rename = "bad_challenge")] + BadChallenge, +} + +// -- Install init arg (captcha disabled). All II init fields are `opt`, so a +// one-field record is a valid subtype. -- + +#[derive(CandidType)] +enum StaticCaptchaTrigger { + CaptchaEnabled, + CaptchaDisabled, +} + +#[derive(CandidType)] +enum CaptchaTrigger { + Dynamic { + threshold_pct: u16, + current_rate_sampling_interval_s: u64, + reference_rate_sampling_interval_s: u64, + }, + Static(StaticCaptchaTrigger), +} + +#[derive(CandidType)] +struct CaptchaConfig { + max_unsolved_captchas: u64, + captcha_trigger: CaptchaTrigger, +} + +#[derive(CandidType)] +struct InternetIdentityInit { + captcha_config: Option, +} + +// ---- II test-anchor helpers (as in the hosted harness) --------------------- + +/// The device pubkey II's test helpers use; not a real DER key, accepted only +/// because the registration captcha is disabled by the install arg. +const DEVICE_PUBKEY: &[u8] = b"test"; + +/// The device principal that authorizes anchor operations. +fn principal_1() -> Principal { + Principal::self_authenticating(DEVICE_PUBKEY) +} + +fn device_data_1() -> DeviceData { + DeviceData { + pubkey: DEVICE_PUBKEY.to_vec(), + alias: "e2e device".to_string(), + credential_id: None, + aaguid: None, + purpose: Purpose::Authentication, + key_type: KeyType::CrossPlatform, + protection: DeviceProtection::Unprotected, + origin: None, + metadata: None, + } +} + +/// Register a fresh anchor (captcha disabled ⇒ the accepted solution is "a"). +async fn register_anchor(pic: &pocket_ic::nonblocking::PocketIc, ii: Principal) -> u64 { + let bytes = pic + .update_call( + ii, + Principal::anonymous(), + "create_challenge", + Encode!().unwrap(), + ) + .await + .expect("create_challenge"); + let challenge = Decode!(&bytes, Challenge).unwrap(); + let attempt = ChallengeResult { + key: challenge.challenge_key, + chars: "a".to_string(), + }; + let bytes = pic + .update_call( + ii, + principal_1(), + "register", + Encode!(&device_data_1(), &attempt, &Option::::None).unwrap(), + ) + .await + .expect("register"); + match Decode!(&bytes, RegisterResponse).unwrap() { + RegisterResponse::Registered { user_number } => user_number, + RegisterResponse::CanisterFull => panic!("register: canister full"), + RegisterResponse::BadChallenge => panic!("register: bad challenge (captcha not disabled?)"), + } +} + +/// Pull `key=value` out of a `&`-separated fragment (values are base64url or +/// plain text here — none contain `&`). +fn field<'a>(blob: &'a str, key: &str) -> Option<&'a str> { + blob.split('&') + .filter_map(|kv| kv.split_once('=')) + .find(|(k, _)| *k == key) + .map(|(_, v)| v) +} + +#[tokio::test] +async fn local_login_end_to_end() { + // Runtime guard: skip cleanly unless the un-fetchable artifacts are provided. + let (Ok(ii_wasm_path), Ok(_)) = (std::env::var("II_WASM"), std::env::var("POCKET_IC_BIN")) + else { + eprintln!( + "skipping local_login_end_to_end: set II_WASM (internet_identity release \ + .wasm.gz) and POCKET_IC_BIN (pocket-ic v15 server) to run it" + ); + return; + }; + let ii_wasm = std::fs::read(&ii_wasm_path).expect("read II_WASM (gz bytes; PocketIC gunzips)"); + + // --- PocketIC + live II --- + let mut pic = pocket_ic::PocketIcBuilder::new() + .with_nns_subnet() + .with_application_subnet() + .build_async() + .await; + // Fresh instances boot at a 2021 mock clock; align to now BEFORE minting + // the registration delegation, or it is already expired. + pic.set_time(SystemTime::now().into()).await; + pic.tick().await; + + let ii = pic.create_canister().await; + pic.add_cycles(ii, 2_000_000_000_000).await; + let init = InternetIdentityInit { + captcha_config: Some(CaptchaConfig { + max_unsolved_captchas: 500, + captcha_trigger: CaptchaTrigger::Static(StaticCaptchaTrigger::CaptchaDisabled), + }), + }; + pic.install_canister(ii, ii_wasm, Encode!(&Some(init)).unwrap(), None) + .await; + + // --- The local binary's composition, exactly as `serve()` wires it, but + // pointed at PocketIC's gateway (the component-3 test configuration: + // endpoint override + fetch_root_key against a loopback replica). --- + let gateway = pic.make_live(None).await; + let agent = imcp2_core::Agent::builder() + .with_url(gateway.as_str()) + .build() + .expect("build agent"); + agent + .fetch_root_key() + .await + .expect("fetch PocketIC root key"); + let identities = Identities::new( + IiInstance { + name: "e2e", + ii_url: gateway.to_string().trim_end_matches('/').to_string(), + ii_canister: ii, + }, + "http://localhost:8000".into(), // management-identity origin; any fixed value + agent, + ); + let slot = SessionSlot::new(); + let driver = LoginDriver::new(identities.clone(), slot.clone(), /* auto_open */ false); + + // --- 1. `authenticate`: fresh session, live listener, the connect link --- + let BeginOutcome::Pending { url, fresh } = driver.begin(false).await.expect("begin") else { + panic!("a fresh driver must start a login flow"); + }; + assert!(fresh); + let fragment = url.split_once('#').expect("connect link fragment").1; + let state = field(fragment, "state").expect("state").to_string(); + let reg_key_x = base64::Engine::decode( + &base64::engine::general_purpose::URL_SAFE_NO_PAD, + field(fragment, "registration_key").expect("registration_key"), + ) + .expect("decode registration_key"); + let (pending_sid, callback_url) = driver.pending_handshake().await.expect("pending"); + assert_eq!(pending_sid, state, "the link's state IS the session id"); + // The link's percent-encoded callback and the listener's own callback are + // the same one value (II matches the allow-list by exact string equality). + let encoded_callback = field(fragment, "callback").expect("callback"); + assert_eq!( + encoded_callback.replace("%3A", ":").replace("%2F", "/"), + callback_url, + "the fragment's callback must be the listener's callback" + ); + let origin = callback_url + .strip_suffix("/callback") + .expect("callback path") + .to_string(); + let http = reqwest::Client::new(); + + // --- 2. The browser lands on the pinned page; II fetches the allow-list --- + let page = http.get(&callback_url).send().await.expect("GET /callback"); + assert_eq!(page.status(), 200); + assert!( + page.headers().get("content-security-policy").is_some(), + "the pinned page ships its CSP" + ); + let wk = http + .get(format!("{origin}{AUTH_CALLBACKS_WELL_KNOWN}")) + .send() + .await + .expect("GET allow-list"); + assert_eq!( + wk.json::().await.unwrap(), + serde_json::json!({ "callbacks": [callback_url] }), + "the #4091 allow-list must declare this callback verbatim" + ); + + // --- 3. The II-side ceremony (browser/frontend role) --- + let anchor = register_anchor(&pic, ii).await; + // Trust an MCP server for the anchor (precondition for `prepare`). + let set_cfg = pic + .update_call( + ii, + principal_1(), + "mcp_set_config", + Encode!( + &anchor, + &McpConfig { + enabled: true, + url: Some("http://localhost:8000/mcp".into()), + } + ) + .unwrap(), + ) + .await + .expect("mcp_set_config call"); + Decode!(&set_cfg, Result<(), String>) + .unwrap() + .expect("mcp_set_config Ok"); + + // prepare: consent (full access, 24h grant) is recorded server-side, + // keyed by the P_reg the canister returns as `user_key`. + const GRANT_TTL_NS: u64 = 24 * 60 * 60 * 1_000_000_000; + let prep_bytes = pic + .update_call( + ii, + principal_1(), + "prepare_mcp_registration_delegation", + Encode!( + &anchor, + ®_key_x, + &Some(Permissions::All), + &Some(GRANT_TTL_NS) + ) + .unwrap(), + ) + .await + .expect("prepare call"); + let prepared = Decode!(&prep_bytes, Result) + .unwrap() + .expect("prepare Ok"); + + // get: the canister-signed hop toward our X. + let get_bytes = pic + .query_call( + ii, + principal_1(), + "get_mcp_registration_delegation", + Encode!( + &anchor, + ®_key_x, + &prepared.user_key, + &prepared.expiration + ) + .unwrap(), + ) + .await + .expect("get call"); + let signed = Decode!(&get_bytes, Result) + .unwrap() + .expect("get Ok"); + assert_eq!( + signed.delegation.pubkey, reg_key_x, + "delegation targets our X" + ); + + // Serialize into the agent-js DelegationChain JSON the redeem parses: hex + // byte fields, HEX-string expiration, publicKey = der(P_reg). + let mut delegation = serde_json::json!({ + "pubkey": hex::encode(&signed.delegation.pubkey), + "expiration": format!("{:x}", signed.delegation.expiration), + }); + if let Some(targets) = &signed.delegation.targets { + delegation["targets"] = + serde_json::Value::from(targets.iter().map(|p| p.to_text()).collect::>()); + } + let chain_json = serde_json::json!({ + "publicKey": hex::encode(&prepared.user_key), + "delegations": [{ "delegation": delegation, "signature": hex::encode(&signed.signature) }], + }) + .to_string(); + + // --- 4. The browser's redeem POST → the LOCAL success contract --- + let resp = http + .post(format!("{origin}/redeem")) + .json(&serde_json::json!({ "state": state, "delegation": chain_json })) + .send() + .await + .expect("POST /redeem"); + assert_eq!(resp.status(), 200, "redeem should succeed against live II"); + assert_eq!( + resp.json::().await.unwrap(), + serde_json::json!({ "done": true }), + "the local redeem answers done (the page renders its terminal state)" + ); + + // --- 5. Signed in: the slot holds this session; status + principal agree --- + assert_eq!( + slot.get(), + Some(state.clone()), + "the session slot is filled" + ); + match driver.status().await { + LoginStatus::SignedIn(g) => { + assert_eq!(g.session_id, state); + assert_eq!( + g.permissions, "all", + "the consent-time access level was recorded" + ); + assert!(g.minutes_left() > 0, "the 24h grant is live"); + assert_eq!( + g.principal, + identities.session_principal(&state).await, + "auth_status reports the real session principal" + ); + } + _ => panic!("a completed redeem must report SignedIn"), + } + // A repeat `authenticate` now reports the live session instead of a link. + match driver.begin(false).await.expect("begin") { + BeginOutcome::AlreadySignedIn(g) => assert_eq!(g.session_id, state), + BeginOutcome::Pending { .. } => panic!("a live grant must be reported, not replaced"), + } + + // --- 6. The transient listener is torn down after the redeem --- + let mut listener_down = false; + for _ in 0..50 { + match http.get(&callback_url).send().await { + Err(_) => { + listener_down = true; + break; + } + Ok(_) => tokio::time::sleep(std::time::Duration::from_millis(100)).await, + } + } + assert!( + listener_down, + "the login listener must shut down once the grant landed" + ); +} diff --git a/crates/imcp2-local/src/main.rs b/crates/imcp2-local/src/main.rs index 53a8a78..b4950ae 100644 --- a/crates/imcp2-local/src/main.rs +++ b/crates/imcp2-local/src/main.rs @@ -33,6 +33,11 @@ mod login; mod server; +// The PocketIC-backed end-to-end login test (the design's local-replica test +// configuration). Compiled only for `cargo test --features e2e`. +#[cfg(all(test, feature = "e2e"))] +mod e2e_local_login; + use imcp2_core::{identities::Identities, skills, IcTools, IiInstance}; use login::SessionSlot; use rmcp::ServiceExt; @@ -107,6 +112,22 @@ async fn main() -> anyhow::Result<()> { tracing::info!(ii = %instance.ii_url, "using Internet Identity instance \"{}\"", instance.name); let identities = Identities::new(instance, management_origin(), agent.clone()); + // The same 60-second session reaper the hosted server runs: expired + // grants, abandoned sign-in keys, and their cached delegations leave + // memory promptly instead of accumulating over a long-lived process. + // (No shutdown plumbing needed: the task ends with the runtime.) + { + let identities = identities.clone(); + tokio::spawn(async move { + let mut tick = tokio::time::interval(std::time::Duration::from_secs(60)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tick.tick().await; + identities.reap_expired_sessions().await; + } + }); + } + // The single-user authentication layer, owned by THIS binary: the login // flow fills (and on re-login replaces) the slot, and the resolver it // hands core answers every tool call's "which session?" from it. From 336a3bf48968dcd29c3788b11cd08d3be5731be4 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 10:19:25 +0000 Subject: [PATCH 10/10] Address review: the session slot is expiry-aware MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reaper removes an expired session's server-side state, but the slot kept resolving its id — session-requiring tools would then bypass their "needs an authenticated session" guidance and surface internal errors (or re-create a keyless session via ensure_session). The slot now stores the grant expiration II returned and reports NO session the moment the grant lapses, so unauthenticated is immediate and accurate — the reaper's 60-second cadence frees memory, it does not gate authentication. auth_status keeps its Expired reporting via the driver's retained Grant. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AqNkpMQiQHzYxC2djTfvBK --- crates/imcp2-local/src/login/routes.rs | 2 +- crates/imcp2-local/src/login/slot.rs | 63 ++++++++++++++++++-------- crates/imcp2-local/src/login/tests.rs | 10 +++- docs/scoping-local-deployment.md | 3 +- 4 files changed, 55 insertions(+), 23 deletions(-) diff --git a/crates/imcp2-local/src/login/routes.rs b/crates/imcp2-local/src/login/routes.rs index 91be4b7..f1ce9f1 100644 --- a/crates/imcp2-local/src/login/routes.rs +++ b/crates/imcp2-local/src/login/routes.rs @@ -240,7 +240,7 @@ async fn redeem(State(ctx): State, Json(body): Json) -> Re Ok(outcome) => { let principal = driver.inner.identities.session_principal(&body.state).await; let mut state = driver.inner.state.lock().await; - driver.inner.slot.set(body.state.clone()); + driver.inner.slot.set(body.state.clone(), outcome.expiration_ns); state.grant = Some(Grant { session_id: body.state.clone(), principal, diff --git a/crates/imcp2-local/src/login/slot.rs b/crates/imcp2-local/src/login/slot.rs index ddcfb16..f5fda4e 100644 --- a/crates/imcp2-local/src/login/slot.rs +++ b/crates/imcp2-local/src/login/slot.rs @@ -1,20 +1,36 @@ //! The [`SessionSlot`] — the local binary's whole authentication state: one -//! shared, replaceable holder for the signed-in session id, read by core +//! shared, replaceable holder for the signed-in session, read by core //! through the injected resolver and written by the login flow. use std::sync::Arc; -/// The single-user session holder — a shared, **replaceable** slot rather -/// than a fixed id. The Internet Identity contract requires a FRESH session -/// key after an `Unauthorized` (a reused id would reuse its keys), so each -/// successful login writes a fresh session id here — same-process -/// reauthentication after grant expiry or revocation, with no client restart. -/// Clones share the slot: the login flow holds a writing handle, and the slot -/// answers `IcTools`' session lookups through [`SessionSlot::resolver`] — the -/// local binary's whole authentication layer, living here rather than in -/// core, which only asks the injected resolver. Starts empty (signed out). +use super::now_ns; + +/// The single-user session holder — a shared, **replaceable**, +/// **expiry-aware** slot rather than a fixed id. The Internet Identity +/// contract requires a FRESH session key after an `Unauthorized` (a reused id +/// would reuse its keys), so each successful login writes a fresh session id +/// here — same-process reauthentication after grant expiry or revocation, +/// with no client restart. Clones share the slot: the login flow holds a +/// writing handle, and the slot answers `IcTools`' session lookups through +/// [`SessionSlot::resolver`] — the local binary's whole authentication layer, +/// living here rather than in core, which only asks the injected resolver. +/// Starts empty (signed out). +/// +/// Expiry-aware: the slot stores the grant's expiration alongside the id and +/// reports **no session** the moment the grant lapses — so tools return their +/// "needs an authenticated session" sign-in guidance instead of acting on a +/// dead session (whose server-side state the reaper also removes, on its own +/// 60-second cadence, which must not be what gates authentication). #[derive(Clone, Debug, Default)] -pub struct SessionSlot(Arc>>); +pub struct SessionSlot(Arc>>); + +#[derive(Debug)] +struct Entry { + session_id: String, + /// Grant expiration, ns since the Unix epoch (from `mcp_register_v2`). + expiration_ns: u64, +} impl SessionSlot { /// A fresh, empty slot (no session — every session-needing tool reports @@ -23,20 +39,29 @@ impl SessionSlot { Self::default() } - /// The current session id, if a login has completed. The critical section - /// is a clone — never held across an await. + /// The current session id, while its grant is live; `None` once it + /// expires. The critical section is a clone — never held across an await. pub fn get(&self) -> Option { - self.0.read().expect("session slot poisoned").clone() + self.0 + .read() + .expect("session slot poisoned") + .as_ref() + .filter(|entry| entry.expiration_ns > now_ns()) + .map(|entry| entry.session_id.clone()) } - /// Replace the session id after a successful login (a fresh id per the II - /// fresh-session-key contract). - pub fn set(&self, session_id: String) { - *self.0.write().expect("session slot poisoned") = Some(session_id); + /// Replace the session after a successful login (a fresh id per the II + /// fresh-session-key contract), with the grant expiration II returned. + pub fn set(&self, session_id: String, expiration_ns: u64) { + *self.0.write().expect("session slot poisoned") = Some(Entry { + session_id, + expiration_ns, + }); } /// This slot as the [`imcp2_core::SessionResolver`] the local binary - /// injects into `IcTools`: every call acts as the one signed-in session. + /// injects into `IcTools`: every call acts as the one signed-in session, + /// or as unauthenticated once the grant lapses. pub fn resolver(&self) -> imcp2_core::SessionResolver { let slot = self.clone(); Arc::new(move |_ctx| slot.get()) diff --git a/crates/imcp2-local/src/login/tests.rs b/crates/imcp2-local/src/login/tests.rs index 0e9eff7..7e01f2e 100644 --- a/crates/imcp2-local/src/login/tests.rs +++ b/crates/imcp2-local/src/login/tests.rs @@ -338,14 +338,20 @@ fn the_session_slot_is_shared_and_replaceable() { let slot = SessionSlot::new(); let reader = slot.clone(); assert_eq!(reader.get(), None, "a fresh slot is signed out"); - slot.set("sess-1".into()); + let hour_from_now = now_ns() + 3_600_000_000_000; + slot.set("sess-1".into(), hour_from_now); assert_eq!(reader.get(), Some("sess-1".into()), "clones share the slot"); - slot.set("sess-2".into()); + slot.set("sess-2".into(), hour_from_now); assert_eq!( reader.get(), Some("sess-2".into()), "a re-login replaces the id" ); + // Expiry-aware: the moment the grant lapses, the slot reports NO session + // — tools return sign-in guidance instead of acting on a dead session, + // regardless of when the 60-second reaper gets to the server-side state. + slot.set("sess-3".into(), now_ns() - 1); + assert_eq!(reader.get(), None, "an expired grant is signed out"); } // The handshake window is exactly HANDSHAKE_TTL: a pending flow older than diff --git a/docs/scoping-local-deployment.md b/docs/scoping-local-deployment.md index 3e6cc5f..2c35576 100644 --- a/docs/scoping-local-deployment.md +++ b/docs/scoping-local-deployment.md @@ -457,7 +457,8 @@ Each binary then owns its authentication end to end, in its own package: redeem. Replaceable, not fixed: the II contract requires a **fresh session key** after `Unauthorized`, so each re-login writes a fresh session id — same-process reauthentication after grant expiry or revocation, with no client restart. The resolver - it injects simply reads the slot. + it injects simply reads the slot — which is expiry-aware: the moment the grant lapses it + reports no session, so tools return sign-in guidance instead of acting on a dead session. The seam is deliberately just a value-returning callback rather than a middleware slot in core: a real authentication middleware needs its transport (the hosted validation is HTTP