supercode_harness/

mcp_oauth.rs

//! P5-2 (COMPOSABLE-HARNESS-DESIGN.md §2 module 15 D7 row 2 "OAuth"; §2.1
//! dep "`model.oauth` → trust-grade token storage" — the same security
//! class applies here): OAuth PROTOCOL support for authenticated remote MCP
//! servers.
//!
//! **Scope, stated plainly.** This module implements:
//! - the OAuth 2.0 Device Authorization Grant (RFC 8628) — the
//!   non-interactive path: no browser/redirect listener needed, just a
//!   `user_code`/`verification_uri` a caller prints and a background poll;
//! - refresh-token exchange.
//!
//! It does **NOT** implement the interactive authorization-code + PKCE +
//! local-redirect-listener browser flow — that needs a UI to open a
//! browser and a local HTTP listener to catch the redirect, which is
//! `tui`'s job (P5 item #4, not yet built). This is a **tui-deferred**
//! citation, not a silent gap: a server that only offers the browser flow
//! (no device-code grant) simply isn't reachable through this module yet.
//!
//! **Token storage is NOT this module's job.** This module only speaks the
//! wire protocol and returns [`McpOAuthTokens`] values — persisting them is
//! a CLI-layer concern (`crates/cli/src/userconfig.rs`'s
//! `save_mcp_oauth_tokens`/`load_mcp_oauth_tokens`), same trust-grade
//! posture (owner-only permissions, user/global-directory-only, never
//! project-readable) as `Config::api_key`/`save_api_key` (§3.2 S13) — this
//! crate never touches a filesystem for a credential.

use std::time::Duration;

use serde::Deserialize;

use crate::error::{Error, Result};

/// The endpoints/identity an MCP server's OAuth device-code flow needs.
/// Carries no token — see the module doc comment.
#[derive(Debug, Clone)]
pub struct OAuthEndpoints {
    /// RFC 8628 device authorization endpoint.
    pub device_authorization_endpoint: String,
    /// Token endpoint (also used for the refresh-token grant).
    pub token_endpoint: String,
    /// OAuth client id.
    pub client_id: String,
    /// Optional scope string.
    pub scope: Option<String>,
}

/// A trust-grade credential pair — see the module doc comment for why this
/// crate never persists one itself.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct McpOAuthTokens {
    /// The bearer access token.
    pub access_token: String,
    /// The refresh token, if the server issued one.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub refresh_token: Option<String>,
    /// Unix-epoch seconds this access token expires at, if known.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expires_at_secs: Option<u64>,
}

impl McpOAuthTokens {
    /// Whether this token is expired (or about to expire within `skew`
    /// seconds) as of `now_secs`. `expires_at_secs: None` (unknown
    /// lifetime) is treated as NOT expired — a server that never told us
    /// when the token dies is assumed valid until it actually fails.
    pub fn is_expired_at(&self, now_secs: u64, skew_secs: u64) -> bool {
        self.expires_at_secs
            .map(|exp| now_secs.saturating_add(skew_secs) >= exp)
            .unwrap_or(false)
    }
}

/// The RFC 8628 device-authorization-response fields this module needs.
#[derive(Debug, Clone)]
pub struct DeviceAuthorization {
    /// Opaque device code the client polls the token endpoint with.
    pub device_code: String,
    /// Short code the USER enters at `verification_uri`.
    pub user_code: String,
    /// The URL the user visits.
    pub verification_uri: String,
    /// A URL that already embeds `user_code`, if the server provided one —
    /// print this instead of `verification_uri` + `user_code` separately
    /// when present.
    pub verification_uri_complete: Option<String>,
    /// How often (seconds) the client should poll the token endpoint.
    pub interval_secs: u64,
    /// How long (seconds) `device_code` remains valid.
    pub expires_in_secs: u64,
}

#[derive(Deserialize)]
struct DeviceAuthResponse {
    device_code: String,
    user_code: String,
    verification_uri: String,
    #[serde(default)]
    verification_uri_complete: Option<String>,
    #[serde(default = "default_interval")]
    interval: u64,
    #[serde(default = "default_expires_in")]
    expires_in: u64,
}
fn default_interval() -> u64 {
    5
}
fn default_expires_in() -> u64 {
    600
}

