From 8eec1c5656e6b2e0d6a6efd4b21ca8bc5a9ab582 Mon Sep 17 00:00:00 2001
From: MathieuHa <mathieu@hanotaux.fr>
Date: Sat, 15 Oct 2022 18:49:41 +0200
Subject: [PATCH] :sparkles: Add Custom header usage to find the user IP

---
 README.md  |  5 +++++
 bouncer.go |  7 ++++++-
 ip/ip.go   | 13 +++----------
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index cc37cfe..7f282e8 100644
--- a/README.md
+++ b/README.md
@@ -74,6 +74,10 @@ At each start of synchronisation, the middleware will wait a random number of se
   - []string
   - default: []
   - List of IPs of trusted Proxies that are in front of traefik (ex: Cloudflare)
+- ForwardedHeadersCustomName
+  - string
+  - default: X-Forwarded-For
+  - Name of the header where the real IP of the client should be retrieved
 
 ### Configuration
 
@@ -123,6 +127,7 @@ http:
           forwardedHeadersTrustedIPs: 
             - 10.0.10.23/32
             - 10.0.20.0/24
+          forwardedHeadersCustomName: X-Custom-Header
 ```
 These are the default values of the plugin except for LapiKey.
 
diff --git a/bouncer.go b/bouncer.go
index 8d1ae7a..3b0c181 100644
--- a/bouncer.go
+++ b/bouncer.go
@@ -28,6 +28,7 @@ const (
 	crowdsecLapiStreamRoute = "v1/decisions/stream"
 	cacheBannedValue        = "t"
 	cacheNoBannedValue      = "f"
+	xForwardedFor           = "X-Forwarded-For"
 )
 
 // Config the plugin configuration.
@@ -40,6 +41,7 @@ type Config struct {
 	UpdateIntervalSeconds      int64    `json:"updateIntervalSeconds,omitempty"`
 	DefaultDecisionSeconds     int64    `json:"defaultDecisionSeconds,omitempty"`
 	ForwardedHeadersTrustedIPs []string `json:"forwardedheaderstrustedips,omitempty"`
+	ForwardedHeadersCustomName string   `json:"forwardedheaderscustomheader,omitempty"`
 }
 
 // CreateConfig creates the default plugin configuration.
@@ -53,6 +55,7 @@ func CreateConfig() *Config {
 		UpdateIntervalSeconds:      60,
 		DefaultDecisionSeconds:     60,
 		ForwardedHeadersTrustedIPs: []string{},
+		ForwardedHeadersCustomName: xForwardedFor,
 	}
 }
 
@@ -70,6 +73,7 @@ type Bouncer struct {
 	crowdsecMode           string
 	updateInterval         int64
 	defaultDecisionTimeout int64
+	customHeader           string
 	poolStrategy           *ip.PoolStrategy
 	client                 *http.Client
 	cache                  *ttl_map.Heap
@@ -96,6 +100,7 @@ func New(ctx context.Context, next http.Handler, config *Config, name string) (h
 		crowdsecHost:           config.CrowdsecLapiHost,
 		crowdsecKey:            config.CrowdsecLapiKey,
 		updateInterval:         config.UpdateIntervalSeconds,
+		customHeader:           config.ForwardedHeadersCustomName,
 		defaultDecisionTimeout: config.DefaultDecisionSeconds,
 		poolStrategy: &ip.PoolStrategy{
 			Checker: checker,
@@ -198,7 +203,7 @@ func contains(source []string, target string) bool {
 
 // It returns the first IP that is not in the pool, or the empty string otherwise.
 func getRemoteIP(bouncer *Bouncer, req *http.Request) (string, error) {
-	remoteIP := bouncer.poolStrategy.GetIP(req)
+	remoteIP := bouncer.poolStrategy.GetIP(req, bouncer.customHeader)
 	if len(remoteIP) != 0 {
 		return remoteIP, nil
 	}
diff --git a/ip/ip.go b/ip/ip.go
index 45617a2..226416a 100644
--- a/ip/ip.go
+++ b/ip/ip.go
@@ -80,16 +80,8 @@ func parseIP(addr string) (net.IP, error) {
 	return userIP, nil
 }
 
-
-
-
-
 // STRATEGY
 
-const (
-	xForwardedFor = "X-Forwarded-For"
-)
-
 // PoolStrategy is a strategy based on an IP Checker.
 // It allows to check whether addresses are in a given pool of IPs.
 type PoolStrategy struct {
@@ -99,12 +91,13 @@ type PoolStrategy struct {
 // GetIP checks the list of Forwarded IPs (most recent first) against the
 // Checker pool of IPs. It returns the first IP that is not in the pool, or the
 // empty string otherwise.
-func (s *PoolStrategy) GetIP(req *http.Request) string {
+func (s *PoolStrategy) GetIP(req *http.Request, customHeader string) string {
 	if s.Checker == nil {
 		return ""
 	}
 
-	xff := req.Header.Get(xForwardedFor)
+	xff := req.Header.Get(customHeader)
+
 	xffs := strings.Split(xff, ",")
 
 	for i := len(xffs) - 1; i >= 0; i-- {