gitlab.com/gitlab-org/labkit@v1.21.0/correlation/inbound_http_test.go

package correlation

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/require"
)

// nolint:gocognit,cyclop
func TestInjectCorrelationID(t *testing.T) {
	tests := []struct {
		name                      string
		opts                      []InboundHandlerOption
		header                    http.Header
		shouldMatch               string
		shouldNotMatch            string
		shouldSetResponseHeader   bool
		expectedResponseHeader    string
		downstreamResponseHeaders http.Header
		remoteAddr                string
		xffHeader                 string
	}{
		{
			name: "without_propagation",
		},
		{
			name: "without_propagation_ignore_incoming_header",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldNotMatch: "123",
		},
		{
			name: "with_propagation_no_incoming_header",
			opts: []InboundHandlerOption{WithPropagation()},
		},
		{
			name: "with_propagation_incoming_header",
			opts: []InboundHandlerOption{WithPropagation()},
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldMatch: "123",
		},
		{
			name: "with_propagation_double_incoming_header",
			opts: []InboundHandlerOption{WithPropagation()},
			header: map[string][]string{
				propagationHeader: {"123", "456"},
			},
			shouldMatch: "123",
		},
		{
			name:                    "with_set_response_header",
			opts:                    []InboundHandlerOption{WithSetResponseHeader()},
			shouldSetResponseHeader: true,
		},
		{
			name:                    "with_set_response_header_and_with_propagation_incoming_header",
			opts:                    []InboundHandlerOption{WithSetResponseHeader(), WithPropagation()},
			shouldSetResponseHeader: true,
		},
		{
			name:                    "with_set_response_header_and_with_propagation_incoming_header_set",
			opts:                    []InboundHandlerOption{WithSetResponseHeader(), WithPropagation()},
			shouldSetResponseHeader: true,
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			expectedResponseHeader: "123",
		},
		{
			name: "mismatching_correlation_ids_without_propagation",
			opts: []InboundHandlerOption{},
			downstreamResponseHeaders: map[string][]string{
				propagationHeader: {"CLIENT_CORRELATION_ID"},
			},
			shouldSetResponseHeader: true,
			expectedResponseHeader:  "CLIENT_CORRELATION_ID",
		},

		{
			name: "mismatching_correlation_ids_with_propagation",
			opts: []InboundHandlerOption{WithSetResponseHeader()},
			downstreamResponseHeaders: map[string][]string{
				propagationHeader: {"CLIENT_CORRELATION_ID"},
			},
			shouldSetResponseHeader: true,
			expectedResponseHeader:  "CLIENT_CORRELATION_ID",
		},
		{
			name:                    "with_set_response_header_and_with_propagation_incoming_header_set",
			opts:                    []InboundHandlerOption{WithSetResponseHeader(), WithPropagation()},
			shouldSetResponseHeader: true,
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			downstreamResponseHeaders: map[string][]string{
				propagationHeader: {"CLIENT_CORRELATION_ID"},
			},
			expectedResponseHeader: "CLIENT_CORRELATION_ID",
		},
		{
			name: "trusted_cidrs_and_propagation",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForXForwardedFor([]string{"192.168.0.1/32"}),
				WithCIDRsTrustedForPropagation([]string{"1.2.3.4/32"}),
				WithPropagation(),
			},
			xffHeader:  "1.2.3.4, 127.0.0.1",
			remoteAddr: "192.168.0.1:1024",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldMatch: "123",
		},
		{
			name: "trusted_xff_cidrs_for_propagation",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForXForwardedFor([]string{"192.168.0.1/32"}),
				WithCIDRsTrustedForPropagation([]string{"192.168.0.1/32"}),
				WithPropagation(),
			},
			xffHeader:  "1.2.3.4, 127.0.0.1",
			remoteAddr: "192.168.0.1:1024",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldMatch: "123",
		},
		{
			name: "untrusted_xff_cidr_propagation",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForXForwardedFor([]string{"1.2.3.4/32"}),
				WithCIDRsTrustedForPropagation([]string{"1.2.3.5/32"}),
				WithPropagation(),
			},
			xffHeader:  "1.2.3.5, 127.0.0.1",
			remoteAddr: "192.168.0.1:1024",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldNotMatch: "123",
		},
		{
			name: "trusted_xff_cidr_untrusted_propagation_cidr",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForXForwardedFor([]string{"192.168.0.1/32"}),
				WithCIDRsTrustedForPropagation([]string{"127.0.0.1/32"}),
				WithPropagation(),
			},
			xffHeader:  "1.2.3.5, 127.0.0.1",
			remoteAddr: "192.168.0.1:1024",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldNotMatch: "123",
		},
		{
			name: "trusted_cidrs_no_propagation",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForXForwardedFor([]string{"192.168.0.1/32"}),
				WithCIDRsTrustedForPropagation([]string{"1.2.3.4/32"}),
			},
			xffHeader:  "1.2.3.4, 127.0.0.1",
			remoteAddr: "192.168.0.1:1024",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldNotMatch: "123",
		},
		{
			name: "trusted_propagation_cidr_remote_addr_propagation",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForPropagation([]string{"1.2.3.4/32"}),
				WithPropagation(),
			},
			remoteAddr: "1.2.3.4:1024",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldMatch: "123",
		},
		{
			name: "no_xff_trusted_cidrs",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForPropagation([]string{"192.168.0.1/8"}),
				WithPropagation(),
			},
			xffHeader:  "1.2.3.4, 127.0.0.1",
			remoteAddr: "192.168.0.1:1024",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldMatch: "123",
		},
		{
			name: "xff_trusted_cidrs_with_propagation_unix_socket",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForXForwardedFor([]string{"192.168.0.1/8", "127.0.0.1/32"}),
				WithCIDRsTrustedForPropagation([]string{"1.2.3.4/32"}),
				WithPropagation(),
			},
			xffHeader:  "1.2.3.4, 127.0.0.1",
			remoteAddr: "@",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldMatch: "123",
		},
		{
			name: "no_xff_with_propagation_unix_socket",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForPropagation([]string{"1.2.3.4/8"}),
				WithPropagation(),
			},
			xffHeader:  "1.2.3.4, 127.0.0.1",
			remoteAddr: "@",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldNotMatch: "123",
		},
		{
			name: "bad_cidrs_ignore_propagation",
			opts: []InboundHandlerOption{
				WithCIDRsTrustedForXForwardedFor([]string{"a.b.c"}),
				WithCIDRsTrustedForPropagation([]string{"x.y.z"}),
				WithPropagation(),
			},
			xffHeader:  "1.2.3.4, 127.0.0.1",
			remoteAddr: "192.168.1.2:9876",
			header: map[string][]string{
				propagationHeader: {"123"},
			},
			shouldNotMatch: "123",
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			invoked := false

			h := InjectCorrelationID(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				invoked = true

				ctx := r.Context()
				correlationID := ExtractFromContext(ctx)
				require.NotNil(t, correlationID, "CorrelationID is missing")
				require.NotEmpty(t, correlationID, "CorrelationID is missing")
				if test.shouldMatch != "" {
					require.Equal(t, test.shouldMatch, correlationID, "CorrelationID should match")
				}

				if test.shouldNotMatch != "" {
					require.NotEqual(t, test.shouldNotMatch, correlationID, "CorrelationID should not match")
				}

				for k, vs := range test.downstreamResponseHeaders {
					for _, v := range vs {
						w.Header().Set(k, v)
					}
				}
			}), test.opts...)

			r := httptest.NewRequest("GET", "http://example.com", nil)
			for k, v := range test.header {
				r.Header[http.CanonicalHeaderKey(k)] = v
			}

			if test.remoteAddr != "" {
				r.RemoteAddr = test.remoteAddr
			}
			if test.xffHeader != "" {
				r.Header.Add("X-Forwarded-For", test.xffHeader)
			}

			w := httptest.NewRecorder()
			h.ServeHTTP(w, r)

			require.True(t, invoked, "handler not executed")

			v, ok := w.Result().Header[http.CanonicalHeaderKey(propagationHeader)]
			require.Equal(t, test.shouldSetResponseHeader, ok, "response header existence mismatch")
			if test.shouldSetResponseHeader {
				require.Len(t, v, 1, "expected exactly one correlation header")
				require.NotEmpty(t, v[0], "expected non-empty correlation header")
			}

			if test.expectedResponseHeader != "" {
				responseHeader := w.Header().Get(propagationHeader)
				require.Equal(t, test.expectedResponseHeader, responseHeader, "response header should match")
			}
		})
	}
}