Skip to content

Commit 78f98b4

Browse files
committed
refactor: Consolidate dangerous pattern checks with contains_in_middle helper
- Unified logic for detecting dangerous patterns into a single function, reducing code complexity. - Simplified safety checks by removing redundant mid-pattern logic. - Added comprehensive test cases to validate `contains_in_middle` functionality.
1 parent 5d5b648 commit 78f98b4

File tree

1 file changed

+59
-19
lines changed
  • crates/chat-cli/src/cli/chat/tools/execute

1 file changed

+59
-19
lines changed

crates/chat-cli/src/cli/chat/tools/execute/mod.rs

Lines changed: 59 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -70,22 +70,9 @@ impl ExecuteCommand {
7070
let Some(args) = shlex::split(&self.command) else {
7171
return true;
7272
};
73-
const DANGEROUS_PATTERNS: &[&str] = &["<(", "$(", "`", ">", "&&", "||", "&", "${", "\n", "\r", "IFS"];
74-
const DANGEROUS_MID_PATTERNS: &[&str] = &[";"];
73+
const DANGEROUS_PATTERNS: &[&str] = &["<(", "$(", "`", ">", "&&", "||", "&", ";", "${", "\n", "\r", "IFS"];
7574

76-
if args
77-
.iter()
78-
.any(|arg| DANGEROUS_PATTERNS.iter().any(|p| arg.contains(p)))
79-
{
80-
return true;
81-
}
82-
83-
// Check for patterns that are dangerous only when they appear in the middle of arguments
84-
if args.iter().any(|arg| {
85-
DANGEROUS_MID_PATTERNS.iter().any(|p| {
86-
arg.contains(p) && !arg.ends_with(p)
87-
})
88-
}) {
75+
if DANGEROUS_PATTERNS.iter().any(|p| contains_in_middle(&self.command, p)) {
8976
return true;
9077
}
9178

@@ -292,13 +279,69 @@ pub fn format_output(output: &str, max_size: usize) -> String {
292279
)
293280
}
294281

282+
fn contains_in_middle(target: &str, pattern: &str) -> bool {
283+
if let Some(pos) = target.find(pattern) {
284+
pos > 0 && pos + pattern.len() < target.len()
285+
} else {
286+
false
287+
}
288+
}
289+
295290
#[cfg(test)]
296291
mod tests {
297292
use std::collections::HashMap;
298293

299294
use super::*;
300295
use crate::cli::agent::ToolSettingTarget;
301296

297+
#[test]
298+
fn test_contains_in_middle() {
299+
let test_cases = &[
300+
// Semicolon tests
301+
("find . -exec cat {} \\; -exec grep pattern {} \\;", ";", true),
302+
("echo hello; echo world", ";", true),
303+
("echo hello", ";", false),
304+
(";echo hello", ";", false),
305+
("echo hello;", ";", false),
306+
// Pipe tests
307+
("echo hello | grep world", "|", true),
308+
("echo hello|grep world", "|", true),
309+
("|echo hello", "|", false),
310+
("echo hello|", "|", false),
311+
// Ampersand tests
312+
("echo hello & echo world", "&", true),
313+
("echo hello&echo world", "&", true),
314+
("&echo hello", "&", false),
315+
("echo hello&", "&", false),
316+
// Greater than tests
317+
("echo hello > file.txt", ">", true),
318+
("echo hello>file.txt", ">", true),
319+
(">echo hello", ">", false),
320+
("echo hello>", ">", false),
321+
// Less than tests
322+
("cat < input.txt", "<", true),
323+
("cat<input.txt", "<", true),
324+
("<cat input.txt", "<", false),
325+
("cat input.txt<", "<", false),
326+
// Dollar sign tests
327+
("echo $HOME test", "$", true),
328+
("echo test$HOME", "$", true),
329+
("$echo test", "$", false),
330+
("echo test$", "$", false),
331+
];
332+
333+
for (target, pattern, expected) in test_cases {
334+
assert_eq!(
335+
contains_in_middle(target, pattern),
336+
*expected,
337+
"expected contains_in_middle('{}', '{}') to be {}",
338+
target,
339+
pattern,
340+
expected
341+
);
342+
}
343+
}
344+
302345
#[test]
303346
fn test_requires_acceptance_for_readonly_commands() {
304347
let cmds = &[
@@ -353,16 +396,13 @@ mod tests {
353396
// `find` with readonly -exec commands (should be allowed)
354397
("find . -name '*.rs' -exec grep -l pattern {} \\;", false),
355398
("find . -type f -exec cat {} \\;", false),
399+
("find . -type f -exec rm {} \\;", true),
356400
("find . -name '*.txt' -exec head {} \\;", false),
357401
("find . -type f -exec ls -l {} \\;", false),
358402
// Multiple -exec commands - mixed readonly and non-readonly
359403
("find . -exec cat {} \\; -exec rm {} \\;", true),
360404
("find . -exec ls {} \\; -exec touch newfile \\;", true),
361405
("find . -exec grep pattern {} \\; -exec chmod 755 {} \\;", true),
362-
// Multiple -exec commands - all readonly (should be allowed)
363-
("find . -exec cat {} \\; -exec grep pattern {} \\;", false),
364-
("find . -exec ls -l {} \\; -exec head {} \\;", false),
365-
("find . -exec echo {} \\; -exec tail {} \\;", false),
366406
(r"find . -${t}exec touch asdf \{\} +", true),
367407
(r"find . -${t:=exec} touch asdf2 \{\} +", true),
368408
// `grep` command arguments

0 commit comments

Comments
 (0)