#[derive(Deserialize)]
struct TokenResponse {
    access_token: String,
    #[serde(default)]
    refresh_token: Option<String>,
    #[serde(default)]
    expires_in: Option<u64>,
}

#[derive(Deserialize)]
struct TokenErrorResponse {
    error: String,
}

/// Step 1 of RFC 8628: request a device/user code pair.
pub async fn start_device_authorization(
    client: &reqwest::Client,
    ep: &OAuthEndpoints,
) -> Result<DeviceAuthorization> {
    let mut form = vec![("client_id", ep.client_id.as_str())];
    if let Some(scope) = &ep.scope {
        form.push(("scope", scope.as_str()));
    }
    let resp = client
        .post(&ep.device_authorization_endpoint)
        .form(&form)
        .send()
        .await
        .map_err(|e| Error::tool("mcp_oauth", format!("device authorization request: {e}")))?;
    if !resp.status().is_success() {
        return Err(Error::tool(
            "mcp_oauth",
            format!("device authorization: http status {}", resp.status()),
        ));
    }
    let body: DeviceAuthResponse = resp.json().await.map_err(|e| {
        Error::tool(
            "mcp_oauth",
            format!("decoding device authorization response: {e}"),
        )
    })?;
    Ok(DeviceAuthorization {
        device_code: body.device_code,
        user_code: body.user_code,
        verification_uri: body.verification_uri,
        verification_uri_complete: body.verification_uri_complete,
        interval_secs: body.interval,
        expires_in_secs: body.expires_in,
    })
}

/// One poll of the token endpoint for a device code — RFC 8628 §3.5. The
/// server replies `authorization_pending` until the user finishes at
/// `verification_uri`; the caller (e.g. `poll_until_authorized`) is
/// expected to sleep `interval_secs` and retry.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DevicePollOutcome {
    /// The user hasn't completed authorization yet — keep polling at the
    /// same interval.
    Pending,
    /// The user hasn't completed authorization yet AND the server wants
    /// polling slowed down (RFC 8628 §3.5) — the next poll should wait
    /// `interval + 5s`, not just `interval`.
    SlowDown,
    /// The user completed authorization; tokens are attached.
    Authorized(McpOAuthTokens),
    /// The user (or the server) denied/cancelled — stop polling.
    Denied,
    /// The device code expired before authorization completed.
    Expired,
}

/// A single token-endpoint poll for the device-code grant.
pub async fn poll_device_token(
    client: &reqwest::Client,
    ep: &OAuthEndpoints,
    device_code: &str,
) -> Result<DevicePollOutcome> {
    let form = [
        ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
        ("device_code", device_code),
        ("client_id", ep.client_id.as_str()),
    ];
    let resp = client
        .post(&ep.token_endpoint)
        .form(&form)
        .send()
        .await
        .map_err(|e| Error::tool("mcp_oauth", format!("token poll: {e}")))?;
    if resp.status().is_success() {
        let body: TokenResponse = resp
            .json()
            .await
            .map_err(|e| Error::tool("mcp_oauth", format!("decoding token response: {e}")))?;
        return Ok(DevicePollOutcome::Authorized(McpOAuthTokens {
            access_token: body.access_token,
            refresh_token: body.refresh_token,
            expires_at_secs: body.expires_in.map(|secs| now_secs() + secs),
        }));
    }
    let body: TokenErrorResponse = resp.json().await.unwrap_or(TokenErrorResponse {
        error: "unknown_error".to_string(),
    });
    match body.error.as_str() {
        "authorization_pending" => Ok(DevicePollOutcome::Pending),
        "slow_down" => Ok(DevicePollOutcome::SlowDown),
        "expired_token" => Ok(DevicePollOutcome::Expired),
        _ => Ok(DevicePollOutcome::Denied),
    }
}

