clawmates/crates/cm-api/tests/oauth_flow.rs

//! The OAuth connect round trip against a REAL local identity provider:
//! a live HTTP server implementing discovery + the token endpoint (the
//! same real-server pattern as the Slack sink). One-time states are
//! replay-proof; the access token ends up broker-held only.

use std::sync::Arc;
use std::time::Duration;

use cm_api::AppState;
use cm_auth::AuthService;
use cm_config::OAuthConfig;
use cm_domain::{Role, User, UserId, Workspace, WorkspaceId};
use cm_llm::ScriptedProvider;
use cm_runtime::{Runtime, RuntimeConfig};
use cm_secrets::{BrokerServer, FileKey};
use serde_json::{json, Value};

/// A minimal real IdP: discovery doc + token endpoint that validates the
/// code and client, then issues a bearer token.
async fn spawn_idp() -> String {
    use axum::extract::State as AxState;
    use axum::routing::{get, post};
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let issuer = format!("http://{addr}");
    let issuer_for_doc = issuer.clone();
    let app = axum::Router::new()
        .route(
            "/.well-known/openid-configuration",
            get(move || {
                let issuer = issuer_for_doc.clone();
                async move {
                    axum::Json(json!({
                        "issuer": issuer,
                        "authorization_endpoint": format!("{issuer}/auth"),
                        "token_endpoint": format!("{issuer}/token"),
                    }))
                }
            }),
        )
        .route(
            "/token",
            post(
                |AxState(()): AxState<()>, axum::Form(form): axum::Form<Value>| async move {
                    if form["grant_type"] == "authorization_code"
                        && form["code"] == "good-code"
                        && form["client_id"] == "clawmates"
                        && form["client_secret"] == "tc-secret"
                    {
                        axum::Json(json!({
                            "access_token": "idp-access-token-42",
                            "token_type": "Bearer",
                        }))
                        .into_response()
                    } else {
                        (
                            axum::http::StatusCode::BAD_REQUEST,
                            axum::Json(json!({"error": "invalid_grant"})),
                        )
                            .into_response()
                    }
                },
            ),
        )
        .with_state(());
    use axum::response::IntoResponse;
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });
    issuer
}

async fn spawn_broker(pool: sqlx::PgPool) -> std::path::PathBuf {
    let dir = std::env::temp_dir().join(format!("tc-oa-{}", uuid::Uuid::now_v7().simple()));
    std::fs::create_dir_all(&dir).unwrap();
    let key_path = dir.join("broker.key");
    FileKey::generate(&key_path).unwrap();
    let key = FileKey::load(&key_path).unwrap();
    let short = uuid::Uuid::now_v7().simple().to_string();
    let socket = std::path::PathBuf::from(format!("/tmp/tco-{}.sock", &short[short.len() - 12..]));
    let server = BrokerServer::new(pool, key, socket.clone());
    tokio::spawn(async move {
        server.serve().await.unwrap();
    });
    for _ in 0..50 {
        if socket.exists() {
            break;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    socket
}

#[tokio::test]
async fn oauth_round_trip_stores_the_token_in_the_broker_only() {
    let pool = cm_testkit::test_pool().await;
    let issuer = spawn_idp().await;
    let socket = spawn_broker(pool.clone()).await;

    let runtime = Runtime::new(
        pool.clone(),
        Arc::new(ScriptedProvider::from_toml("").unwrap()),
        RuntimeConfig::basic("scripted", 1024),
    );
    let oauth = OAuthConfig {
        issuer_url: Some(issuer.clone()),
        client_id: Some("clawmates".into()),
        client_secret: Some("tc-secret".into()),
        redirect_base: Some("http://127.0.0.1:9".into()), // shape only
    };
    let app = cm_api::router(
        AppState::new(pool.clone(), runtime)
            .with_broker(socket)
            .with_oauth(oauth),
    );
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });
    let base = format!("http://{addr}");
    let client = reqwest::Client::builder()
        .redirect(reqwest::redirect::Policy::none())
        .build()
        .unwrap();

    // Seed and login.
    let ws = Workspace {
        id: WorkspaceId::new(),
        name: "Acme".into(),
        plan: "team".into(),
    };
    cm_db::repo::workspaces::insert(&pool, &ws).await.unwrap();
    let owner = User {
        id: UserId::new(),
        workspace_id: ws.id,
        email: format!("{}@acme.test", UserId::new()),
        role: Role::Owner,
        display_name: "Owner".into(),
        created_at: time::OffsetDateTime::UNIX_EPOCH,
    };
    cm_db::repo::users::insert(&pool, &owner).await.unwrap();
    AuthService::new(pool.clone())
        .set_password(owner.id, "pw")
        .await
        .unwrap();
    let token = client
        .post(format!("{base}/api/auth/login"))
        .json(&json!({"email": owner.email, "password": "pw"}))
        .send()
        .await
        .unwrap()
        .json::<Value>()
        .await
        .unwrap()["token"]
        .as_str()
        .unwrap()
        .to_owned();
    let claw: Value = client
        .post(format!("{base}/api/claws"))
        .bearer_auth(&token)
        .json(&json!({"name": "Scout", "job_title": "Analyst"}))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    let claw_id = claw["id"].as_str().unwrap();

    // Sanity: the IdP discovery endpoint is reachable.
    let doc: Value = reqwest::get(format!("{issuer}/.well-known/openid-configuration"))
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    assert!(doc["token_endpoint"].is_string(), "discovery doc: {doc}");

    // Start: authorize URL points at the IdP with our state.
    let start_res = client
        .post(format!("{base}/api/apps/oauth/start"))
        .bearer_auth(&token)
        .json(&json!({"clawId": claw_id, "provider": "notion"}))
        .send()
        .await
        .unwrap();
    let status = start_res.status();
    let start: Value = start_res.json().await.unwrap();
    assert_eq!(status, 200, "start failed: {start}");
    let authorize_url = start["authorize_url"].as_str().unwrap();
    assert!(authorize_url.starts_with(&format!("{issuer}/auth?")));
    assert!(authorize_url.contains("client_id=clawmates"));
    let oauth_state = start["state"].as_str().unwrap();

    // The IdP redirects back with a code: exchange succeeds, connection
    // exists, token is broker-held only.
    let callback = client
        .get(format!(
            "{base}/api/apps/oauth/callback?code=good-code&state={oauth_state}"
        ))
        .send()
        .await
        .unwrap();
    assert_eq!(callback.status(), 303, "redirects back to the panel");

    let directory: Value = client
        .get(format!("{base}/api/apps?clawId={claw_id}"))
        .bearer_auth(&token)
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    let notion = directory
        .as_array()
        .unwrap()
        .iter()
        .find(|a| a["id"] == "notion")
        .unwrap();
    assert_eq!(notion["connected"], true);

    // The raw token never landed in Postgres outside the encrypted store.
    let leaked = sqlx::query_scalar::<_, i64>(
        "SELECT count(*) FROM secrets WHERE convert_from(ciphertext, 'UTF8') LIKE '%idp-access-token%'",
    )
    .fetch_one(&pool)
    .await
    .unwrap_or(0);
    assert_eq!(leaked, 0, "token must be encrypted at rest");

    // Replayed state is refused.
    let replay = client
        .get(format!(
            "{base}/api/apps/oauth/callback?code=good-code&state={oauth_state}"
        ))
        .send()
        .await
        .unwrap();
    assert_eq!(replay.status(), 404);

    // Forged state is refused.
    let forged = client
        .get(format!(
            "{base}/api/apps/oauth/callback?code=good-code&state=forged"
        ))
        .send()
        .await
        .unwrap();
    assert_eq!(forged.status(), 404);

    // A bad code fails the exchange (fresh state, IdP rejects).
    let start2: Value = client
        .post(format!("{base}/api/apps/oauth/start"))
        .bearer_auth(&token)
        .json(&json!({"clawId": claw_id, "provider": "linear"}))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    let bad = client
        .get(format!(
            "{base}/api/apps/oauth/callback?code=bad-code&state={}",
            start2["state"].as_str().unwrap()
        ))
        .send()
        .await
        .unwrap();
    assert_eq!(bad.status(), 500);
}

