feat: more elaborate XFF compute (#350)

* feat: more elaborate XFF compute

#328 followup

now featuring configuration and
defaults that shouldn't break most
setups.

fixes #344

* refactor: obvious condition eval order optimization

* feat: add StripLLU implementation

* chore: I'm sorry it's 7 AM

* test: add test environment for unix socket serving

Signed-off-by: Xe Iaso <me@xeiaso.net>

* test(unix-socket-xff): comment out the shell script more

Signed-off-by: Xe Iaso <me@xeiaso.net>

* fix(internal): fix logic bug in XFF computation, add tests

Signed-off-by: Xe Iaso <me@xeiaso.net>

* fix(internal): prevent panic in local testing

Signed-off-by: Xe Iaso <me@xeiaso.net>

* fix(internal): shuffle around return values to flow better

Signed-off-by: Xe Iaso <me@xeiaso.net>

---------

Signed-off-by: Xe Iaso <me@xeiaso.net>
Co-authored-by: Xe Iaso <me@xeiaso.net>

2 files changed, 268 insertions, 22 deletions

diff --git a/internal/headers.go b/internal/headers.go
index 4516b40..639d2b1 100644
--- a/internal/headers.go
+++ b/internal/headers.go
@@ -1,15 +1,29 @@
 package internal
 
 import (
+	"errors"
+	"fmt"
 	"log/slog"
 	"net"
 	"net/http"
+	"net/netip"
 	"strings"
 
 	"github.com/TecharoHQ/anubis"
 	"github.com/sebest/xff"
 )
 
+// TODO: move into config
+type XFFComputePreferences struct {
+	StripPrivate  bool
+	StripLoopback bool
+	StripCGNAT    bool
+	StripLLU      bool
+	Flatten       bool
+}
+
+var CGNat = netip.MustParsePrefix("100.64.0.0/10")
+
 // UnchangingCache sets the Cache-Control header to cache a response for 1 year if
 // and only if the application is compiled in "release" mode by Docker.
 func UnchangingCache(next http.Handler) http.Handler {
@@ -71,40 +85,106 @@ func XForwardedForUpdate(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		defer next.ServeHTTP(w, r)
 
-		remoteIP, _, err := net.SplitHostPort(r.RemoteAddr)
+		pref := XFFComputePreferences{
+			StripPrivate:  true,
+			StripLoopback: true,
+			StripCGNAT:    true,
+			Flatten:       true,
+			StripLLU:      true,
+		}
+
+		remoteAddr := r.RemoteAddr
+		origXFFHeader := r.Header.Get("X-Forwarded-For")
 
-		if parsedRemoteIP := net.ParseIP(remoteIP); parsedRemoteIP != nil && parsedRemoteIP.IsLoopback() {
-			// anubis is likely deployed behind a local reverse proxy
-			// pass header as-is to not break existing applications
+		if remoteAddr == "@" {
+			// remote is a unix socket
+			// do not touch chain
 			return
 		}
 
+		xffHeaderString, err := computeXFFHeader(remoteAddr, origXFFHeader, pref)
 		if err != nil {
-			slog.Warn("The default format of request.RemoteAddr should be IP:Port", "remoteAddr", r.RemoteAddr)
+			slog.Debug("computing X-Forwarded-For header failed", "err", err)
 			return
 		}
-		if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
-			forwardedList := strings.Split(",", xff)
-			forwardedList = append(forwardedList, remoteIP)
-			// this behavior is equivalent to
-			// ingress-nginx "compute-full-forwarded-for"
-			// https://kubernetes.github.io/ingress-nginx/user-guide/nginx-configuration/configmap/#compute-full-forwarded-for
-			//
-			// this would be the correct place to strip and/or flatten this list
-			//
-			// strip - iterate backwards and eliminate configured trusted IPs
-			// flatten - only return the last element to avoid spoofing confusion
-			//
-			// many applications handle this in different ways, but
-			// generally they'd be expected to do these two things on
-			// their own end to find the first non-spoofed IP
-			r.Header.Set("X-Forwarded-For", strings.Join(forwardedList, ","))
+
+		if len(xffHeaderString) == 0 {
+			r.Header.Del("X-Forwarded-For")
 		} else {
-			r.Header.Set("X-Forwarded-For", remoteIP)
+			r.Header.Set("X-Forwarded-For", xffHeaderString)
 		}
 	})
 }
 
+var (
+	ErrCantSplitHostParse = errors.New("internal: unable to net.SplitHostParse")
+	ErrCantParseRemoteIP  = errors.New("internal: unable to parse remote IP")
+)
+
+func computeXFFHeader(remoteAddr string, origXFFHeader string, pref XFFComputePreferences) (string, error) {
+	remoteIP, _, err := net.SplitHostPort(remoteAddr)
+	if err != nil {
+		return "", fmt.Errorf("%w: %w", ErrCantSplitHostParse, err)
+	}
+	parsedRemoteIP, err := netip.ParseAddr(remoteIP)
+	if err != nil {
+		return "", fmt.Errorf("%w: %w", ErrCantParseRemoteIP, err)
+	}
+
+	origForwardedList := make([]string, 0, 4)
+	if origXFFHeader != "" {
+		origForwardedList = strings.Split(origXFFHeader, ",")
+	}
+	origForwardedList = append(origForwardedList, parsedRemoteIP.String())
+	forwardedList := make([]string, 0, len(origForwardedList))
+	// this behavior is equivalent to
+	// ingress-nginx "compute-full-forwarded-for"
+	// https://kubernetes.github.io/ingress-nginx/user-guide/nginx-configuration/configmap/#compute-full-forwarded-for
+	//
+	// this would be the correct place to strip and/or flatten this list
+	//
+	// strip - iterate backwards and eliminate configured trusted IPs
+	// flatten - only return the last element to avoid spoofing confusion
+	//
+	// many applications handle this in different ways, but
+	// generally they'd be expected to do these two things on
+	// their own end to find the first non-spoofed IP
+	for i := len(origForwardedList) - 1; i >= 0; i-- {
+		segmentIP, err := netip.ParseAddr(origForwardedList[i])
+		if err != nil {
+			// can't assess this element, so the remainder of the chain
+			// can't be trusted. not a fatal error, since anyone can
+			// spoof an XFF header
+			slog.Debug("failed to parse XFF segment", "err", err)
+			break
+		}
+		if pref.StripPrivate && segmentIP.IsPrivate() {
+			continue
+		}
+		if pref.StripLoopback && segmentIP.IsLoopback() {
+			continue
+		}
+		if pref.StripLLU && segmentIP.IsLinkLocalUnicast() {
+			continue
+		}
+		if pref.StripCGNAT && CGNat.Contains(segmentIP) {
+			continue
+		}
+		forwardedList = append([]string{segmentIP.String()}, forwardedList...)
+	}
+	var xffHeaderString string
+	if len(forwardedList) == 0 {
+		xffHeaderString = ""
+		return xffHeaderString, nil
+	}
+	if pref.Flatten {
+		xffHeaderString = forwardedList[len(forwardedList)-1]
+	} else {
+		xffHeaderString = strings.Join(forwardedList, ",")
+	}
+	return xffHeaderString, nil
+}
+
 // NoStoreCache sets the Cache-Control header to no-store for the response.
 func NoStoreCache(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
diff --git a/internal/xff_test.go b/internal/xff_test.go
new file mode 100644
index 0000000..26bb0e6
--- /dev/null
+++ b/internal/xff_test.go
@@ -0,0 +1,166 @@
+package internal
+
+import (
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestXForwardedForUpdateIgnoreUnix(t *testing.T) {
+	var remoteAddr = ""
+	var xff = ""
+
+	h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		remoteAddr = r.RemoteAddr
+		xff = r.Header.Get("X-Forwarded-For")
+		w.WriteHeader(http.StatusOK)
+	})
+
+	r := httptest.NewRequest(http.MethodGet, "/", nil)
+
+	r.RemoteAddr = "@"
+
+	w := httptest.NewRecorder()
+
+	XForwardedForUpdate(h).ServeHTTP(w, r)
+
+	if r.RemoteAddr != remoteAddr {
+		t.Errorf("wanted remoteAddr to be %s, got: %s", r.RemoteAddr, remoteAddr)
+	}
+
+	if xff != "" {
+		t.Error("handler added X-Forwarded-For when it should not have")
+	}
+}
+
+func TestXForwardedForUpdateAddToChain(t *testing.T) {
+	var xff = ""
+	const expected = "1.1.1.1"
+
+	h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		xff = r.Header.Get("X-Forwarded-For")
+		w.WriteHeader(http.StatusOK)
+	})
+
+	srv := httptest.NewServer(XForwardedForUpdate(h))
+
+	r, err := http.NewRequest(http.MethodGet, srv.URL, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r.Header.Set("X-Forwarded-For", "1.1.1.1,10.20.30.40")
+
+	if _, err := srv.Client().Do(r); err != nil {
+		t.Fatal(err)
+	}
+
+	if xff != expected {
+		t.Logf("expected: %s", expected)
+		t.Logf("got:      %s", xff)
+		t.Error("X-Forwarded-For header was not what was expected")
+	}
+}
+
+func TestComputeXFFHeader(t *testing.T) {
+	for _, tt := range []struct {
+		name          string
+		remoteAddr    string
+		origXFFHeader string
+		pref          XFFComputePreferences
+		result        string
+		err           error
+	}{
+		{
+			name:          "StripPrivate",
+			remoteAddr:    "127.0.0.1:80",
+			origXFFHeader: "1.1.1.1,10.0.0.1",
+			pref: XFFComputePreferences{
+				StripPrivate: true,
+			},
+			result: "1.1.1.1,127.0.0.1",
+		},
+		{
+			name:          "StripLoopback",
+			remoteAddr:    "127.0.0.1:80",
+			origXFFHeader: "1.1.1.1,10.0.0.1,127.0.0.1",
+			pref: XFFComputePreferences{
+				StripLoopback: true,
+			},
+			result: "1.1.1.1,10.0.0.1",
+		},
+		{
+			name:          "StripCGNAT",
+			remoteAddr:    "100.64.0.1:80",
+			origXFFHeader: "1.1.1.1,10.0.0.1,100.64.0.1",
+			pref: XFFComputePreferences{
+				StripCGNAT: true,
+			},
+			result: "1.1.1.1,10.0.0.1",
+		},
+		{
+			name:          "StripLinkLocalUnicastIPv4",
+			remoteAddr:    "169.254.0.1:80",
+			origXFFHeader: "1.1.1.1,10.0.0.1,169.254.0.1",
+			pref: XFFComputePreferences{
+				StripLLU: true,
+			},
+			result: "1.1.1.1,10.0.0.1",
+		},
+		{
+			name:          "StripLinkLocalUnicastIPv6",
+			remoteAddr:    "169.254.0.1:80",
+			origXFFHeader: "1.1.1.1,10.0.0.1,fe80::",
+			pref: XFFComputePreferences{
+				StripLLU: true,
+			},
+			result: "1.1.1.1,10.0.0.1",
+		},
+		{
+			name:          "Flatten",
+			remoteAddr:    "127.0.0.1:80",
+			origXFFHeader: "1.1.1.1,10.0.0.1,fe80::,100.64.0.1,169.254.0.1",
+			pref: XFFComputePreferences{
+				StripPrivate:  true,
+				StripLoopback: true,
+				StripCGNAT:    true,
+				StripLLU:      true,
+				Flatten:       true,
+			},
+			result: "1.1.1.1",
+		},
+		{
+			name:       "invalid-ip-port",
+			remoteAddr: "fe80::",
+			err:        ErrCantSplitHostParse,
+		},
+		{
+			name:       "invalid-remote-ip",
+			remoteAddr: "anubis:80",
+			err:        ErrCantParseRemoteIP,
+		},
+		{
+			name:       "no-xff-dont-panic",
+			remoteAddr: "127.0.0.1:80",
+			pref: XFFComputePreferences{
+				StripPrivate:  true,
+				StripLoopback: true,
+				StripCGNAT:    true,
+				StripLLU:      true,
+				Flatten:       true,
+			},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := computeXFFHeader(tt.remoteAddr, tt.origXFFHeader, tt.pref)
+			if err != nil && !errors.Is(err, tt.err) {
+				t.Errorf("computeXFFHeader got the wrong error, wanted %v but got: %v", tt.err, err)
+			}
+
+			if result != tt.result {
+				t.Errorf("computeXFFHeader returned the wrong result, wanted %q but got: %q", tt.result, result)
+			}
+		})
+	}
+}