/// The whole non-interactive device-code flow: start authorization, invoke
/// `on_prompt` exactly once with the [`DeviceAuthorization`] (so the caller
/// can print `verification_uri`/`user_code` for the user), then poll until
/// authorized/denied/expired — bounded by `expires_in_secs`, sleeping
/// `interval_secs` between attempts (never faster, per RFC 8628's
/// `slow_down` semantics — a `slow_down` response widens the interval by a
/// further 5s, same as the spec recommends).
pub async fn run_device_flow(
    client: &reqwest::Client,
    ep: &OAuthEndpoints,
    on_prompt: impl FnOnce(&DeviceAuthorization),
) -> Result<McpOAuthTokens> {
    let auth = start_device_authorization(client, ep).await?;
    on_prompt(&auth);
    let deadline = now_secs() + auth.expires_in_secs;
    let mut interval = auth.interval_secs.max(1);
    loop {
        tokio::time::sleep(Duration::from_secs(interval)).await;
        match poll_device_token(client, ep, &auth.device_code).await? {
            DevicePollOutcome::Authorized(tokens) => return Ok(tokens),
            DevicePollOutcome::Pending => {
                if now_secs() >= deadline {
                    return Err(Error::tool(
                        "mcp_oauth",
                        "device code expired while polling",
                    ));
                }
            }
            DevicePollOutcome::SlowDown => {
                interval += 5;
                if now_secs() >= deadline {
                    return Err(Error::tool(
                        "mcp_oauth",
                        "device code expired while polling",
                    ));
                }
            }
            DevicePollOutcome::Denied => {
                return Err(Error::tool("mcp_oauth", "authorization was denied"));
            }
            DevicePollOutcome::Expired => {
                return Err(Error::tool("mcp_oauth", "device code expired"));
            }
        }
    }
}

/// Exchange a refresh token for a new access token.
pub async fn refresh_token(
    client: &reqwest::Client,
    ep: &OAuthEndpoints,
    refresh_token: &str,
) -> Result<McpOAuthTokens> {
    let form = [
        ("grant_type", "refresh_token"),
        ("refresh_token", refresh_token),
        ("client_id", ep.client_id.as_str()),
    ];
    let resp = client
        .post(&ep.token_endpoint)
        .form(&form)
        .send()
        .await
        .map_err(|e| Error::tool("mcp_oauth", format!("refresh request: {e}")))?;
    if !resp.status().is_success() {
        return Err(Error::tool(
            "mcp_oauth",
            format!("refresh: http status {}", resp.status()),
        ));
    }
    let body: TokenResponse = resp
        .json()
        .await
        .map_err(|e| Error::tool("mcp_oauth", format!("decoding refresh response: {e}")))?;
    Ok(McpOAuthTokens {
        access_token: body.access_token,
        // A server that omits `refresh_token` on refresh means "reuse the
        // same one" per RFC 6749 §6 — the caller (which already has the OLD
        // refresh token) is responsible for keeping it if this is `None`.
        refresh_token: body.refresh_token,
        expires_at_secs: body.expires_in.map(|secs| now_secs() + secs),
    })
}

fn now_secs() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

/// Build the `Authorization: Bearer <token>` header value for a stored
/// [`McpOAuthTokens`] — the shape [`crate::mcp::McpClient::connect_http`]/
/// `connect_sse`'s `headers` map expects.
pub fn bearer_header(tokens: &McpOAuthTokens) -> (String, String) {
    (
        "Authorization".to_string(),
        format!("Bearer {}", tokens.access_token),
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn is_expired_at_treats_unknown_lifetime_as_not_expired() {
        let t = McpOAuthTokens {
            access_token: "x".into(),
            refresh_token: None,
            expires_at_secs: None,
        };
        assert!(!t.is_expired_at(u64::MAX / 2, 0));
    }

    #[test]
    fn is_expired_at_honors_skew() {
        let t = McpOAuthTokens {
            access_token: "x".into(),
            refresh_token: None,
            expires_at_secs: Some(1000),
        };
        assert!(!t.is_expired_at(900, 30));
        assert!(t.is_expired_at(980, 30)); // within skew of expiry
        assert!(t.is_expired_at(1000, 0));
    }

    #[test]
    fn bearer_header_has_the_expected_shape() {
        let t = McpOAuthTokens {
            access_token: "secret123".into(),
            refresh_token: None,
            expires_at_secs: None,
        };
        let (name, value) = bearer_header(&t);
        assert_eq!(name, "Authorization");
        assert_eq!(value, "Bearer secret123");
    }
}