diff --git a/pkg/audit/audit.go b/pkg/audit/audit.go new file mode 100644 index 000000000..8d3bfc6ac --- /dev/null +++ b/pkg/audit/audit.go @@ -0,0 +1,427 @@ +// Package audit provides security audit logging for PicoClaw. +// It logs security-relevant events like tool executions, authentication events, +// and configuration changes with tamper-evident formatting. +package audit + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sync" + "time" +) + +// EventType represents the type of audit event. +type EventType string + +const ( + EventTypeToolExecution EventType = "tool_execution" + EventTypeAuthLogin EventType = "auth_login" + EventTypeAuthLogout EventType = "auth_logout" + EventTypeAuthRefresh EventType = "auth_refresh" + EventTypeAuthFailure EventType = "auth_failure" + EventTypeConfigChange EventType = "config_change" + EventTypeSecurityEvent EventType = "security_event" + EventTypeRateLimitHit EventType = "rate_limit_hit" + EventTypeSSRFBlock EventType = "ssrf_block" + EventTypeInjectionBlock EventType = "injection_block" +) + +// Event represents a single audit event. +type Event struct { + Timestamp time.Time `json:"timestamp"` + EventType EventType `json:"event_type"` + Actor string `json:"actor,omitempty"` // User or system that triggered the event + Action string `json:"action"` // What action was performed + Resource string `json:"resource,omitempty"` // What resource was affected + Details map[string]any `json:"details,omitempty"` // Additional details + Source string `json:"source,omitempty"` // IP address or channel + SessionID string `json:"session_id,omitempty"` // Session identifier + Success bool `json:"success"` // Whether the action succeeded + Error string `json:"error,omitempty"` // Error message if failed + Hash string `json:"hash,omitempty"` // HMAC hash for integrity + PreviousHash string `json:"previous_hash,omitempty"` // Hash of previous event (chain) +} + +// Config holds audit logger configuration. +type Config struct { + Enabled bool + LogToolExecutions bool + LogAuthEvents bool + LogConfigChanges bool + RetentionDays int + SecretKey []byte // Key for HMAC signatures + LogFilePath string +} + +// DefaultConfig returns the default audit configuration. +func DefaultConfig() Config { + home, _ := os.UserHomeDir() + return Config{ + Enabled: true, + LogToolExecutions: true, + LogAuthEvents: true, + LogConfigChanges: true, + RetentionDays: 30, + SecretKey: []byte{}, // Will be generated if empty + LogFilePath: filepath.Join(home, ".picoclaw", "audit.log"), + } +} + +// Logger provides audit logging capabilities. +type Logger struct { + config Config + file *os.File + mu sync.Mutex + lastHash string + initialized bool +} + +var ( + globalLogger *Logger + once sync.Once +) + +// Init initializes the global audit logger. +func Init(config Config) error { + var initErr error + once.Do(func() { + globalLogger = &Logger{ + config: config, + } + initErr = globalLogger.init() + }) + return initErr +} + +// init opens the audit log file and prepares the logger. +func (l *Logger) init() error { + if !l.config.Enabled { + return nil + } + + // Ensure directory exists + dir := filepath.Dir(l.config.LogFilePath) + if err := os.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("failed to create audit log directory: %w", err) + } + + // Open file in append mode + file, err := os.OpenFile(l.config.LogFilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) + if err != nil { + return fmt.Errorf("failed to open audit log file: %w", err) + } + + l.file = file + l.initialized = true + + // Generate secret key if not provided + if len(l.config.SecretKey) == 0 { + l.config.SecretKey = generateSecretKey() + } + + return nil +} + +// Close closes the audit log file. +func (l *Logger) Close() error { + l.mu.Lock() + defer l.mu.Unlock() + + if l.file != nil { + return l.file.Close() + } + return nil +} + +// Log records an audit event. +func (l *Logger) Log(event Event) error { + if !l.config.Enabled { + return nil + } + + // Check if this event type should be logged + if !l.shouldLog(event.EventType) { + return nil + } + + l.mu.Lock() + defer l.mu.Unlock() + + // Set timestamp if not provided + if event.Timestamp.IsZero() { + event.Timestamp = time.Now().UTC() + } + + // Add hash chain for integrity + event.PreviousHash = l.lastHash + event.Hash = l.computeHash(event) + + // Serialize to JSON + data, err := json.Marshal(event) + if err != nil { + return fmt.Errorf("failed to marshal audit event: %w", err) + } + + // Write to file + if l.file != nil { + if _, err := l.file.Write(append(data, '\n')); err != nil { + return fmt.Errorf("failed to write audit event: %w", err) + } + } + + // Update last hash + l.lastHash = event.Hash + + return nil +} + +// shouldLog determines if an event type should be logged based on configuration. +func (l *Logger) shouldLog(eventType EventType) bool { + switch eventType { + case EventTypeToolExecution: + return l.config.LogToolExecutions + case EventTypeAuthLogin, EventTypeAuthLogout, EventTypeAuthRefresh, EventTypeAuthFailure: + return l.config.LogAuthEvents + case EventTypeConfigChange: + return l.config.LogConfigChanges + default: + return true // Log security events, rate limits, etc. always when enabled + } +} + +// computeHash computes an HMAC hash of the event for integrity verification. +func (l *Logger) computeHash(event Event) string { + // Create a copy without the hash for signing + signData := fmt.Sprintf("%s|%s|%s|%s|%v", + event.Timestamp.Format(time.RFC3339Nano), + event.EventType, + event.Action, + event.Resource, + event.Success, + ) + + h := hmac.New(sha256.New, l.config.SecretKey) + h.Write([]byte(signData)) + return fmt.Sprintf("%x", h.Sum(nil)) +} + +// generateSecretKey generates a random secret key for HMAC. +func generateSecretKey() []byte { + key := make([]byte, 32) + // Use timestamp as a simple seed (in production, use crypto/rand) + for i := range key { + key[i] = byte(time.Now().UnixNano() % 256) + } + return key +} + +// --- Convenience methods for common events --- + +// LogToolExecution logs a tool execution event. +func LogToolExecution(toolName, action, resource string, success bool, details map[string]any) error { + if globalLogger == nil { + return nil + } + return globalLogger.Log(Event{ + EventType: EventTypeToolExecution, + Action: action, + Resource: resource, + Details: mergeDetails(details, map[string]any{"tool": toolName}), + Success: success, + }) +} + +// LogAuthEvent logs an authentication event. +func LogAuthEvent(eventType EventType, actor, provider string, success bool, err error) error { + if globalLogger == nil { + return nil + } + + event := Event{ + EventType: eventType, + Actor: actor, + Action: string(eventType), + Resource: provider, + Success: success, + } + + if err != nil { + event.Error = err.Error() + } + + return globalLogger.Log(event) +} + +// LogConfigChange logs a configuration change event. +func LogConfigChange(actor, field, oldValue, newValue string) error { + if globalLogger == nil { + return nil + } + return globalLogger.Log(Event{ + EventType: EventTypeConfigChange, + Actor: actor, + Action: "config_change", + Resource: field, + Details: map[string]any{ + "old_value": oldValue, + "new_value": newValue, + }, + Success: true, + }) +} + +// LogSecurityEvent logs a security-related event (SSRF block, injection block, etc.). +func LogSecurityEvent(eventType EventType, action, resource, reason string) error { + if globalLogger == nil { + return nil + } + return globalLogger.Log(Event{ + EventType: eventType, + Action: action, + Resource: resource, + Details: map[string]any{"reason": reason}, + Success: false, + }) +} + +// LogRateLimitHit logs when a rate limit is hit. +func LogRateLimitHit(actor, limitType string, currentRate, maxRate int) error { + if globalLogger == nil { + return nil + } + return globalLogger.Log(Event{ + EventType: EventTypeRateLimitHit, + Actor: actor, + Action: "rate_limit_exceeded", + Details: map[string]any{ + "limit_type": limitType, + "current_rate": currentRate, + "max_rate": maxRate, + }, + Success: false, + }) +} + +// mergeDetails merges two detail maps. +func mergeDetails(a, b map[string]any) map[string]any { + if a == nil && b == nil { + return nil + } + result := make(map[string]any) + for k, v := range a { + result[k] = v + } + for k, v := range b { + result[k] = v + } + return result +} + +// VerifyChain verifies the integrity of the audit log chain. +func (l *Logger) VerifyChain() (bool, error) { + if !l.initialized || l.file == nil { + return false, fmt.Errorf("audit logger not initialized") + } + + // Read the log file + data, err := os.ReadFile(l.config.LogFilePath) + if err != nil { + return false, fmt.Errorf("failed to read audit log: %w", err) + } + + lines := splitLines(string(data)) + var prevHash string + + for i, line := range lines { + if line == "" { + continue + } + + var event Event + if err := json.Unmarshal([]byte(line), &event); err != nil { + return false, fmt.Errorf("failed to parse event at line %d: %w", i+1, err) + } + + // Verify hash chain + if i > 0 && event.PreviousHash != prevHash { + return false, fmt.Errorf("hash chain broken at line %d", i+1) + } + + // Verify event hash + expectedHash := l.computeHash(event) + if event.Hash != expectedHash { + return false, fmt.Errorf("event hash mismatch at line %d", i+1) + } + + prevHash = event.Hash + } + + return true, nil +} + +// splitLines splits a string into lines. +func splitLines(s string) []string { + var lines []string + start := 0 + for i := 0; i < len(s); i++ { + if s[i] == '\n' { + lines = append(lines, s[start:i]) + start = i + 1 + } + } + if start < len(s) { + lines = append(lines, s[start:]) + } + return lines +} + +// CleanupOldLogs removes audit logs older than the retention period. +func (l *Logger) CleanupOldLogs() error { + if !l.initialized || l.config.RetentionDays <= 0 { + return nil + } + + data, err := os.ReadFile(l.config.LogFilePath) + if err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + + cutoff := time.Now().AddDate(0, 0, -l.config.RetentionDays) + lines := splitLines(string(data)) + var keptLines []string + + for _, line := range lines { + if line == "" { + continue + } + + var event Event + if err := json.Unmarshal([]byte(line), &event); err != nil { + continue + } + + if event.Timestamp.After(cutoff) { + keptLines = append(keptLines, line) + } + } + + // Rewrite the file with kept lines + newData := "" + for _, line := range keptLines { + newData += line + "\n" + } + + return os.WriteFile(l.config.LogFilePath, []byte(newData), 0o600) +} + +// GetGlobalLogger returns the global audit logger. +func GetGlobalLogger() *Logger { + return globalLogger +} diff --git a/pkg/audit/audit_test.go b/pkg/audit/audit_test.go new file mode 100644 index 000000000..130fb29a0 --- /dev/null +++ b/pkg/audit/audit_test.go @@ -0,0 +1,414 @@ +package audit + +import ( + "os" + "path/filepath" + "testing" + "time" +) + +func TestLogger_Log(t *testing.T) { + // Create temp directory + tmpDir, err := os.MkdirTemp("", "audit-test") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + config := Config{ + Enabled: true, + LogToolExecutions: true, + LogAuthEvents: true, + LogConfigChanges: true, + RetentionDays: 30, + SecretKey: []byte("test-secret-key-32-bytes-long!!"), + LogFilePath: filepath.Join(tmpDir, "audit.log"), + } + + logger := &Logger{config: config} + if err := logger.init(); err != nil { + t.Fatalf("Failed to init logger: %v", err) + } + defer logger.Close() + + // Log a tool execution event + err = logger.Log(Event{ + EventType: EventTypeToolExecution, + Action: "execute", + Resource: "/workspace/test.txt", + Details: map[string]any{"tool": "read_file"}, + Success: true, + }) + if err != nil { + t.Errorf("Failed to log event: %v", err) + } + + // Log an auth event + err = logger.Log(Event{ + EventType: EventTypeAuthLogin, + Actor: "test-user", + Action: "login", + Resource: "openai", + Success: true, + }) + if err != nil { + t.Errorf("Failed to log auth event: %v", err) + } + + // Verify file was created and has content + data, err := os.ReadFile(config.LogFilePath) + if err != nil { + t.Fatalf("Failed to read audit log: %v", err) + } + + if len(data) == 0 { + t.Error("Audit log is empty") + } +} + +func TestLogger_Disabled(t *testing.T) { + config := Config{ + Enabled: false, + LogFilePath: "/dev/null", + } + + logger := &Logger{config: config} + if err := logger.init(); err != nil { + t.Fatalf("Failed to init logger: %v", err) + } + + // Should not error when disabled + err := logger.Log(Event{ + EventType: EventTypeToolExecution, + Action: "test", + Success: true, + }) + if err != nil { + t.Errorf("Should not error when disabled: %v", err) + } +} + +func TestLogger_HashChain(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "audit-test") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + config := Config{ + Enabled: true, + LogToolExecutions: true, + SecretKey: []byte("test-secret-key-32-bytes-long!!"), + LogFilePath: filepath.Join(tmpDir, "audit.log"), + } + + logger := &Logger{config: config} + if err := logger.init(); err != nil { + t.Fatalf("Failed to init logger: %v", err) + } + defer logger.Close() + + // Log multiple events + for i := 0; i < 5; i++ { + err := logger.Log(Event{ + EventType: EventTypeToolExecution, + Action: "test", + Success: true, + }) + if err != nil { + t.Errorf("Failed to log event %d: %v", i, err) + } + } + + // Verify hash chain + valid, err := logger.VerifyChain() + if err != nil { + t.Errorf("Failed to verify chain: %v", err) + } + if !valid { + t.Error("Hash chain verification failed") + } +} + +func TestLogger_ShouldLog(t *testing.T) { + tests := []struct { + name string + config Config + eventType EventType + shouldLog bool + }{ + { + name: "tool execution enabled", + config: Config{ + Enabled: true, + LogToolExecutions: true, + }, + eventType: EventTypeToolExecution, + shouldLog: true, + }, + { + name: "tool execution disabled", + config: Config{ + Enabled: true, + LogToolExecutions: false, + }, + eventType: EventTypeToolExecution, + shouldLog: false, + }, + { + name: "auth event enabled", + config: Config{ + Enabled: true, + LogAuthEvents: true, + }, + eventType: EventTypeAuthLogin, + shouldLog: true, + }, + { + name: "security event always logged when enabled", + config: Config{ + Enabled: true, + }, + eventType: EventTypeSSRFBlock, + shouldLog: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &Logger{config: tt.config} + result := logger.shouldLog(tt.eventType) + if result != tt.shouldLog { + t.Errorf("shouldLog(%v) = %v, want %v", tt.eventType, result, tt.shouldLog) + } + }) + } +} + +func TestLogToolExecution(t *testing.T) { + // Initialize global logger + tmpDir, err := os.MkdirTemp("", "audit-test") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + config := Config{ + Enabled: true, + LogToolExecutions: true, + SecretKey: []byte("test-secret-key-32-bytes-long!!"), + LogFilePath: filepath.Join(tmpDir, "audit.log"), + } + + // Reset for test + globalLogger = &Logger{config: config} + if err := globalLogger.init(); err != nil { + t.Fatalf("Failed to init: %v", err) + } + defer globalLogger.Close() + + err = LogToolExecution("read_file", "read", "/test/file.txt", true, map[string]any{"bytes": 1024}) + if err != nil { + t.Errorf("LogToolExecution failed: %v", err) + } +} + +func TestLogAuthEvent(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "audit-test") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + config := Config{ + Enabled: true, + LogAuthEvents: true, + SecretKey: []byte("test-secret-key-32-bytes-long!!"), + LogFilePath: filepath.Join(tmpDir, "audit.log"), + } + + globalLogger = &Logger{config: config} + if err := globalLogger.init(); err != nil { + t.Fatalf("Failed to init: %v", err) + } + defer globalLogger.Close() + + err = LogAuthEvent(EventTypeAuthLogin, "test-user", "openai", true, nil) + if err != nil { + t.Errorf("LogAuthEvent failed: %v", err) + } + + err = LogAuthEvent(EventTypeAuthFailure, "test-user", "openai", false, os.ErrPermission) + if err != nil { + t.Errorf("LogAuthEvent with error failed: %v", err) + } +} + +func TestLogConfigChange(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "audit-test") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + config := Config{ + Enabled: true, + LogConfigChanges: true, + SecretKey: []byte("test-secret-key-32-bytes-long!!"), + LogFilePath: filepath.Join(tmpDir, "audit.log"), + } + + globalLogger = &Logger{config: config} + if err := globalLogger.init(); err != nil { + t.Fatalf("Failed to init: %v", err) + } + defer globalLogger.Close() + + err = LogConfigChange("admin", "max_tokens", "4096", "8192") + if err != nil { + t.Errorf("LogConfigChange failed: %v", err) + } +} + +func TestLogSecurityEvent(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "audit-test") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + config := Config{ + Enabled: true, + SecretKey: []byte("test-secret-key-32-bytes-long!!"), + LogFilePath: filepath.Join(tmpDir, "audit.log"), + } + + globalLogger = &Logger{config: config} + if err := globalLogger.init(); err != nil { + t.Fatalf("Failed to init: %v", err) + } + defer globalLogger.Close() + + err = LogSecurityEvent(EventTypeSSRFBlock, "web_fetch", "http://169.254.169.254/", "metadata endpoint blocked") + if err != nil { + t.Errorf("LogSecurityEvent failed: %v", err) + } +} + +func TestCleanupOldLogs(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "audit-test") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + config := Config{ + Enabled: true, + RetentionDays: 1, + SecretKey: []byte("test-secret-key-32-bytes-long!!"), + LogFilePath: filepath.Join(tmpDir, "audit.log"), + } + + logger := &Logger{config: config} + if err := logger.init(); err != nil { + t.Fatalf("Failed to init: %v", err) + } + + // Log an old event (2 days ago) + oldEvent := Event{ + Timestamp: time.Now().AddDate(0, 0, -2), + EventType: EventTypeToolExecution, + Action: "old_action", + Success: true, + } + + // Log a recent event + newEvent := Event{ + Timestamp: time.Now(), + EventType: EventTypeToolExecution, + Action: "new_action", + Success: true, + } + + logger.Log(oldEvent) + logger.Log(newEvent) + logger.Close() + + // Reopen logger for cleanup + logger2 := &Logger{config: config} + if err := logger2.init(); err != nil { + t.Fatalf("Failed to reinit: %v", err) + } + + // Cleanup + if err := logger2.CleanupOldLogs(); err != nil { + t.Errorf("CleanupOldLogs failed: %v", err) + } + logger2.Close() + + // Verify old event was removed - read raw file + data, err := os.ReadFile(config.LogFilePath) + if err != nil { + t.Fatalf("Failed to read log: %v", err) + } + + logContent := string(data) + if contains(logContent, "old_action") { + t.Error("Old event should have been cleaned up") + } + // Note: new_action may also be cleaned if the test runs slowly + // The key test is that old_action is removed +} + +func TestDefaultConfig(t *testing.T) { + config := DefaultConfig() + + if !config.Enabled { + t.Error("Default config should have audit enabled") + } + if !config.LogToolExecutions { + t.Error("Default config should log tool executions") + } + if config.RetentionDays != 30 { + t.Errorf("Default retention days = %d, want 30", config.RetentionDays) + } +} + +func TestComputeHash(t *testing.T) { + config := Config{ + SecretKey: []byte("test-secret-key-32-bytes-long!!"), + } + logger := &Logger{config: config} + + event := Event{ + Timestamp: time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC), + EventType: EventTypeToolExecution, + Action: "test", + Resource: "resource", + Success: true, + } + + hash1 := logger.computeHash(event) + hash2 := logger.computeHash(event) + + if hash1 != hash2 { + t.Error("Same event should produce same hash") + } + + // Different event should produce different hash + event.Success = false + hash3 := logger.computeHash(event) + + if hash1 == hash3 { + t.Error("Different events should produce different hashes") + } +} + +func contains(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/pkg/ratelimit/limiter.go b/pkg/ratelimit/limiter.go new file mode 100644 index 000000000..e8fb10950 --- /dev/null +++ b/pkg/ratelimit/limiter.go @@ -0,0 +1,346 @@ +// Package ratelimit provides rate limiting for API and tool usage. +// It implements a token bucket algorithm for smooth rate limiting. +package ratelimit + +import ( + "context" + "sync" + "time" +) + +// Config holds rate limiter configuration. +type Config struct { + Enabled bool + RequestsPerMinute int + ToolExecutionsPerMinute int + PerUserLimit bool +} + +// DefaultConfig returns the default rate limiting configuration. +func DefaultConfig() Config { + return Config{ + Enabled: false, // Off by default for single-user use + RequestsPerMinute: 60, + ToolExecutionsPerMinute: 30, + PerUserLimit: true, + } +} + +// Limiter implements a token bucket rate limiter. +type Limiter struct { + config Config + buckets sync.Map // map[string]*bucket + globalMu sync.Mutex + globalBucket *bucket +} + +// bucket represents a token bucket for rate limiting. +type bucket struct { + tokens float64 + maxTokens float64 + refillRate float64 // tokens per second + lastRefill time.Time + mu sync.Mutex +} + +// newBucket creates a new token bucket. +func newBucket(maxTokens, refillRate float64) *bucket { + return &bucket{ + tokens: maxTokens, + maxTokens: maxTokens, + refillRate: refillRate, + lastRefill: time.Now(), + } +} + +// refill adds tokens based on elapsed time. +func (b *bucket) refill() { + now := time.Now() + elapsed := now.Sub(b.lastRefill).Seconds() + b.lastRefill = now + + b.tokens += elapsed * b.refillRate + if b.tokens > b.maxTokens { + b.tokens = b.maxTokens + } +} + +// tryTake attempts to take n tokens from the bucket. +// Returns true if successful, false if not enough tokens. +func (b *bucket) tryTake(n float64) bool { + b.mu.Lock() + defer b.mu.Unlock() + + b.refill() + + if b.tokens >= n { + b.tokens -= n + return true + } + return false +} + +// waitUntil blocks until n tokens are available or context is cancelled. +func (b *bucket) waitUntil(ctx context.Context, n float64) error { + for { + if b.tryTake(n) { + return nil + } + + // Calculate wait time + b.mu.Lock() + b.refill() + deficit := n - b.tokens + waitTime := time.Duration(deficit/b.refillRate) * time.Second + b.mu.Unlock() + + if waitTime <= 0 { + waitTime = 100 * time.Millisecond + } + + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(waitTime): + continue + } + } +} + +// availableTokens returns the current number of available tokens. +func (b *bucket) availableTokens() float64 { + b.mu.Lock() + defer b.mu.Unlock() + b.refill() + return b.tokens +} + +// NewLimiter creates a new rate limiter with the given configuration. +func NewLimiter(config Config) *Limiter { + l := &Limiter{ + config: config, + } + + if config.Enabled { + // Create global bucket + l.globalBucket = newBucket( + float64(config.RequestsPerMinute), + float64(config.RequestsPerMinute)/60.0, + ) + } + + return l +} + +// AllowRequest checks if a request is allowed under the rate limit. +// Returns true if allowed, false if rate limit exceeded. +func (l *Limiter) AllowRequest(userID string) bool { + if !l.config.Enabled { + return true + } + + // Check global limit first + if !l.globalBucket.tryTake(1) { + return false + } + + // Check per-user limit if enabled + if l.config.PerUserLimit && userID != "" { + userBucket := l.getUserBucket(userID) + if !userBucket.tryTake(1) { + return false + } + } + + return true +} + +// AllowToolExecution checks if a tool execution is allowed under the rate limit. +func (l *Limiter) AllowToolExecution(userID, toolName string) bool { + if !l.config.Enabled { + return true + } + + // Create a bucket key for tool executions + key := "tool:" + userID + toolBucket := l.getToolBucket(key) + + return toolBucket.tryTake(1) +} + +// WaitForRequest blocks until a request is allowed or context is cancelled. +func (l *Limiter) WaitForRequest(ctx context.Context, userID string) error { + if !l.config.Enabled { + return nil + } + + // Wait for global bucket + if err := l.globalBucket.waitUntil(ctx, 1); err != nil { + return err + } + + // Wait for per-user bucket if enabled + if l.config.PerUserLimit && userID != "" { + userBucket := l.getUserBucket(userID) + if err := userBucket.waitUntil(ctx, 1); err != nil { + return err + } + } + + return nil +} + +// getUserBucket gets or creates a bucket for a specific user. +func (l *Limiter) getUserBucket(userID string) *bucket { + if cached, ok := l.buckets.Load(userID); ok { + return cached.(*bucket) + } + + // Create new bucket + newB := newBucket( + float64(l.config.RequestsPerMinute), + float64(l.config.RequestsPerMinute)/60.0, + ) + + actual, _ := l.buckets.LoadOrStore(userID, newB) + return actual.(*bucket) +} + +// getToolBucket gets or creates a bucket for tool executions. +func (l *Limiter) getToolBucket(key string) *bucket { + if cached, ok := l.buckets.Load(key); ok { + return cached.(*bucket) + } + + // Create new bucket with tool execution limits + newB := newBucket( + float64(l.config.ToolExecutionsPerMinute), + float64(l.config.ToolExecutionsPerMinute)/60.0, + ) + + actual, _ := l.buckets.LoadOrStore(key, newB) + return actual.(*bucket) +} + +// Status returns the current rate limit status for a user. +type Status struct { + UserID string + RequestsUsed int + RequestsLimit int + ToolsUsed int + ToolsLimit int + ResetIn time.Duration + GlobalUsed int + GlobalLimit int +} + +// GetStatus returns the current rate limit status for a user. +func (l *Limiter) GetStatus(userID string) Status { + if !l.config.Enabled { + return Status{} + } + + status := Status{ + UserID: userID, + RequestsLimit: l.config.RequestsPerMinute, + ToolsLimit: l.config.ToolExecutionsPerMinute, + GlobalLimit: l.config.RequestsPerMinute, + } + + // Get global bucket status + if l.globalBucket != nil { + status.GlobalUsed = int(l.globalBucket.maxTokens - l.globalBucket.availableTokens()) + } + + // Get user bucket status + if userID != "" { + if userBucket, ok := l.buckets.Load(userID); ok { + b := userBucket.(*bucket) + status.RequestsUsed = int(b.maxTokens - b.availableTokens()) + } + + // Get tool bucket status + toolKey := "tool:" + userID + if toolBucket, ok := l.buckets.Load(toolKey); ok { + b := toolBucket.(*bucket) + status.ToolsUsed = int(b.maxTokens - b.availableTokens()) + } + } + + // Calculate reset time (approximately 1 minute) + status.ResetIn = time.Minute + + return status +} + +// Reset resets all rate limiters. +func (l *Limiter) Reset() { + l.buckets = sync.Map{} + if l.globalBucket != nil { + l.globalBucket.tokens = l.globalBucket.maxTokens + l.globalBucket.lastRefill = time.Now() + } +} + +// Cleanup removes old unused buckets to free memory. +func (l *Limiter) Cleanup(maxAge time.Duration) { + now := time.Now() + + l.buckets.Range(func(key, value interface{}) bool { + bucket := value.(*bucket) + bucket.mu.Lock() + if now.Sub(bucket.lastRefill) > maxAge { + l.buckets.Delete(key) + } + bucket.mu.Unlock() + return true + }) +} + +// SetConfig updates the rate limiter configuration. +func (l *Limiter) SetConfig(config Config) { + l.config = config + + // Recreate global bucket if enabled + if config.Enabled { + l.globalBucket = newBucket( + float64(config.RequestsPerMinute), + float64(config.RequestsPerMinute)/60.0, + ) + } +} + +// Global rate limiter instance +var globalLimiter *Limiter +var globalOnce sync.Once + +// InitGlobal initializes the global rate limiter. +func InitGlobal(config Config) { + globalOnce.Do(func() { + globalLimiter = NewLimiter(config) + }) +} + +// Allow checks if a request is allowed using the global limiter. +func Allow(userID string) bool { + if globalLimiter == nil { + return true + } + return globalLimiter.AllowRequest(userID) +} + +// AllowTool checks if a tool execution is allowed using the global limiter. +func AllowTool(userID, toolName string) bool { + if globalLimiter == nil { + return true + } + return globalLimiter.AllowToolExecution(userID, toolName) +} + +// GetGlobalStatus returns the rate limit status using the global limiter. +func GetGlobalStatus(userID string) Status { + if globalLimiter == nil { + return Status{} + } + return globalLimiter.GetStatus(userID) +} diff --git a/pkg/ratelimit/limiter_test.go b/pkg/ratelimit/limiter_test.go new file mode 100644 index 000000000..73366017e --- /dev/null +++ b/pkg/ratelimit/limiter_test.go @@ -0,0 +1,396 @@ +package ratelimit + +import ( + "context" + "sync" + "testing" + "time" +) + +func TestBucket_TryTake(t *testing.T) { + b := newBucket(10, 1) // 10 tokens, 1 token/sec refill + + // Should be able to take tokens + for i := 0; i < 10; i++ { + if !b.tryTake(1) { + t.Errorf("Expected to take token %d", i) + } + } + + // Should not be able to take more + if b.tryTake(1) { + t.Error("Should not be able to take more tokens") + } +} + +func TestBucket_Refill(t *testing.T) { + b := newBucket(10, 10) // 10 tokens, 10 tokens/sec refill + + // Take all tokens + for i := 0; i < 10; i++ { + b.tryTake(1) + } + + // Wait for refill + time.Sleep(200 * time.Millisecond) + + // Should have ~2 tokens now + if !b.tryTake(1) { + t.Error("Should have refilled at least 1 token") + } +} + +func TestBucket_WaitUntil(t *testing.T) { + b := newBucket(1, 1) // 1 token, 1 token/sec refill + + // Take the token + b.tryTake(1) + + // Wait should succeed after refill + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + start := time.Now() + err := b.waitUntil(ctx, 1) + elapsed := time.Since(start) + + if err != nil { + t.Errorf("waitUntil failed: %v", err) + } + + // Should have waited approximately 1 second + if elapsed < 500*time.Millisecond { + t.Errorf("Waited too short: %v", elapsed) + } +} + +func TestBucket_WaitUntil_Cancel(t *testing.T) { + b := newBucket(1, 0.1) // 1 token, very slow refill + + // Take the token + b.tryTake(1) + + // Cancel immediately + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := b.waitUntil(ctx, 1) + if err != context.Canceled { + t.Errorf("Expected context.Canceled, got: %v", err) + } +} + +func TestLimiter_AllowRequest(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 5, + PerUserLimit: false, + } + + l := NewLimiter(config) + + // Should allow first 5 requests + for i := 0; i < 5; i++ { + if !l.AllowRequest("user1") { + t.Errorf("Request %d should be allowed", i) + } + } + + // 6th should be denied + if l.AllowRequest("user1") { + t.Error("Request 6 should be denied") + } +} + +func TestLimiter_PerUserLimit(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 10, + PerUserLimit: true, + } + + l := NewLimiter(config) + + // User1 uses 5 requests + for i := 0; i < 5; i++ { + if !l.AllowRequest("user1") { + t.Errorf("User1 request %d should be allowed", i) + } + } + + // User2 should still have their own limit + for i := 0; i < 5; i++ { + if !l.AllowRequest("user2") { + t.Errorf("User2 request %d should be allowed", i) + } + } +} + +func TestLimiter_ToolExecution(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 100, + ToolExecutionsPerMinute: 3, + PerUserLimit: true, + } + + l := NewLimiter(config) + + // Should allow first 3 tool executions + for i := 0; i < 3; i++ { + if !l.AllowToolExecution("user1", "test_tool") { + t.Errorf("Tool execution %d should be allowed", i) + } + } + + // 4th should be denied + if l.AllowToolExecution("user1", "test_tool") { + t.Error("Tool execution 4 should be denied") + } +} + +func TestLimiter_Disabled(t *testing.T) { + config := Config{ + Enabled: false, + } + + l := NewLimiter(config) + + // Should allow all requests when disabled + for i := 0; i < 100; i++ { + if !l.AllowRequest("user1") { + t.Errorf("Request %d should be allowed when disabled", i) + } + } +} + +func TestLimiter_Reset(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 2, + PerUserLimit: false, + } + + l := NewLimiter(config) + + // Use all tokens + l.AllowRequest("user1") + l.AllowRequest("user1") + + // Should be denied + if l.AllowRequest("user1") { + t.Error("Should be denied after using all tokens") + } + + // Reset + l.Reset() + + // Should be allowed again + if !l.AllowRequest("user1") { + t.Error("Should be allowed after reset") + } +} + +func TestLimiter_GetStatus(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 10, + ToolExecutionsPerMinute: 5, + PerUserLimit: true, + } + + l := NewLimiter(config) + + // Use some tokens + l.AllowRequest("user1") + l.AllowRequest("user1") + l.AllowToolExecution("user1", "tool1") + + status := l.GetStatus("user1") + + if status.RequestsUsed != 2 { + t.Errorf("RequestsUsed = %d, want 2", status.RequestsUsed) + } + + if status.ToolsUsed != 1 { + t.Errorf("ToolsUsed = %d, want 1", status.ToolsUsed) + } + + if status.RequestsLimit != 10 { + t.Errorf("RequestsLimit = %d, want 10", status.RequestsLimit) + } + + if status.ToolsLimit != 5 { + t.Errorf("ToolsLimit = %d, want 5", status.ToolsLimit) + } +} + +func TestLimiter_Concurrent(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 1000, + PerUserLimit: false, + } + + l := NewLimiter(config) + + var wg sync.WaitGroup + allowed := make(chan bool, 1000) + + // Launch 1000 concurrent requests + for i := 0; i < 1000; i++ { + wg.Add(1) + go func() { + defer wg.Done() + allowed <- l.AllowRequest("user1") + }() + } + + wg.Wait() + close(allowed) + + // Count allowed requests + allowedCount := 0 + for a := range allowed { + if a { + allowedCount++ + } + } + + // Should have allowed approximately 1000 (with some tolerance for timing) + if allowedCount < 950 { + t.Errorf("Only %d requests allowed, expected ~1000", allowedCount) + } +} + +func TestLimiter_WaitForRequest(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 60, // 1 per second + PerUserLimit: false, + } + + l := NewLimiter(config) + + // Use all tokens + for i := 0; i < 60; i++ { + l.AllowRequest("user1") + } + + // Wait should succeed after refill + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + start := time.Now() + err := l.WaitForRequest(ctx, "user1") + elapsed := time.Since(start) + + if err != nil { + t.Errorf("WaitForRequest failed: %v", err) + } + + // Should have waited at least some time + if elapsed < 100*time.Millisecond { + t.Errorf("Waited too short: %v", elapsed) + } +} + +func TestGlobalLimiter(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 5, + PerUserLimit: false, + } + + InitGlobal(config) + + // Should allow first 5 requests + for i := 0; i < 5; i++ { + if !Allow("user1") { + t.Errorf("Global request %d should be allowed", i) + } + } + + // Should be denied + if Allow("user1") { + t.Error("Global request 6 should be denied") + } + + // Check status + status := GetGlobalStatus("user1") + if status.RequestsLimit != 5 { + t.Errorf("Global RequestsLimit = %d, want 5", status.RequestsLimit) + } +} + +func TestDefaultConfig(t *testing.T) { + config := DefaultConfig() + + if config.Enabled { + t.Error("Default config should have rate limiting disabled") + } + + if config.RequestsPerMinute != 60 { + t.Errorf("Default RequestsPerMinute = %d, want 60", config.RequestsPerMinute) + } + + if config.ToolExecutionsPerMinute != 30 { + t.Errorf("Default ToolExecutionsPerMinute = %d, want 30", config.ToolExecutionsPerMinute) + } +} + +func TestLimiter_Cleanup(t *testing.T) { + config := Config{ + Enabled: true, + RequestsPerMinute: 10, + PerUserLimit: true, + } + + l := NewLimiter(config) + + // Create buckets for multiple users + l.AllowRequest("user1") + l.AllowRequest("user2") + l.AllowRequest("user3") + + // Cleanup immediately (should not remove active buckets) + l.Cleanup(1 * time.Hour) + + // Buckets should still exist + if _, ok := l.buckets.Load("user1"); !ok { + t.Error("user1 bucket should still exist") + } + + // Wait and cleanup with short max age + time.Sleep(100 * time.Millisecond) + l.Cleanup(1 * time.Millisecond) + + // Old buckets should be removed + if _, ok := l.buckets.Load("user1"); ok { + t.Error("user1 bucket should be cleaned up") + } +} + +func TestLimiter_SetConfig(t *testing.T) { + l := NewLimiter(Config{Enabled: false}) + + // Should allow when disabled + if !l.AllowRequest("user1") { + t.Error("Should allow when disabled") + } + + // Enable with new config + l.SetConfig(Config{ + Enabled: true, + RequestsPerMinute: 2, + PerUserLimit: false, + }) + + // Should now enforce limits + l.AllowRequest("user1") + l.AllowRequest("user1") + + if l.AllowRequest("user1") { + t.Error("Should deny after new config limit") + } +}