From 3bfbd1d704b81da0cfa5bb2bbf7767e9e90f6077 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Mazeau?= <francois.mazeau@datadoghq.com>
Date: Tue, 7 Jun 2022 16:17:27 +0200
Subject: [PATCH] contrib/internal/httptrace: add tests for env var
 configuration

---
 contrib/internal/httptrace/httptrace_test.go | 23 ++++++--
 contrib/internal/httptrace/options.go        | 11 +++-
 contrib/internal/httptrace/options_test.go   | 58 ++++++++++++++++++++
 3 files changed, 84 insertions(+), 8 deletions(-)
 create mode 100644 contrib/internal/httptrace/options_test.go

diff --git a/contrib/internal/httptrace/httptrace_test.go b/contrib/internal/httptrace/httptrace_test.go
index a75a51bd8b..25e5a85c1e 100644
--- a/contrib/internal/httptrace/httptrace_test.go
+++ b/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
 }
diff --git a/contrib/internal/httptrace/options.go b/contrib/internal/httptrace/options.go
index 7e85167f1b..28581b34ac 100644
--- a/contrib/internal/httptrace/options.go
+++ b/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),
+	}
 }
diff --git a/contrib/internal/httptrace/options_test.go b/contrib/internal/httptrace/options_test.go
new file mode 100644
index 0000000000..3cd264dfc6
--- /dev/null
+++ b/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)
+		}
+	}
+}