feature: add support for custom client IP header

.traefik.yml

@@ -8,6 +8,7 @@ testData:
   # This will cause the plugin to not attempt to load the database file.
   enabled: false
   # databaseFilePath: IP2LOCATION-LITE-DB1.IPV6.BIN
+  # ipHeader: CF-Connecting-IP
   # allowedCountries: [ "CH", "DE" ]
   # blockedCountries: [ "RU" ]
   # defaultAllow: false

README.md

@@ -47,6 +47,8 @@ http:
           enabled: true
           # Path to ip2location database file
           databaseFilePath: /plugins-local/src/github.com/nscuro/traefik-plugin-geoblock/IP2LOCATION-LITE-DB1.IPV6.BIN
+          # Optional trusted header to use for the client IP instead of X-Forwarded-For / X-Real-IP
+          ipHeader: CF-Connecting-IP
           # Whitelist of countries to allow (ISO 3166-1 alpha-2)
           allowedCountries: [ "AT", "CH", "DE" ]
           # Blocklist of countries to block (ISO 3166-1 alpha-2)
@@ -62,3 +64,5 @@ http:
           # Add CIDR to be blacklisted, even if in an allowed country or IP block
           blockedIPBlocks: ["66.249.64.5/32"]
 ```
+
+If the optional `ipHeader` parameter is configured, the plugin reads the client IP from that header and does not use `X-Forwarded-For` or `X-Real-IP` for that request. This is useful when a trusted upstream service such as a WAF (eg. Cloudflare) specifies the client IP in a custom header.

plugin.go

@@ -18,6 +18,7 @@ import (
 type Config struct {
 	Enabled              bool     // Enable this plugin?
 	DatabaseFilePath     string   // Path to ip2location database file
+	IPHeader             string   // Optional header to trust for the client IP address
 	AllowedCountries     []string // Whitelist of countries to allow (ISO 3166-1 alpha-2)
 	BlockedCountries     []string // Blocklist of countries to be blocked (ISO 3166-1 alpha-2)
 	DefaultAllow         bool     // If source matches neither blocklist nor whitelist, should it be allowed through?
@@ -37,6 +38,7 @@ type Plugin struct {
 	name                 string
 	db                   *ip2location.DB
 	enabled              bool
+	ipHeader             string
 	allowedCountries     []string
 	blockedCountries     []string
 	defaultAllow         bool
@@ -94,6 +96,7 @@ func New(_ context.Context, next http.Handler, cfg *Config, name string) (http.H
 		name:                 name,
 		db:                   db,
 		enabled:              cfg.Enabled,
+		ipHeader:             cfg.IPHeader,
 		allowedCountries:     cfg.AllowedCountries,
 		blockedCountries:     cfg.BlockedCountries,
 		defaultAllow:         cfg.DefaultAllow,
@@ -128,21 +131,28 @@ func (p Plugin) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	p.next.ServeHTTP(rw, req)
 }
 
-// GetRemoteIPs collects the remote IPs from the X-Forwarded-For and X-Real-IP headers.
+// GetRemoteIPs collects the remote IPs from a configured client IP header, or
+// falls back to X-Forwarded-For and X-Real-IP when no custom header is configured.
 func (p Plugin) GetRemoteIPs(req *http.Request) []string {
+	if p.ipHeader != "" {
+		return uniqueIPs(req.Header.Values(p.ipHeader))
+	}
+
+	return uniqueIPs([]string{
+		req.Header.Get("x-forwarded-for"),
+		req.Header.Get("x-real-ip"),
+	})
+}
+
+func uniqueIPs(headerValues []string) []string {
 	uniqIPs := make(map[string]struct{})
 
-	if xff := req.Header.Get("x-forwarded-for"); xff != "" {
-		for _, ip := range strings.Split(xff, ",") {
-			ip = strings.TrimSpace(ip)
-			if ip == "" {
-				continue
-			}
-			uniqIPs[ip] = struct{}{}
+	for _, headerValue := range headerValues {
+		if headerValue == "" {
+			continue
 		}
-	}
-	if xri := req.Header.Get("x-real-ip"); xri != "" {
-		for _, ip := range strings.Split(xri, ",") {
+
+		for _, ip := range strings.Split(headerValue, ",") {
 			ip = strings.TrimSpace(ip)
 			if ip == "" {
 				continue

plugin_test.go

@@ -203,6 +203,66 @@ func TestPlugin_ServeHTTP(t *testing.T) {
 
 		testRequest(t, "Default allow false", cfg, "8.8.4.4", http.StatusForbidden)
 	})
+
+	t.Run("CustomIPHeader", func(t *testing.T) {
+		cfg := &Config{
+			Enabled:              true,
+			DatabaseFilePath:     dbFilePath,
+			IPHeader:             "CF-Connecting-IP",
+			AllowedCountries:     []string{"US"},
+			DisallowedStatusCode: http.StatusForbidden,
+		}
+
+		plugin, err := New(context.TODO(), &noopHandler{}, cfg, pluginName)
+		if err != nil {
+			t.Errorf("expected no error, but got: %v", err)
+		}
+
+		req := httptest.NewRequest(http.MethodGet, "/foobar", nil)
+		req.Header.Set("CF-Connecting-IP", "1.1.1.1")
+		req.Header.Set("X-Real-IP", "185.5.82.105")
+
+		rr := httptest.NewRecorder()
+		plugin.ServeHTTP(rr, req)
+
+		if rr.Code != http.StatusTeapot {
+			t.Errorf("expected status code %d, but got: %d", http.StatusTeapot, rr.Code)
+		}
+	})
+}
+
+func TestPlugin_GetRemoteIPs(t *testing.T) {
+	t.Run("UsesConfiguredHeader", func(t *testing.T) {
+		plugin := Plugin{ipHeader: "CF-Connecting-IP"}
+
+		req := httptest.NewRequest(http.MethodGet, "/foobar", nil)
+		req.Header.Add("CF-Connecting-IP", "1.1.1.1")
+		req.Header.Add("CF-Connecting-IP", "8.8.8.8, 1.1.1.1")
+		req.Header.Set("X-Forwarded-For", "185.5.82.105")
+		req.Header.Set("X-Real-IP", "185.5.82.105")
+
+		ips := plugin.GetRemoteIPs(req)
+
+		if len(ips) != 2 {
+			t.Fatalf("expected 2 IPs, but got %d", len(ips))
+		}
+
+		expected := map[string]struct{}{
+			"1.1.1.1": {},
+			"8.8.8.8": {},
+		}
+
+		for _, ip := range ips {
+			if _, ok := expected[ip]; !ok {
+				t.Fatalf("unexpected IP %q returned", ip)
+			}
+			delete(expected, ip)
+		}
+
+		if len(expected) != 0 {
+			t.Fatalf("missing IPs: %v", expected)
+		}
+	})
 }
 
 func testRequest(t *testing.T, testName string, cfg *Config, ip string, expectedStatus int) {