enh(security): Improve whitelisting and threshold applications

config.example.yml

@@ -74,3 +74,8 @@ security:
     auto_block_enabled: true
     auto_block_threshold: 50
     auto_block_duration: 24h0m0s
+    # Only list proxies you control. Forwarded client IPs are honored solely
+    # when the connecting peer matches; an empty list ignores them, which
+    # prevents X-Forwarded-For / CF-Connecting-IP spoofing.
+    trusted_proxies: []
+    trust_cf_header: false

internal/api/config_handlers.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"fmt"
 	"net/http"
 	"strings"
 
@@ -53,23 +54,61 @@ func (s *Server) updateConfigKey(c *gin.Context) {
 	}
 
 	applied := false
+	var applyErr error
 	if apply, ok := s.runtimeAppliers()[key]; ok {
-		apply(s)
-		applied = true
+		applyErr = apply(s)
+		applied = applyErr == nil
 	}
 
 	entry, _ := config.Get(s.config, key)
-	c.JSON(http.StatusOK, gin.H{
+	resp := gin.H{
 		"entry":   entry,
 		"applied": applied,
-	})
+	}
+	if applyErr != nil {
+		resp["apply_error"] = applyErr.Error()
+	}
+	c.JSON(http.StatusOK, resp)
 }
 
