Merge pull request #83 from Bestigor89/fix/xff-country-asn-blocking

fix: use X-Forwarded-For for country/ASN blocking behind proxies

Author: fabriziosalmi

handler.go

@@ -328,7 +328,8 @@ func (m *Middleware) handlePhase(w http.ResponseWriter, r *http.Request, phase i
 		// Whitelisting
 		if m.CountryWhitelist.Enabled {
 			m.logger.Debug("Starting country whitelisting phase")
-			allowed, err := m.isCountryInList(r.RemoteAddr, m.CountryWhitelist.CountryList, m.CountryWhitelist.geoIP)
+			clientIP := getClientIP(r)
+			allowed, err := m.isCountryInList(clientIP, m.CountryWhitelist.CountryList, m.CountryWhitelist.geoIP)
 			if err != nil {
 				m.logRequest(zapcore.ErrorLevel, "Failed to check country whitelist",
 					r,
@@ -360,7 +361,8 @@ func (m *Middleware) handlePhase(w http.ResponseWriter, r *http.Request, phase i
 		// ASN Blocking
 		if m.BlockASNs.Enabled {
 			m.logger.Debug("Starting ASN blocking phase")
-			blocked, err := m.geoIPHandler.IsASNInList(r.RemoteAddr, m.BlockASNs.BlockedASNs, m.BlockASNs.geoIP)
+			clientIP := getClientIP(r)
+			blocked, err := m.geoIPHandler.IsASNInList(clientIP, m.BlockASNs.BlockedASNs, m.BlockASNs.geoIP)
 			if err != nil {
 				m.logRequest(zapcore.ErrorLevel, "Failed to check ASN blocking",
 					r,
@@ -377,7 +379,7 @@ func (m *Middleware) handlePhase(w http.ResponseWriter, r *http.Request, phase i
 					return
 				}
 			} else if blocked {
-				asnInfo := m.geoIPHandler.GetASN(r.RemoteAddr, m.BlockASNs.geoIP)
+				asnInfo := m.geoIPHandler.GetASN(clientIP, m.BlockASNs.geoIP)
 				m.blockRequest(w, r, state, http.StatusForbidden, "asn_block", "asn_block_rule",
 					zap.String("message", "Request blocked by ASN"),
 					zap.String("asn", asnInfo),
@@ -394,7 +396,8 @@ func (m *Middleware) handlePhase(w http.ResponseWriter, r *http.Request, phase i
 		// Blacklisting
 		if m.CountryBlacklist.Enabled {
 			m.logger.Debug("Starting country blacklisting phase")
-			blocked, err := m.isCountryInList(r.RemoteAddr, m.CountryBlacklist.CountryList, m.CountryBlacklist.geoIP)
+			clientIP := getClientIP(r)
+			blocked, err := m.isCountryInList(clientIP, m.CountryBlacklist.CountryList, m.CountryBlacklist.geoIP)
 			if err != nil {
 				m.logRequest(zapcore.ErrorLevel, "Failed to check country blacklisting",
 					r,

helpers.go

@@ -2,6 +2,7 @@ package caddywaf
 
 import (
 	"net"
+	"net/http"
 	"os"
 	"strings"
 )
@@ -45,3 +46,14 @@ func extractIP(remoteAddr string) string {
 	}
 	return host
 }
+
+// getClientIP returns the real client IP, checking X-Forwarded-For header first.
+func getClientIP(r *http.Request) string {
+	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
+		ips := strings.Split(xff, ",")
+		if len(ips) > 0 {
+			return strings.TrimSpace(ips[0])
+		}
+	}
+	return r.RemoteAddr
+}

helpers_test.go

@@ -1,6 +1,7 @@
 package caddywaf
 
 import (
+	"net/http"
 	"os"
 	"testing"
 )
@@ -48,3 +49,57 @@ func TestFileExists(t *testing.T) {
 		})
 	}
 }
+
+func TestGetClientIP(t *testing.T) {
+	tests := []struct {
+		name       string
+		remoteAddr string
+		xff        string
+		want       string
+	}{
+		{
+			name:       "no X-Forwarded-For, use RemoteAddr",
+			remoteAddr: "192.168.1.1:12345",
+			xff:        "",
+			want:       "192.168.1.1:12345",
+		},
+		{
+			name:       "single IP in X-Forwarded-For",
+			remoteAddr: "10.0.0.1:12345",
+			xff:        "203.0.113.50",
+			want:       "203.0.113.50",
+		},
+		{
+			name:       "multiple IPs in X-Forwarded-For",
+			remoteAddr: "10.0.0.1:12345",
+			xff:        "203.0.113.50, 70.41.3.18, 150.172.238.178",
+			want:       "203.0.113.50",
+		},
+		{
+			name:       "X-Forwarded-For with spaces",
+			remoteAddr: "10.0.0.1:12345",
+			xff:        "  203.0.113.50  ,  70.41.3.18  ",
+			want:       "203.0.113.50",
+		},
+		{
+			name:       "IPv6 in X-Forwarded-For",
+			remoteAddr: "10.0.0.1:12345",
+			xff:        "2001:db8::1",
+			want:       "2001:db8::1",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req, _ := http.NewRequest("GET", "/", nil)
+			req.RemoteAddr = tt.remoteAddr
+			if tt.xff != "" {
+				req.Header.Set("X-Forwarded-For", tt.xff)
+			}
+
+			if got := getClientIP(req); got != tt.want {
+				t.Errorf("getClientIP() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}