clawmates/crates/tc-secrets/tests/broker.rs

//! Broker tests against the real wire: a real unix socket, real Postgres,
//! real crypto, and a real local HTTP receiver standing in for the external
//! service (its base URL is configuration — the same mechanism air-gapped
//! installs use).

use std::sync::Arc;

use axum::extract::State;
use axum::routing::post;
use serde_json::json;
use tc_domain::{
    AccessPolicy, Agent, AgentId, AgentStatus, GatedCategory, Role, User, UserId, Workspace,
    WorkspaceId,
};
use tc_safety::{approvals, Decision, NewApproval};
use tc_secrets::{BrokerClient, BrokerError, BrokerServer, FileKey};
use tokio::sync::Mutex;

struct Seed {
    workspace: Workspace,
    owner: User,
    agent: Agent,
    run_id: uuid::Uuid,
}

async fn seeded(pool: &sqlx::PgPool) -> Seed {
    let workspace = Workspace {
        id: WorkspaceId::new(),
        name: "Acme".into(),
        plan: "team".into(),
    };
    tc_db::repo::workspaces::insert(pool, &workspace)
        .await
        .unwrap();
    let owner = User {
        id: UserId::new(),
        workspace_id: workspace.id,
        email: format!("{}@acme.test", UserId::new()),
        role: Role::Owner,
        display_name: "Owner".into(),
        created_at: time::OffsetDateTime::UNIX_EPOCH,
    };
    tc_db::repo::users::insert(pool, &owner).await.unwrap();
    let agent = Agent {
        id: AgentId::new(),
        workspace_id: workspace.id,
        name: "Scout".into(),
        job_title: "Analyst".into(),
        system_prompt: String::new(),
        avatar: String::new(),
        accent: String::new(),
        wallpaper: String::new(),
        managed_by: owner.id,
        status: AgentStatus::Online,
    };
    tc_db::repo::agents::insert(pool, &agent, &AccessPolicy::default())
        .await
        .unwrap();
    let session = tc_db::repo::sessions::create(pool, agent.id, workspace.id, "Chat")
        .await
        .unwrap();
    let run_id = tc_db::repo::runs::create(pool, session.id).await.unwrap();
    Seed {
        workspace,
        owner,
        agent,
        run_id,
    }
}

/// A real HTTP receiver that records the Authorization headers it sees.
async fn spawn_receiver() -> (String, Arc<Mutex<Vec<String>>>) {
    let seen: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
    let state = seen.clone();
    let app =
        axum::Router::new()
            .route(
                "/hook",
                post(
                    |State(seen): State<Arc<Mutex<Vec<String>>>>,
                     headers: axum::http::HeaderMap| async move {
                        let auth = headers
                            .get("authorization")
                            .and_then(|v| v.to_str().ok())
                            .unwrap_or("")
                            .to_owned();
                        seen.lock().await.push(auth);
                        axum::Json(json!({"received": true}))
                    },
                ),
            )
            .with_state(state);
    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();
    });
    (format!("http://{addr}"), seen)
}

async fn broker_pair(pool: sqlx::PgPool, dir: &std::path::Path) -> BrokerClient {
    let key_path = dir.join("broker.key");
    FileKey::generate(&key_path).unwrap();
    let key = FileKey::load(&key_path).unwrap();
    // Unix socket paths are limited to ~104 bytes on macOS; keep it short.
    // Use the RANDOM tail of the uuid — the v7 prefix is a timestamp and
    // collides across tests started in the same millisecond.
    let short = uuid::Uuid::now_v7().simple().to_string();
    let socket = std::path::PathBuf::from(format!("/tmp/tcb-{}.sock", &short[short.len() - 12..]));
    let server = BrokerServer::new(pool, key, socket.clone());
    tokio::spawn(async move {
        server.serve().await.unwrap();
    });
    // The socket appears asynchronously.
    for _ in 0..50 {
        if socket.exists() {
            break;
        }
        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
    }
    BrokerClient::connect(&socket).await.unwrap()
}

async fn approved_approval(pool: &sqlx::PgPool, seed: &Seed) -> uuid::Uuid {
    let approval = approvals::create(
        pool,
        NewApproval {
            workspace_id: seed.workspace.id,
            run_id: seed.run_id,
            session_key: "agent:x-claw-0:session:y:z".into(),
            action_type: "http.request".into(),
            category: GatedCategory::OutboundMessage,
            payload: json!({}),
            preview: json!({}),
            requested_by_agent: seed.agent.id,
            taint_sources: vec![],
            expires_at: None,
        },
    )
    .await
    .unwrap();
    approvals::decide(pool, approval.id, seed.owner.id, Decision::Approve)
        .await
        .unwrap();
    approval.id
}

