refactor: enhance IP blacklist handling and add panic recovery in middleware

Author: fabriziosalmi

blacklist.go

@@ -65,12 +65,25 @@ func (bl *BlacklistLoader) LoadDNSBlacklistFromFile(path string, dnsBlacklist ma
 }
 
 func (m *Middleware) isIPBlacklisted(ip string) bool {
-	if m.ipBlacklist.Contains(netip.MustParseAddr(ip)) {
-		m.muIPBlacklistMetrics.Lock()                            // Acquire lock before accessing shared counter
-		m.IPBlacklistBlockCount++                                // Increment the counter
-		m.muIPBlacklistMetrics.Unlock()                          // Release lock after accessing counter
-		m.logger.Debug("IP blacklist hit", zap.String("ip", ip)) // Keep existing debug log
-		return true                                              // Indicate that the IP is blacklisted
+	// Extract IP address without port
+	cleanIP := extractIP(ip, m.logger)
+
+	addr, err := netip.ParseAddr(cleanIP)
+	if err != nil {
+		m.logger.Warn("Failed to parse IP address for blacklist check",
+			zap.String("ip", ip),
+			zap.String("clean_ip", cleanIP),
+			zap.Error(err),
+		)
+		return false
+	}
+
+	if m.ipBlacklist.Contains(addr) {
+		m.muIPBlacklistMetrics.Lock()                                 // Acquire lock before accessing shared counter
+		m.IPBlacklistBlockCount++                                     // Increment the counter
+		m.muIPBlacklistMetrics.Unlock()                               // Release lock after accessing counter
+		m.logger.Debug("IP blacklist hit", zap.String("ip", cleanIP)) // Keep existing debug log
+		return true                                                   // Indicate that the IP is blacklisted
 	}
 	return false // Indicate that the IP is NOT blacklisted
 }

caddywaf.go

@@ -82,7 +82,8 @@ func parseCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error)
 
 func (m *Middleware) Provision(ctx caddy.Context) error {
 	m.logger = ctx.Logger(m)
-	m.ruleCache = NewRuleCache() // Initialize RuleCache
+	m.ruleCache = NewRuleCache()   // Initialize RuleCache
+	m.Rules = make(map[int][]Rule) // Initialize Rules map to prevent nil pointer panic
 
 	// Set default log severity if not provided
 	if m.LogSeverity == "" {
@@ -466,7 +467,16 @@ func (m *Middleware) loadIPBlacklist(path string, blacklistMap iptrie.Trie) erro
 
 	// Convert the map to CIDRTrie
 	for ip := range blacklist {
-		blacklistMap.Insert(netip.MustParsePrefix(ip), nil)
+		// Add /32 suffix if the IP doesn't have CIDR notation
+		if !strings.Contains(ip, "/") {
+			ip = ip + "/32"
+		}
+		prefix, err := netip.ParsePrefix(ip)
+		if err != nil {
+			m.logger.Warn("Skipping invalid IP in blacklist", zap.String("ip", ip), zap.Error(err))
+			continue
+		}
+		blacklistMap.Insert(prefix, nil)
 	}
 	return nil
 }

handler.go

@@ -22,6 +22,20 @@ type (
 func (m *Middleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
 	logID := uuid.New().String()
 
+	// Add panic recovery to catch and log panics
+	defer func() {
+		if rec := recover(); rec != nil {
+			m.logger.Error("PANIC in ServeHTTP",
+				zap.String("log_id", logID),
+				zap.Any("panic", rec),
+				zap.Stack("stack"),
+			)
+			// Return 500 error to client
+			w.WriteHeader(http.StatusInternalServerError)
+			w.Write([]byte("Internal Server Error"))
+		}
+	}()
+
 	m.logRequestStart(r, logID)
 
 	// Propagate log ID within the request context for logging
@@ -152,7 +166,14 @@ func (m *Middleware) handleResponseBodyPhase(recorder *responseRecorder, r *http
 	}
 	m.logger.Debug("Response body captured for Phase 4 analysis", zap.String("log_id", logID))
 
-	for _, rule := range m.Rules[4] {
+	// Check if rules exist for Phase 4 before iterating
+	rules, ok := m.Rules[4]
+	if !ok || len(rules) == 0 {
+		m.logger.Debug("No rules found for Phase 4")
+		return
+	}
+
+	for _, rule := range rules {
 		if rule.regex.MatchString(body) {
 			if m.processRuleMatch(recorder, r, &rule, body, state) {
 				return