Allow specifying header source for client IP

internal/cmd/deploy.go

@@ -64,6 +64,7 @@ func newDeployCommand() *deployCommand {
 	deployCommand.cmd.Flags().StringSliceVar(&deployCommand.args.TargetOptions.LogResponseHeaders, "log-response-header", nil, "Additional response header to log (may be specified multiple times)")
 	deployCommand.cmd.Flags().StringSliceVar(&deployCommand.args.ServiceOptions.ExcludeMetricsPaths, "exclude-metrics-path", nil, "Request path(s) to exclude from Prometheus metrics (may be specified multiple times)")
 	deployCommand.cmd.Flags().BoolVar(&deployCommand.args.TargetOptions.ForwardHeaders, "forward-headers", false, "Forward X-Forwarded headers to target (default false if TLS enabled; otherwise true)")
+	deployCommand.cmd.Flags().StringVar(&deployCommand.args.ServiceOptions.ClientIPHeader, "client-ip-header", "", "Request header containing the original client IP; used to populate X-Forwarded-For when present")
 	deployCommand.cmd.Flags().BoolVar(&deployCommand.args.TargetOptions.ScopeCookiePaths, "scope-cookie-paths", false, "Scope cookie paths to match path prefix")
 
 	deployCommand.cmd.MarkFlagRequired("target")

internal/server/client_ip_middleware.go

@@ -0,0 +1,24 @@
+package server
+
+import (
+	"net/http"
+)
+
+type ClientIPMiddleware struct {
+	headerName string
+	next       http.Handler
+}
+
+func WithClientIPMiddleware(headerName string, next http.Handler) http.Handler {
+	return &ClientIPMiddleware{
+		headerName: http.CanonicalHeaderKey(headerName),
+		next:       next,
+	}
+}
+
+func (h *ClientIPMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if clientIP := r.Header.Get(h.headerName); clientIP != "" {
+		r.Header.Set("X-Forwarded-For", clientIP)
+	}
+	h.next.ServeHTTP(w, r)
+}

internal/server/logging_middleware_test.go

@@ -98,3 +98,26 @@ func TestMiddleware_LoggingMiddleware(t *testing.T) {
 	assert.Equal(t, "HTTP/1.1", logline.Proto)
 	assert.Equal(t, "http", logline.Scheme)
 }
+
+func TestMiddleware_LoggingMiddlewareLogsClientIPHeaderAsRemoteAddr(t *testing.T) {
+	out := &strings.Builder{}
+	logger := slog.New(slog.NewJSONHandler(out, nil))
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
+
+	middleware := WithLoggingMiddleware(logger, 80, 443, WithClientIPMiddleware("True-Client-IP", handler))
+
+	req := httptest.NewRequest("GET", "http://app.example.com/", nil)
+	req.Header.Set("X-Forwarded-For", "10.10.10.10")
+	req.Header.Set("True-Client-IP", "203.0.113.7")
+
+	middleware.ServeHTTP(httptest.NewRecorder(), req)
+
+	logline := struct {
+		RemoteAddr string `json:"remote_addr"`
+	}{}
+
+	err := json.NewDecoder(strings.NewReader(out.String())).Decode(&logline)
+	require.NoError(t, err)
+
+	assert.Equal(t, "203.0.113.7", logline.RemoteAddr)
+}

internal/server/service.go

@@ -94,6 +94,7 @@ type ServiceOptions struct {
 	WriterAffinityTimeout       time.Duration `json:"writer_affinity_timeout"`
 	ReadTargetsAcceptWebsockets bool          `json:"read_targets_accept_websockets"`
 	ExcludeMetricsPaths         []string      `json:"exclude_metrics_paths"`
+	ClientIPHeader              string        `json:"client_ip_header"`
 }
 
 func (so *ServiceOptions) ShouldExcludeMetrics(r *http.Request) bool {
@@ -456,6 +457,10 @@ func (s *Service) createMiddleware(options ServiceOptions, certManager CertManag
 		handler = certManager.HTTPHandler(handler)
 	}
 
+	if options.ClientIPHeader != "" {
+		handler = WithClientIPMiddleware(options.ClientIPHeader, handler)
+	}
+
 	return handler, nil
 }
 

internal/server/service_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -25,6 +26,49 @@ func TestService_ServeRequest(t *testing.T) {
 	require.Equal(t, http.StatusOK, w.Result().StatusCode)
 }
 
+func TestService_ClientIPHeaderRewritesXForwardedFor(t *testing.T) {
+	var xForwardedFor, trueClientIP string
+
+	serviceOptions := defaultServiceOptions
+	serviceOptions.ClientIPHeader = "True-Client-IP"
+
+	targetOptions := defaultTargetOptions
+	targetOptions.ForwardHeaders = true
+
+	service := testCreateServiceWithHandler(t, serviceOptions, targetOptions,
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path != defaultHealthCheckConfig.Path {
+				xForwardedFor = r.Header.Get("X-Forwarded-For")
+				trueClientIP = r.Header.Get("True-Client-IP")
+			}
+		}))
+
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req.Header.Set("True-Client-IP", "203.0.113.9")
+	req.Header.Set("X-Forwarded-For", "6.6.6.6")
+
+	clientIP, _, err := net.SplitHostPort(req.RemoteAddr)
+	require.NoError(t, err)
+
+	w := httptest.NewRecorder()
+	service.ServeHTTP(w, req)
+
+	require.Equal(t, http.StatusOK, w.Result().StatusCode)
+	require.Equal(t, "203.0.113.9, "+clientIP, xForwardedFor)
+	require.Equal(t, "203.0.113.9", trueClientIP)
+
+	// Without the trusted header, the client-supplied X-Forwarded-For is
+	// forwarded unmodified, as usual.
+	req = httptest.NewRequest(http.MethodGet, "http://example.com/", nil)
+	req.Header.Set("X-Forwarded-For", "6.6.6.6")
+
+	w = httptest.NewRecorder()
+	service.ServeHTTP(w, req)
+
+	require.Equal(t, http.StatusOK, w.Result().StatusCode)
+	require.Equal(t, "6.6.6.6, "+clientIP, xForwardedFor)
+}
+
 func TestService_RedirectToHTTPSWhenTLSRequired(t *testing.T) {
 	service := testCreateService(t, ServiceOptions{Hosts: []string{"example.com"}, TLSEnabled: true, TLSRedirect: true}, defaultTargetOptions)