git.iko.soy/crush
on main help sign in
refactor: improve command blocking system and fix test isolation

- Simplify command blocking logic by using utility functions instead of complex closures - Add sudo to banned commands list - Move command blocking from bash tool to shell layer for better separation of concerns - Add comprehensive tests for command blocking functionality - Fix test isolation by using temporary directories to prevent npm package files from polluting source tree - Remove redundant command validation logic from bash tool

commitf99f50427a16e36701abf1e965dd1c605c42557f authorKujtim Hoxha <kujtimii.h@gmail.com> dateFri 11 Jul 2025, 00:21 (1 year ago) parentde3b815
3 files changed+213−15
compare with main
unifiedsplit
internal/llm/tools/bash.go+16−8 internal/shell/command_block_test.go+123 internal/shell/shell.go+74−7
internal/llm/tools/bash.go+16−8
@@ -43,7 +43,7 @@ const (
4343var bannedCommands = []string{
4444 "alias", "curl", "curlie", "wget", "axel", "aria2c",
4545 "nc", "telnet", "lynx", "w3m", "links", "httpie", "xh",
46− "http-prompt", "chrome", "firefox", "safari",
46+ "http-prompt", "chrome", "firefox", "safari", "sudo",
4747}
4848
4949// getSafeReadOnlyCommands returns platform-appropriate safe commands
@@ -244,7 +244,22 @@ Important:
244244- Never update git config`, bannedCommandsStr, MaxOutputLength)
245245}
246246
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+
247258func 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+
248263 return &bashTool{
249264 permissions: permission,
250265 workingDir: workingDir,
@@ -289,13 +304,6 @@ func (b *bashTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
289304 return NewTextErrorResponse("missing command"), nil
290305 }
291306
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−
299307 isSafeReadOnly := false
300308 cmdLower := strings.ToLower(params.Command)
301309
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{}
4444
4545func (noopLogger) InfoPersist(msg string, keysAndValues ...interface{}) {}
4646
47+// CommandBlockFunc is a function that determines if a command should be blocked
48+type CommandBlockFunc func(args []string) bool
49+
4750// Shell provides cross-platform shell execution with optional state persistence
4851type Shell struct {
49− env []string
50− cwd string
51− mu sync.Mutex
52− logger Logger
52+ env []string
53+ cwd string
54+ mu sync.Mutex
55+ logger Logger
56+ blockFuncs []CommandBlockFunc
5357}
5458
5559// Options for creating a new shell
@@ -57,6 +61,7 @@ type Options struct {
5761 WorkingDir string
5862 Env []string
5963 Logger Logger
64+ BlockFuncs []CommandBlockFunc
6065}
6166
6267// NewShell creates a new shell instance with the given options
@@ -81,9 +86,10 @@ func NewShell(opts *Options) *Shell {
8186 }
8287
8388 return &Shell{
84− cwd: cwd,
85− env: env,
86− logger: logger,
89+ cwd: cwd,
90+ env: env,
91+ logger: logger,
92+ blockFuncs: opts.BlockFuncs,
8793 }
8894}
8995
@@ -152,6 +158,13 @@ func (s *Shell) SetEnv(key, value string) {
152158 s.env = append(s.env, keyPrefix+value)
153159}
154160
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+
155168// Windows-specific commands that should use native shell
156169var windowsNativeCommands = map[string]bool{
157170 "dir": true,
@@ -203,6 +216,59 @@ func (s *Shell) determineShellType(command string) ShellType {
203216 return ShellTypePOSIX
204217}
205218
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+
206272// execWindows executes commands using native Windows shells (cmd.exe or PowerShell)
207273func (s *Shell) execWindows(ctx context.Context, command string, shell string) (string, string, error) {
208274 var cmd *exec.Cmd
@@ -291,6 +357,7 @@ func (s *Shell) execPOSIX(ctx context.Context, command string) (string, string,
291357 interp.Interactive(false),
292358 interp.Env(expand.ListEnviron(s.env...)),
293359 interp.Dir(s.cwd),
360+ interp.ExecHandlers(s.createCommandBlockHandler()),
294361 )
295362 if err != nil {
296363 return "", "", fmt.Errorf("could not run command: %w", err)