From a56b30d96ca01943146e82ffe50aad87a6ccdc1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Mazeau?= Date: Tue, 7 Jun 2022 11:56:15 +0200 Subject: [PATCH] contrib/internal/httptrace: handle DD_CLIENT_IP_HEADER env var --- contrib/internal/httptrace/httptrace.go | 22 ++++++- contrib/internal/httptrace/httptrace_test.go | 23 +++++-- contrib/internal/httptrace/options.go | 23 +++++++ contrib/internal/httptrace/options_test.go | 63 ++++++++++++++++++++ 4 files changed, 123 insertions(+), 8 deletions(-) create mode 100644 contrib/internal/httptrace/options.go create mode 100644 contrib/internal/httptrace/options_test.go diff --git a/contrib/internal/httptrace/httptrace.go b/contrib/internal/httptrace/httptrace.go index 24fc379a1d..5eb89fffa6 100644 --- a/contrib/internal/httptrace/httptrace.go +++ b/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, " ") diff --git a/contrib/internal/httptrace/httptrace_test.go b/contrib/internal/httptrace/httptrace_test.go index ebe72fc6cc..4dfd6c02b5 100644 --- a/contrib/internal/httptrace/httptrace_test.go +++ b/contrib/internal/httptrace/httptrace_test.go @@ -32,10 +32,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 { @@ -146,6 +147,18 @@ func genIPTestCases() []IPTestCase { expectedIP: netaddr.MustParseIP(ipv4Global), headers: map[string]string{"X-fOrWaRdEd-FoR": ipv4Global}, }, + { + 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 @@ -158,7 +171,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()) }) } diff --git a/contrib/internal/httptrace/options.go b/contrib/internal/httptrace/options.go new file mode 100644 index 0000000000..1823772171 --- /dev/null +++ b/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), + } +} diff --git a/contrib/internal/httptrace/options_test.go b/contrib/internal/httptrace/options_test.go new file mode 100644 index 0000000000..2e33ba3879 --- /dev/null +++ b/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) + } + } +}