clawmates/crates/cm-api/src/routes/oauth.rs

//! OAuth authorization-code connects (§7.8 [+] and §10 MCP-OAuth).
//!
//! `start` records a one-time state and returns the authorization URL;
//! `callback` consumes the state (CAS delete — replay-proof), exchanges
//! the code at the issuer's REAL token endpoint, hands the access token to
//! the secret broker, and records the connection. The token, like every
//! credential, is never readable again outside the broker.

use axum::extract::{Query, State};
use axum::response::Redirect;
use axum::Json;
use cm_db::repo::audit::Actor;
use cm_domain::{AgentId, UserId, WorkspaceId};
use serde::Deserialize;
use serde_json::{json, Value};
use uuid::Uuid;

use crate::routes::claws::workspace_agent;
use crate::{ApiError, AppState, Authed};

/// Fetches the issuer's discovery document for its endpoints.
async fn discover(issuer_url: &str) -> Result<(String, String), ApiError> {
    let doc: Value = reqwest::get(format!(
        "{}/.well-known/openid-configuration",
        issuer_url.trim_end_matches('/')
    ))
    .await
    .map_err(|_| ApiError::Internal)?
    .json()
    .await
    .map_err(|_| ApiError::Internal)?;
    let authorize = doc["authorization_endpoint"]
        .as_str()
        .ok_or(ApiError::Internal)?
        .to_owned();
    let token = doc["token_endpoint"]
        .as_str()
        .ok_or(ApiError::Internal)?
        .to_owned();
    Ok((authorize, token))
}

#[derive(Deserialize)]
pub struct StartRequest {
    #[serde(rename = "clawId")]
    claw_id: AgentId,
    provider: String,
    /// "oauth" (configured IdP) or "mcp_oauth" (custom issuer below).
    #[serde(rename = "authType", default = "default_auth_type")]
    auth_type: String,
    /// Required for mcp_oauth: the MCP server's issuer.
    #[serde(rename = "issuerUrl")]
    issuer_url: Option<String>,
}

fn default_auth_type() -> String {
    "oauth".into()
}

/// POST /api/apps/oauth/start → { authorize_url }
pub async fn start(
    State(state): State<AppState>,
    Authed(user): Authed,
    Json(body): Json<StartRequest>,
) -> Result<Json<Value>, ApiError> {
    let agent = workspace_agent(&state, &user, body.claw_id).await?;
    let issuer = match body.auth_type.as_str() {
        "mcp_oauth" => body.issuer_url.clone().ok_or(ApiError::Conflict)?,
        "oauth" => state.oauth.issuer_url.clone().ok_or(ApiError::Conflict)?,
        _ => return Err(ApiError::Conflict),
    };
    let client_id = state.oauth.client_id.clone().ok_or(ApiError::Conflict)?;
    let redirect_base = state
        .oauth
        .redirect_base
        .clone()
        .ok_or(ApiError::Conflict)?;
    let (authorize_endpoint, _) = discover(&issuer).await?;

    let oauth_state = Uuid::now_v7().simple().to_string();
    sqlx::query!(
        "INSERT INTO oauth_states
            (state, workspace_id, user_id, agent_id, provider, auth_type,
             issuer_url, expires_at)
         VALUES ($1, $2, $3, $4, $5, $6, $7, now() + interval '10 minutes')",
        oauth_state,
        user.workspace_id.as_uuid(),
        user.user_id.as_uuid(),
        agent.id.as_uuid(),
        body.provider,
        body.auth_type,
        issuer,
    )
    .execute(&state.pool)
    .await
    .map_err(|_| ApiError::Internal)?;

    let redirect_uri = format!("{redirect_base}/api/apps/oauth/callback");
    let authorize_url = format!(
        "{authorize_endpoint}?response_type=code&client_id={}&redirect_uri={}&state={oauth_state}&scope=openid",
        urlencoding::encode(&client_id),
        urlencoding::encode(&redirect_uri),
    );
    Ok(Json(
        json!({ "authorize_url": authorize_url, "state": oauth_state }),
    ))
}

#[derive(Deserialize)]
pub struct CallbackQuery {
    code: String,
    state: String,
}

/// GET /api/apps/oauth/callback?code=&state= — the IdP redirect target.
/// Unauthenticated by nature; trust comes from the one-time state.
pub async fn callback(
    State(state): State<AppState>,
    Query(query): Query<CallbackQuery>,
) -> Result<Redirect, ApiError> {
    // Consume the state exactly once (replays and forgeries both 404).
    let pending = sqlx::query!(
        "DELETE FROM oauth_states
         WHERE state = $1 AND expires_at > now()
         RETURNING workspace_id, user_id, agent_id, provider, auth_type, issuer_url",
        query.state,
    )
    .fetch_optional(&state.pool)
    .await
    .map_err(|_| ApiError::Internal)?
    .ok_or(ApiError::NotFound)?;

    let client_id = state.oauth.client_id.clone().ok_or(ApiError::Internal)?;
    let redirect_base = state
        .oauth
        .redirect_base
        .clone()
        .ok_or(ApiError::Internal)?;
    let (_, token_endpoint) = discover(&pending.issuer_url).await?;

    // Exchange the code at the REAL token endpoint.
    let mut form = vec![
        ("grant_type", "authorization_code".to_owned()),
        ("code", query.code.clone()),
        ("client_id", client_id),
        (
            "redirect_uri",
            format!("{redirect_base}/api/apps/oauth/callback"),
        ),
    ];
    if let Some(secret) = state.oauth.client_secret.clone() {
        form.push(("client_secret", secret));
    }
    let token_response: Value = reqwest::Client::new()
        .post(&token_endpoint)
        .form(&form)
        .send()
        .await
        .map_err(|_| ApiError::Internal)?
        .json()
        .await
        .map_err(|_| ApiError::Internal)?;
    let access_token = token_response["access_token"]
        .as_str()
        .ok_or(ApiError::Internal)?;

    // The token goes straight to the broker; only the ref persists.
    let socket = state.broker_socket.as_ref().ok_or(ApiError::Internal)?;
    let mut broker = cm_secrets::BrokerClient::connect(socket)
        .await
        .map_err(|_| ApiError::Internal)?;
    let workspace_id = WorkspaceId::from(pending.workspace_id);
    let secret_ref = broker
        .store_secret(
            workspace_id,
            &format!("{}_oauth_token", pending.provider),
            access_token,
        )
        .await
        .map_err(|_| ApiError::Internal)?;
    let connection = cm_db::repo::connections::insert(
        &state.pool,
        workspace_id,
        Some(AgentId::from(pending.agent_id)),
        &pending.provider,
        &pending.auth_type,
        secret_ref,
    )
    .await?;
    cm_db::repo::audit::append(
        &state.pool,
        workspace_id,
        Actor::User(UserId::from(pending.user_id)),
        "app.connected",
        "app_connection",
        &connection.id.to_string(),
        json!({"provider": pending.provider, "auth_type": pending.auth_type}),
    )
    .await?;

    // Back to the claw's Add Apps panel.
    Ok(Redirect::to(&format!(
        "/claws/{}?app=apps",
        pending.agent_id
    )))
}