gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-09-01 00:59:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(agent): block thread_id-based context pollution across users (#760)

Browse Source 
* 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:
pikaxinge
2026-03-11 16:52:31 -07:00
committed by GitHub
parent c8cac0925d
commit 2094d6e30d
10 changed files with 448 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent/agent_loop.rs
+3 -1
View File
@@ -803,7 +803,9 @@ impl Agent {
thread_id = %external_thread_id,
"Hydrating thread from DB"
);
self.maybe_hydrate_thread(message, external_thread_id).await;
if let Some(rejection) = self.maybe_hydrate_thread(message, external_thread_id).await {
return Ok(Some(format!("Error: {}", rejection)));
}
}
// Resolve session and thread
src/agent/thread_ops.rs
+164 -26
View File
@@ -23,6 +23,14 @@ use crate::error::Error;
use crate::llm::{ChatMessage, ToolCall};
use crate::tools::redact_params;
const FORGED_THREAD_ID_ERROR: &str = "Invalid or unauthorized thread ID.";
fn requires_preexisting_uuid_thread(channel: &str) -> bool {
// Gateway-style channels send server-issued conversation UUIDs.
// Unknown UUIDs should be rejected instead of silently creating a new thread.
matches!(channel, "gateway" | "test")
}
impl Agent {
/// Hydrate a historical thread from DB into memory if not already present.
///
@@ -37,11 +45,11 @@ impl Agent {
&self,
message: &IncomingMessage,
external_thread_id: &str,
) {
) -> Option<String> {
// Only hydrate UUID-shaped thread IDs (web gateway uses UUIDs)
let thread_uuid = match Uuid::parse_str(external_thread_id) {
Ok(id) => id,
Err(_) => return,
Err(_) => return None,
};
// Check if already in memory
@@ -52,7 +60,7 @@ impl Agent {
{
let sess = session.lock().await;
if sess.threads.contains_key(&thread_uuid) {
return;
return None;
}
}
@@ -61,6 +69,62 @@ impl Agent {
let msg_count;
if let Some(store) = self.store() {
// Never hydrate history from a conversation UUID that isn't owned
// by the current authenticated user.
let owned = match store
.conversation_belongs_to_user(thread_uuid, &message.user_id)
.await
{
Ok(v) => v,
Err(e) => {
tracing::warn!(
"Failed to verify conversation ownership for hydration {}: {}",
thread_uuid,
e
);
if requires_preexisting_uuid_thread(&message.channel) {
return Some(FORGED_THREAD_ID_ERROR.to_string());
}
return None;
}
};
if !owned {
let exists = match store.get_conversation_metadata(thread_uuid).await {
Ok(Some(_)) => true,
Ok(None) => false,
Err(e) => {
tracing::warn!(
"Failed to inspect conversation metadata for hydration {}: {}",
thread_uuid,
e
);
if requires_preexisting_uuid_thread(&message.channel) {
return Some(FORGED_THREAD_ID_ERROR.to_string());
}
return None;
}
};
if requires_preexisting_uuid_thread(&message.channel) {
tracing::warn!(
user = %message.user_id,
channel = %message.channel,
thread_id = %thread_uuid,
exists,
"Rejected message for unavailable thread id"
);
return Some(FORGED_THREAD_ID_ERROR.to_string());
}
tracing::warn!(
user = %message.user_id,
thread_id = %thread_uuid,
exists,
"Skipped hydration for thread id not owned by sender"
);
return None;
}
let db_messages = store
.list_conversation_messages(thread_uuid)
.await
@@ -104,6 +168,8 @@ impl Agent {
thread_uuid,
msg_count
);
None
}
pub(super) async fn process_user_input(
@@ -303,8 +369,13 @@ impl Agent {
thread_id = %thread_id,
"Persisting user message to DB"
);
self.persist_user_message(thread_id, &message.user_id, effective_content)
.await;
self.persist_user_message(
thread_id,
&message.channel,
&message.user_id,
effective_content,
)
.await;
tracing::debug!(
message_id = %message.id,
@@ -386,10 +457,21 @@ impl Agent {
.await;
// Persist tool calls then assistant response (user message already persisted at turn start)
self.persist_tool_calls(thread_id, &message.user_id, turn_number, &tool_calls)
.await;
self.persist_assistant_response(thread_id, &message.user_id, &response)
.await;
self.persist_tool_calls(
thread_id,
&message.channel,
&message.user_id,
turn_number,
&tool_calls,
)
.await;
self.persist_assistant_response(
thread_id,
&message.channel,
&message.user_id,
&response,
)
.await;
Ok(SubmissionResult::response(response))
}
@@ -423,6 +505,41 @@ impl Agent {
}
}
/// Ensure a thread UUID is writable for `(channel, user_id)`.
///
/// Returns `false` for foreign/unowned conversation IDs or DB errors.
async fn ensure_writable_conversation(
&self,
store: &Arc<dyn crate::db::Database>,
thread_id: Uuid,
channel: &str,
user_id: &str,
) -> bool {
match store
.ensure_conversation(thread_id, channel, user_id, None)
.await
{
Ok(true) => true,
Ok(false) => {
tracing::warn!(
user = %user_id,
channel = %channel,
thread_id = %thread_id,
"Rejected write for unavailable thread id"
);
false
}
Err(e) => {
tracing::warn!(
"Failed to ensure writable conversation {}: {}",
thread_id,
e
);
false
}
}
}
/// Persist the user message to the DB at turn start (before the agentic loop).
///
/// This ensures the user message is durable even if the process crashes
@@ -430,6 +547,7 @@ impl Agent {
pub(super) async fn persist_user_message(
&self,
thread_id: Uuid,
channel: &str,
user_id: &str,
user_input: &str,
) {
@@ -438,11 +556,10 @@ impl Agent {
None => return,
};
if let Err(e) = store
.ensure_conversation(thread_id, "gateway", user_id, None)
if !self
.ensure_writable_conversation(&store, thread_id, channel, user_id)
.await
{
tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e);
return;
}
@@ -462,6 +579,7 @@ impl Agent {
pub(super) async fn persist_assistant_response(
&self,
thread_id: Uuid,
channel: &str,
user_id: &str,
response: &str,
) {
@@ -470,11 +588,10 @@ impl Agent {
None => return,
};
if let Err(e) = store
.ensure_conversation(thread_id, "gateway", user_id, None)
if !self
.ensure_writable_conversation(&store, thread_id, channel, user_id)
.await
{
tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e);
return;
}
@@ -494,6 +611,7 @@ impl Agent {
pub(super) async fn persist_tool_calls(
&self,
thread_id: Uuid,
channel: &str,
user_id: &str,
turn_number: usize,
tool_calls: &[crate::agent::session::TurnToolCall],
@@ -543,11 +661,10 @@ impl Agent {
}
};
if let Err(e) = store
.ensure_conversation(thread_id, "gateway", user_id, None)
if !self
.ensure_writable_conversation(&store, thread_id, channel, user_id)
.await
{
tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e);
return;
}
@@ -1214,10 +1331,21 @@ impl Agent {
.map(|t| (t.turn_number, t.tool_calls.clone()))
.unwrap_or_default();
// User message already persisted at turn start; save tool calls then assistant response
self.persist_tool_calls(thread_id, &message.user_id, turn_number, &tool_calls)
.await;
self.persist_assistant_response(thread_id, &message.user_id, &response)
.await;
self.persist_tool_calls(
thread_id,
&message.channel,
&message.user_id,
turn_number,
&tool_calls,
)
.await;
self.persist_assistant_response(
thread_id,
&message.channel,
&message.user_id,
&response,
)
.await;
let _ = self
.channels
.send_status(
@@ -1270,8 +1398,13 @@ impl Agent {
thread.clear_pending_approval();
thread.complete_turn(&rejection);
// User message already persisted at turn start; save rejection response
self.persist_assistant_response(thread_id, &message.user_id, &rejection)
.await;
self.persist_assistant_response(
thread_id,
&message.channel,
&message.user_id,
&rejection,
)
.await;
}
}
@@ -1309,8 +1442,13 @@ impl Agent {
thread.enter_auth_mode(ext_name.clone());
thread.complete_turn(&instructions);
// User message already persisted at turn start; save auth instructions
self.persist_assistant_response(thread_id, &message.user_id, &instructions)
.await;
self.persist_assistant_response(
thread_id,
&message.channel,
&message.user_id,
&instructions,
)
.await;
}
}
let _ = self