diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index dfbef9fbc..424bebeb9 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -55,6 +55,32 @@ func NewAgentInstance( toolsRegistry.Register(tools.NewEditFileTool(workspace, restrict)) toolsRegistry.Register(tools.NewAppendFileTool(workspace, restrict)) + // Configure security middleware from config + if cfg != nil { + sec := cfg.Tools.Security + mw := tools.NewToolMiddleware() + for toolName, pc := range sec.ToolPolicies { + enabled := true + if pc.Enabled != nil { + enabled = *pc.Enabled + } + maxArg := pc.MaxArgSize + if maxArg == 0 { + maxArg = sec.DefaultMaxArgSize + } + maxCalls := pc.MaxCallsPerMin + if maxCalls == 0 { + maxCalls = sec.DefaultMaxCallsPerMin + } + mw.SetPolicy(toolName, tools.ToolPolicy{ + Enabled: enabled, + MaxArgSize: maxArg, + MaxCallsPerMin: maxCalls, + }) + } + toolsRegistry.Middleware = mw + } + sessionsDir := filepath.Join(workspace, "sessions") sessionsManager := session.NewSessionManager(sessionsDir) diff --git a/pkg/config/config.go b/pkg/config/config.go index 036021e49..73e6780af 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -452,11 +452,24 @@ type ExecConfig struct { CustomDenyPatterns []string `json:"custom_deny_patterns" env:"PICOCLAW_TOOLS_EXEC_CUSTOM_DENY_PATTERNS"` } +type ToolPolicyConfig struct { + Enabled *bool `json:"enabled,omitempty"` + MaxArgSize int `json:"max_arg_size,omitempty" env:"MAX_ARG_SIZE"` + MaxCallsPerMin int `json:"max_calls_per_min,omitempty" env:"MAX_CALLS_PER_MIN"` +} + +type SecurityConfig struct { + DefaultMaxArgSize int `json:"default_max_arg_size" env:"PICOCLAW_TOOLS_SECURITY_DEFAULT_MAX_ARG_SIZE"` + DefaultMaxCallsPerMin int `json:"default_max_calls_per_min" env:"PICOCLAW_TOOLS_SECURITY_DEFAULT_MAX_CALLS_PER_MIN"` + ToolPolicies map[string]ToolPolicyConfig `json:"tool_policies,omitempty"` +} + type ToolsConfig struct { - Web WebToolsConfig `json:"web"` - Cron CronToolsConfig `json:"cron"` - Exec ExecConfig `json:"exec"` - Skills SkillsToolsConfig `json:"skills"` + Web WebToolsConfig `json:"web"` + Cron CronToolsConfig `json:"cron"` + Exec ExecConfig `json:"exec"` + Skills SkillsToolsConfig `json:"skills"` + Security SecurityConfig `json:"security"` } type SkillsToolsConfig struct { diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 7654326e7..bf041f2ed 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -290,6 +290,15 @@ func DefaultConfig() *Config { Exec: ExecConfig{ EnableDenyPatterns: true, }, + Security: SecurityConfig{ + DefaultMaxArgSize: 100000, + ToolPolicies: map[string]ToolPolicyConfig{ + "exec": {MaxCallsPerMin: 30}, + "spawn": {MaxCallsPerMin: 5}, + "web_fetch": {MaxCallsPerMin: 20}, + "web_search": {MaxCallsPerMin: 10}, + }, + }, Skills: SkillsToolsConfig{ Registries: SkillsRegistriesConfig{ ClawHub: ClawHubRegistryConfig{ diff --git a/pkg/tools/middleware.go b/pkg/tools/middleware.go new file mode 100644 index 000000000..a91acda68 --- /dev/null +++ b/pkg/tools/middleware.go @@ -0,0 +1,93 @@ +package tools + +import ( + "fmt" + "sync" + "time" +) + +// ToolPolicy defines per-tool security constraints. +type ToolPolicy struct { + MaxArgSize int // Max total size of all args in bytes (0 = unlimited) + MaxCallsPerMin int // Rate limit: calls per minute (0 = unlimited) + Enabled bool // Whether the tool is allowed to execute +} + +// ToolMiddleware provides pre-execution security checks for tool calls. +type ToolMiddleware struct { + policies map[string]ToolPolicy + limiters map[string]*rateBucket + mu sync.RWMutex + nowFunc func() time.Time +} + +// NewToolMiddleware creates a new middleware with default settings. +func NewToolMiddleware() *ToolMiddleware { + return &ToolMiddleware{ + policies: make(map[string]ToolPolicy), + limiters: make(map[string]*rateBucket), + nowFunc: time.Now, + } +} + +// SetPolicy configures the security policy for a specific tool. +func (m *ToolMiddleware) SetPolicy(toolName string, policy ToolPolicy) { + m.mu.Lock() + defer m.mu.Unlock() + m.policies[toolName] = policy + if policy.MaxCallsPerMin > 0 { + m.limiters[toolName] = newRateBucket(policy.MaxCallsPerMin, m.nowFunc) + } else { + delete(m.limiters, toolName) + } +} + +// Check validates a tool call against its policy. +// Returns nil if allowed, or an error describing why it was blocked. +func (m *ToolMiddleware) Check(toolName string, args map[string]any) error { + m.mu.RLock() + policy, hasPolicy := m.policies[toolName] + limiter := m.limiters[toolName] + m.mu.RUnlock() + + if !hasPolicy { + return nil // no policy = allow + } + + if !policy.Enabled { + return fmt.Errorf("tool %q is disabled by policy", toolName) + } + + // Input size validation + if policy.MaxArgSize > 0 { + totalSize := estimateArgSize(args) + if totalSize > policy.MaxArgSize { + return fmt.Errorf("tool %q input too large (%d bytes, max %d)", toolName, totalSize, policy.MaxArgSize) + } + } + + // Rate limiting + if limiter != nil && !limiter.Allow() { + return fmt.Errorf("tool %q rate limited (max %d/min)", toolName, policy.MaxCallsPerMin) + } + + return nil +} + +// estimateArgSize calculates the approximate size of tool arguments in bytes. +func estimateArgSize(args map[string]any) int { + total := 0 + for _, v := range args { + switch val := v.(type) { + case string: + total += len(val) + case float64: + total += 8 + case bool: + total += 1 + default: + total += 64 // estimate for complex types + } + } + return total +} diff --git a/pkg/tools/middleware_test.go b/pkg/tools/middleware_test.go new file mode 100644 index 000000000..f7526e15f --- /dev/null +++ b/pkg/tools/middleware_test.go @@ -0,0 +1,183 @@ +package tools + +import ( + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestMiddleware_NoPolicy_Allows(t *testing.T) { + mw := NewToolMiddleware() + err := mw.Check("any_tool", map[string]any{"key": "value"}) + assert.NoError(t, err) +} + +func TestMiddleware_Disabled_Blocks(t *testing.T) { + mw := NewToolMiddleware() + mw.SetPolicy("dangerous_tool", ToolPolicy{Enabled: false}) + + err := mw.Check("dangerous_tool", map[string]any{}) + require.Error(t, err) + assert.Contains(t, err.Error(), "disabled by policy") +} + +func TestMiddleware_Enabled_Allows(t *testing.T) { + mw := NewToolMiddleware() + mw.SetPolicy("safe_tool", ToolPolicy{Enabled: true}) + + err := mw.Check("safe_tool", map[string]any{"key": "value"}) + assert.NoError(t, err) +} + +func TestMiddleware_ArgSizeLimit_Blocks(t *testing.T) { + mw := NewToolMiddleware() + mw.SetPolicy("exec", ToolPolicy{Enabled: true, MaxArgSize: 100}) + + // Small args — should pass + err := mw.Check("exec", map[string]any{"cmd": "ls"}) + assert.NoError(t, err) + + // Large args — should block + largeCmd := strings.Repeat("a", 200) + err = mw.Check("exec", map[string]any{"cmd": largeCmd}) + require.Error(t, err) + assert.Contains(t, err.Error(), "input too large") +} + +func TestMiddleware_RateLimit_AllowsThenBlocks(t *testing.T) { + now := time.Now() + mw := NewToolMiddleware() + mw.nowFunc = func() time.Time { return now } + mw.SetPolicy("web_fetch", ToolPolicy{Enabled: true, MaxCallsPerMin: 3}) + + args := map[string]any{"url": "https://example.com"} + + // First 3 calls should pass + for i := 0; i < 3; i++ { + err := mw.Check("web_fetch", args) + assert.NoError(t, err, "call %d should be allowed", i+1) + } + + // 4th call should be blocked + err := mw.Check("web_fetch", args) + require.Error(t, err) + assert.Contains(t, err.Error(), "rate limited") +} + +func TestMiddleware_RateLimit_ResetsAfterWindow(t *testing.T) { + now := time.Now() + currentTime := now + mw := NewToolMiddleware() + mw.nowFunc = func() time.Time { return currentTime } + mw.SetPolicy("exec", ToolPolicy{Enabled: true, MaxCallsPerMin: 2}) + + args := map[string]any{"cmd": "ls"} + + assert.NoError(t, mw.Check("exec", args)) + assert.NoError(t, mw.Check("exec", args)) + assert.Error(t, mw.Check("exec", args), "should be blocked at limit") + + // Advance past window + currentTime = now.Add(61 * time.Second) + assert.NoError(t, mw.Check("exec", args), "should be allowed after window reset") +} + +func TestMiddleware_DifferentTools_IndependentPolicies(t *testing.T) { + mw := NewToolMiddleware() + mw.SetPolicy("exec", ToolPolicy{Enabled: true, MaxCallsPerMin: 1}) + mw.SetPolicy("web_fetch", ToolPolicy{Enabled: true, MaxCallsPerMin: 1}) + + // Use up exec limit + assert.NoError(t, mw.Check("exec", map[string]any{"cmd": "ls"})) + assert.Error(t, mw.Check("exec", map[string]any{"cmd": "pwd"})) + + // web_fetch should still have its own limit + assert.NoError(t, mw.Check("web_fetch", map[string]any{"url": "https://example.com"})) + assert.Error(t, mw.Check("web_fetch", map[string]any{"url": "https://other.com"})) +} + +func TestMiddleware_UnknownTool_Allowed(t *testing.T) { + mw := NewToolMiddleware() + mw.SetPolicy("exec", ToolPolicy{Enabled: false}) + + // Tool without a policy should be allowed + err := mw.Check("read_file", map[string]any{"path": "/tmp/file"}) + assert.NoError(t, err) +} + +func TestMiddleware_ConcurrentAccess(t *testing.T) { + now := time.Now() + mw := NewToolMiddleware() + mw.nowFunc = func() time.Time { return now } + mw.SetPolicy("exec", ToolPolicy{Enabled: true, MaxCallsPerMin: 50}) + + var wg sync.WaitGroup + allowed := 0 + var mu sync.Mutex + + for i := 0; i < 100; i++ { + wg.Add(1) + go func() { + defer wg.Done() + err := mw.Check("exec", map[string]any{"cmd": "ls"}) + if err == nil { + mu.Lock() + allowed++ + mu.Unlock() + } + }() + } + wg.Wait() + + assert.Equal(t, 50, allowed, "exactly 50 calls should be allowed") +} + +func TestEstimateArgSize(t *testing.T) { + tests := []struct { + name string + args map[string]any + expected int + }{ + { + name: "empty args", + args: map[string]any{}, + expected: 0, + }, + { + name: "string arg", + args: map[string]any{"cmd": "hello"}, + expected: 5, + }, + { + name: "float64 arg", + args: map[string]any{"count": float64(42)}, + expected: 8, + }, + { + name: "bool arg", + args: map[string]any{"verbose": true}, + expected: 1, + }, + { + name: "mixed args", + args: map[string]any{"cmd": "ls -la", "count": float64(5), "verbose": true}, + expected: 6 + 8 + 1, // "ls -la" + float64 + bool + }, + { + name: "unknown type", + args: map[string]any{"data": []int{1, 2, 3}}, + expected: 64, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := estimateArgSize(tt.args) + assert.Equal(t, tt.expected, got) + }) + } +} diff --git a/pkg/tools/ratelimit.go b/pkg/tools/ratelimit.go new file mode 100644 index 000000000..4cad95eaf --- /dev/null +++ b/pkg/tools/ratelimit.go @@ -0,0 +1,49 @@ +package tools + +import ( + "sync" + "time" +) + +// rateBucket is a sliding-window rate limiter that tracks timestamps +// of recent calls within a 1-minute window. Zero external dependencies. +type rateBucket struct { + mu sync.Mutex + max int + calls []time.Time + nowFunc func() time.Time +} + +func newRateBucket(maxPerMinute int, nowFunc func() time.Time) *rateBucket { + return &rateBucket{ + max: maxPerMinute, + calls: make([]time.Time, 0, maxPerMinute), + nowFunc: nowFunc, + } +} + +// Allow returns true if the call is within the rate limit. +func (rb *rateBucket) Allow() bool { + rb.mu.Lock() + defer rb.mu.Unlock() + + now := rb.nowFunc() + cutoff := now.Add(-time.Minute) + + // Prune expired entries + valid := 0 + for _, t := range rb.calls { + if t.After(cutoff) { + rb.calls[valid] = t + valid++ + } + } + rb.calls = rb.calls[:valid] + + if len(rb.calls) >= rb.max { + return false + } + + rb.calls = append(rb.calls, now) + return true +} diff --git a/pkg/tools/ratelimit_test.go b/pkg/tools/ratelimit_test.go new file mode 100644 index 000000000..c698ebd67 --- /dev/null +++ b/pkg/tools/ratelimit_test.go @@ -0,0 +1,92 @@ +package tools + +import ( + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestRateBucket_AllowsUnderLimit(t *testing.T) { + now := time.Now() + rb := newRateBucket(5, func() time.Time { return now }) + + for i := 0; i < 5; i++ { + assert.True(t, rb.Allow(), "call %d should be allowed", i+1) + } +} + +func TestRateBucket_BlocksOverLimit(t *testing.T) { + now := time.Now() + rb := newRateBucket(3, func() time.Time { return now }) + + for i := 0; i < 3; i++ { + assert.True(t, rb.Allow()) + } + assert.False(t, rb.Allow(), "4th call should be blocked") + assert.False(t, rb.Allow(), "5th call should also be blocked") +} + +func TestRateBucket_ResetsAfterWindow(t *testing.T) { + now := time.Now() + currentTime := now + rb := newRateBucket(2, func() time.Time { return currentTime }) + + // Fill the bucket + assert.True(t, rb.Allow()) + assert.True(t, rb.Allow()) + assert.False(t, rb.Allow(), "should be blocked at limit") + + // Advance time past the 1-minute window + currentTime = now.Add(61 * time.Second) + + // Should be allowed again + assert.True(t, rb.Allow(), "should be allowed after window expires") + assert.True(t, rb.Allow(), "second call after reset should be allowed") + assert.False(t, rb.Allow(), "should be blocked again at limit") +} + +func TestRateBucket_PartialExpiry(t *testing.T) { + now := time.Now() + currentTime := now + rb := newRateBucket(3, func() time.Time { return currentTime }) + + // Make 3 calls + assert.True(t, rb.Allow()) + currentTime = now.Add(10 * time.Second) + assert.True(t, rb.Allow()) + currentTime = now.Add(20 * time.Second) + assert.True(t, rb.Allow()) + assert.False(t, rb.Allow(), "should be blocked at limit") + + // Advance 50s — first call (at t=0) expires, others (t=10, t=20) still valid + currentTime = now.Add(61 * time.Second) + + // One slot freed up + assert.True(t, rb.Allow(), "should allow one more after partial expiry") +} + +func TestRateBucket_Concurrent(t *testing.T) { + now := time.Now() + rb := newRateBucket(100, func() time.Time { return now }) + + var wg sync.WaitGroup + allowed := int32(0) + var mu sync.Mutex + + for i := 0; i < 200; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if rb.Allow() { + mu.Lock() + allowed++ + mu.Unlock() + } + }() + } + wg.Wait() + + assert.Equal(t, int32(100), allowed, "exactly 100 calls should be allowed") +} diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index 6ecb8ae7c..3beae1aed 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -11,8 +11,9 @@ import ( ) type ToolRegistry struct { - tools map[string]Tool - mu sync.RWMutex + tools map[string]Tool + mu sync.RWMutex + Middleware *ToolMiddleware } func NewToolRegistry() *ToolRegistry { @@ -77,6 +78,18 @@ func (r *ToolRegistry) ExecuteWithContext( }) } + // Security middleware check (pre-execution) + if r.Middleware != nil { + if err := r.Middleware.Check(name, args); err != nil { + logger.WarnCF("tool", "Tool blocked by middleware", + map[string]any{ + "tool": name, + "reason": err.Error(), + }) + return ErrorResult(fmt.Sprintf("Blocked: %v", err)) + } + } + start := time.Now() result := tool.Execute(ctx, args) duration := time.Since(start)