feat: add platform key usage to mcp auth

2026-03-31 15:48:09 +03:00
parent 84f4437ce0
commit 563e17c3df
11 changed files with 363 additions and 30 deletions
@@ -49,19 +49,22 @@ mod tests {
    };

    use axum::{Json, Router, http::header, routing::post};
+    use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
    use crank_adapter_grpc::test_support as grpc_test_support;
    use crank_core::{
        Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, DescriptorId,
        ExecutionConfig, GraphqlOperationType, GraphqlTarget, GrpcTarget, HttpMethod, Operation,
-        OperationId, OperationStatus, Protocol, RestTarget, Target, ToolDescription, WorkspaceId,
+        OperationId, OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyScope,
+        PlatformApiKeyStatus, Protocol, RestTarget, Target, ToolDescription, WorkspaceId,
    };
    use crank_mapping::{MappingRule, MappingSet};
    use crank_registry::{
-        CreateAgentRequest, ListInvocationLogsQuery, PostgresRegistry, PublishAgentRequest,
-        PublishRequest,
+        CreateAgentRequest, CreatePlatformApiKeyRequest, ListInvocationLogsQuery, PostgresRegistry,
+        PublishAgentRequest, PublishRequest,
    };
    use crank_schema::{Schema, SchemaKind};
    use serde_json::{Value, json};
+    use sha2::{Digest, Sha256};
    use sqlx::{Executor, postgres::PgPoolOptions};
    use tokio::net::TcpListener;