-func (s *Server) runtimeAppliers() map[string]func(*Server) {
-	return map[string]func(*Server){
-		"cleanup.timeout": func(srv *Server) {
+func (s *Server) runtimeAppliers() map[string]func(*Server) error {
+	applyDetectorThresholds := func(srv *Server) error {
+		if srv.securityManager == nil {
+			return nil
+		}
+		srv.securityManager.SetDetectorThresholds(
+			srv.config.Security.RateThreshold,
+			srv.config.Security.NotFoundThreshold,
+			srv.config.Security.AuthFailureThreshold,
+			srv.config.Security.UniquePathsThreshold,
+			srv.config.Security.RepeatedHitsThreshold,
+			srv.config.Security.DetectionWindow,
+		)
+		return nil
+	}
+	regenerateSecurityScripts := func(srv *Server) error {
+		if !srv.config.Security.Enabled {
+			return nil
+		}
+		if !srv.infraManager.IsNginxRunning() {
+			return fmt.Errorf("value saved but nginx is not running; regenerate security scripts to apply")
+		}
+		_, err := srv.infraManager.RefreshSecurityScripts()
+		return err
+	}
+	return map[string]func(*Server) error{
+		"cleanup.timeout": func(srv *Server) error {
 			srv.manager.SetCleanupTimeout(srv.config.Cleanup.Timeout)
+			return nil
 		},
+		"security.rate_threshold":          applyDetectorThresholds,
+		"security.not_found_threshold":     applyDetectorThresholds,
+		"security.auth_failure_threshold":  applyDetectorThresholds,
+		"security.unique_paths_threshold":  applyDetectorThresholds,
+		"security.repeated_hits_threshold": applyDetectorThresholds,
+		"security.detection_window":        applyDetectorThresholds,
+		"security.trusted_proxies":         regenerateSecurityScripts,
+		"security.trust_cf_header":         regenerateSecurityScripts,
 	}
 }
 

internal/api/config_handlers_test.go

@@ -0,0 +1,31 @@
+package api
+
+import (
+	"testing"
+
+	"github.com/flatrun/agent/pkg/config"
+)
+
+func TestRuntimeConfigKeysAdvertisesTrustKeys(t *testing.T) {
+	server := &Server{config: &config.Config{}}
+	keys := server.runtimeConfigKeys()
+
+	for _, key := range []string{"security.trusted_proxies", "security.trust_cf_header"} {
+		if !keys[key] {
+			t.Errorf("expected %q to be advertised as a runtime config key", key)
+		}
+	}
+}
+
+func TestTrustKeyApplierNoOpWhenSecurityDisabled(t *testing.T) {
+	server := &Server{config: &config.Config{}}
+	server.config.Security.Enabled = false
+
+	apply := server.runtimeAppliers()["security.trusted_proxies"]
+	if apply == nil {
+		t.Fatal("expected an applier for security.trusted_proxies")
+	}
+	if err := apply(server); err != nil {
+		t.Fatalf("applier should be a no-op when security is disabled, got: %v", err)
+	}
+}

internal/infra/manager.go

@@ -379,7 +379,7 @@ func (m *Manager) SetNginxRealtimeCaptureWithStatus(enabled bool) (map[string]in
 			result["agent_ip"] = agentIP
 			result["agent_port"] = agentPort
 
-			securityLua, err := templates.GetNginxSecurityLuaWithConfig(agentIP, agentPort, m.config.Security.InternalAPIToken)
+			securityLua, err := templates.GetNginxSecurityLuaWithConfig(agentIP, agentPort, m.config.Security.InternalAPIToken, m.config.Security.TrustedProxies, m.config.Security.TrustCFHeader)
 			if err != nil {
 				errors = append(errors, fmt.Sprintf("failed to get security.lua template: %v", err))
 			} else {
@@ -1245,7 +1245,7 @@ func (m *Manager) RefreshSecurityScripts() (*RefreshSecurityScriptsResult, error
 	}
 
 	// Generate and write security.lua with injected IP
-	securityLua, err := templates.GetNginxSecurityLuaWithConfig(agentIP, agentPort, m.config.Security.InternalAPIToken)
+	securityLua, err := templates.GetNginxSecurityLuaWithConfig(agentIP, agentPort, m.config.Security.InternalAPIToken, m.config.Security.TrustedProxies, m.config.Security.TrustCFHeader)
 	if err != nil {
 		result.Errors = append(result.Errors, fmt.Sprintf("failed to generate security.lua: %v", err))
 		result.Success = false

internal/security/manager.go

@@ -11,6 +11,9 @@ type Manager struct {
 	detector        *Detector
 	deploymentsPath string
 	mu              sync.RWMutex
+
+	wlMu    sync.RWMutex
+	wlCache *whitelistCache
 }
 
 func NewManager(deploymentsPath string) (*Manager, error) {
@@ -57,6 +60,14 @@ func (m *Manager) IngestEvent(event *IngestEvent, autoBlockDuration time.Duratio
 
 	result := &IngestResult{}
 
+	whitelisted, err := m.IsRequestWhitelisted(event.SourceIP, event.RequestPath)
+	if err != nil {
+		return nil, err
+	}
+	if whitelisted {
+		return result, nil
+	}
+
 	// Check if IP is blocked - if so, don't process
 	blocked, err := m.db.IsIPBlocked(event.SourceIP)
 	if err != nil {
@@ -180,11 +191,19 @@ func (m *Manager) GetWhitelist() ([]WhitelistEntry, error) {
 }
 
 func (m *Manager) AddWhitelistEntry(value, entryType, reason string) (int64, error) {
-	return m.db.AddWhitelistEntry(value, entryType, reason, false)
+	id, err := m.db.AddWhitelistEntry(value, entryType, reason, false)
+	if err == nil {
+		m.invalidateWhitelistCache()
+	}
+	return id, err
 }
 
 func (m *Manager) RemoveWhitelistEntry(id int64) error {
-	return m.db.RemoveWhitelistEntry(id)
+	err := m.db.RemoveWhitelistEntry(id)
+	if err == nil {
+		m.invalidateWhitelistCache()
+	}
+	return err
 }
 
 func (m *Manager) IsWhitelisted(value string) (bool, error) {
@@ -196,6 +215,9 @@ func (m *Manager) AddDockerGatewayToWhitelist(gatewayIP string) error {
 		return nil
 	}
 	_, err := m.db.AddWhitelistEntry(gatewayIP, "ip", "Docker gateway", true)
+	if err == nil {
+		m.invalidateWhitelistCache()
+	}
 	return err
 }
 

internal/security/manager_test.go

@@ -0,0 +1,186 @@
+package security
+
+import (
+	"testing"
+	"time"
+)
+
+func newTestManager(t *testing.T) *Manager {
+	t.Helper()
+	m, err := NewManager(t.TempDir())
+	if err != nil {
+		t.Fatalf("NewManager: %v", err)
+	}
+	t.Cleanup(func() { m.Close() })
+	return m
+}
+
+func ingestAuthFailures(t *testing.T, m *Manager, ip, path string, n int) *IngestResult {
+	t.Helper()
+	var last *IngestResult
+	for i := 0; i < n; i++ {
+		var err error
+		last, err = m.IngestEvent(&IngestEvent{
+			SourceIP:      ip,
+			RequestPath:   path,
+			RequestMethod: "GET",
+			StatusCode:    401,
+			UserAgent:     "Mozilla/5.0",
+		}, time.Hour)
+		if err != nil {
+			t.Fatalf("IngestEvent: %v", err)
+		}
+	}
+	return last
+}
+
+func TestIngestEventAutoBlocksOnRepeatedAuthFailures(t *testing.T) {
+	m := newTestManager(t)
+
+	result := ingestAuthFailures(t, m, "203.0.113.10", "/api/v1/stats", 5)
+
+	if !result.AutoBlocked {
+		t.Fatal("expected IP to be auto-blocked after repeated auth failures")
+	}
+	blocked, err := m.IsIPBlocked("203.0.113.10")
+	if err != nil {
+		t.Fatalf("IsIPBlocked: %v", err)
+	}
+	if !blocked {
+		t.Fatal("expected IP to be in blocked list")
+	}
+}
+
+func TestIngestEventSkipsWhitelistedIP(t *testing.T) {
+	m := newTestManager(t)
+
+	if _, err := m.AddWhitelistEntry("203.0.113.7", "ip", "test"); err != nil {
+		t.Fatalf("AddWhitelistEntry: %v", err)
+	}
+
+	result := ingestAuthFailures(t, m, "203.0.113.7", "/api/v1/stats", 20)
+
+	if result.Event != nil {
+		t.Fatal("expected no event for whitelisted IP")
+	}
+	if result.AutoBlocked {
+		t.Fatal("expected whitelisted IP to never be auto-blocked")
+	}
+	blocked, err := m.IsIPBlocked("203.0.113.7")
+	if err != nil {
+		t.Fatalf("IsIPBlocked: %v", err)
+	}
+	if blocked {
+		t.Fatal("whitelisted IP must not be blocked")
+	}
+}
+
+func TestIngestEventSkipsIPInWhitelistedCIDR(t *testing.T) {
+	m := newTestManager(t)
+
+	if _, err := m.AddWhitelistEntry("198.51.100.0/24", "cidr", "test range"); err != nil {
+		t.Fatalf("AddWhitelistEntry: %v", err)
+	}
+
+	result := ingestAuthFailures(t, m, "198.51.100.20", "/api/v1/stats", 20)
+
+	if result.Event != nil || result.AutoBlocked {
+		t.Fatal("expected IP inside whitelisted CIDR to be skipped")
+	}
+}
+
+func TestIngestEventSkipsSeededPrivateNetworks(t *testing.T) {
+	m := newTestManager(t)
+
+	for _, ip := range []string{"127.0.0.1", "10.1.2.3", "172.18.0.5", "192.168.1.50"} {
+		result := ingestAuthFailures(t, m, ip, "/api/v1/stats", 20)
+		if result.Event != nil || result.AutoBlocked {
+			t.Fatalf("expected default-whitelisted IP %s to be skipped", ip)
+		}
+	}
+}
+
+func TestIngestEventSkipsWhitelistedPathPrefix(t *testing.T) {
+	m := newTestManager(t)
+
+	result, err := m.IngestEvent(&IngestEvent{
+		SourceIP:      "203.0.113.99",
+		RequestPath:   "/api/health",
+		RequestMethod: "GET",
+		StatusCode:    500,
+		UserAgent:     "Mozilla/5.0",
+	}, time.Hour)
+	if err != nil {
+		t.Fatalf("IngestEvent: %v", err)
+	}
+
+	if result.Event != nil || result.AutoBlocked {
+		t.Fatal("expected request to whitelisted path to be skipped")
+	}
+}
+
+func TestWhitelistCacheInvalidatedOnMutation(t *testing.T) {
+	m := newTestManager(t)
+
+	got, err := m.IsRequestWhitelisted("203.0.113.40", "/api/v1/stats")
+	if err != nil {
+		t.Fatalf("IsRequestWhitelisted: %v", err)
+	}
+	if got {
+		t.Fatal("IP unexpectedly whitelisted before adding entry")
+	}
+
+	id, err := m.AddWhitelistEntry("203.0.113.40", "ip", "test")
+	if err != nil {
+		t.Fatalf("AddWhitelistEntry: %v", err)
+	}
+	got, err = m.IsRequestWhitelisted("203.0.113.40", "/api/v1/stats")
+	if err != nil {
+		t.Fatalf("IsRequestWhitelisted: %v", err)
+	}
+	if !got {
+		t.Fatal("expected entry added after cache build to be honored")
+	}
+
+	if err := m.RemoveWhitelistEntry(id); err != nil {
+		t.Fatalf("RemoveWhitelistEntry: %v", err)
+	}
+	got, err = m.IsRequestWhitelisted("203.0.113.40", "/api/v1/stats")
+	if err != nil {
+		t.Fatalf("IsRequestWhitelisted: %v", err)
+	}
+	if got {
+		t.Fatal("expected removed entry to stop matching")
+	}
+}
+
+func TestIsRequestWhitelisted(t *testing.T) {
+	m := newTestManager(t)
+
+	if _, err := m.AddWhitelistEntry("2001:db8::/32", "cidr", "test v6"); err != nil {
+		t.Fatalf("AddWhitelistEntry: %v", err)
+	}
+
+	cases := []struct {
+		ip   string
+		path string
+		want bool
+	}{
+		{"127.0.0.1", "/anything", true},
+		{"10.255.0.1", "/anything", true},
+		{"2001:db8::1", "/anything", true},
+		{"203.0.113.5", "/api/_internal/blocked-ips", true},
+		{"203.0.113.5", "/wp-login.php", false},
+		{"not-an-ip", "/wp-login.php", false},
+	}
+
+	for _, tc := range cases {
+		got, err := m.IsRequestWhitelisted(tc.ip, tc.path)
+		if err != nil {
+			t.Fatalf("IsRequestWhitelisted(%s, %s): %v", tc.ip, tc.path, err)
+		}
+		if got != tc.want {
+			t.Errorf("IsRequestWhitelisted(%s, %s) = %v, want %v", tc.ip, tc.path, got, tc.want)
+		}
+	}
+}