ibkr-agent-gateway 0.5.2

//! Remote MCP HTTP auth enforcement.

use super::{
    http_server::HttpMcpResponse, oauth_metadata::protected_resource_metadata,
    session::HttpMcpSessionIds,
};
use crate::internal::auth::{OAuthContextInput, RemoteAuthContext, remote_auth_context_from_input};
use crate::internal::config::RemoteMcpConfig;
use crate::internal::domain::{ErrorCode, GatewayError};
use crate::internal::oauth::{Jwks, OAuthIssuerConfig, PreparedOAuthVerifier};
use serde_json::json;
use std::collections::BTreeMap;
use time::OffsetDateTime;

/// Prepared remote MCP auth state for repeated HTTP requests.
#[derive(Clone, Debug)]
pub struct RemoteMcpAuthVerifier {
    verifier: PreparedOAuthVerifier,
}

impl RemoteMcpAuthVerifier {
    /// Builds a prepared verifier from fail-closed remote MCP config and JWKS.
    pub fn new(config: &RemoteMcpConfig, jwks: &Jwks) -> Result<Self, GatewayError> {
        let oauth_config = oauth_issuer_config(config)?;
        let verifier = PreparedOAuthVerifier::new(oauth_config, jwks)?;
        Ok(Self { verifier })
    }
}

/// Authorizes a remote MCP request before tool execution.
pub fn authorize_remote_request(
    config: &RemoteMcpConfig,
    jwks: &Jwks,
    headers: &BTreeMap<String, String>,
    required_scope: &str,
) -> Result<RemoteAuthContext, HttpMcpResponse> {
    let verifier = RemoteMcpAuthVerifier::new(config, jwks)
        .map_err(|error| auth_error_response(config, error))?;
    authorize_remote_request_with_verifier(config, &verifier, headers, required_scope)
}

/// Authorizes a remote MCP request using precompiled auth state.
pub fn authorize_remote_request_with_verifier(
    config: &RemoteMcpConfig,
    verifier: &RemoteMcpAuthVerifier,
    headers: &BTreeMap<String, String>,
    required_scope: &str,
) -> Result<RemoteAuthContext, HttpMcpResponse> {
    authorize_remote_request_with_optional_scope(config, verifier, headers, Some(required_scope))
}

/// Authorizes a remote MCP request without requiring one specific scope.
///
/// Use this for JSON-RPC methods such as `initialize` and `tools/list`, where
/// the server needs the validated token's whole granted-scope set before it can
/// decide which tools to show.
pub fn authorize_remote_request_without_required_scope(
    config: &RemoteMcpConfig,
    verifier: &RemoteMcpAuthVerifier,
    headers: &BTreeMap<String, String>,
) -> Result<RemoteAuthContext, HttpMcpResponse> {
    authorize_remote_request_with_optional_scope(config, verifier, headers, None)
}

fn authorize_remote_request_with_optional_scope(
    config: &RemoteMcpConfig,
    verifier: &RemoteMcpAuthVerifier,
    headers: &BTreeMap<String, String>,
    required_scope: Option<&str>,
) -> Result<RemoteAuthContext, HttpMcpResponse> {
    let token =
        bearer_token(headers).ok_or_else(|| auth_error_response(config, missing_token()))?;
    let validated = verifier
        .verifier
        .validate_bearer_jwt(token, required_scope, OffsetDateTime::now_utc())
        .map_err(|error| auth_error_response(config, error))?;
    let session_ids = HttpMcpSessionIds::from_headers(headers);
    let expires_at = OffsetDateTime::from_unix_timestamp(validated.claims.exp)
        .map_err(|_| auth_error_response(config, invalid_time()))?;

    remote_auth_context_from_input(OAuthContextInput {
        subject: validated.claims.sub,
        issuer: validated.claims.iss,
        audience: validated.audience,
        scopes: validated.granted_scopes,
        expires_at,
        token_id_hash: validated.token_id_hash,
        request_id: session_ids.request_id,
        session_id: session_ids.session_id,
    })
    .map_err(|error| auth_error_response(config, error))
}

fn oauth_issuer_config(config: &RemoteMcpConfig) -> Result<OAuthIssuerConfig, GatewayError> {
    let issuer = config.issuer.as_ref().ok_or_else(|| {
        GatewayError::new(
            ErrorCode::ConfigInvalid,
            "Remote MCP issuer is not configured",
            false,
            Some("Configure remote_mcp.issuer".to_string()),
        )
    })?;
    let jwks_url = config.jwks_url.as_ref().ok_or_else(|| {
        GatewayError::new(
            ErrorCode::ConfigInvalid,
            "Remote MCP JWKS URL is not configured",
            false,
            Some("Configure remote_mcp.jwks_url".to_string()),
        )
    })?;
    let token_id_hmac_secret = config
        .token_id_hmac_secret
        .as_deref()
        .filter(|secret| !secret.trim().is_empty())
        .ok_or_else(|| {
            GatewayError::new(
                ErrorCode::ConfigInvalid,
                "Remote MCP token id HMAC secret is not configured",
                false,
                Some("Configure remote_mcp.token_id_hmac_secret from a secret store".to_string()),
            )
        })?;

    Ok(OAuthIssuerConfig {
        issuer: issuer.to_string(),
        jwks_url: jwks_url.to_string(),
        audiences: config.audiences.clone(),
        allowed_scopes: config.allowed_scopes.clone(),
        clock_skew_seconds: config.clock_skew_seconds,
        metadata_url: config.metadata_url.as_ref().map(ToString::to_string),
        token_id_hmac_secret: token_id_hmac_secret.as_bytes().to_vec(),
    })
}

fn bearer_token(headers: &BTreeMap<String, String>) -> Option<&str> {
    let value = headers
        .iter()
        .find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
        .map(|(_, value)| value.as_str())?;
    value
        .strip_prefix("Bearer ")
        .or_else(|| value.strip_prefix("bearer "))
        .filter(|token| !token.trim().is_empty())
}

fn auth_error_response(config: &RemoteMcpConfig, error: GatewayError) -> HttpMcpResponse {
    let status = match error.code {
        ErrorCode::AuthMissingScope => 403,
        ErrorCode::AuthTokenMissing
        | ErrorCode::AuthTokenInvalid
        | ErrorCode::AuthTokenExpired
        | ErrorCode::AuthInvalidIssuer
        | ErrorCode::AuthInvalidAudience => 401,
        ErrorCode::ConfigInvalid => 500,
        ErrorCode::BrokerBackendUnavailable => 503,
        _ => 401,
    };
    let mut response = HttpMcpResponse::json(
        status,
        json!({
            "error": error.code,
            "message": auth_response_message(status),
            "oauth_protected_resource": protected_resource_metadata(config)
        }),
    );
    if status == 401 {
        response.headers.insert(
            "www-authenticate".to_string(),
            "Bearer resource_metadata=\"/.well-known/oauth-protected-resource\"".to_string(),
        );
    }
    response
}

const fn auth_response_message(status: u16) -> &'static str {
    match status {
        401 => "Authentication failed",
        403 => "Authorization failed",
        503 => "Authentication service unavailable",
        _ => "Remote MCP authentication could not be completed",
    }
}

fn missing_token() -> GatewayError {
    GatewayError::new(
        ErrorCode::AuthTokenMissing,
        "Bearer token is required",
        false,
        Some("Send Authorization: Bearer <token>".to_string()),
    )
}

fn invalid_time() -> GatewayError {
    GatewayError::new(
        ErrorCode::AuthTokenInvalid,
        "JWT expiry timestamp is invalid",
        false,
        Some("Use a token with a valid exp claim".to_string()),
    )
}