mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-09-01 09:09:19 +00:00
fix(agent): block thread_id-based context pollution across users (#760)
* fix(agent): prevent forged thread UUID context/write contamination * fix(agent): close thread_id race and reject forged UUID hydration * fix(ci): satisfy clippy and fmt checks after rebase
This commit is contained in:
@@ -0,0 +1,183 @@
|
||||
//! E2E regression test: forged thread IDs must not cross user boundaries.
|
||||
//!
|
||||
//! Demonstrates that a client cannot provide another user's conversation UUID
|
||||
//! and get that history hydrated into prompt context or written into.
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
mod support;
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use ironclaw::channels::{IncomingMessage, OutgoingResponse};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::support::test_rig::TestRigBuilder;
|
||||
use crate::support::trace_llm::{LlmTrace, TraceResponse, TraceStep};
|
||||
|
||||
fn assert_safe_thread_rejection(response: &OutgoingResponse) {
|
||||
let msg = response.content.to_lowercase();
|
||||
assert!(
|
||||
msg.contains("thread") && (msg.contains("invalid") || msg.contains("unauthorized")),
|
||||
"expected safe thread-id rejection response, got: {}",
|
||||
response.content
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forged_existing_foreign_thread_id_is_rejected_without_hydration_or_persistence() {
|
||||
let trace = LlmTrace::single_turn(
|
||||
"thread-id-isolation",
|
||||
"attacker turn",
|
||||
vec![TraceStep {
|
||||
request_hint: None,
|
||||
response: TraceResponse::Text {
|
||||
content: "safe response".to_string(),
|
||||
input_tokens: 12,
|
||||
output_tokens: 4,
|
||||
},
|
||||
expected_tool_results: Vec::new(),
|
||||
}],
|
||||
);
|
||||
|
||||
let rig = TestRigBuilder::new().with_trace(trace).build().await;
|
||||
|
||||
let foreign_thread_id = Uuid::new_v4();
|
||||
let marker = format!("FOREIGN-MARKER-{}", Uuid::new_v4());
|
||||
let store = rig.database();
|
||||
assert!(
|
||||
store
|
||||
.ensure_conversation(foreign_thread_id, "gateway", "victim-user", None)
|
||||
.await
|
||||
.expect("failed to create victim conversation"),
|
||||
"test setup failed: victim conversation was not created"
|
||||
);
|
||||
store
|
||||
.add_conversation_message(
|
||||
foreign_thread_id,
|
||||
"user",
|
||||
&format!("victim-only secret marker: {marker}"),
|
||||
)
|
||||
.await
|
||||
.expect("failed to seed victim conversation message");
|
||||
|
||||
let before_messages = store
|
||||
.list_conversation_messages(foreign_thread_id)
|
||||
.await
|
||||
.expect("failed to read victim conversation before forged send");
|
||||
assert!(
|
||||
before_messages.iter().any(|m| m.content.contains(&marker)),
|
||||
"test setup failed: victim marker message missing"
|
||||
);
|
||||
let before_len = before_messages.len();
|
||||
|
||||
let forged = IncomingMessage::new("test", "test-user", "attacker turn")
|
||||
.with_thread(foreign_thread_id.to_string());
|
||||
rig.send_incoming(forged).await;
|
||||
let responses = rig.wait_for_responses(1, Duration::from_secs(20)).await;
|
||||
assert_eq!(
|
||||
responses.len(),
|
||||
1,
|
||||
"expected one assistant response for forged-thread request"
|
||||
);
|
||||
assert_safe_thread_rejection(&responses[0]);
|
||||
|
||||
let captured = rig.captured_llm_requests();
|
||||
assert!(
|
||||
captured.is_empty(),
|
||||
"forged thread-id request should be rejected before any LLM call"
|
||||
);
|
||||
let prompt_dump = captured
|
||||
.iter()
|
||||
.flat_map(|req| req.iter().map(|m| m.content.as_str()))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
assert!(
|
||||
!prompt_dump.contains(&marker),
|
||||
"forged thread_id leaked foreign marker into LLM prompt context: {prompt_dump}"
|
||||
);
|
||||
|
||||
let after_messages = store
|
||||
.list_conversation_messages(foreign_thread_id)
|
||||
.await
|
||||
.expect("failed to read victim conversation after forged send");
|
||||
assert_eq!(
|
||||
after_messages.len(),
|
||||
before_len,
|
||||
"forged thread_id wrote new messages into victim conversation"
|
||||
);
|
||||
assert!(
|
||||
after_messages
|
||||
.iter()
|
||||
.all(|m| m.content != "attacker turn" && m.content != "safe response"),
|
||||
"forged request content was persisted to victim conversation"
|
||||
);
|
||||
|
||||
rig.shutdown();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forged_nonexistent_thread_id_is_rejected_and_followup_request_still_works() {
|
||||
let trace = LlmTrace::single_turn(
|
||||
"thread-id-isolation-nonexistent",
|
||||
"real follow-up turn",
|
||||
vec![TraceStep {
|
||||
request_hint: None,
|
||||
response: TraceResponse::Text {
|
||||
content: "safe response".to_string(),
|
||||
input_tokens: 12,
|
||||
output_tokens: 4,
|
||||
},
|
||||
expected_tool_results: Vec::new(),
|
||||
}],
|
||||
);
|
||||
|
||||
let rig = TestRigBuilder::new().with_trace(trace).build().await;
|
||||
|
||||
let forged_thread_id = Uuid::new_v4();
|
||||
let store = rig.database();
|
||||
|
||||
let forged = IncomingMessage::new("test", "test-user", "attacker turn")
|
||||
.with_thread(forged_thread_id.to_string());
|
||||
rig.send_incoming(forged).await;
|
||||
let responses = rig.wait_for_responses(1, Duration::from_secs(20)).await;
|
||||
assert_eq!(
|
||||
responses.len(),
|
||||
1,
|
||||
"expected one response for forged nonexistent-thread request"
|
||||
);
|
||||
assert_safe_thread_rejection(&responses[0]);
|
||||
assert!(
|
||||
rig.captured_llm_requests().is_empty(),
|
||||
"forged nonexistent thread-id request should be rejected before any LLM call"
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.get_conversation_metadata(forged_thread_id)
|
||||
.await
|
||||
.expect("get metadata for forged thread id")
|
||||
.is_none(),
|
||||
"forged nonexistent thread id must not create a conversation row"
|
||||
);
|
||||
|
||||
rig.send_message("real follow-up turn").await;
|
||||
let responses = rig.wait_for_responses(2, Duration::from_secs(20)).await;
|
||||
assert_eq!(
|
||||
responses.len(),
|
||||
2,
|
||||
"expected follow-up response after rejection"
|
||||
);
|
||||
assert_eq!(
|
||||
responses[1].content, "safe response",
|
||||
"follow-up valid request should still be handled normally"
|
||||
);
|
||||
assert_eq!(
|
||||
rig.captured_llm_requests().len(),
|
||||
1,
|
||||
"only follow-up request should reach LLM"
|
||||
);
|
||||
|
||||
rig.shutdown();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user