From fa86be031078f515055e263370d7704325948f4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ask=20Bj=C3=B8rn=20Hansen?= <ask@develooper.com>
Date: Sun, 12 Nov 2023 15:53:34 -0800
Subject: [PATCH] Use Fastly IPs

---
 go.mod           |  2 +-
 go.sum           |  2 ++
 server/server.go | 26 +++++++++++++++++++++++++-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index e0f4534..f9b70cd 100644
--- a/go.mod
+++ b/go.mod
@@ -11,7 +11,7 @@ require (
 	github.com/samber/slog-echo v1.6.0
 	github.com/spf13/cobra v1.7.0
 	github.com/stretchr/testify v1.8.4
-	go.ntppool.org/common v0.2.4
+	go.ntppool.org/common v0.2.5-0.20231112235121-2bff6d8ef307
 	go.opentelemetry.io/contrib/instrumentation/github.com/labstack/echo/otelecho v0.45.0
 	go.opentelemetry.io/otel v1.19.0
 	go.opentelemetry.io/otel/trace v1.19.0
diff --git a/go.sum b/go.sum
index a2cc102..9408408 100644
--- a/go.sum
+++ b/go.sum
@@ -131,6 +131,8 @@ github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9dec
 go.mongodb.org/mongo-driver v1.11.4/go.mod h1:PTSz5yu21bkT/wXpkS7WR5f0ddqw5quethTUn9WM+2g=
 go.ntppool.org/common v0.2.4 h1:OqKR1OHYayv6AsERAR8RYKdOEigJqXBpqkGWlaGF3+Q=
 go.ntppool.org/common v0.2.4/go.mod h1:kYshXIaeI13tj6CSW56KHkcwp0lJbM8bFCe3tm3BZEQ=
+go.ntppool.org/common v0.2.5-0.20231112235121-2bff6d8ef307 h1:bJPpvb3aP3sIdO/ptxH9Jqhksk0+c5qQBSa/xHLhscc=
+go.ntppool.org/common v0.2.5-0.20231112235121-2bff6d8ef307/go.mod h1:kYshXIaeI13tj6CSW56KHkcwp0lJbM8bFCe3tm3BZEQ=
 go.opentelemetry.io/contrib/instrumentation/github.com/labstack/echo/otelecho v0.45.0 h1:JJCIHAxGCB5HM3NxeIwFjHc087Xwk96TG9kaZU6TAec=
 go.opentelemetry.io/contrib/instrumentation/github.com/labstack/echo/otelecho v0.45.0/go.mod h1:Px9kH7SJ+NhsgWRtD/eMcs15Tyt4uL3rM7X54qv6pfA=
 go.opentelemetry.io/contrib/propagators/b3 v1.20.0 h1:Yty9Vs4F3D6/liF1o6FNt0PvN85h/BJJ6DQKJ3nrcM0=
diff --git a/server/server.go b/server/server.go
index dfa4b16..83d880c 100644
--- a/server/server.go
+++ b/server/server.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"os"
 
 	"golang.org/x/sync/errgroup"
 
@@ -19,6 +20,7 @@ import (
 	"go.ntppool.org/common/logger"
 	"go.ntppool.org/common/metricsserver"
 	"go.ntppool.org/common/tracing"
+	"go.ntppool.org/common/xff/fastlyxff"
 
 	chdb "go.ntppool.org/data-api/chdb"
 	"go.ntppool.org/data-api/ntpdb"
@@ -81,6 +83,29 @@ func (srv *Server) Run() error {
 	})
 
 	e := echo.New()
+
+	trustOptions := []echo.TrustOption{
+		echo.TrustLoopback(true),
+		echo.TrustLinkLocal(false),
+		echo.TrustPrivateNet(true),
+	}
+
+	if fileName := os.Getenv("FASTLY_IPS"); len(fileName) > 0 {
+		xff, err := fastlyxff.New(fileName)
+		if err != nil {
+			return err
+		}
+		cdnTrustRanges, err := xff.EchoTrustOption()
+		if err != nil {
+			return err
+		}
+		trustOptions = append(trustOptions, cdnTrustRanges...)
+	} else {
+		log.Warn("Fastly IPs not configured (FASTLY_IPS)")
+	}
+
+	e.IPExtractor = echo.ExtractIPFromXFFHeader(trustOptions...)
+
 	e.Use(otelecho.Middleware("data-api"))
 	e.Use(slogecho.New(log))
 
@@ -106,7 +131,6 @@ func (srv *Server) Run() error {
 	})
 
 	e.GET("/api/usercc", srv.userCountryData)
-
 	e.GET("/api/server/dns/answers/:server", srv.dnsAnswers)
 
 	g.Go(func() error {