contrib/internal/httptrace: handle DD_CLIENT_IP_HEADER env var

contrib/internal/httptrace/httptrace.go

@@ -37,7 +37,7 @@ func StartRequestSpan(r *http.Request, opts ...ddtrace.StartSpanOption) (tracer.
 			tracer.Tag("http.host", r.Host),
 		}, opts...)
 	}
-	if ip := getClientIP(r.RemoteAddr, r.Header); ip.IsValid() {
+	if ip := getClientIP(r.RemoteAddr, r.Header, cfg.ipHeader); ip.IsValid() {
 		opts = append(opts, tracer.Tag(ext.HTTPClientIP, ip.String()))
 	}
 	if spanctx, err := tracer.Extract(tracer.HTTPHeadersCarrier(r.Header)); err == nil {
@@ -74,10 +74,26 @@ var (
 	ipv6SpecialNetworks = []*netaddr.IPPrefix{
 		ippref("fec0::/10"), // site local
 	}
-	ipHeaders = []string{"x-forwarded-for", "x-real-ip", "x-client-ip", "x-forwarded", "x-cluster-client-ip", "forwarded-for", "forwarded", "via", "true-client-ip"}
+	ipHeaders = []string{
+		"x-forwarded-for",
+		"x-real-ip",
+		"x-client-ip",
+		"x-forwarded",
+		"x-cluster-client-ip",
+		"forwarded-for",
+		"forwarded",
+		"via",
+		"true-client-ip",
+	}
 )
 
-func getClientIP(remoteAddr string, headers http.Header) netaddr.IP {
+// getClientIP uses the request headers to resolve the client IP. If a specific header to check is provided through
+// DD_CLIENT_IP_HEADER, then only this header is checked.
+func getClientIP(remoteAddr string, headers http.Header, clientIPHeader string) netaddr.IP {
+	ipHeaders := ipHeaders
+	if len(clientIPHeader) > 0 {
+		ipHeaders = []string{clientIPHeader}
+	}
 	check := func(value string) netaddr.IP {
 		for _, ip := range strings.Split(value, ",") {
 			ipStr := strings.Trim(ip, " ")

contrib/internal/httptrace/httptrace_test.go

@@ -31,10 +31,11 @@ func TestStartRequestSpan(t *testing.T) {
 }
 
 type IPTestCase struct {
-	name       string
-	remoteAddr string
-	headers    map[string]string
-	expectedIP netaddr.IP
+	name         string
+	remoteAddr   string
+	headers      map[string]string
+	expectedIP   netaddr.IP
+	userIPHeader string
 }
 
 func genIPTestCases() []IPTestCase {
@@ -135,11 +136,24 @@ func genIPTestCases() []IPTestCase {
 			expectedIP: netaddr.MustParseIP(ipv6Global),
 		},
 	}, tcs...)
-	// No headers
-	tcs = append(tcs, IPTestCase{
-		name:       "no-headers",
-		expectedIP: netaddr.IP{},
-	})
+	tcs = append([]IPTestCase{
+		{
+			name:       "no-headers",
+			expectedIP: netaddr.IP{},
+		},
+		{
+			name:         "user-header",
+			expectedIP:   netaddr.MustParseIP(ipv4Global),
+			headers:      map[string]string{"x-forwarded-for": ipv6Global, "custom-header": ipv4Global},
+			userIPHeader: "custom-header",
+		},
+		{
+			name:         "user-header-not-found",
+			expectedIP:   netaddr.IP{},
+			headers:      map[string]string{"x-forwarded-for": ipv4Global},
+			userIPHeader: "custom-header",
+		},
+	}, tcs...)
 
 	return tcs
 }
@@ -151,7 +165,7 @@ func TestIPHeaders(t *testing.T) {
 			for k, v := range tc.headers {
 				header.Add(k, v)
 			}
-			require.Equal(t, tc.expectedIP.String(), getClientIP(tc.remoteAddr, header).String())
+			require.Equal(t, tc.expectedIP.String(), getClientIP(tc.remoteAddr, header, tc.userIPHeader).String())
 		})
 	}
 

contrib/internal/httptrace/options.go

@@ -0,0 +1,23 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2022 Datadog, Inc.
+
+package httptrace
+
+import "os"
+
+type config struct {
+	ipHeader string
+}
+
+var (
+	clientIPHeaderEnvVar = "DD_CLIENT_IP_HEADER"
+	cfg                  = newConfig()
+)
+
+func newConfig() *config {
+	return &config{
+		ipHeader: os.Getenv(clientIPHeaderEnvVar),
+	}
+}

contrib/internal/httptrace/options_test.go

@@ -0,0 +1,63 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2022 Datadog, Inc.
+
+package httptrace
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestConfig(t *testing.T) {
+	t.Run("config", func(t *testing.T) {
+		t.Run("client-ip-header-unset", func(t *testing.T) {
+			cfg := newConfig()
+			require.Empty(t, cfg.ipHeader)
+
+		})
+		t.Run("client-ip-header-empty", func(t *testing.T) {
+			restore := cleanEnv()
+			err := os.Setenv(clientIPHeaderEnvVar, "")
+			require.NoError(t, err)
+			cfg := newConfig()
+			require.Empty(t, cfg.ipHeader)
+			defer restore()
+
+		})
+		t.Run("client-ip-header-set", func(t *testing.T) {
+			restore := cleanEnv()
+			err := os.Setenv(clientIPHeaderEnvVar, "custom-header")
+			require.NoError(t, err)
+			cfg := newConfig()
+			require.Equal(t, "custom-header", cfg.ipHeader)
+			defer restore()
+
+		})
+	})
+}
+
+func cleanEnv() func() {
+	val := os.Getenv(clientIPHeaderEnvVar)
+	if err := os.Unsetenv(clientIPHeaderEnvVar); err != nil {
+		panic(err)
+	}
+	return func() {
+		restoreEnv(clientIPHeaderEnvVar, val)
+	}
+}
+
+func restoreEnv(key, value string) {
+	if value != "" {
+		if err := os.Setenv(key, value); err != nil {
+			panic(err)
+		}
+	} else {
+		if err := os.Unsetenv(key); err != nil {
+			panic(err)
+		}
+	}
+}