fix(shell): block `go test -exec` and ensure it works using equals

Andrey Nering created 2 months ago

Change summary

internal/llm/tools/bash.go           |  3 ++
internal/shell/command_block_test.go | 36 ++++++++++++++++++++++++++++++
internal/shell/shell.go              |  7 +++++
3 files changed, 45 insertions(+), 1 deletion(-)

Detailed changes

internal/llm/tools/bash.go

@@ -298,6 +298,9 @@ func blockFuncs() []shell.BlockFunc {
 		shell.ArgumentsBlocker("pnpm", []string{"add"}, []string{"--global"}),
 		shell.ArgumentsBlocker("pnpm", []string{"add"}, []string{"-g"}),
 		shell.ArgumentsBlocker("yarn", []string{"global", "add"}, nil),
+
+		// `go test -exec` can run arbitrary commands
+		shell.ArgumentsBlocker("go", []string{"test"}, []string{"-exec"}),
 	}
 }
 

internal/shell/command_block_test.go

@@ -231,6 +231,24 @@ func TestArgumentsBlocker(t *testing.T) {
 			input:       []string{"pacman", "-Q", "package"},
 			shouldBlock: false,
 		},
+
+		// `go test -exec`
+		{
+			name:        "go test exec",
+			cmd:         "go",
+			args:        []string{"test"},
+			flags:       []string{"-exec"},
+			input:       []string{"go", "test", "-exec", "bash -c 'echo hello'"},
+			shouldBlock: true,
+		},
+		{
+			name:        "go test exec",
+			cmd:         "go",
+			args:        []string{"test"},
+			flags:       []string{"-exec"},
+			input:       []string{"go", "test", `-exec="bash -c 'echo hello'"`},
+			shouldBlock: true,
+		},
 	}
 
 	for _, tt := range tests {
@@ -329,6 +347,24 @@ func TestSplitArgsFlags(t *testing.T) {
 			wantArgs:  []string{"package"},
 			wantFlags: []string{"-S"},
 		},
+		{
+			name:      "flag with equals sign",
+			input:     []string{"-exec=bash", "package"},
+			wantArgs:  []string{"package"},
+			wantFlags: []string{"-exec"},
+		},
+		{
+			name:      "long flag with equals sign",
+			input:     []string{"--config=/path/to/config", "run"},
+			wantArgs:  []string{"run"},
+			wantFlags: []string{"--config"},
+		},
+		{
+			name:      "flag with complex value",
+			input:     []string{`-exec="bash -c 'echo hello'"`, "test"},
+			wantArgs:  []string{"test"},
+			wantFlags: []string{"-exec"},
+		},
 	}
 
 	for _, tt := range tests {

internal/shell/shell.go

@@ -197,7 +197,12 @@ func splitArgsFlags(parts []string) (args []string, flags []string) {
 	flags = make([]string, 0, len(parts))
 	for _, part := range parts {
 		if strings.HasPrefix(part, "-") {
-			flags = append(flags, part)
+			// Extract flag name before '=' if present
+			flag := part
+			if idx := strings.IndexByte(part, '='); idx != -1 {
+				flag = part[:idx]
+			}
+			flags = append(flags, flag)
 		} else {
 			args = append(args, part)
 		}