contrib/internal/httptrace: add tests for env var configuration

contrib/internal/httptrace/httptrace_test.go

@@ -136,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
 }

contrib/internal/httptrace/options.go

@@ -6,8 +6,13 @@ type config struct {
 	ipHeader string
 }
 
-var cfg = config{}
+var (
+	clientIPHeaderEnvVar = "DD_CLIENT_IP_HEADER"
+	cfg                  = newConfig()
+)
 
-func init() {
-	cfg.ipHeader = os.Getenv("DD_CLIENT_IP_HEADER")
+func newConfig() *config {
+	return &config{
+		ipHeader: os.Getenv(clientIPHeaderEnvVar),
+	}
 }

contrib/internal/httptrace/options_test.go

@@ -0,0 +1,58 @@
+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)
+		}
+	}
+}