#[tokio::test]
async fn secrets_are_encrypted_at_rest_and_never_returned() {
    let pool = tc_testkit::test_pool().await;
    let seed = seeded(&pool).await;
    let dir = std::env::temp_dir().join(format!("tc-broker-{}", uuid::Uuid::now_v7()));
    std::fs::create_dir_all(&dir).unwrap();
    let mut client = broker_pair(pool.clone(), &dir).await;

    let secret_id = client
        .store_secret(seed.workspace.id, "api_key", "sk-super-secret-token")
        .await
        .unwrap();

    // At rest: ciphertext only, never the plaintext bytes.
    let ciphertext: Vec<u8> = sqlx::query_scalar("SELECT ciphertext FROM secrets WHERE id = $1")
        .bind(secret_id)
        .fetch_one(&pool)
        .await
        .unwrap();
    let raw = String::from_utf8_lossy(&ciphertext);
    assert!(!raw.contains("sk-super-secret-token"));

    // No protocol operation returns plaintext — metadata only.
    let kind = client.secret_kind(secret_id).await.unwrap();
    assert_eq!(kind, "api_key");
}

#[tokio::test]
async fn capability_requires_an_unconsumed_grant() {
    let pool = tc_testkit::test_pool().await;
    let seed = seeded(&pool).await;
    let dir = std::env::temp_dir().join(format!("tc-broker-{}", uuid::Uuid::now_v7()));
    std::fs::create_dir_all(&dir).unwrap();
    let mut client = broker_pair(pool.clone(), &dir).await;
    let (receiver_url, seen) = spawn_receiver().await;

    let secret_id = client
        .store_secret(seed.workspace.id, "api_key", "sk-live-12345")
        .await
        .unwrap();

    // Pending approval (no grant yet): the broker refuses outright.
    let pending = approvals::create(
        &pool,
        NewApproval {
            workspace_id: seed.workspace.id,
            run_id: seed.run_id,
            session_key: "k".into(),
            action_type: "http.request".into(),
            category: GatedCategory::OutboundMessage,
            payload: json!({}),
            preview: json!({}),
            requested_by_agent: seed.agent.id,
            taint_sources: vec![],
            expires_at: None,
        },
    )
    .await
    .unwrap();
    let refused = client
        .invoke_http(
            pending.id,
            secret_id,
            &format!("{receiver_url}/hook"),
            json!({}),
        )
        .await;
    assert!(matches!(refused, Err(BrokerError::GrantRefused)));
    assert!(seen.lock().await.is_empty(), "nothing may execute");

    // Approved: the call executes WITH the credential injected; the
    // credential itself never crosses back over the socket.
    let approval_id = approved_approval(&pool, &seed).await;
    let status = client
        .invoke_http(
            approval_id,
            secret_id,
            &format!("{receiver_url}/hook"),
            json!({}),
        )
        .await
        .unwrap();
    assert_eq!(status, 200);
    let observed = seen.lock().await;
    assert_eq!(observed.as_slice(), ["Bearer sk-live-12345"]);
    drop(observed);

    // The grant is single-use: replay refused, no second call.
    let replay = client
        .invoke_http(
            approval_id,
            secret_id,
            &format!("{receiver_url}/hook"),
            json!({}),
        )
        .await;
    assert!(matches!(replay, Err(BrokerError::GrantRefused)));
    assert_eq!(seen.lock().await.len(), 1);
}

#[tokio::test]
async fn non_http_urls_are_rejected() {
    let pool = tc_testkit::test_pool().await;
    let seed = seeded(&pool).await;
    let dir = std::env::temp_dir().join(format!("tc-broker-{}", uuid::Uuid::now_v7()));
    std::fs::create_dir_all(&dir).unwrap();
    let mut client = broker_pair(pool.clone(), &dir).await;

    let secret_id = client
        .store_secret(seed.workspace.id, "api_key", "sk-x")
        .await
        .unwrap();
    let approval_id = approved_approval(&pool, &seed).await;
    let refused = client
        .invoke_http(approval_id, secret_id, "file:///etc/passwd", json!({}))
        .await;
    assert!(matches!(refused, Err(BrokerError::Invalid(_))));
}

#[test]
fn key_files_round_trip_and_reject_corruption() {
    let dir = std::env::temp_dir().join(format!("tc-key-{}", uuid::Uuid::now_v7()));
    std::fs::create_dir_all(&dir).unwrap();
    let path = dir.join("broker.key");
    FileKey::generate(&path).unwrap();
    let key = FileKey::load(&path).unwrap();

    let sealed = key.seal(b"attack at dawn").unwrap();
    assert_ne!(sealed.ciphertext, b"attack at dawn");
    let opened = key.open(&sealed).unwrap();
    assert_eq!(opened, b"attack at dawn");

    // Bit-flip is detected (AEAD), not silently decrypted.
    let mut tampered = sealed;
    tampered.ciphertext[0] ^= 0xff;
    assert!(key.open(&tampered).is_err());
}