git.iko.soy/crush
on main help sign in
refactor: improve command blocking system and fix test isolation
f99f504Kujtim Hoxha
internal/llm/tools/bash.go+16−8
compare with main
unifiedsplit
@@ -43,7 +43,7 @@ const (
43var bannedCommands = []string{
43var bannedCommands = []string{
44 "alias", "curl", "curlie", "wget", "axel", "aria2c",
44 "alias", "curl", "curlie", "wget", "axel", "aria2c",
45 "nc", "telnet", "lynx", "w3m", "links", "httpie", "xh",
45 "nc", "telnet", "lynx", "w3m", "links", "httpie", "xh",
46− "http-prompt", "chrome", "firefox", "safari",
46+ "http-prompt", "chrome", "firefox", "safari", "sudo",
47}
47}
48
48
49// getSafeReadOnlyCommands returns platform-appropriate safe commands
49// getSafeReadOnlyCommands returns platform-appropriate safe commands
@@ -244,7 +244,22 @@ Important:
244- Never update git config`, bannedCommandsStr, MaxOutputLength)
244- Never update git config`, bannedCommandsStr, MaxOutputLength)
245}
245}
246
246
247+func createCommandBlockFuncs() []shell.CommandBlockFunc {
248+ return []shell.CommandBlockFunc{
249+ shell.CreateSimpleCommandBlocker(bannedCommands),
250+ shell.CreateSubCommandBlocker([][]string{
251+ {"brew", "install"},
252+ {"npm", "install", "-g"},
253+ {"npm", "install", "--global"},
254+ }),
255+ }
256+}
257+
247func NewBashTool(permission permission.Service, workingDir string) BaseTool {
258func NewBashTool(permission permission.Service, workingDir string) BaseTool {
259+ // Set up command blocking on the persistent shell
260+ persistentShell := shell.GetPersistentShell(workingDir)
261+ persistentShell.SetBlockFuncs(createCommandBlockFuncs())
262+
248 return &bashTool{
263 return &bashTool{
249 permissions: permission,
264 permissions: permission,
250 workingDir: workingDir,
265 workingDir: workingDir,
@@ -289,13 +304,6 @@ func (b *bashTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
289 return NewTextErrorResponse("missing command"), nil
304 return NewTextErrorResponse("missing command"), nil
290 }
305 }
291
306
292− baseCmd := strings.Fields(params.Command)[0]
293− for _, banned := range bannedCommands {
294− if strings.EqualFold(baseCmd, banned) {
295− return NewTextErrorResponse(fmt.Sprintf("command '%s' is not allowed", baseCmd)), nil
296− }
297− }
298−
299 isSafeReadOnly := false
307 isSafeReadOnly := false
300 cmdLower := strings.ToLower(params.Command)
308 cmdLower := strings.ToLower(params.Command)
301
309
internal/shell/command_block_test.gonew file+123
@@ -0,0 +1,123 @@
1+package shell
2+
3+import (
4+ "context"
5+ "os"
6+ "strings"
7+ "testing"
8+)
9+
10+func TestCommandBlocking(t *testing.T) {
11+ tests := []struct {
12+ name string
13+ blockFuncs []CommandBlockFunc
14+ command string
15+ shouldBlock bool
16+ }{
17+ {
18+ name: "block simple command",
19+ blockFuncs: []CommandBlockFunc{
20+ func(args []string) bool {
21+ return len(args) > 0 && args[0] == "curl"
22+ },
23+ },
24+ command: "curl https://example.com",
25+ shouldBlock: true,
26+ },
27+ {
28+ name: "allow non-blocked command",
29+ blockFuncs: []CommandBlockFunc{
30+ func(args []string) bool {
31+ return len(args) > 0 && args[0] == "curl"
32+ },
33+ },
34+ command: "echo hello",
35+ shouldBlock: false,
36+ },
37+ {
38+ name: "block subcommand",
39+ blockFuncs: []CommandBlockFunc{
40+ func(args []string) bool {
41+ return len(args) >= 2 && args[0] == "brew" && args[1] == "install"
42+ },
43+ },
44+ command: "brew install wget",
45+ shouldBlock: true,
46+ },
47+ {
48+ name: "allow different subcommand",
49+ blockFuncs: []CommandBlockFunc{
50+ func(args []string) bool {
51+ return len(args) >= 2 && args[0] == "brew" && args[1] == "install"
52+ },
53+ },
54+ command: "brew list",
55+ shouldBlock: false,
56+ },
57+ {
58+ name: "block npm global install with -g",
59+ blockFuncs: []CommandBlockFunc{
60+ CreateSubCommandBlocker([][]string{
61+ {"npm", "install", "-g"},
62+ {"npm", "install", "--global"},
63+ }),
64+ },
65+ command: "npm install -g typescript",
66+ shouldBlock: true,
67+ },
68+ {
69+ name: "block npm global install with --global",
70+ blockFuncs: []CommandBlockFunc{
71+ CreateSubCommandBlocker([][]string{
72+ {"npm", "install", "-g"},
73+ {"npm", "install", "--global"},
74+ }),
75+ },
76+ command: "npm install --global typescript",
77+ shouldBlock: true,
78+ },
79+ {
80+ name: "allow npm local install",
81+ blockFuncs: []CommandBlockFunc{
82+ CreateSubCommandBlocker([][]string{
83+ {"npm", "install", "-g"},
84+ {"npm", "install", "--global"},
85+ }),
86+ },
87+ command: "npm install typescript",
88+ shouldBlock: false,
89+ },
90+ }
91+
92+ for _, tt := range tests {
93+ t.Run(tt.name, func(t *testing.T) {
94+ // Create a temporary directory for each test
95+ tmpDir, err := os.MkdirTemp("", "shell-test-*")
96+ if err != nil {
97+ t.Fatalf("Failed to create temp dir: %v", err)
98+ }
99+ defer os.RemoveAll(tmpDir)
100+
101+ shell := NewShell(&Options{
102+ WorkingDir: tmpDir,
103+ BlockFuncs: tt.blockFuncs,
104+ })
105+
106+ _, _, err = shell.Exec(context.Background(), tt.command)
107+
108+ if tt.shouldBlock {
109+ if err == nil {
110+ t.Errorf("Expected command to be blocked, but it was allowed")
111+ } else if !strings.Contains(err.Error(), "not allowed for security reasons") {
112+ t.Errorf("Expected security error, got: %v", err)
113+ }
114+ } else {
115+ // For non-blocked commands, we might get other errors (like command not found)
116+ // but we shouldn't get the security error
117+ if err != nil && strings.Contains(err.Error(), "not allowed for security reasons") {
118+ t.Errorf("Command was unexpectedly blocked: %v", err)
119+ }
120+ }
121+ })
122+ }
123+}
internal/shell/shell.go+74−7
@@ -44,12 +44,16 @@ type noopLogger struct{}
44
44
45func (noopLogger) InfoPersist(msg string, keysAndValues ...interface{}) {}
45func (noopLogger) InfoPersist(msg string, keysAndValues ...interface{}) {}
46
46
47+// CommandBlockFunc is a function that determines if a command should be blocked
48+type CommandBlockFunc func(args []string) bool
49+
47// Shell provides cross-platform shell execution with optional state persistence
50// Shell provides cross-platform shell execution with optional state persistence
48type Shell struct {
51type Shell struct {
49− env []string
52+ env []string
50− cwd string
53+ cwd string
51− mu sync.Mutex
54+ mu sync.Mutex
52− logger Logger
55+ logger Logger
56+ blockFuncs []CommandBlockFunc
53}
57}
54
58
55// Options for creating a new shell
59// Options for creating a new shell
@@ -57,6 +61,7 @@ type Options struct {
57 WorkingDir string
61 WorkingDir string
58 Env []string
62 Env []string
59 Logger Logger
63 Logger Logger
64+ BlockFuncs []CommandBlockFunc
60}
65}
61
66
62// NewShell creates a new shell instance with the given options
67// NewShell creates a new shell instance with the given options
@@ -81,9 +86,10 @@ func NewShell(opts *Options) *Shell {
81 }
86 }
82
87
83 return &Shell{
88 return &Shell{
84− cwd: cwd,
89+ cwd: cwd,
85− env: env,
90+ env: env,
86− logger: logger,
91+ logger: logger,
92+ blockFuncs: opts.BlockFuncs,
87 }
93 }
88}
94}
89
95
@@ -152,6 +158,13 @@ func (s *Shell) SetEnv(key, value string) {
152 s.env = append(s.env, keyPrefix+value)
158 s.env = append(s.env, keyPrefix+value)
153}
159}
154
160
161+// SetBlockFuncs sets the command block functions for the shell
162+func (s *Shell) SetBlockFuncs(blockFuncs []CommandBlockFunc) {
163+ s.mu.Lock()
164+ defer s.mu.Unlock()
165+ s.blockFuncs = blockFuncs
166+}
167+
155// Windows-specific commands that should use native shell
168// Windows-specific commands that should use native shell
156var windowsNativeCommands = map[string]bool{
169var windowsNativeCommands = map[string]bool{
157 "dir": true,
170 "dir": true,
@@ -203,6 +216,59 @@ func (s *Shell) determineShellType(command string) ShellType {
203 return ShellTypePOSIX
216 return ShellTypePOSIX
204}
217}
205
218
219+// CreateSimpleCommandBlocker creates a CommandBlockFunc that blocks exact command matches
220+func CreateSimpleCommandBlocker(bannedCommands []string) CommandBlockFunc {
221+ bannedSet := make(map[string]bool)
222+ for _, cmd := range bannedCommands {
223+ bannedSet[cmd] = true
224+ }
225+
226+ return func(args []string) bool {
227+ if len(args) == 0 {
228+ return false
229+ }
230+ return bannedSet[args[0]]
231+ }
232+}
233+
234+// CreateSubCommandBlocker creates a CommandBlockFunc that blocks specific subcommands
235+func CreateSubCommandBlocker(blockedSubCommands [][]string) CommandBlockFunc {
236+ return func(args []string) bool {
237+ for _, blocked := range blockedSubCommands {
238+ if len(args) >= len(blocked) {
239+ match := true
240+ for i, part := range blocked {
241+ if args[i] != part {
242+ match = false
243+ break
244+ }
245+ }
246+ if match {
247+ return true
248+ }
249+ }
250+ }
251+ return false
252+ }
253+}
254+func (s *Shell) createCommandBlockHandler() func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {
255+ return func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc {
256+ return func(ctx context.Context, args []string) error {
257+ if len(args) == 0 {
258+ return next(ctx, args)
259+ }
260+
261+ for _, blockFunc := range s.blockFuncs {
262+ if blockFunc(args) {
263+ return fmt.Errorf("command '%s' is not allowed for security reasons", strings.Join(args, " "))
264+ }
265+ }
266+
267+ return next(ctx, args)
268+ }
269+ }
270+}
271+
206// execWindows executes commands using native Windows shells (cmd.exe or PowerShell)
272// execWindows executes commands using native Windows shells (cmd.exe or PowerShell)
207func (s *Shell) execWindows(ctx context.Context, command string, shell string) (string, string, error) {
273func (s *Shell) execWindows(ctx context.Context, command string, shell string) (string, string, error) {
208 var cmd *exec.Cmd
274 var cmd *exec.Cmd
@@ -291,6 +357,7 @@ func (s *Shell) execPOSIX(ctx context.Context, command string) (string, string,
291 interp.Interactive(false),
357 interp.Interactive(false),
292 interp.Env(expand.ListEnviron(s.env...)),
358 interp.Env(expand.ListEnviron(s.env...)),
293 interp.Dir(s.cwd),
359 interp.Dir(s.cwd),
360+ interp.ExecHandlers(s.createCommandBlockHandler()),
294 )
361 )
295 if err != nil {
362 if err != nil {
296 return "", "", fmt.Errorf("could not run command: %w", err)
363 return "", "", fmt.Errorf("could not run command: %w", err)