bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: harden streaming transport and ownership checks

Browse Source
This commit is contained in:
a.tolmachev
2026-04-11 01:29:01 +03:00
parent 1a4f0ea6f3
commit 2770c5935f
3 changed files with 348 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/admin-api/src/service.rs
+12 -4
View File
@@ -2246,10 +2246,18 @@ impl AdminService {
let created_at = now_string().map_err(|error| RuntimeError::SecretCrypto {
details: error.to_string(),
})?;
let expires_at =
add_millis(&created_at, 300_000).map_err(|error| RuntimeError::SecretCrypto {
details: error.to_string(),
})?;
let streaming = runtime.execution_config.streaming.as_ref().ok_or_else(|| {
RuntimeError::MissingStreamingConfig {
operation_id: runtime.operation_id.as_str().to_owned(),
}
})?;
let expires_at = add_millis(
&created_at,
streaming.max_session_lifetime_ms.unwrap_or(300_000),
)
.map_err(|error| RuntimeError::SecretCrypto {
details: error.to_string(),
})?;
let job = AsyncJobHandle {
id: crank_core::AsyncJobId::new(new_prefixed_id("job")),
workspace_id: workspace_id.clone(),
apps/mcp-server/src/app.rs
+107 -3
View File
@@ -889,6 +889,16 @@ async fn handle_session_poll_call(
Err(error) => return internal_jsonrpc_error(message, error),
};
if !stream_session_belongs_to_tool(&loaded, &tool) {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"stream_session_not_found",
format!("stream session {} was not found", control.session_id),
);
}
if loaded.is_expired(&now) {
let _ = state
.registry
@@ -1008,7 +1018,7 @@ async fn handle_session_stop_call(
session: &SessionState,
message: &Value,
response_mode: ResponseMode,
_tool: PublishedAgentTool,
tool: PublishedAgentTool,
arguments: Value,
) -> Response {
let control: SessionControlArgs = match serde_json::from_value(arguments.clone()) {
@@ -1024,6 +1034,34 @@ async fn handle_session_stop_call(
}
};
let loaded = match state
.registry
.get_stream_session(&StreamSessionId::new(control.session_id.clone()))
.await
{
Ok(Some(session_record)) => session_record,
Ok(None) => {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"stream_session_not_found",
format!("stream session {} was not found", control.session_id),
);
}
Err(error) => return internal_jsonrpc_error(message, error),
};
if !stream_session_belongs_to_tool(&loaded, &tool) {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"stream_session_not_found",
format!("stream session {} was not found", control.session_id),
);
}
match state
.registry
.close_stream_session(
@@ -1075,7 +1113,10 @@ async fn handle_async_job_start_call(
progress: json!({ "pct": 0 }),
result: None,
error: None,
expires_at: Some(add_millis(&now, 300_000)),
expires_at: Some(add_millis(
&now,
streaming.max_session_lifetime_ms.unwrap_or(300_000),
)),
created_at: now.clone(),
updated_at: now.clone(),
finished_at: None,
@@ -1222,6 +1263,16 @@ async fn handle_async_job_status_call(
Err(error) => return internal_jsonrpc_error(message, error),
};
if !async_job_belongs_to_tool(&job, &_tool) {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"async_job_not_found",
format!("async job {} was not found", control.job_id),
);
}
success_tool_response(
message,
response_mode,
@@ -1275,6 +1326,16 @@ async fn handle_async_job_result_call(
Err(error) => return internal_jsonrpc_error(message, error),
};
if !async_job_belongs_to_tool(&job, &_tool) {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"async_job_not_found",
format!("async job {} was not found", control.job_id),
);
}
match job.status {
JobStatus::Completed => success_tool_response(
message,
@@ -1330,6 +1391,34 @@ async fn handle_async_job_cancel_call(
}
};
match state
.registry
.get_async_job(&AsyncJobId::new(control.job_id.clone()))
.await
{
Ok(Some(job)) => {
if !async_job_belongs_to_tool(&job, &_tool) {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"async_job_not_found",
format!("async job {} was not found", control.job_id),
);
}
}
Ok(None) => {
return tool_error_response(
message,
response_mode,
&session.protocol_version,
"async_job_not_found",
format!("async job {} was not found", control.job_id),
);
}
Err(error) => return internal_jsonrpc_error(message, error),
};
match state
.registry
.cancel_async_job(&AsyncJobId::new(control.job_id.clone()), &now_rfc3339())
@@ -1588,7 +1677,7 @@ fn negotiate_post_response_mode(headers: &HeaderMap) -> Result<ResponseMode, Sta
}
}
if saw_json && saw_sse {
if saw_json || saw_sse {
return preferred.ok_or(StatusCode::NOT_ACCEPTABLE);
}
@@ -1620,6 +1709,21 @@ fn validate_get_accept_header(headers: &HeaderMap) -> Result<(), StatusCode> {
Err(StatusCode::NOT_ACCEPTABLE)
}
fn stream_session_belongs_to_tool(
session_record: &StreamSession,
tool: &PublishedAgentTool,
) -> bool {
session_record.workspace_id == tool.workspace_id
&& session_record.agent_id.as_ref() == Some(&tool.agent_id)
&& session_record.operation_id == tool.operation.id
}
fn async_job_belongs_to_tool(job: &AsyncJobHandle, tool: &PublishedAgentTool) -> bool {
job.workspace_id == tool.workspace_id
&& job.agent_id.as_ref() == Some(&tool.agent_id)
&& job.operation_id == tool.operation.id
}
fn protocol_version_from_headers(headers: &HeaderMap) -> Result<String, StatusCode> {
let Some(version) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else {
return Ok(DEFAULT_PROTOCOL_VERSION.to_owned());
apps/mcp-server/src/main.rs
+229
View File
@@ -568,6 +568,47 @@ mod tests {
assert_eq!(after_delete.status(), reqwest::StatusCode::NOT_FOUND);
}
#[tokio::test]
async fn initialize_accepts_json_only_response_negotiation() {
let registry = test_registry().await;
publish_agent_with_bindings(&registry, "sales-json-accept", vec![]).await;
let api_key =
create_platform_api_key(&registry, "mcp-json-accept", &[PlatformApiKeyScope::Read])
.await;
let base_url = spawn_mcp_server(build_test_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-json-accept"))
.header(header::ACCEPT, "application/json")
.header(header::AUTHORIZATION, format!("Bearer {api_key}"))
.json(&json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2025-11-25"
}
}))
.send()
.await
.unwrap();
assert_eq!(response.status(), reqwest::StatusCode::OK);
assert_eq!(
response
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap(),
"application/json"
);
}
#[tokio::test]
async fn rejects_get_with_protocol_version_mismatch() {
let registry = test_registry().await;
@@ -989,6 +1030,194 @@ mod tests {
);
}
#[tokio::test]
async fn rejects_cross_agent_session_poll() {
let registry = test_registry().await;
let server_addr = grpc_test_support::spawn_unary_echo_server().await;
let operation_a = test_grpc_session_operation(&server_addr, "echo_stream_session_a");
let operation_b = test_grpc_session_operation(&server_addr, "echo_stream_session_b");
for operation in [&operation_a, &operation_b] {
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_a, "sales-session-a").await;
publish_agent_for_operation(&registry, &operation_b, "sales-session-b").await;
let api_key = create_platform_api_key(
&registry,
"mcp-cross-session",
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
)
.await;
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
))
.await;
let client = reqwest::Client::new();
let agent_a_url = agent_mcp_url(&base_url, "sales-session-a");
let agent_b_url = agent_mcp_url(&base_url, "sales-session-b");
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 start_response = post_jsonrpc(
&client,
&agent_a_url,
&api_key,
Some(&session_a),
json!({
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "echo_stream_session_a_start",
"arguments": { "message": "hello" }
}
}),
)
.await;
let foreign_session_id = start_response["result"]["structuredContent"]["session_id"]
.as_str()
.unwrap()
.to_owned();
let poll_response = post_jsonrpc(
&client,
&agent_b_url,
&api_key,
Some(&session_b),
json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": {
"name": "echo_stream_session_b_poll",
"arguments": { "session_id": foreign_session_id }
}
}),
)
.await;
assert_eq!(
poll_response["result"]["structuredContent"]["error"]["code"],
json!("stream_session_not_found")
);
assert!(
poll_response["result"]["structuredContent"]["error"]["message"]
.as_str()
.is_some_and(|message| message.starts_with("stream session "))
);
}
#[tokio::test]
async fn rejects_cross_agent_async_job_access() {
let registry = test_registry().await;
let upstream_base_url = spawn_upstream_server().await;
let operation_a = test_rest_async_job_operation(&upstream_base_url, "crm_async_a");
let operation_b = test_rest_async_job_operation(&upstream_base_url, "crm_async_b");
for operation in [&operation_a, &operation_b] {
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_a, "sales-async-a").await;
publish_agent_for_operation(&registry, &operation_b, "sales-async-b").await;
let api_key = create_platform_api_key(
&registry,
"mcp-cross-async",
&[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
)
.await;
let base_url = spawn_mcp_server(build_test_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
))
.await;
let client = reqwest::Client::new();
let agent_a_url = agent_mcp_url(&base_url, "sales-async-a");
let agent_b_url = agent_mcp_url(&base_url, "sales-async-b");
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 start_response = post_jsonrpc(
&client,
&agent_a_url,
&api_key,
Some(&session_a),
json!({
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": {
"name": "crm_async_a_start",
"arguments": { "email": "user@example.com" }
}
}),
)
.await;
let foreign_job_id = start_response["result"]["structuredContent"]["job_id"]
.as_str()
.unwrap()
.to_owned();
let status_response = post_jsonrpc(
&client,
&agent_b_url,
&api_key,
Some(&session_b),
json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": {
"name": "crm_async_b_status",
"arguments": { "job_id": foreign_job_id }
}
}),
)
.await;
assert_eq!(
status_response["result"]["structuredContent"]["error"]["code"],
json!("async_job_not_found")
);
assert!(
status_response["result"]["structuredContent"]["error"]["message"]
.as_str()
.is_some_and(|message| message.starts_with("async job "))
);
}
#[tokio::test]
async fn exposes_and_runs_async_job_tool_family_via_mcp() {
let registry = test_registry().await;