diff --git a/.gitignore b/.gitignore index eb45ca5..a186ff2 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,6 @@ candid-wasm/target # native deploy build artifact /build-out + +# runtime OAuth client registrations (OAUTH_CLIENTS_FILE) +/oauth-clients.json diff --git a/README.md b/README.md index ec0f219..32a6b6e 100644 --- a/README.md +++ b/README.md @@ -96,11 +96,11 @@ Add the server to Claude Code (replace the URL with wherever it's hosted): claude mcp add --transport http ic-poc https://YOUR-HOST/mcp ``` -Then run `/mcp` → **ic-poc** → authenticate: the client obtains a device code and -opens a verification link that launches **Internet Identity**'s `/mcp` handshake; -you sign in once, and the tools become available — that single login registers the -connection's session key with II as a time-boxed grant. (Any MCP client that -supports the OAuth 2.0 device authorization grant, RFC 8628, works.) +Then run `/mcp` → **ic-poc** → authenticate: the client sends the browser to +**Internet Identity**'s `/mcp` handshake; you sign in once, II registers the +connection's session key as a time-boxed grant and returns you to the client, and +the tools become available. Redirect-based clients (e.g. Claude.ai) use the +authorization-code flow; clients that poll can use the device grant (RFC 8628). ## Run @@ -149,35 +149,48 @@ curl -s "${H[@]}" -d '{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"n # => ("Internet Computer") ``` -## Auth (OAuth 2.0 device grant, login via Internet Identity) +## Auth (OAuth 2.1, login via Internet Identity) `/mcp` is gated by a bearer token. II's `/mcp` handshake has **no redirect back to this server** — the II tab makes two background `fetch()` POSTs to our callback and then finishes on its own — so a classic authorization-code -`redirect_uri` can't be delivered. The MCP client's login is therefore modelled -as the **OAuth 2.0 Device Authorization Grant (RFC 8628)**: the client gets a -`device_code` + a verification link, the user opens it (which launches II's -handshake), and the client polls the token endpoint until the grant is live. +`redirect_uri` can't be delivered *by II*. We bridge that gap and support **two +flows**, so any MCP client works: + +- **Authorization code (finish redirect)** — for redirect-based clients (e.g. + Claude.ai). `/oauth/authorize` redirects to II's handshake; the key-request + response carries a `finish_url` on our origin, so after `mcp_register` II + navigates the browser to `/oauth/finish`, which confirms registration, mints a + PKCE-bound code, and 302s to the client's `redirect_uri?code=…&state=…`. +- **Device grant (RFC 8628)** — for clients that can poll: `device_code` + + `verification_uri`, and the client polls `/oauth/token` until the grant is live. Endpoints: -- `GET /.well-known/oauth-authorization-server` — AS metadata (advertises the - device-code grant and the `device_authorization_endpoint`) +- `GET /.well-known/oauth-authorization-server` — AS metadata (advertises both + the `authorization_code` and `device_code` grants) - `GET /.well-known/oauth-protected-resource` — points clients at the AS - `POST /oauth/register` — dynamic client registration (RFC 7591); `redirect_uris` - are optional (the device grant has none) and persisted to `OAUTH_CLIENTS_FILE` -- `POST /oauth/device_authorization` — mints a `device_code`/`user_code` and a - `verification_uri` (RFC 8628 §3.2) + are stored (device-only clients may register none) and persisted to + `OAUTH_CLIENTS_FILE`; requested `grant_types` are honoured +- `GET /oauth/authorize` — validates the client + redirect, requires PKCE, then + redirects to II's handshake +- `GET /oauth/finish` — II navigates here after registration; confirms the grant + is live, then 302s to `redirect_uri?code=…&state=…` +- `POST /oauth/device_authorization` — device-grant entry: mints a `device_code`/ + `user_code` + `verification_uri` (RFC 8628 §3.2) - `GET /oauth/device` — the user opens this; it launches II's `/mcp` handshake with the connect `state` (= session id) in the URL fragment, `ttl` in **seconds** - `POST /oauth/connect/callback` — serves II's **two cross-origin JSON POSTs**: a key request `{state}` → `{public_key}` (a fresh session keypair, minted here), and a completion notification `{state, expiration}` → mark the grant live -- `POST /oauth/token` — the client polls with the `device_code` - (`authorization_pending`/`slow_down`) until the token is issued +- `POST /oauth/token` — exchanges an authorization `code` (PKCE), or polls with a + `device_code` (`authorization_pending`/`slow_down`), for the access token -Unauthenticated `/mcp` requests get `401` with a `WWW-Authenticate` header -pointing at the resource metadata, as the MCP spec expects. +"Grant is live" means II's completion POST arrived **or** a signed +`mcp_get_accounts` now succeeds (belt-and-suspenders, since the completion POST is +best-effort). Unauthenticated `/mcp` requests get `401` with a `WWW-Authenticate` +header pointing at the resource metadata, as the MCP spec expects. **The server is passive during the handshake and holds no key at link time.** On the key-request POST it generates a fresh per-connection Ed25519 keypair and @@ -189,10 +202,10 @@ calls `mcp_register`.** The issued access token is bound to the session key's principal (`self_authenticating(session_pubkey)`), which is exactly the identity the grant is bound to. -**PKCE (S256)** is bound to the poll when the client supplies a `code_challenge`; -device codes live 600s, access tokens 1h. Treat any `Unauthorized` from II as -"session over → reconnect": the server surfaces a reconnect message and does not -retry. +**PKCE (S256)** is required for the authorization-code flow and honoured for the +device grant; auth codes live 120s, device codes 600s, access tokens 1h. Treat any +`Unauthorized` from II as "session over → reconnect": the server surfaces a +reconnect message and does not retry. Set the public base URL (used in the discovery docs, as the MCP origin, and as the management identity's derivation origin) with `PUBLIC_URL`. The Internet Identity diff --git a/src/auth.rs b/src/auth.rs index 0511d37..c59d755 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1,24 +1,30 @@ //! OAuth 2.1 authorization server for the MCP endpoint, with **Internet Identity** //! as the login mechanism, using II's session-key registration handshake. //! -//! II's new `/mcp` handshake has NO redirect back to this server — the II tab -//! runs two background `fetch()` POSTs to our callback and then "finishes on its -//! own" (see `docs/mcp-server-guide.md` / dfinity/internet-identity#4086). A -//! classic auth-code `redirect_uri` therefore can't be delivered. We model the -//! MCP client's login as the **OAuth 2.0 Device Authorization Grant (RFC 8628)**: -//! the client asks `/oauth/device_authorization` for a `device_code` + a -//! `verification_uri`, the user opens that URI (which launches II's `/mcp` -//! handshake), and the client polls `/oauth/token` until the grant is live. +//! II's `/mcp` handshake registers the session key under the user's own auth and, +//! if we ask, navigates the browser back to a `finish_url` on our origin (see +//! `docs/mcp-server-guide.md` §2a/§2c / dfinity/internet-identity#4086). We use +//! that hook to serve two flows, so any MCP client works: //! -//! Connect handshake (Phase 1b): our `/oauth/connect/callback` serves the two -//! cross-origin JSON POSTs II makes — a key request `{state}` → `{public_key}` -//! (a fresh session keypair minted per connection) and a completion notification -//! `{state, expiration}` → mark the grant live. We never receive or verify a -//! delegation chain, and never call `mcp_register` (II's frontend does, under the -//! user's own authentication). +//! * **Authorization code (finish redirect)** — for redirect-based clients +//! (e.g. Claude.ai). `/oauth/authorize` redirects to II's handshake; the +//! key-request response carries a `finish_url`, so after `mcp_register` II +//! navigates the browser to `/oauth/finish`, which confirms registration, +//! mints a PKCE-bound code, and 302s to the client's `redirect_uri`. +//! * **Device authorization grant (RFC 8628)** — for clients that can poll the +//! token endpoint: `/oauth/device_authorization` → `device_code` + +//! `verification_uri`; the user opens it (launching the same II handshake, no +//! `finish_url`); the client polls `/oauth/token` until the grant is live. //! -//! Implemented: dynamic client registration, PKCE (S256) bound to the poll, -//! short-lived device codes, 1h access tokens, session-key-bound principal. +//! Connect handshake (Phase 1b): `/oauth/connect/callback` serves the two +//! cross-origin JSON POSTs II makes — a key request `{state}` → `{public_key +//! [, finish_url]}` (a fresh session keypair minted per connection) and a +//! completion notification `{state, expiration}` → mark the grant live. We never +//! receive or verify a delegation chain, and never call `mcp_register` (II's +//! frontend does, under the user's own authentication). +//! +//! Implemented: dynamic client registration, PKCE (S256) enforced, short-lived +//! codes/device codes, 1h access tokens, session-key-bound principal. use std::{ collections::HashMap, @@ -43,15 +49,16 @@ use uuid::Uuid; use crate::identities::Identities; -/// How long a device-authorization request (and its pending II handshake) stays -/// valid before the user must restart. -const DEVICE_CODE_TTL: Duration = Duration::from_secs(600); +/// How long an authorization request (auth-code or device) and its pending II +/// handshake stay valid before the user must restart. +const CONNECT_TTL: Duration = Duration::from_secs(600); +/// Lifetime of a minted authorization code (auth-code flow). +const CODE_TTL: Duration = Duration::from_secs(120); /// Minimum seconds a client should wait between token polls (RFC 8628 `interval`). const POLL_INTERVAL_SECS: u64 = 5; /// Access-token lifetime (also the II grant's default, 1h). const TOKEN_TTL: Duration = Duration::from_secs(3600); -/// `ttl` (seconds) requested for the II grant. Omitting would default to 3600; -/// we send it explicitly. Clamped by II to [600, 2592000]. +/// `ttl` (seconds) requested for the II grant. Clamped by II to [600, 2592000]. const GRANT_TTL_SECS: u64 = 3600; /// RFC 8628 device-code grant type. @@ -62,8 +69,9 @@ pub fn base_url() -> String { std::env::var("PUBLIC_URL").unwrap_or_else(|_| "http://localhost:8000".to_string()) } -/// A registered OAuth client (RFC 7591). Redirect URIs are recorded for -/// completeness but unused: the device grant has no redirect. +/// A registered OAuth client (RFC 7591): the redirect URIs it declared. The +/// auth-code flow only redirects a code to one of these (exact match), so the +/// server is not an open redirector and needs no hardcoded host allowlist. #[derive(Clone, Debug, Serialize, Deserialize)] struct ClientReg { #[serde(default)] @@ -72,7 +80,7 @@ struct ClientReg { /// File the dynamic client registrations are persisted to. RFC 7591 clients are /// long-lived (they cache their `client_id`), so registrations must survive a -/// restart — unlike device codes/tokens, which are short-lived and stay in +/// restart — unlike codes/tokens/connects, which are short-lived and stay in /// memory. Override with `OAUTH_CLIENTS_FILE`. fn clients_file() -> String { std::env::var("OAUTH_CLIENTS_FILE").unwrap_or_else(|_| "oauth-clients.json".to_string()) @@ -106,37 +114,86 @@ fn persist_clients(clients: &HashMap) { } } +/// Acceptance rule for a redirect (OAuth 2.1): the client must be REGISTERED, +/// and the requested redirect must either exactly match a registered URI, or be +/// a loopback URI matching a registered loopback URI on everything but the +/// port. RFC 8252 §7.3 requires the any-port latitude — native clients bind an +/// ephemeral loopback port at runtime, so the exact port can't be registered — +/// but registration itself is still required, so every client that can receive +/// a code is on record (DCR is open, so this is an audit trail, not vetting). +fn redirect_allowed(reg: Option<&ClientReg>, redirect_uri: &str) -> bool { + let Some(reg) = reg else { return false }; + reg.redirect_uris + .iter() + .any(|u| u == redirect_uri || loopback_match(u, redirect_uri)) +} + +/// Whether `requested` is a loopback redirect matching the registered loopback +/// URI `registered` on scheme, host, path, and query — any port (RFC 8252 §7.3). +/// Both sides must independently pass [`is_loopback_redirect`], so a registered +/// hosted URI grants no loopback latitude and look-alike hosts are rejected. +fn loopback_match(registered: &str, requested: &str) -> bool { + if !is_loopback_redirect(registered) || !is_loopback_redirect(requested) { + return false; + } + let (Ok(a), Ok(b)) = (url::Url::parse(registered), url::Url::parse(requested)) else { + return false; + }; + a.host_str() == b.host_str() && a.path() == b.path() && a.query() == b.query() +} + #[derive(Clone)] pub struct AuthStore { clients: Arc>>, tokens: Arc>>, - /// Device-authorization grants in flight, keyed by `session_id` (which is - /// also the II connect `state`). Holds the `device_code`/`user_code` used by - /// the token poll and verification page. + /// Auth-code connects in flight, keyed by `session_id` (= the II connect + /// `state`). Serve the poll bridge for redirect-based clients. + authz: Arc>>, + /// Minted authorization codes awaiting exchange at `/oauth/token`. + codes: Arc>>, + /// Device-authorization grants in flight, keyed by `session_id`. devices: Arc>>, /// Shared with the MCP tools: the session's backend key / grant expiration /// live here (keyed by `session_id`) for the tools to sign with. identities: Identities, } +/// An auth-code connect awaiting the user's II handshake (poll bridge). +#[derive(Clone, Debug)] +struct AuthzPending { + client_id: String, + redirect_uri: String, + /// The OAuth client's own `state`, echoed back on the final redirect. + client_state: String, + code_challenge: Option, + created: Instant, + /// Grant confirmed live (completion POST or a signed-call fallback). + live: bool, + /// The authorization code minted once the grant is live (idempotent polls). + code: Option, +} + +/// A minted authorization code awaiting exchange. +#[derive(Clone, Debug)] +struct CodeGrant { + client_id: String, + code_challenge: Option, + session_id: String, + created: Instant, +} + /// A device-authorization grant awaiting the user's II handshake. #[derive(Clone, Debug)] struct DeviceAuth { device_code: String, user_code: String, - /// The connection's session id; also the II connect `state`. session_id: String, #[allow(dead_code)] client_id: Option, scope: Option, - /// PKCE challenge (S256), if the client supplied one; verified at token time. code_challenge: Option, created: Instant, - /// Verification page visited and II launched. - launched: bool, - /// Grant confirmed live (completion POST or a signed-call fallback). live: bool, - /// Last token poll, for RFC 8628 `slow_down` enforcement. last_poll: Option, } @@ -153,17 +210,207 @@ impl AuthStore { Self { clients: Arc::new(RwLock::new(load_clients())), tokens: Arc::default(), + authz: Arc::default(), + codes: Arc::default(), devices: Arc::default(), identities, } } + /// Whether `redirect_uri` is acceptable for `client_id`: the client must be + /// registered, and the redirect must match a registered URI (exactly, or + /// port-agnostically for loopback per RFC 8252 §7.3). + async fn validate_client(&self, client_id: &str, redirect_uri: &str) -> bool { + redirect_allowed(self.clients.read().await.get(client_id), redirect_uri) + } + /// The verified principal + session id behind a bearer token, if valid. pub async fn session_for_token(&self, token: &str) -> Option<(String, String)> { let tokens = self.tokens.read().await; let info = tokens.get(token)?; (info.created.elapsed() < info.ttl).then(|| (info.principal.clone(), info.session_id.clone())) } + + /// Whether this connection's grant is live: the completion POST arrived, or + /// (best-effort fallback) a signed `mcp_get_accounts` now succeeds. Takes no + /// lock across the network call. + async fn grant_live(&self, session_id: &str, flag: bool) -> bool { + flag || self.identities.grant_is_live(session_id).await + } +} + +// ---- Authorization code (poll bridge) ---------------------------------- + +#[derive(Debug, Deserialize)] +pub struct AuthorizeQuery { + #[serde(default)] + response_type: Option, + client_id: String, + redirect_uri: String, + #[serde(default)] + state: Option, + #[serde(default)] + code_challenge: Option, + #[allow(dead_code)] + #[serde(default)] + code_challenge_method: Option, + #[allow(dead_code)] + #[serde(default)] + scope: Option, + #[allow(dead_code)] + #[serde(default)] + resource: Option, +} + +/// GET /oauth/authorize — the redirect-based entry point. Validates the client +/// and PKCE, records a pending connect, and redirects the browser to II's `/mcp` +/// handshake; II navigates back to our `finish_url` once it registers the key. +pub async fn authorize(State(store): State, Query(q): Query) -> Response { + // Only the authorization-code response type is supported. + match q.response_type.as_deref() { + Some("code") => {} + Some(_) => return oauth_err(StatusCode::BAD_REQUEST, "unsupported_response_type", "only response_type=code"), + None => return oauth_err(StatusCode::BAD_REQUEST, "invalid_request", "response_type=code required"), + } + if !store.validate_client(&q.client_id, &q.redirect_uri).await { + return oauth_err(StatusCode::BAD_REQUEST, "invalid_request", "unknown client_id / redirect_uri"); + } + // OAuth 2.1: PKCE is required for public clients. + let Some(code_challenge) = q.code_challenge.clone() else { + return oauth_err(StatusCode::BAD_REQUEST, "invalid_request", "code_challenge (PKCE S256) required"); + }; + if let Some(m) = &q.code_challenge_method { + if m != "S256" { + return oauth_err(StatusCode::BAD_REQUEST, "invalid_request", "only code_challenge_method=S256 is supported"); + } + } + + let session_id = format!("sess-{}", Uuid::new_v4()); + store.authz.write().await.insert( + session_id.clone(), + AuthzPending { + client_id: q.client_id.clone(), + redirect_uri: q.redirect_uri.clone(), + client_state: q.state.clone().unwrap_or_default(), + code_challenge: Some(code_challenge), + created: Instant::now(), + live: false, + code: None, + }, + ); + + // Redirect the browser to II's handshake. II navigates back to our + // `finish_url` (returned in the key-request response) once it registers. + js_redirect(&ii_mcp_url(&session_id)) +} + +#[derive(Debug, Deserialize)] +pub struct FinishQuery { + /// The pending-auth id (= session id) carried by `finish_url`. + id: String, + /// Retry counter, so the "finishing" reload is bounded. + #[serde(default)] + r: u32, +} + +/// GET /oauth/finish — II navigates the browser here after registering the +/// session key (this is the `finish_url` returned in the key-request response). +/// Arrival is NOT proof of registration: confirm it (the completion POST flag, or +/// a signed `mcp_get_accounts` that returns `Ok`), then mint the authorization +/// code and 302 to the client's `redirect_uri` with `code` + the client's `state`. +pub async fn finish(State(store): State, Query(q): Query) -> Response { + // Snapshot without holding the lock across the network probe. + let snap = { + let authz = store.authz.read().await; + authz.get(&q.id).map(|a| { + ( + a.created.elapsed() >= CONNECT_TTL, + a.live, + a.code.clone(), + a.client_id.clone(), + a.redirect_uri.clone(), + a.client_state.clone(), + a.code_challenge.clone(), + ) + }) + }; + let Some((expired, live_flag, existing_code, client_id, redirect_uri, client_state, code_challenge)) = snap else { + return connect_error("unknown or already-used connect request — restart from your client"); + }; + if expired { + return connect_error("connect request expired — restart from your client"); + } + // Idempotent: if the code was already minted, just redirect again. + if let Some(code) = existing_code { + return redirect_302(&build_redirect(&redirect_uri, &code, &client_state)); + } + // Confirm registration (best-effort signed-call fallback, no lock held). If it + // isn't confirmable yet (a race with II's registration/propagation), reload + // shortly — bounded so we don't loop forever. + if !store.grant_live(&q.id, live_flag).await { + if q.r >= 8 { + return connect_error("could not confirm the connection with Internet Identity — reconnect and try again"); + } + return finishing_page(&q.id, q.r + 1); + } + + // Reserve the code under the `authz` lock, then insert into `codes` AFTER + // releasing it — never hold one map's lock while awaiting the other's, so the + // lock order is consistent with `token_authorization_code` (no deadlock). + let fresh = format!("mcp-code-{}", Uuid::new_v4()); + let (code, newly_minted) = { + let mut authz = store.authz.write().await; + let Some(a) = authz.get_mut(&q.id) else { + return connect_error("connect request vanished — restart from your client"); + }; + match &a.code { + Some(existing) => (existing.clone(), false), + None => { + a.code = Some(fresh.clone()); + (fresh, true) + } + } + }; + if newly_minted { + store.codes.write().await.insert( + code.clone(), + CodeGrant { + client_id, + code_challenge, + session_id: q.id.clone(), + created: Instant::now(), + }, + ); + } + tracing::info!(session_id = %q.id, "grant confirmed; issued authorization code"); + redirect_302(&build_redirect(&redirect_uri, &code, &client_state)) +} + +/// A tiny self-reloading page shown while we wait for II's registration to become +/// confirmable, then it re-hits `/oauth/finish` (bounded by the retry counter). +fn finishing_page(id: &str, next_try: u32) -> Response { + let url = js_escape(&format!("/oauth/finish?id={}&r={}", urlencoding::encode(id), next_try)); + Html(format!( + "Finishing…\ + \ +