@@ -96,6 +99,12 @@ mod tests {
            .await
            .unwrap();
        publish_agent_for_operation(&registry, &operation, "sales-rest").await;
+        let api_key = create_platform_api_key(
+            &registry,
+            "mcp-rest",
+            &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
+        )
+        .await;

        let base_url = spawn_mcp_server(build_app(
            registry.clone(),
@@ -105,11 +114,12 @@ mod tests {
        .await;
        let client = reqwest::Client::new();
        let mcp_url = agent_mcp_url(&base_url, "sales-rest");
-        let initialized_session = initialize_session(&client, &mcp_url).await;
+        let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;

        let tools = post_jsonrpc(
            &client,
            &mcp_url,
+            &api_key,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
@@ -122,6 +132,7 @@ mod tests {
        let call_result = post_jsonrpc(
            &client,
            &mcp_url,
+            &api_key,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
@@ -165,6 +176,12 @@ mod tests {
        );
        assert_eq!(logs[0].log.status, crank_core::InvocationStatus::Ok);
        assert_eq!(logs[0].log.tool_name, "crm_create_lead");
+
+        let keys = registry
+            .list_platform_api_keys(&test_workspace_id())
+            .await
+            .unwrap();
+        assert!(keys[0].api_key.last_used_at.is_some());
    }

    #[tokio::test]
@@ -188,6 +205,12 @@ mod tests {
            .await
            .unwrap();
        publish_agent_for_operation(&registry, &operation, "sales-graphql").await;
+        let api_key = create_platform_api_key(
+            &registry,
+            "mcp-graphql",
+            &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
+        )
+        .await;

        let base_url = spawn_mcp_server(build_app(
            registry,
@@ -197,11 +220,12 @@ mod tests {
        .await;
        let client = reqwest::Client::new();
        let mcp_url = agent_mcp_url(&base_url, "sales-graphql");
-        let initialized_session = initialize_session(&client, &mcp_url).await;
+        let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;

        let call_result = post_jsonrpc(
            &client,
            &mcp_url,
+            &api_key,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
@@ -245,6 +269,12 @@ mod tests {
            .await
            .unwrap();
        publish_agent_for_operation(&registry, &operation, "sales-grpc").await;
+        let api_key = create_platform_api_key(
+            &registry,
+            "mcp-grpc",
+            &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
+        )
+        .await;

        let base_url = spawn_mcp_server(build_app(
            registry,
@@ -254,11 +284,12 @@ mod tests {
        .await;
        let client = reqwest::Client::new();
        let mcp_url = agent_mcp_url(&base_url, "sales-grpc");
-        let initialized_session = initialize_session(&client, &mcp_url).await;
+        let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;

        let call_result = post_jsonrpc(
            &client,
            &mcp_url,
+            &api_key,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
@@ -284,6 +315,9 @@ mod tests {
    #[tokio::test]
    async fn requires_initialized_notification_before_tool_methods() {
        let registry = test_registry().await;
+        publish_agent_with_bindings(&registry, "sales-init", vec![]).await;
+        let api_key =
+            create_platform_api_key(&registry, "mcp-init", &[PlatformApiKeyScope::Read]).await;
        let base_url = spawn_mcp_server(build_app(
            registry,
            Duration::from_millis(0),
@@ -295,6 +329,7 @@ mod tests {
        let initialize_response = client
            .post(&mcp_url)
            .header(header::ACCEPT, "application/json, text/event-stream")
+            .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
            .json(&json!({
                "jsonrpc": "2.0",
                "id": 1,
@@ -316,6 +351,7 @@ mod tests {
        let tools_list = client
            .post(&mcp_url)
            .header(header::ACCEPT, "application/json, text/event-stream")
+            .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
            .header("MCP-Session-Id", session_id)
            .header("MCP-Protocol-Version", "2025-11-25")
            .json(&json!({
@@ -346,11 +382,18 @@ mod tests {
        .await;
        let client = reqwest::Client::new();
        let mcp_url = agent_mcp_url(&base_url, "sales-refresh");
-        let initialized_session = initialize_session(&client, &mcp_url).await;
+        let api_key = create_platform_api_key(
+            &registry,
+            "mcp-refresh",
+            &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
+        )
+        .await;
+        let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;

        let before_publish = post_jsonrpc(
            &client,
            &mcp_url,
+            &api_key,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
@@ -381,6 +424,7 @@ mod tests {
        let after_publish = post_jsonrpc(
            &client,
            &mcp_url,
+            &api_key,
            Some(&initialized_session),
            json!({
                "jsonrpc": "2.0",
@@ -434,6 +478,12 @@ mod tests {
            vec![binding_for_operation(&operation_b)],
        )
        .await;
+        let api_key = create_platform_api_key(
+            &registry,
+            "mcp-agents",
+            &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
+        )
+        .await;

        let base_url = spawn_mcp_server(build_app(
            registry,
@@ -444,12 +494,13 @@ mod tests {
        let client = reqwest::Client::new();
        let agent_a_url = agent_mcp_url(&base_url, "sales-a");
        let agent_b_url = agent_mcp_url(&base_url, "sales-b");
-        let session_a = initialize_session(&client, &agent_a_url).await;
-        let session_b = initialize_session(&client, &agent_b_url).await;
+        let session_a = initialize_session(&client, &agent_a_url, &api_key).await;
+        let session_b = initialize_session(&client, &agent_b_url, &api_key).await;

        let tools_a = post_jsonrpc(
            &client,
            &agent_a_url,
+            &api_key,
            Some(&session_a),
            json!({
                "jsonrpc": "2.0",
@@ -462,6 +513,7 @@ mod tests {
        let tools_b = post_jsonrpc(
            &client,
            &agent_b_url,
+            &api_key,
            Some(&session_b),
            json!({
                "jsonrpc": "2.0",
@@ -476,10 +528,97 @@ mod tests {
        assert_eq!(tools_b["result"]["tools"][0]["name"], "crm_update_lead");
    }

-    async fn initialize_session(client: &reqwest::Client, mcp_url: &str) -> String {
+    #[tokio::test]
+    async fn rejects_initialize_without_platform_api_key() {
+        let registry = test_registry().await;
+        publish_agent_with_bindings(&registry, "sales-auth", vec![]).await;
+        let base_url = spawn_mcp_server(build_app(
+            registry,
+            Duration::from_millis(0),
+            Some("https://crank.example.com".to_owned()),
+        ))
+        .await;
+        let client = reqwest::Client::new();
+        let response = client
+            .post(agent_mcp_url(&base_url, "sales-auth"))
+            .header(header::ACCEPT, "application/json, text/event-stream")
+            .json(&json!({
+                "jsonrpc": "2.0",
+                "id": 1,
+                "method": "initialize",
+                "params": {
+                    "protocolVersion": "2025-11-25"
+                }
+            }))
+            .send()
+            .await
+            .unwrap();
+
+        assert_eq!(response.status(), reqwest::StatusCode::UNAUTHORIZED);
+    }
+
+    #[tokio::test]
+    async fn rejects_tool_call_with_read_only_platform_api_key() {
+        let registry = test_registry().await;
+        let upstream_base_url = spawn_upstream_server().await;
+        let operation = test_operation(&upstream_base_url, "crm_read_only");
+
+        registry
+            .create_operation(&test_workspace_id(), &operation, Some("alice"))
+            .await
+            .unwrap();
+        registry
+            .publish_operation(PublishRequest {
+                workspace_id: &test_workspace_id(),
+                operation_id: &operation.id,
+                version: 1,
+                published_at: "2026-03-26T10:00:00Z",
+                published_by: Some("alice"),
+            })
+            .await
+            .unwrap();
+        publish_agent_for_operation(&registry, &operation, "sales-read-only").await;
+        let api_key =
+            create_platform_api_key(&registry, "mcp-read", &[PlatformApiKeyScope::Read]).await;
+
+        let base_url = spawn_mcp_server(build_app(
+            registry,
+            Duration::from_millis(0),
+            Some("https://crank.example.com".to_owned()),
+        ))
+        .await;
+        let client = reqwest::Client::new();
+        let mcp_url = agent_mcp_url(&base_url, "sales-read-only");
+        let initialized_session = initialize_session(&client, &mcp_url, &api_key).await;
+        let response = client
+            .post(&mcp_url)
+            .header(header::ACCEPT, "application/json, text/event-stream")
+            .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
+            .header("MCP-Session-Id", initialized_session)
+            .header("MCP-Protocol-Version", "2025-11-25")
+            .json(&json!({
+                "jsonrpc": "2.0",
+                "id": 3,
+                "method": "tools/call",
+                "params": {
+                    "name": "crm_read_only",
+                    "arguments": {
+                        "email": "user@example.com"
+                    }
+                }
+            }))
+            .send()
+            .await
+            .unwrap();
+
+        assert_eq!(response.status(), reqwest::StatusCode::FORBIDDEN);
+    }
+
+    async fn initialize_session(client: &reqwest::Client, mcp_url: &str, api_key: &str) -> String {
        let initialize_response = client
            .post(mcp_url)
            .header(header::ACCEPT, "application/json, text/event-stream")
+            .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
            .json(&json!({
                "jsonrpc": "2.0",
                "id": 1,
@@ -502,6 +641,7 @@ mod tests {
        let initialized_response = client
            .post(mcp_url)
            .header(header::ACCEPT, "application/json, text/event-stream")
+            .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
            .header("MCP-Session-Id", &session_id)
            .header("MCP-Protocol-Version", "2025-11-25")
            .json(&json!({
@@ -521,12 +661,14 @@ mod tests {
    async fn post_jsonrpc(
        client: &reqwest::Client,
        mcp_url: &str,
+        api_key: &str,
        session_id: Option<&str>,
        payload: Value,
    ) -> Value {
        let mut request = client
            .post(mcp_url)
            .header(header::ACCEPT, "application/json, text/event-stream")
+            .header(header::AUTHORIZATION, format!("Bearer {api_key}"))
            .header("MCP-Protocol-Version", "2025-11-25");

        if let Some(session_id) = session_id {
@@ -543,6 +685,39 @@ mod tests {
            .unwrap()
    }

+    async fn create_platform_api_key(
+        registry: &PostgresRegistry,
+        name: &str,
+        scopes: &[PlatformApiKeyScope],
+    ) -> String {
+        let secret = format!("crk_{}_{}", name, uuid::Uuid::now_v7().simple());
+        let api_key = PlatformApiKey {
+            id: PlatformApiKeyId::new(format!("pk_{name}")),
+            workspace_id: test_workspace_id(),
+            name: name.to_owned(),
+            prefix: secret.chars().take(16).collect(),
+            scopes: scopes.to_vec(),
+            status: PlatformApiKeyStatus::Active,
+            created_at: "2026-03-26T10:00:00Z".to_owned(),
+            last_used_at: None,
+        };
+
+        registry
+            .create_platform_api_key(CreatePlatformApiKeyRequest {
+                api_key: &api_key,
+                secret_hash: &hash_access_secret(&secret),
+            })
+            .await
+            .unwrap();
+
+        secret
+    }
+
+    fn hash_access_secret(secret: &str) -> String {
+        let digest = Sha256::digest(secret.as_bytes());
+        URL_SAFE_NO_PAD.encode(digest)
+    }
+
    async fn spawn_upstream_server() -> String {
        let app = Router::new().route("/crm/leads", post(create_lead));
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();