feat: to tackle proxies, reads xforwardedfor first then remoteaddr

Author: Christ Borg

.traefik.yml

@@ -8,3 +8,4 @@ import: github.com/traefik-plugins/traefikgeoip2
 
 testData:
   dbPath: 'GeoLite2-Country.mmdb'
+  preferXForwardedForHeader: false

middleware.go

@@ -21,7 +21,8 @@ func ResetLookup() {
 
 // Config the plugin configuration.
 type Config struct {
-	DBPath string `json:"dbPath,omitempty"`
+	DBPath                    string `json:"dbPath,omitempty"`
+	PreferXForwardedForHeader bool
 }
 
 // CreateConfig creates the default plugin configuration.
@@ -33,17 +34,19 @@ func CreateConfig() *Config {
 
 // TraefikGeoIP2 a traefik geoip2 plugin.
 type TraefikGeoIP2 struct {
-	next http.Handler
-	name string
+	next                      http.Handler
+	name                      string
+	preferXForwardedForHeader bool
 }
 
 // New created a new TraefikGeoIP2 plugin.
 func New(_ context.Context, next http.Handler, cfg *Config, name string) (http.Handler, error) {
 	if _, err := os.Stat(cfg.DBPath); err != nil {
 		log.Printf("[geoip2] DB not found: db=%s, name=%s, err=%v", cfg.DBPath, name, err)
 		return &TraefikGeoIP2{
-			next: next,
-			name: name,
+			next:                      next,
+			name:                      name,
+			preferXForwardedForHeader: cfg.PreferXForwardedForHeader,
 		}, nil
 	}
 
@@ -68,8 +71,9 @@ func New(_ context.Context, next http.Handler, cfg *Config, name string) (http.H
 	}
 
 	return &TraefikGeoIP2{
-		next: next,
-		name: name,
+		next:                      next,
+		name:                      name,
+		preferXForwardedForHeader: cfg.PreferXForwardedForHeader,
 	}, nil
 }
 
@@ -83,12 +87,7 @@ func (mw *TraefikGeoIP2) ServeHTTP(reqWr http.ResponseWriter, req *http.Request)
 		return
 	}
 
-	ipStr := req.RemoteAddr
-	tmp, _, err := net.SplitHostPort(ipStr)
-	if err == nil {
-		ipStr = tmp
-	}
-
+	ipStr := getClientIP(req, mw.preferXForwardedForHeader)
 	res, err := lookup(net.ParseIP(ipStr))
 	if err != nil {
 		log.Printf("[geoip2] Unable to find: ip=%s, err=%v", ipStr, err)
@@ -106,3 +105,22 @@ func (mw *TraefikGeoIP2) ServeHTTP(reqWr http.ResponseWriter, req *http.Request)
 
 	mw.next.ServeHTTP(reqWr, req)
 }
+
+func getClientIP(req *http.Request, preferXForwardedForHeader bool) string {
+	if preferXForwardedForHeader {
+		// Check X-Forwarded-For header first
+		forwardedFor := req.Header.Get("X-Forwarded-For")
+		if forwardedFor != "" {
+			ips := strings.Split(forwardedFor, ",")
+			return strings.TrimSpace(ips[0])
+		}
+	}
+
+	// If X-Forwarded-For is not present or retrieval is not enabled, fallback to RemoteAddr
+	remoteAddr := req.RemoteAddr
+	tmp, _, err := net.SplitHostPort(remoteAddr)
+	if err == nil {
+		remoteAddr = tmp
+	}
+	return remoteAddr
+}

middleware_test.go

@@ -11,8 +11,9 @@ import (
 )
 
 const (
-	ValidIP       = "188.193.88.199"
-	ValidIPNoCity = "20.1.184.61"
+	ValidIP          = "188.193.88.199"
+	ValidAlternateIP = "188.193.88.200"
+	ValidIPNoCity    = "20.1.184.61"
 )
 
 func TestGeoIPConfig(t *testing.T) {
@@ -123,6 +124,43 @@ func TestGeoIPFromRemoteAddr(t *testing.T) {
 	assertHeader(t, req, mw.IPAddressHeader, "qwerty")
 }
 
+func TestGeoIPFromXForwardedFor(t *testing.T) {
+	mwCfg := mw.CreateConfig()
+	mwCfg.DBPath = "./GeoLite2-City.mmdb"
+	mwCfg.PreferXForwardedForHeader = true
+
+	next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {})
+	mw.ResetLookup()
+	instance, _ := mw.New(context.TODO(), next, mwCfg, "traefik-geoip2")
+
+	req := httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	req.RemoteAddr = fmt.Sprintf("%s:9999", ValidIP)
+	req.Header.Set("X-Forwarded-For", ValidAlternateIP)
+	instance.ServeHTTP(httptest.NewRecorder(), req)
+	assertHeader(t, req, mw.CountryHeader, "DE")
+	assertHeader(t, req, mw.RegionHeader, "BY")
+	assertHeader(t, req, mw.CityHeader, "Munich")
+	assertHeader(t, req, mw.IPAddressHeader, ValidAlternateIP)
+
+	req = httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	req.RemoteAddr = fmt.Sprintf("%s:9999", ValidIP)
+	req.Header.Set("X-Forwarded-For", ValidAlternateIP+",188.193.88.100")
+	instance.ServeHTTP(httptest.NewRecorder(), req)
+	assertHeader(t, req, mw.CountryHeader, "DE")
+	assertHeader(t, req, mw.RegionHeader, "BY")
+	assertHeader(t, req, mw.CityHeader, "Munich")
+	assertHeader(t, req, mw.IPAddressHeader, ValidAlternateIP)
+
+	req = httptest.NewRequest(http.MethodGet, "http://localhost", nil)
+	req.RemoteAddr = ValidIP + ":9999"
+	req.Header.Set("X-Forwarded-For", "qwerty")
+	instance.ServeHTTP(httptest.NewRecorder(), req)
+	assertHeader(t, req, mw.CountryHeader, mw.Unknown)
+	assertHeader(t, req, mw.RegionHeader, mw.Unknown)
+	assertHeader(t, req, mw.CityHeader, mw.Unknown)
+	assertHeader(t, req, mw.IPAddressHeader, "qwerty")
+}
+
 func TestGeoIPCountryDBFromRemoteAddr(t *testing.T) {
 	mwCfg := mw.CreateConfig()
 	mwCfg.DBPath = "./GeoLite2-Country.mmdb"