Finishing sign-in…

\ + " + )) + .into_response() +} + +/// A 302 to an absolute URL (used for the top-level hop back to the OAuth client). +fn redirect_302(url: &str) -> Response { + (StatusCode::FOUND, [(axum::http::header::LOCATION, url.to_string())]).into_response() +} + +fn build_redirect(redirect_uri: &str, code: &str, client_state: &str) -> String { + let sep = if redirect_uri.contains('?') { '&' } else { '?' }; + let mut r = format!("{redirect_uri}{sep}code={}", urlencoding::encode(code)); + if !client_state.is_empty() { + r.push_str(&format!("&state={}", urlencoding::encode(client_state))); + } + r } // ---- Device authorization: start the II handshake ----------------------- @@ -182,13 +429,9 @@ pub struct DeviceAuthzForm { /// POST /oauth/device_authorization (RFC 8628 §3.1–3.2) — mint a device code and /// point the user at the verification URI that launches II's `/mcp` handshake. -pub async fn device_authorization( - State(store): State, - Form(req): Form, -) -> Response { - // We only verify PKCE with S256; reject any other method up front (rather - // than storing the challenge and failing confusingly with S256 at token - // time), and require a challenge whenever a method is named. +pub async fn device_authorization(State(store): State, Form(req): Form) -> Response { + // We only verify PKCE with S256; reject any other method up front, and + // require a challenge whenever a method is named. if let Some(method) = &req.code_challenge_method { if method != "S256" { return oauth_err(StatusCode::BAD_REQUEST, "invalid_request", "only code_challenge_method=S256 is supported"); @@ -200,7 +443,6 @@ pub async fn device_authorization( let session_id = format!("sess-{}", Uuid::new_v4()); let device_code = format!("dc-{}", Uuid::new_v4()); - // A short, human-typable code for the verification page. let user_code = Uuid::new_v4().simple().to_string()[..8].to_uppercase(); store.devices.write().await.insert( @@ -213,7 +455,6 @@ pub async fn device_authorization( scope: req.scope.clone().filter(|s| !s.is_empty()), code_challenge: req.code_challenge.clone(), created: Instant::now(), - launched: false, live: false, last_poll: None, }, @@ -227,14 +468,12 @@ pub async fn device_authorization( "user_code": user_code, "verification_uri": verification_uri, "verification_uri_complete": verification_uri_complete, - "expires_in": DEVICE_CODE_TTL.as_secs(), + "expires_in": CONNECT_TTL.as_secs(), "interval": POLL_INTERVAL_SECS, })) .into_response() } -// ---- Verification page: launch II's /mcp handshake ---------------------- - #[derive(Debug, Deserialize)] pub struct DeviceVerifyQuery { #[serde(default)] @@ -251,33 +490,16 @@ pub async fn device_verify(State(store): State, Query(q): Query { - d.launched = true; - d.session_id.clone() - } - None => return connect_error("unknown or expired code — restart the connection from your client"), - } + let devices = store.devices.read().await; + devices + .values() + .find(|d| d.user_code == user_code && d.created.elapsed() < CONNECT_TTL) + .map(|d| d.session_id.clone()) }; - - // Launch II's `/mcp` handshake. Everything is in the URL fragment (never sent - // to II's servers): the callback on our origin, the single-use `state` (= the - // session id), and the requested grant `ttl` in SECONDS. NO key material is - // put in the link — the session key is minted inside the key-request callback. - let base = base_url(); - let callback = format!("{base}/oauth/connect/callback"); - let ii_mcp_url = format!( - "{ii}/mcp#callback={cb}&state={st}&ttl={ttl}", - ii = crate::identities::ii_url(), - cb = urlencoding::encode(&callback), - st = urlencoding::encode(&session_id), - ttl = GRANT_TTL_SECS, - ); - js_redirect(&ii_mcp_url) + match session_id { + Some(sid) => js_redirect(&ii_mcp_url(&sid)), + None => connect_error("unknown or expired code — restart the connection from your client"), + } } /// A minimal manual code-entry page for the bare `verification_uri`. @@ -295,11 +517,27 @@ fn device_code_form() -> Response { .into_response() } +/// Build II's `/mcp` handshake URL for a connection. Everything is in the URL +/// fragment (never sent to II's servers): the callback on our origin, the +/// single-use `state` (= session id), and the requested grant `ttl` in SECONDS. +/// NO key material is put in the link. +fn ii_mcp_url(session_id: &str) -> String { + let base = base_url(); + let callback = format!("{base}/oauth/connect/callback"); + format!( + "{ii}/mcp#callback={cb}&state={st}&ttl={ttl}", + ii = crate::identities::ii_url(), + cb = urlencoding::encode(&callback), + st = urlencoding::encode(session_id), + ttl = GRANT_TTL_SECS, + ) +} + // ---- Connect callback: II's two cross-origin JSON POSTs ----------------- #[derive(Debug, Deserialize)] pub struct ConnectCallback { - /// The single-use connect state (= session id) set at device-authorization. + /// The single-use connect state (= session id). state: String, /// Present only on the completion notification; a decimal string of ns since /// the epoch (u64 ns overflows JSON numbers, so it is a string on the wire). @@ -318,42 +556,32 @@ pub async fn connect_callback(State(store): State, Json(body): Json { - { - let devices = store.devices.read().await; - match devices.get(&body.state) { - Some(d) if d.created.elapsed() < DEVICE_CODE_TTL => {} - _ => { - return (StatusCode::FORBIDDEN, Json(json!({ "error": "invalid_state" }))).into_response() - } - } + let is_authcode = store.connect_known_authz(&body.state).await; + if !is_authcode && !store.connect_known_device(&body.state).await { + return (StatusCode::FORBIDDEN, Json(json!({ "error": "invalid_state" }))).into_response(); } let public_key = store.identities.session_pubkey_b64(&body.state).await; - (StatusCode::OK, Json(json!({ "public_key": public_key }))).into_response() + let mut resp = json!({ "public_key": public_key }); + if is_authcode { + resp["finish_url"] = + json!(format!("{}/oauth/finish?id={}", base_url(), urlencoding::encode(&body.state))); + } + (StatusCode::OK, Json(resp)).into_response() } - // (b) Completion notification — best-effort. Update the grant if the + // (b) Completion notification — best-effort. Mark the grant live if the // connection is still known; tolerate a missing/expired state (e.g. the // grant was already consumed via the signed-call fallback) with a 2xx, so // a late completion POST doesn't make II treat an otherwise-successful // connect as failed. Never create a session for an unknown state. Some(exp) => { - let known = { - let mut devices = store.devices.write().await; - match devices.get_mut(&body.state) { - Some(d) => { - d.live = true; - true - } - None => false, - } - }; + let known = store.mark_live(&body.state).await; if known { match exp.trim().parse::() { Ok(exp_ns) => store.identities.set_grant_expiration(&body.state, exp_ns).await, - // A malformed expiration is non-fatal: the grant may still be - // live (registration happened under the user's auth); a signed - // call is the source of truth. Err(_) => tracing::warn!("connect completion had unparseable expiration"), } } @@ -362,20 +590,57 @@ pub async fn connect_callback(State(store): State, Json(body): Json bool { + self.authz + .read() + .await + .get(state) + .is_some_and(|a| a.created.elapsed() < CONNECT_TTL) + } + + /// Whether `state` names a known, unexpired device pending connect. + async fn connect_known_device(&self, state: &str) -> bool { + self.devices + .read() + .await + .get(state) + .is_some_and(|d| d.created.elapsed() < CONNECT_TTL) + } + + /// Mark the pending connect live in whichever flow owns `state`; returns + /// whether it was known. + async fn mark_live(&self, state: &str) -> bool { + let mut known = false; + if let Some(a) = self.authz.write().await.get_mut(state) { + a.live = true; + known = true; + } + if let Some(d) = self.devices.write().await.get_mut(state) { + d.live = true; + known = true; + } + known + } +} + /// Top-level redirect via a script-initiated navigation (`location.replace`) /// rather than an HTTP `Location` header, so the II `/mcp` URL's fragment (`#…`) /// is preserved (a `Location` redirect drops it in some clients). fn js_redirect(url: &str) -> Response { - let safe = url - .replace('\\', "\\\\") - .replace('"', "\\\"") - .replace('<', "\\x3c"); Html(format!( - "" + "", + js_escape(url) )) .into_response() } +/// Escape a string for embedding inside a double-quoted JS string literal. +fn js_escape(s: &str) -> String { + s.replace('\\', "\\\\").replace('"', "\\\"").replace('<', "\\x3c") +} + fn connect_error(message: &str) -> Response { let safe = message.replace('<', "<"); ( @@ -387,12 +652,14 @@ fn connect_error(message: &str) -> Response { .into_response() } -// ---- Token: poll for the device grant ----------------------------------- +// ---- Token: exchange a code, or poll a device grant --------------------- #[derive(Debug, Deserialize)] pub struct TokenForm { grant_type: String, #[serde(default)] + code: String, + #[serde(default)] device_code: String, #[serde(default)] client_id: String, @@ -400,24 +667,49 @@ pub struct TokenForm { code_verifier: Option, } -/// POST /oauth/token (RFC 8628 §3.4–3.5) — the client polls here with the device -/// code until the user finishes the II handshake, then receives the access token. +/// POST /oauth/token — `authorization_code` (poll-bridge clients) or the RFC 8628 +/// `device_code` grant. pub async fn token(State(store): State, Form(req): Form) -> Response { - if req.grant_type != DEVICE_GRANT_TYPE { - return oauth_err(StatusCode::BAD_REQUEST, "unsupported_grant_type", "only the device_code grant"); + match req.grant_type.as_str() { + "authorization_code" => token_authorization_code(store, req).await, + g if g == DEVICE_GRANT_TYPE => token_device_code(store, req).await, + _ => oauth_err(StatusCode::BAD_REQUEST, "unsupported_grant_type", "authorization_code or device_code"), } +} + +async fn token_authorization_code(store: AuthStore, req: TokenForm) -> Response { + let grant = match store.codes.write().await.remove(&req.code) { + Some(g) if g.created.elapsed() < CODE_TTL => g, + Some(_) => return oauth_err(StatusCode::BAD_REQUEST, "invalid_grant", "code expired"), + None => return oauth_err(StatusCode::BAD_REQUEST, "invalid_grant", "unknown or used code"), + }; + if !req.client_id.is_empty() && req.client_id != grant.client_id { + return oauth_err(StatusCode::BAD_REQUEST, "invalid_client", "client_id mismatch"); + } + // Enforce PKCE (a challenge is always stored by `authorize`). + if let Some(challenge) = &grant.code_challenge { + let verifier = match &req.code_verifier { + Some(v) => v, + None => return oauth_err(StatusCode::BAD_REQUEST, "invalid_grant", "code_verifier required"), + }; + if &pkce_s256(verifier) != challenge { + return oauth_err(StatusCode::BAD_REQUEST, "invalid_grant", "PKCE verification failed"); + } + } + store.authz.write().await.remove(&grant.session_id); + issue_token(&store, &grant.session_id, None).await +} + +async fn token_device_code(store: AuthStore, req: TokenForm) -> Response { if req.device_code.is_empty() { return oauth_err(StatusCode::BAD_REQUEST, "invalid_request", "device_code required"); } - - // Locate the device grant, enforce expiry + slow_down, and snapshot what we - // need. Done means: completion POST arrived, or a signed call now succeeds. let (session_id, code_challenge, scope) = { let mut devices = store.devices.write().await; let Some(d) = devices.values_mut().find(|d| d.device_code == req.device_code) else { return oauth_err(StatusCode::BAD_REQUEST, "invalid_grant", "unknown or used device_code"); }; - if d.created.elapsed() >= DEVICE_CODE_TTL { + if d.created.elapsed() >= CONNECT_TTL { return oauth_err(StatusCode::BAD_REQUEST, "expired_token", "device_code expired — restart"); } if !req.client_id.is_empty() { @@ -427,7 +719,6 @@ pub async fn token(State(store): State, Form(req): Form) - } } } - // slow_down if polling faster than the advertised interval. if let Some(last) = d.last_poll { if last.elapsed() < Duration::from_secs(POLL_INTERVAL_SECS) { return oauth_err(StatusCode::BAD_REQUEST, "slow_down", "poll no faster than the interval"); @@ -437,24 +728,18 @@ pub async fn token(State(store): State, Form(req): Form) - (d.session_id.clone(), d.code_challenge.clone(), d.scope.clone()) }; - // Determine liveness. Prefer the completion-POST flag; fall back to a cheap - // signed call (the completion POST is best-effort and may never arrive). - let live = { - let flag = store - .devices - .read() - .await - .values() - .find(|d| d.device_code == req.device_code) - .map(|d| d.live) - .unwrap_or(false); - flag || store.identities.grant_is_live(&session_id).await - }; - if !live { + let live_flag = store + .devices + .read() + .await + .values() + .find(|d| d.device_code == req.device_code) + .map(|d| d.live) + .unwrap_or(false); + if !store.grant_live(&session_id, live_flag).await { return oauth_err(StatusCode::BAD_REQUEST, "authorization_pending", "waiting for Internet Identity"); } - // Enforce PKCE if the client bound one at device-authorization time. if let Some(challenge) = &code_challenge { let verifier = match &req.code_verifier { Some(v) => v, @@ -465,12 +750,15 @@ pub async fn token(State(store): State, Form(req): Form) - } } - // Consume the device grant (single use) and issue the token, bound to the - // session key's principal (self_authenticating(session_pubkey)). store.devices.write().await.retain(|_, d| d.device_code != req.device_code); + issue_token(&store, &session_id, scope).await +} + +/// Mint + store an access token bound to the session key's principal. +async fn issue_token(store: &AuthStore, session_id: &str, scope: Option) -> Response { let principal = store .identities - .session_principal(&session_id) + .session_principal(session_id) .await .unwrap_or_else(|| "unknown".to_string()); @@ -479,12 +767,12 @@ pub async fn token(State(store): State, Form(req): Form) - access_token.clone(), TokenInfo { principal: principal.clone(), - session_id, + session_id: session_id.to_string(), created: Instant::now(), ttl: TOKEN_TTL, }, ); - tracing::info!(%principal, "issued MCP access token (device grant)"); + tracing::info!(%principal, "issued MCP access token"); let mut resp = json!({ "access_token": access_token, @@ -510,32 +798,42 @@ pub struct RegisterRequest { client_name: Option, #[serde(default)] redirect_uris: Vec, + #[serde(default)] + grant_types: Vec, } -/// POST /oauth/register — the device grant needs no redirect, so `redirect_uris` -/// is optional; any supplied are recorded but unused. +/// POST /oauth/register (RFC 7591). `redirect_uris` are stored for the auth-code +/// flow; a device-only client may register without any. pub async fn register(State(store): State, Json(req): Json) -> Response { let client_id = format!("client-{}", Uuid::new_v4()); let snapshot = { let mut clients = store.clients.write().await; - clients.insert( - client_id.clone(), - ClientReg { - redirect_uris: req.redirect_uris.clone(), - }, - ); + clients.insert(client_id.clone(), ClientReg { redirect_uris: req.redirect_uris.clone() }); clients.clone() }; tokio::task::spawn_blocking(move || persist_clients(&snapshot)).await.ok(); + // Honour the requested grant types (intersected with what we support); fall + // back to both if the client didn't ask for any. + let supported = ["authorization_code", DEVICE_GRANT_TYPE]; + let granted: Vec = if req.grant_types.is_empty() { + supported.iter().map(|s| s.to_string()).collect() + } else { + req.grant_types + .iter() + .filter(|g| supported.contains(&g.as_str())) + .cloned() + .collect() + }; + // Public client (PKCE, no secret): OMIT client_secret entirely (returning // null breaks clients that validate it as a string). let mut resp = json!({ "client_id": client_id, "redirect_uris": req.redirect_uris, "token_endpoint_auth_method": "none", - "grant_types": [DEVICE_GRANT_TYPE], - "response_types": [], + "grant_types": granted, + "response_types": ["code"], }); if let Some(name) = req.client_name { resp["client_name"] = json!(name); @@ -550,10 +848,12 @@ pub async fn authorization_server_metadata() -> Response { let base = base_url(); Json(json!({ "issuer": base, + "authorization_endpoint": format!("{base}/oauth/authorize"), "device_authorization_endpoint": format!("{base}/oauth/device_authorization"), "token_endpoint": format!("{base}/oauth/token"), "registration_endpoint": format!("{base}/oauth/register"), - "grant_types_supported": [DEVICE_GRANT_TYPE], + "response_types_supported": ["code"], + "grant_types_supported": ["authorization_code", DEVICE_GRANT_TYPE], "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["none"], })) @@ -614,6 +914,21 @@ pub async fn require_token(State(store): State, mut request: Request< } } +/// An `http://` loopback redirect (any port), matched on the parsed **host** so +/// look-alikes can't slip through. Parsing (not `strip_prefix`) is what defends +/// against authority tricks: `http://localhost.evil.com`, `http://localhost@evil.com`, +/// and the userinfo-with-port form `http://localhost:1234@evil.com` all parse to +/// host `evil.com` (or carry userinfo) and are rejected. +fn is_loopback_redirect(redirect_uri: &str) -> bool { + let Ok(url) = url::Url::parse(redirect_uri) else { + return false; + }; + url.scheme() == "http" + && url.username().is_empty() + && url.password().is_none() + && matches!(url.host_str(), Some("localhost" | "127.0.0.1" | "[::1]")) +} + fn oauth_err(status: StatusCode, error: &str, desc: &str) -> Response { (status, Json(json!({ "error": error, "error_description": desc }))).into_response() } @@ -623,7 +938,7 @@ pub type _JsonValue = Value; #[cfg(test)] mod tests { - use super::pkce_s256; + use super::{build_redirect, is_loopback_redirect, pkce_s256, redirect_allowed, ClientReg}; /// RFC 7636 Appendix B test vector. #[test] @@ -632,4 +947,60 @@ mod tests { let expected = "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"; assert_eq!(pkce_s256(verifier), expected); } + + #[test] + fn redirect_requires_registration() { + let reg = ClientReg { + redirect_uris: vec!["https://claude.ai/api/mcp/auth_callback".to_string()], + }; + // Hosted redirects: exact registered match only. + assert!(redirect_allowed(Some(®), "https://claude.ai/api/mcp/auth_callback")); + assert!(!redirect_allowed(Some(®), "https://claude.ai/api/mcp/auth_callback/x")); + // Unregistered clients get nothing — not even loopback. + assert!(!redirect_allowed(None, "https://claude.ai/api/mcp/auth_callback")); + assert!(!redirect_allowed(None, "http://127.0.0.1:51000/callback")); + assert!(!redirect_allowed(None, "http://[::1]:8080/cb")); + } + + /// A registered loopback redirect matches at ANY port (RFC 8252 §7.3 — the + /// client binds an ephemeral port each run), but host and path must match, + /// and a registered hosted URI grants no loopback latitude. + #[test] + fn registered_loopback_matches_any_port() { + let reg = ClientReg { + redirect_uris: vec!["http://localhost:54321/callback".to_string()], + }; + assert!(redirect_allowed(Some(®), "http://localhost:54321/callback")); + assert!(redirect_allowed(Some(®), "http://localhost:61832/callback")); + assert!(redirect_allowed(Some(®), "http://localhost/callback")); + // Different path or host (even another loopback host): rejected. + assert!(!redirect_allowed(Some(®), "http://localhost:61832/other")); + assert!(!redirect_allowed(Some(®), "http://127.0.0.1:61832/callback")); + // Look-alike hosts fail is_loopback_redirect on the requested side. + assert!(!redirect_allowed(Some(®), "http://localhost.evil.com:54321/callback")); + // A registered HOSTED uri gives no loopback latitude. + let hosted = ClientReg { + redirect_uris: vec!["https://claude.ai/cb".to_string()], + }; + assert!(!redirect_allowed(Some(&hosted), "http://localhost:1234/cb")); + } + + #[test] + fn loopback_rejects_lookalikes() { + assert!(is_loopback_redirect("http://127.0.0.1:51000/callback")); + assert!(is_loopback_redirect("http://[::1]:8080/cb")); + assert!(!is_loopback_redirect("http://localhost.evil.com/cb")); + assert!(!is_loopback_redirect("http://localhost@evil.com/cb")); + assert!(!is_loopback_redirect("http://localhost:1234@evil.com/cb")); + assert!(!is_loopback_redirect("https://evil.com/cb")); + } + + #[test] + fn build_redirect_encodes_code_and_state() { + let r = build_redirect("https://claude.ai/cb", "mcp-code-1", "abc/def"); + assert_eq!(r, "https://claude.ai/cb?code=mcp-code-1&state=abc%2Fdef"); + // Appends with & when the redirect already has a query. + let r2 = build_redirect("https://x.test/cb?foo=1", "c", ""); + assert_eq!(r2, "https://x.test/cb?foo=1&code=c"); + } } diff --git a/src/main.rs b/src/main.rs index 9efec37..83cbfea 100644 --- a/src/main.rs +++ b/src/main.rs @@ -949,9 +949,12 @@ async fn main() -> anyhow::Result<()> { let store = auth::AuthStore::new(identities.clone()); // /mcp is gated by a bearer token issued after Internet Identity login. + // `route_layer` (not `layer`) applies the gate ONLY to matched /mcp routes, so + // unmatched paths fall through to a 404 instead of the token gate's 401 — and + // the `/oauth/*` and `/.well-known/*` routes stay exempt from the bearer check. let protected_mcp = axum::Router::new() .nest_service("/mcp", mcp) - .layer(axum::middleware::from_fn_with_state( + .route_layer(axum::middleware::from_fn_with_state( store.clone(), auth::require_token, )); @@ -983,6 +986,8 @@ async fn main() -> anyhow::Result<()> { "/.well-known/oauth-protected-resource/mcp", axum::routing::get(auth::protected_resource_metadata), ) + .route("/oauth/authorize", axum::routing::get(auth::authorize)) + .route("/oauth/finish", axum::routing::get(auth::finish)) .route("/oauth/device_authorization", axum::routing::post(auth::device_authorization)) .route("/oauth/device", axum::routing::get(auth::device_verify)) .route("/oauth/connect/callback", axum::routing::post(auth::connect_callback))