fix: prevent prompt injection via JSON extraction (#355)

Merges #416
This commit is contained in:
Argenis 2026-02-16 14:17:24 -05:00 committed by GitHub
parent dc5a85c85c
commit bff0507132
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -255,6 +255,15 @@ fn parse_tool_calls_from_json_value(value: &serde_json::Value) -> Vec<ParsedTool
calls calls
} }
/// Extract JSON values from a string.
///
/// # Security Warning
///
/// This function extracts ANY JSON objects/arrays from the input. It MUST only
/// be used on content that is already trusted to be from the LLM, such as
/// content inside `<invoke>` tags where the LLM has explicitly indicated intent
/// to make a tool call. Do NOT use this on raw user input or content that
/// could contain prompt injection payloads.
fn extract_json_values(input: &str) -> Vec<serde_json::Value> { fn extract_json_values(input: &str) -> Vec<serde_json::Value> {
let mut values = Vec::new(); let mut values = Vec::new();
let trimmed = input.trim(); let trimmed = input.trim();
@ -353,14 +362,13 @@ fn parse_tool_calls(response: &str) -> (String, Vec<ParsedToolCall>) {
} }
} }
if calls.is_empty() { // SECURITY: We do NOT fall back to extracting arbitrary JSON from the response
for value in extract_json_values(response) { // here. That would enable prompt injection attacks where malicious content
let parsed_calls = parse_tool_calls_from_json_value(&value); // (e.g., in emails, files, or web pages) could include JSON that mimics a
if !parsed_calls.is_empty() { // tool call. Tool calls MUST be explicitly wrapped in either:
calls.extend(parsed_calls); // 1. OpenAI-style JSON with a "tool_calls" array
} // 2. ZeroClaw <invoke>...</invoke> tags
} // This ensures only the LLM's intentional tool calls are executed.
}
// Remaining text after last tool call // Remaining text after last tool call
if !remaining.trim().is_empty() { if !remaining.trim().is_empty() {
@ -1246,18 +1254,16 @@ I will now call the tool with this payload:
} }
#[test] #[test]
fn parse_tool_calls_handles_raw_tool_json_without_tags() { fn parse_tool_calls_rejects_raw_tool_json_without_tags() {
// SECURITY: Raw JSON without explicit wrappers should NOT be parsed
// This prevents prompt injection attacks where malicious content
// could include JSON that mimics a tool call.
let response = r#"Sure, creating the file now. let response = r#"Sure, creating the file now.
{"name": "file_write", "arguments": {"path": "hello.py", "content": "print('hello')"}}"#; {"name": "file_write", "arguments": {"path": "hello.py", "content": "print('hello')"}}"#;
let (text, calls) = parse_tool_calls(response); let (text, calls) = parse_tool_calls(response);
assert!(text.contains("Sure, creating the file now.")); assert!(text.contains("Sure, creating the file now."));
assert_eq!(calls.len(), 1); assert_eq!(calls.len(), 0, "Raw JSON without wrappers should not be parsed");
assert_eq!(calls[0].name, "file_write");
assert_eq!(
calls[0].arguments.get("path").unwrap().as_str().unwrap(),
"hello.py"
);
} }
#[test] #[test]