From 0f0601c567b19e888b42e46bc88fc8614c8f4695 Mon Sep 17 00:00:00 2001
From: Egor Aristov <ef@efprojects.com>
Date: Mon, 10 Feb 2025 18:38:42 +0300
Subject: [PATCH] better ip detection

---
 cmd/webserver/webserver.go | 35 ++++++++++++++++++++++++++---------
 internal/config/config.go  |  1 +
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/cmd/webserver/webserver.go b/cmd/webserver/webserver.go
index 75f3a8d..616245d 100644
--- a/cmd/webserver/webserver.go
+++ b/cmd/webserver/webserver.go
@@ -52,15 +52,7 @@ func main() {
 	e.Use(middleware.Logger())
 	e.Use(middleware.Recover())
 
-	var trustOptions []echo.TrustOption
-	for _, ipRange := range slices.Concat(IpRanges, cfg.TrustedIpRanges) {
-		_, network, err := net.ParseCIDR(ipRange)
-		if err != nil {
-			log.Panicf("Invalid ip range: %s", ipRange)
-		}
-		trustOptions = append(trustOptions, echo.TrustIPRange(network))
-	}
-	e.IPExtractor = echo.ExtractIPFromXFFHeader(trustOptions...)
+	setIPExtractor(e, cfg)
 
 	e.StaticFS("/", echo.MustSubFS(wizard_vue.EmbedFS, wizard_vue.FSPrefix))
 
@@ -85,3 +77,28 @@ func main() {
 		e.Logger.Errorf("failed to shutdown server: %v", err)
 	}
 }
+
+func setIPExtractor(e *echo.Echo, cfg config.Config) {
+	if len(cfg.RealIpHeader) > 0 {
+		// Real ip header
+		e.IPExtractor = func(req *http.Request) string {
+			if len(req.Header.Get(cfg.RealIpHeader)) > 0 {
+				return req.Header.Get(cfg.RealIpHeader)
+			}
+			// fallback
+			ra, _, _ := net.SplitHostPort(req.RemoteAddr)
+			return ra
+		}
+	} else {
+		// X-Forwarded-For with trusted ip ranges
+		var trustOptions []echo.TrustOption
+		for _, ipRange := range slices.Concat(IpRanges, cfg.TrustedIpRanges) {
+			_, network, err := net.ParseCIDR(ipRange)
+			if err != nil {
+				log.Panicf("Invalid ip range: %s", ipRange)
+			}
+			trustOptions = append(trustOptions, echo.TrustIPRange(network))
+		}
+		e.IPExtractor = echo.ExtractIPFromXFFHeader(trustOptions...)
+	}
+}
diff --git a/internal/config/config.go b/internal/config/config.go
index ca8f803..34c0a07 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -23,6 +23,7 @@ type Config struct {
 	RateLimitBurst int     `env:"RATE_LIMIT_BURST" env-default:"10" validate:"number,gte=0"`
 	// IP ranges of reverse proxies for correct real ip detection (cidr format, sep. by comma)
 	TrustedIpRanges []string `env:"TRUSTED_IP_RANGES" env-default:"" validate:"omitempty,dive,cidr"`
+	RealIpHeader    string   `env:"REAL_IP_HEADER" env-default:"" validate:"omitempty"`
 }
 
 func Read() (Config, error) {