#[tokio::test]
async fn mcp_oauth_uses_the_custom_issuer() {
    let pool = cm_testkit::test_pool().await;
    let issuer = spawn_idp().await;
    let socket = spawn_broker(pool.clone()).await;
    let runtime = Runtime::new(
        pool.clone(),
        Arc::new(ScriptedProvider::from_toml("").unwrap()),
        RuntimeConfig::basic("scripted", 1024),
    );
    // No default issuer configured: only mcp_oauth with an explicit issuer
    // can start.
    let oauth = OAuthConfig {
        issuer_url: None,
        client_id: Some("clawmates".into()),
        client_secret: Some("tc-secret".into()),
        redirect_base: Some("http://127.0.0.1:9".into()),
    };
    let app = cm_api::router(
        AppState::new(pool.clone(), runtime)
            .with_broker(socket)
            .with_oauth(oauth),
    );
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });
    let base = format!("http://{addr}");
    let client = reqwest::Client::new();

    let ws = Workspace {
        id: WorkspaceId::new(),
        name: "Acme".into(),
        plan: "team".into(),
    };
    cm_db::repo::workspaces::insert(&pool, &ws).await.unwrap();
    let owner = User {
        id: UserId::new(),
        workspace_id: ws.id,
        email: format!("{}@acme.test", UserId::new()),
        role: Role::Owner,
        display_name: "Owner".into(),
        created_at: time::OffsetDateTime::UNIX_EPOCH,
    };
    cm_db::repo::users::insert(&pool, &owner).await.unwrap();
    AuthService::new(pool.clone())
        .set_password(owner.id, "pw")
        .await
        .unwrap();
    let token = client
        .post(format!("{base}/api/auth/login"))
        .json(&json!({"email": owner.email, "password": "pw"}))
        .send()
        .await
        .unwrap()
        .json::<Value>()
        .await
        .unwrap()["token"]
        .as_str()
        .unwrap()
        .to_owned();
    let claw: Value = client
        .post(format!("{base}/api/claws"))
        .bearer_auth(&token)
        .json(&json!({"name": "Scout", "job_title": "Analyst"}))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    let claw_id = claw["id"].as_str().unwrap();

    // Plain oauth refuses (no configured issuer)...
    let refused = client
        .post(format!("{base}/api/apps/oauth/start"))
        .bearer_auth(&token)
        .json(&json!({"clawId": claw_id, "provider": "notion"}))
        .send()
        .await
        .unwrap();
    assert_eq!(refused.status(), 409);

    // ...while mcp_oauth with the custom issuer proceeds.
    let start: Value = client
        .post(format!("{base}/api/apps/oauth/start"))
        .bearer_auth(&token)
        .json(&json!({
            "clawId": claw_id,
            "provider": "my-mcp-server",
            "authType": "mcp_oauth",
            "issuerUrl": issuer,
        }))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    assert!(start["authorize_url"]
        .as_str()
        .unwrap()
        .starts_with(&format!("{issuer}/auth?")));
}