package server

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"maps"
	"net"
	"net/http"
	"net/url"
	"path"
	"slices"
	"strconv"
	"strings"
	"testing"

	"buf.build/go/protovalidate"
	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab"
	gapi "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kubernetes_api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/kubernetes_api/rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/cache"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/featureflag"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/httpz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/jwttool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/prototool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/matcher"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_cache"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_event_tracker"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_gitlab"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_kubernetes_api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_modserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_otel"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_prototool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_usage_metrics"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testlogger"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/version"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/agentcfg"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tool/grpcz"
	otelmetric "go.opentelemetry.io/otel/metric"
	metricnoop "go.opentelemetry.io/otel/metric/noop"
	"go.opentelemetry.io/otel/propagation"
	tracesdk "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	"go.uber.org/mock/gomock"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/anypb"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/utils/ptr"
)

const (
	jobToken            = "asdfgasdfxadf"
	personalAccessToken = "glpat-asdfgasdfxadf"
	requestPath         = "/api/bla"
	requestPayload      = "asdfndaskjfadsbfjsadhvfjhavfjasvf"
	responsePayload     = `{"key":"value"}`
	testHost            = "any-host-will-do.local"
)

var (
	fakeGitLabReleasesList = []string{
		"18.0",
		"17.11", "17.10", "17.9", "17.8", "17.7", "17.6", "17.5", "17.4", "17.3", "17.2", "17.1", "17.0",
		"16.11", "16.10", "16.9", "16.8", "16.7", "16.6", "16.5", "16.4", "16.3", "16.2", "16.1", "16.0",
	}
	queryParamName              string
	queryParamValue             string
	validWebSocketTokenSecret   = []byte("test-secret")
	invalidWebSocketTokenSecret = []byte("invalid-test-secret")
)

func init() {
	key := make([]byte, 256)
	for i := range 256 {
		key[i] = byte(i)
	}
	val := slices.Clone(key)
	slices.Reverse(val)
	queryParamName = string(key)
	queryParamValue = string(val)
}

func TestProxy_CORSPreflight(t *testing.T) {
	_, _, client, req, _, _, _, _, _, _, _, _, _, _, _, _, _ := setupProxyWithHandler(t, "/", func(w http.ResponseWriter, r *http.Request) {
		t.Fail() // unexpected invocation
	})

	req.Method = http.MethodOptions

	// set CORS headers
	req.Header.Set("Access-Control-Request-Method", "POST")
	req.Header.Set("Access-Control-Request-Headers", "Content-Type, Cookie, X-Csrf-Token, Gitlab-Agent-Id")

	resp, err := client.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()

	assert.Equal(t, http.StatusOK, resp.StatusCode)

	assertCORSHeaders(t, resp.Header)
	assertCORSPreflightHeaders(t, resp.Header)
}

func TestProxy_OriginIsNotAllowed(t *testing.T) {
	methods := []string{
		http.MethodGet,
		http.MethodHead,
		http.MethodPost,
		http.MethodPut,
		http.MethodPatch,
		http.MethodDelete,
		http.MethodConnect,
		http.MethodOptions,
		http.MethodTrace,
	}
	for _, method := range methods {
		t.Run(method, func(t *testing.T) {
			_, _, client, req, _, _, _, _, _, _, _, _, _, _, _, _, _ := setupProxyWithHandler(t, "/", func(w http.ResponseWriter, r *http.Request) {
				t.Fail() // unexpected invocation
			})

			req.Method = method
			req.Header.Set("Origin", "https://not-allowed.example.com")

			resp, err := client.Do(req)
			require.NoError(t, err)
			defer resp.Body.Close()

			assert.Equal(t, http.StatusForbidden, resp.StatusCode)
		})
	}
}

func assertCORSHeaders(t *testing.T, header http.Header) {
	// Assert headers used for CORS preflight and normal requests
	assert.Equal(t, []string{"kas.gitlab.example.com"}, header.Values("Access-Control-Allow-Origin"))
	assert.Equal(t, []string{"true"}, header.Values("Access-Control-Allow-Credentials"))
	assert.Equal(t, []string{"Origin"}, header.Values("Vary"))
}

func assertCORSPreflightHeaders(t *testing.T, header http.Header) {
	// Assert CORS preflight response headers
	assert.Equal(t, []string{"Content-Type, Cookie, X-Csrf-Token, Gitlab-Agent-Id"}, header.Values("Access-Control-Allow-Headers"))
	assert.Equal(t, []string{"GET, HEAD, POST, PUT, DELETE, CONNECT, OPTIONS, TRACE, PATCH"}, header.Values("Access-Control-Allow-Methods"))
	assert.Equal(t, []string{"86400"}, header.Values("Access-Control-Max-Age"))
}

func TestProxy_AuthorizationErrors(t *testing.T) {
	tests := []struct {
		name            string
		auth            []string
		cookie          *string
		agentIDHeader   []string
		csrfTokenHeader []string
		webSocketToken  *string
		message         string
	}{
		{
			name:    "missing credentials",
			message: "GitLab Agent Server: Unauthorized: no valid credentials provided",
		},
		{
			name:    "job token: multiple headers",
			auth:    []string{"a", "b"},
			message: "GitLab Agent Server: Unauthorized: Authorization header: expecting a single header, got 2",
		},
		{
			name:    "job token: invalid format1",
			auth:    []string{"Token asdfadsf"},
			message: "GitLab Agent Server: Unauthorized: Authorization header: expecting Bearer token",
		},
		{
			name:    "job token: invalid format2",
			auth:    []string{"Bearer asdfadsf"},
			message: "GitLab Agent Server: Unauthorized: Authorization header: invalid value",
		},
		{
			name:    "job token: invalid agent id",
			auth:    []string{"Bearer ci:asdf:as"},
			message: `GitLab Agent Server: Unauthorized: Authorization header: failed to parse: strconv.ParseInt: parsing "asdf": invalid syntax`,
		},
		{
			name:    "job token: empty token",
			auth:    []string{"Bearer ci:1:"},
			message: "GitLab Agent Server: Unauthorized: Authorization header: empty token",
		},
		{
			name:    "job token: unknown token type",
			auth:    []string{"Bearer blabla:1:asd"},
			message: "GitLab Agent Server: Unauthorized: Authorization header: unknown token type",
		},
		{
			name:    "cookie: empty string",
			cookie:  new(string),
			message: "GitLab Agent Server: Unauthorized: _gitlab_kas cookie value must not be empty",
		},
		{
			name:    "cookie: missing agent ID header",
			cookie:  ptr.To("the cookie"),
			message: "GitLab Agent Server: Unauthorized: exactly one agent id must be provided either in the \"Gitlab-Agent-Id\" header or \"gitlab-agent-id\" query parameter",
		},
		{
			name:    "cookie: missing agent ID and type header",
			cookie:  ptr.To("the cookie"),
			message: "GitLab Agent Server: Unauthorized: exactly one agent id must be provided either in the \"Gitlab-Agent-Id\" header or \"gitlab-agent-id\" query parameter",
		},
		{
			name:          "cookie: multiple agent ID header values",
			cookie:        ptr.To("the cookie"),
			agentIDHeader: []string{"a", "b"},
			message:       "GitLab Agent Server: Unauthorized: Gitlab-Agent-Id header must have exactly one value",
		},
		{
			name:          "cookie: invalid agent ID value",
			cookie:        ptr.To("the cookie"),
			agentIDHeader: []string{"abcd"},
			message:       "GitLab Agent Server: Unauthorized: agent id in request: invalid value: \"abcd\"",
		},
		{
			name:          "cookie: missing CSRF token header",
			cookie:        ptr.To("the cookie"),
			agentIDHeader: []string{"1234"},
			message:       "GitLab Agent Server: Unauthorized: exactly one CSRF token must be provided either in the \"X-Csrf-Token\" header or \"gitlab-csrf-token\" query parameter",
		},
		{
			name:            "cookie: multiple CSRF token header values",
			cookie:          ptr.To("the cookie"),
			agentIDHeader:   []string{"1234"},
			csrfTokenHeader: []string{"x", "y"},
			message:         "GitLab Agent Server: Unauthorized: X-Csrf-Token header must have exactly one value",
		},
		{
			name:    "personal access token: invalid agent id",
			auth:    []string{"Bearer pat:asdf:as"},
			message: `GitLab Agent Server: Unauthorized: Authorization header: failed to parse: strconv.ParseInt: parsing "asdf": invalid syntax`,
		},
		{
			name:    "personal access token token: empty token",
			auth:    []string{"Bearer pat:1:"},
			message: "GitLab Agent Server: Unauthorized: Authorization header: empty token",
		},
		{
			name:           "websocket token: empty token",
			webSocketToken: ptr.To(encodeTokenWithPrefix("")),
			message:        "GitLab Agent Server: Unauthorized: token is malformed: token contains an invalid number of segments",
		},
		{
			name:           "websocket token: invalid token",
			webSocketToken: ptr.To(encodeTokenWithPrefix(generateInvalidTestWebSocketToken(t, "/"))),
			message:        "GitLab Agent Server: Unauthorized: token signature is invalid: signature is invalid",
		},
		{
			name:           "websocket token: URL (endpoint) mismatch",
			webSocketToken: ptr.To(encodeTokenWithPrefix(generateValidTestWebSocketToken(t, "/another-path", nil))),
			message:        "GitLab Agent Server: Unauthorized: token has invalid claims: token has invalid endpoint",
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			_, _, client, req, _, _, _, _, _, _, _, _, _, _, mockRequestsTotalCounter, _, mockRequestsErrorCounter := setupProxyWithHandler(t, "/", func(w http.ResponseWriter, r *http.Request) {
				t.Fail() // unexpected invocation
			})
			mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockRequestsErrorCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			if len(tc.auth) > 0 {
				req.Header[httpz.AuthorizationHeader] = tc.auth
			}
			if tc.cookie != nil {
				req.AddCookie(&http.Cookie{Name: gitLabKASCookieName, Value: *tc.cookie})
			}
			if len(tc.agentIDHeader) > 0 {
				req.Header[httpz.GitlabAgentIDHeader] = tc.agentIDHeader
			}
			if len(tc.csrfTokenHeader) > 0 {
				req.Header[httpz.CSRFTokenHeader] = tc.csrfTokenHeader
			}
			if tc.webSocketToken != nil {
				req.Header[httpz.ConnectionHeader] = []string{"Upgrade"}
				req.Header[httpz.UpgradeHeader] = []string{"WebSocket"}
				req.Header[httpz.SecWebSocketProtocolHeader] = []string{*tc.webSocketToken}
			}
			resp, err := client.Do(req)
			require.NoError(t, err)
			defer resp.Body.Close()
			assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
			expected := metav1.Status{
				TypeMeta: metav1.TypeMeta{
					Kind:       "Status",
					APIVersion: "v1",
				},
				Status: metav1.StatusFailure,
				Reason: metav1.StatusReasonUnauthorized,
				Code:   http.StatusUnauthorized,
			}
			actualStatus := readStatus(t, resp)
			assert.Truef(t, strings.HasPrefix(actualStatus.Message, tc.message), "Message %q does not have prefix %q", actualStatus.Message, tc.message)
			assert.Empty(t, cmp.Diff(expected, actualStatus, cmpopts.IgnoreFields(metav1.Status{}, "Message")))
		})
	}
}

func TestProxy_AllowedAgentsError(t *testing.T) {
	tests := []struct {
		allowedAgentsHTTPStatus int
		expectedHTTPStatus      int
		message                 string
		captureErr              bool
	}{
		{
			allowedAgentsHTTPStatus: http.StatusUnauthorized, // token is invalid
			expectedHTTPStatus:      http.StatusUnauthorized,
			message:                 "GitLab Agent Server: Unauthorized: CI job token: HTTP status code: 401 for path /api/v4/job/allowed_agents",
		},
		{
			allowedAgentsHTTPStatus: http.StatusForbidden, // token is forbidden
			expectedHTTPStatus:      http.StatusForbidden,
			message:                 "GitLab Agent Server: Forbidden: CI job token: HTTP status code: 403 for path /api/v4/job/allowed_agents",
		},
		{
			allowedAgentsHTTPStatus: http.StatusNotFound, // agent is not found
			expectedHTTPStatus:      http.StatusNotFound,
			message:                 "GitLab Agent Server: Not found: agents for CI job token: HTTP status code: 404 for path /api/v4/job/allowed_agents",
		},
		{
			allowedAgentsHTTPStatus: http.StatusBadGateway, // some weird error
			expectedHTTPStatus:      http.StatusInternalServerError,
			message:                 "GitLab Agent Server: Failed to get allowed agents for CI job token: HTTP status code: 502 for path /api/v4/job/allowed_agents",
			captureErr:              true,
		},
	}
	for _, tc := range tests {
		t.Run(strconv.Itoa(tc.allowedAgentsHTTPStatus), func(t *testing.T) {
			api, _, client, req, _, _, _, _, _, _, _, _, _, _, mockRequestsTotalCounter, _, mockRequestsErrorCounter := setupProxyWithHandler(t, "/", func(w http.ResponseWriter, r *http.Request) {
				assertToken(t, r)
				w.WriteHeader(tc.allowedAgentsHTTPStatus)
			})
			mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockRequestsErrorCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			if tc.captureErr {
				api.EXPECT().
					HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(),
						matcher.ErrorEq(fmt.Sprintf("HTTP status code: %d for path /api/v4/job/allowed_agents", tc.allowedAgentsHTTPStatus)),
						logz.AgentKey(testhelpers.AgentkKey1))
			}
			req.Header.Set(httpz.AuthorizationHeader, fmt.Sprintf("Bearer %s:%d:%s", tokenTypeCI, testhelpers.AgentkKey1.ID, jobToken))
			resp, err := client.Do(req)
			require.NoError(t, err)
			defer resp.Body.Close()
			assert.Equal(t, tc.expectedHTTPStatus, resp.StatusCode)
			expected := metav1.Status{
				TypeMeta: metav1.TypeMeta{
					Kind:       "Status",
					APIVersion: "v1",
				},
				Status: metav1.StatusFailure,
				Reason: code2reason[int32(tc.expectedHTTPStatus)],
				Code:   int32(tc.expectedHTTPStatus),
			}
			actualStatus := readStatus(t, resp)
			assert.True(t, strings.HasPrefix(actualStatus.Message, tc.message))
			assert.Empty(t, cmp.Diff(expected, actualStatus, cmpopts.IgnoreFields(metav1.Status{}, "Message")))
		})
	}
}

func TestProxy_AuthorizeProxyUserError(t *testing.T) {
	tests := []struct {
		authorizeProxyUserHTTPStatus    int
		authorizeProxyUserContentType   string
		authorizeProxyUserFailureReason []byte
		expectedHTTPStatus              int
		expectedErrReason               string
		message                         string
		captureErr                      bool
	}{
		{
			authorizeProxyUserHTTPStatus:    http.StatusUnauthorized, // invalid credentials
			authorizeProxyUserContentType:   "application/json",
			authorizeProxyUserFailureReason: []byte(`{"message": "unauthorized"}`),
			expectedHTTPStatus:              http.StatusUnauthorized,
			expectedErrReason:               "unauthorized",
			message:                         "GitLab Agent Server: Unauthorized",
		},
		{
			authorizeProxyUserHTTPStatus:    http.StatusForbidden, // user has no access to agent
			authorizeProxyUserContentType:   "application/json",
			authorizeProxyUserFailureReason: []byte(`{"message": "forbidden"}`),
			expectedHTTPStatus:              http.StatusUnauthorized,
			expectedErrReason:               "forbidden",
			message:                         "GitLab Agent Server: Unauthorized",
		},
		{
			authorizeProxyUserHTTPStatus:    http.StatusNotFound, // user or agent not found
			authorizeProxyUserContentType:   "application/json",
			authorizeProxyUserFailureReason: []byte(`{"message": "not found"}`),
			expectedHTTPStatus:              http.StatusUnauthorized,
			expectedErrReason:               "not found",
			message:                         "GitLab Agent Server: Unauthorized",
		},
		{
			authorizeProxyUserHTTPStatus:    http.StatusBadRequest, // token invalid
			authorizeProxyUserContentType:   "application/json",
			authorizeProxyUserFailureReason: []byte(`{"message": "unable to get user from request data"}`),
			expectedHTTPStatus:              http.StatusUnauthorized,
			expectedErrReason:               "Unable to get user from request data",
			message:                         "GitLab Agent Server: Unauthorized",
		},
		{
			authorizeProxyUserHTTPStatus:  http.StatusBadGateway, // some weird error
			authorizeProxyUserContentType: "text/plain",
			expectedHTTPStatus:            http.StatusInternalServerError,
			expectedErrReason:             "<unknown reason: expected application/json content type, but got text/plain>",
			message:                       "GitLab Agent Server: Failed to authorize user session",
			captureErr:                    true,
		},
	}
	for _, tc := range tests {
		t.Run(strconv.Itoa(tc.authorizeProxyUserHTTPStatus), func(t *testing.T) {
			api, _, client, req, _, _, _, _, _, _, _, _, _, _, mockRequestsTotalCounter, _, mockRequestsErrorCounter := setupProxyWithHandler(t, "/", func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set(httpz.ContentTypeHeader, tc.authorizeProxyUserContentType)
				w.WriteHeader(tc.authorizeProxyUserHTTPStatus)
				w.Write(tc.authorizeProxyUserFailureReason)
			})
			mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockRequestsErrorCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			if tc.captureErr {
				api.EXPECT().
					HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(),
						matcher.ErrorEq(fmt.Sprintf("HTTP status code: %d for path /api/v4/internal/kubernetes/authorize_proxy_user with reason %s", tc.authorizeProxyUserHTTPStatus, tc.expectedErrReason)),
						logz.AgentKey(testhelpers.AgentkKey1))
			}
			setExpectedSessionCookieParamsInHeader(req)
			resp, err := client.Do(req)
			require.NoError(t, err)
			defer resp.Body.Close()
			assert.Equal(t, tc.expectedHTTPStatus, resp.StatusCode)
			expected := metav1.Status{
				TypeMeta: metav1.TypeMeta{
					Kind:       "Status",
					APIVersion: "v1",
				},
				Status: metav1.StatusFailure,
				Reason: code2reason[int32(tc.expectedHTTPStatus)],
				Code:   int32(tc.expectedHTTPStatus),
			}
			assertCORSHeaders(t, resp.Header)
			actualStatus := readStatus(t, resp)
			assert.True(t, strings.HasPrefix(actualStatus.Message, tc.message))
			assert.Empty(t, cmp.Diff(expected, actualStatus, cmpopts.IgnoreFields(metav1.Status{}, "Message")))
		})
	}
}

func TestProxy_NoExpectedURLPathPrefix(t *testing.T) {
	_, _, client, req, _, _, _, _, _, _, _, _, _, _, mockRequestsTotalCounter, _, mockRequestsErrorCounter := setupProxyWithHandler(t, "/bla/", configCIAccessGitLabHandler(t, nil, nil))
	mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
	mockRequestsErrorCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
	req.URL.Path = requestPath
	req.Header.Set(httpz.AuthorizationHeader, fmt.Sprintf("Bearer %s:%d:%s", tokenTypeCI, testhelpers.AgentkKey1.ID, jobToken))
	resp, err := client.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
	expected := metav1.Status{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Status",
			APIVersion: "v1",
		},
		Status:  metav1.StatusFailure,
		Message: "Correlation ID: ",
		Reason:  metav1.StatusReasonBadRequest,
		Code:    http.StatusBadRequest,
	}
	actualStatus := readStatus(t, resp)
	assert.True(t, strings.HasPrefix(actualStatus.Message, "GitLab Agent Server: Bad request: URL does not start with expected prefix. Trace ID: "))
	assert.Empty(t, cmp.Diff(expected, actualStatus, cmpopts.IgnoreFields(metav1.Status{}, "Message")))
}

func TestProxy_ForbiddenAgentID(t *testing.T) {
	_, _, client, req, _, _, _, _, _, _, _, _, _, _, mockRequestsTotalCounter, _, mockRequestsErrorCounter := setupProxy(t)
	mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
	mockRequestsErrorCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
	req.Header.Set(httpz.AuthorizationHeader, fmt.Sprintf("Bearer %s:%d:%s", tokenTypeCI, 15 /* disallowed id */, jobToken))
	resp, err := client.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	assert.Equal(t, http.StatusForbidden, resp.StatusCode)
	expected := metav1.Status{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Status",
			APIVersion: "v1",
		},
		Status: metav1.StatusFailure,
		Reason: metav1.StatusReasonForbidden,
		Code:   http.StatusForbidden,
	}
	actualStatus := readStatus(t, resp)
	assert.True(t, strings.HasPrefix(actualStatus.Message, "GitLab Agent Server: Forbidden: agentKey is not allowed. Trace ID: "))
	assert.Empty(t, cmp.Diff(expected, actualStatus, cmpopts.IgnoreFields(metav1.Status{}, "Message")))
}

func TestProxy_CIAccessHappyPath(t *testing.T) {
	tests := []struct {
		name          string
		urlPathPrefix string
		config        *gapi.Configuration
		env           *gapi.Environment
		expectedExtra *rpc.HeaderExtra
	}{
		{
			name:          "no prefix",
			urlPathPrefix: "/",
		},
		{
			name:          "with prefix",
			urlPathPrefix: "/bla/",
		},
		{
			name:          "impersonate agent",
			urlPathPrefix: "/",
			config: &gapi.Configuration{
				AccessAs: &agentcfg.CiAccessAsCF{
					As: &agentcfg.CiAccessAsCF_Agent{
						Agent: &agentcfg.CiAccessAsAgentCF{},
					},
				},
			},
		},
		{
			name:          "impersonate",
			urlPathPrefix: "/",
			config: &gapi.Configuration{
				AccessAs: &agentcfg.CiAccessAsCF{
					As: &agentcfg.CiAccessAsCF_Impersonate{
						Impersonate: &agentcfg.CiAccessAsImpersonateCF{
							Username: "user1",
							Groups:   []string{"g1", "g2"},
							Uid:      "uid",
							Extra: []*agentcfg.ExtraKeyValCF{
								{
									Key: "k1",
									Val: []string{"v1", "v2"},
								},
							},
						},
					},
				},
			},
			expectedExtra: &rpc.HeaderExtra{
				ImpConfig: &kubernetes_api.ImpersonationConfig{
					Username: "user1",
					Groups:   []string{"g1", "g2"},
					Uid:      "uid",
					Extra: []*kubernetes_api.ExtraKeyVal{
						{
							Key: "k1",
							Val: []string{"v1", "v2"},
						},
					},
				},
			},
		},
		{
			name:          "impersonate ci job no env",
			urlPathPrefix: "/",
			config: &gapi.Configuration{
				AccessAs: &agentcfg.CiAccessAsCF{
					As: &agentcfg.CiAccessAsCF_CiJob{
						CiJob: &agentcfg.CiAccessAsCiJobCF{},
					},
				},
			},
			expectedExtra: &rpc.HeaderExtra{
				ImpConfig: &kubernetes_api.ImpersonationConfig{
					Username: "gitlab:ci_job:1",
					Groups:   []string{"gitlab:ci_job", "gitlab:group:6", "gitlab:project:3"},
					Extra: []*kubernetes_api.ExtraKeyVal{
						{
							Key: "agent.gitlab.com/id",
							Val: []string{"123"},
						},
						{
							Key: "agent.gitlab.com/config_project_id",
							Val: []string{"5"},
						},
						{
							Key: "agent.gitlab.com/project_id",
							Val: []string{"3"},
						},
						{
							Key: "agent.gitlab.com/ci_pipeline_id",
							Val: []string{"2"},
						},
						{
							Key: "agent.gitlab.com/ci_job_id",
							Val: []string{"1"},
						},
						{
							Key: "agent.gitlab.com/username",
							Val: []string{"testuser"},
						},
					},
				},
			},
		},
		{
			name:          "impersonate ci job prod env",
			urlPathPrefix: "/",
			config: &gapi.Configuration{
				AccessAs: &agentcfg.CiAccessAsCF{
					As: &agentcfg.CiAccessAsCF_CiJob{
						CiJob: &agentcfg.CiAccessAsCiJobCF{},
					},
				},
			},
			env: &gapi.Environment{
				Slug: "prod",
				Tier: "production",
			},
			expectedExtra: &rpc.HeaderExtra{
				ImpConfig: &kubernetes_api.ImpersonationConfig{
					Username: "gitlab:ci_job:1",
					Groups:   []string{"gitlab:ci_job", "gitlab:group:6", "gitlab:group_env_tier:6:production", "gitlab:project:3", "gitlab:project_env:3:prod", "gitlab:project_env_tier:3:production"},
					Extra: []*kubernetes_api.ExtraKeyVal{
						{
							Key: "agent.gitlab.com/id",
							Val: []string{"123"},
						},
						{
							Key: "agent.gitlab.com/config_project_id",
							Val: []string{"5"},
						},
						{
							Key: "agent.gitlab.com/project_id",
							Val: []string{"3"},
						},
						{
							Key: "agent.gitlab.com/ci_pipeline_id",
							Val: []string{"2"},
						},
						{
							Key: "agent.gitlab.com/ci_job_id",
							Val: []string{"1"},
						},
						{
							Key: "agent.gitlab.com/username",
							Val: []string{"testuser"},
						},
						{
							Key: "agent.gitlab.com/environment_slug",
							Val: []string{"prod"},
						},
						{
							Key: "agent.gitlab.com/environment_tier",
							Val: []string{"production"},
						},
					},
				},
			},
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			_, mrClient, client, req, requestCount, mockCIAccessRequestCounter, mockCIAccessAgentsCounter, mockCIAccessEventTracker, _, _, _, _, _, _, mockRequestsTotalCounter, mockRequestsSuccessCounter, _ := setupProxyWithHandler(t, tc.urlPathPrefix, configCIAccessGitLabHandler(t, tc.config, tc.env))
			mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockRequestsSuccessCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockCIAccessRequestCounter.EXPECT().Inc()
			mockCIAccessAgentsCounter.EXPECT().Add(testhelpers.AgentkKey1.ID)
			mockCIAccessEventTracker.EXPECT().EmitEvent(K8sAPIProxyRequestsEvent{UserID: testhelpers.UserID, ProjectID: testhelpers.ConfigProjectID})
			testProxyHappyPath(t, setExpectedJobToken, tc.expectedExtra, mrClient, client, req, requestCount, false)
		})
	}
}

func setExpectedJobToken(req *http.Request) {
	req.Header.Set(httpz.AuthorizationHeader, fmt.Sprintf("Bearer %s:%d:%s", tokenTypeCI, testhelpers.AgentkKey1.ID, jobToken))
}

func TestProxy_PreferAuthorizationHeaderOverSessionCookie(t *testing.T) {
	var expectedExtra *rpc.HeaderExtra // nil
	setJobTokenAndCookie := func(req *http.Request) {
		setExpectedJobToken(req)
		setExpectedSessionCookieParamsInHeader(req)
	}

	_, mrClient, client, req, requestCount, mockCIAccessRequestCounter, mockCIAccessAgentsCounter, mockCIAccessEventTracker, _, _, _, _, _, _, mockRequestsTotalCounter, mockRequestsSuccessCounter, _ := setupProxyWithHandler(t, "/", configCIAccessGitLabHandler(t, nil, nil))
	mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
	mockRequestsSuccessCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
	mockCIAccessRequestCounter.EXPECT().Inc()
	mockCIAccessAgentsCounter.EXPECT().Add(testhelpers.AgentkKey1.ID)
	mockCIAccessEventTracker.EXPECT().EmitEvent(K8sAPIProxyRequestsEvent{UserID: testhelpers.UserID, ProjectID: testhelpers.ConfigProjectID})
	testProxyHappyPath(t, setJobTokenAndCookie, expectedExtra, mrClient, client, req, requestCount, false)
}

func TestProxy_UserAccessHappyPath(t *testing.T) {
	authenticationMethods := []struct {
		name          string
		setParamsFunc func(r *http.Request)
	}{
		{name: "auth in header", setParamsFunc: setExpectedSessionCookieParamsInHeader},
		{name: "auth in params", setParamsFunc: setExpectedSessionCookieParamsInQuery},
	}
	testcases := []struct {
		name          string
		urlPathPrefix string
		auth          *gapi.AuthorizeProxyUserAPIResponse
		expectedExtra *rpc.HeaderExtra
	}{
		{
			name:          "impersonate agent",
			urlPathPrefix: "/",
			auth: &gapi.AuthorizeProxyUserAPIResponse{
				AccessAs: &gapi.AccessAsProxyAuthorization{
					AccessAs: &gapi.AccessAsProxyAuthorization_Agent{
						Agent: &gapi.AccessAsAgentAuthorization{},
					},
				},
			},
		},
		{
			name:          "impersonate user",
			urlPathPrefix: "/",
			auth: &gapi.AuthorizeProxyUserAPIResponse{
				AccessAs: &gapi.AccessAsProxyAuthorization{
					AccessAs: &gapi.AccessAsProxyAuthorization_User{
						User: &gapi.AccessAsUserAuthorization{
							Projects: []*gapi.ProjectAccessCF{
								{
									Id:    1234,
									Roles: []string{"guest", "developer"},
								},
							},
						},
					},
				},
			},
			expectedExtra: &rpc.HeaderExtra{
				ImpConfig: &kubernetes_api.ImpersonationConfig{
					Username: "gitlab:user:testuser",
					Groups: []string{
						"gitlab:user",
						"gitlab:project_role:1234:guest",
						"gitlab:project_role:1234:developer",
					},
					Extra: []*kubernetes_api.ExtraKeyVal{
						{
							Key: "agent.gitlab.com/id",
							Val: []string{"123"},
						},
						{
							Key: "agent.gitlab.com/username",
							Val: []string{"testuser"},
						},
						{
							Key: "agent.gitlab.com/access_type",
							Val: []string{"session_cookie"},
						},
						{
							Key: "agent.gitlab.com/config_project_id",
							Val: []string{"5"},
						},
					},
				},
			},
		},
	}
	for _, atc := range authenticationMethods {
		for _, tc := range testcases {
			t.Run(fmt.Sprintf("%s-%s", atc.name, tc.name), func(t *testing.T) {
				t.Parallel()
				tc.auth = proto.CloneOf(tc.auth) // clone to avoid data race
				tc.auth.Agent = &gapi.AuthorizedAgentForUser{
					Id:            testhelpers.AgentkKey1.ID,
					ConfigProject: &gapi.ConfigProject{Id: 5},
				}
				tc.auth.User = &gapi.User{
					Id:       testhelpers.UserID,
					Username: "testuser",
				}
				_, mrClient, client, req, requestCount, _, _, _, mockUserAccessRequestCounter, mockUserAccessAgentsCounter, mockUserAccessEventTracker, _, _, _, mockRequestsTotalCounter, mockRequestsSuccessCounter, _ := setupProxyWithHandler(t, tc.urlPathPrefix, configUserAccessGitLabHandler(t, tc.auth))
				mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
				mockRequestsSuccessCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
				mockUserAccessRequestCounter.EXPECT().Inc()
				mockUserAccessAgentsCounter.EXPECT().Add(testhelpers.AgentkKey1.ID)
				mockUserAccessEventTracker.EXPECT().EmitEvent(K8sAPIProxyRequestsEvent{UserID: testhelpers.UserID, ProjectID: testhelpers.ConfigProjectID})
				testProxyHappyPath(t, atc.setParamsFunc, tc.expectedExtra, mrClient, client, req, requestCount, false)
			})
		}
	}
}

func TestProxy_StatusResponses(t *testing.T) {
	userHandler := func(t *testing.T) func(http.ResponseWriter, *http.Request) {
		auth := &gapi.AuthorizeProxyUserAPIResponse{
			Agent: &gapi.AuthorizedAgentForUser{
				Id:            testhelpers.AgentkKey1.ID,
				ConfigProject: &gapi.ConfigProject{Id: 5},
			},
			User: &gapi.User{
				Id:       testhelpers.UserID,
				Username: "testuser",
			},
			AccessAs: &gapi.AccessAsProxyAuthorization{
				AccessAs: &gapi.AccessAsProxyAuthorization_Agent{
					Agent: &gapi.AccessAsAgentAuthorization{},
				},
			},
		}
		return configUserAccessGitLabHandler(t, auth)
	}

	req := &prototool.HttpRequest{
		Method: http.MethodPost,
		Header: []*prototool.HeaderKV{
			mock_prototool.NewHeaderKV("Accept-Encoding", "gzip"), // added by the Go client
			mock_prototool.NewHeaderKV(httpz.UserAgentHeader, "test-agent"),
			mock_prototool.NewHeaderKV("Content-Length", strconv.Itoa(len(requestPayload))), // added by the Go client
			mock_prototool.NewHeaderKV(httpz.ViaHeader, "gRPC/1.0 sv1"),
			mock_prototool.NewHeaderKV(httpz.OriginHeader, "kas.gitlab.example.com"),
			mock_prototool.NewHeaderKV(mock_prototool.ValidHeaderName(), mock_prototool.ValidHeaderValue()),
		},
		UrlPath: requestPath,
		Query: []*prototool.QueryKV{
			mock_prototool.NewQueryKV(queryParamName, queryParamValue),
		},
	}

	type testcase struct {
		name                  string
		authn                 func(*http.Request)
		handler               func(*testing.T) func(http.ResponseWriter, *http.Request)
		request               *prototool.HttpRequest
		response              *prototool.HttpResponse
		expectedErr           string
		expectedStatus        metav1.Status
		expectedStatusCode    int
		expectedMessagePrefix string
	}
	var testcases []testcase
	addRedirectTC := func(status int, authType string, authn func(*http.Request)) testcase {
		return testcase{
			name:    fmt.Sprintf("%s/redirect/%d", authType, status),
			authn:   authn,
			handler: userHandler,
			request: req,
			response: &prototool.HttpResponse{
				StatusCode: int32(status),
				Status:     http.StatusText(status),
				Header: []*prototool.HeaderKV{
					mock_prototool.NewHeaderKV("Date", "NOW!"),
					mock_prototool.NewHeaderKV("Location", "https://example.com/"),
				},
			},
			expectedErr: "redirects are not allowed",
			expectedStatus: metav1.Status{
				TypeMeta: metav1.TypeMeta{
					Kind:       "Status",
					APIVersion: "v1",
				},
				Status: metav1.StatusFailure,
				Reason: metav1.StatusReasonInternalError,
				Code:   http.StatusBadGateway,
			},
			expectedStatusCode:    http.StatusBadGateway,
			expectedMessagePrefix: "GitLab Agent Server: HTTP->gRPC: bad server response: redirects are not allowed",
		}
	}

	for status := 300; status < 400; status++ {
		testcases = append(testcases,
			addRedirectTC(status, "auth in header", setExpectedSessionCookieParamsInHeader),
			addRedirectTC(status, "auth in params", setExpectedSessionCookieParamsInQuery),
		)
	}

	addBadContentTypeTC := func(contentType []string, errMsg string, authType string, authn func(*http.Request)) testcase {
		return testcase{
			name:    fmt.Sprintf("%s/%s", authType, contentType),
			authn:   authn,
			handler: userHandler,
			request: req,
			response: &prototool.HttpResponse{
				StatusCode: http.StatusOK,
				Status:     http.StatusText(http.StatusOK),
				Header: []*prototool.HeaderKV{
					mock_prototool.NewHeaderKV("Date", "NOW!"),
					mock_prototool.NewHeaderKV(httpz.ContentTypeHeader, contentType...),
				},
			},
			expectedErr: errMsg,
			expectedStatus: metav1.Status{
				TypeMeta: metav1.TypeMeta{
					Kind:       "Status",
					APIVersion: "v1",
				},
				Status: metav1.StatusFailure,
				Reason: metav1.StatusReasonInternalError,
				Code:   http.StatusBadGateway,
			},
			expectedStatusCode:    http.StatusBadGateway,
			expectedMessagePrefix: "GitLab Agent Server: HTTP->gRPC: bad server response: " + errMsg + ". Trace ID: ",
		}
	}

	ctCases := [][]string{
		{"text/html"},
		{"text/html; encoding=utf8"},
		// two headers. Invalid, but we must ensure our code does not get confused somehow
		{"application/json", "text/html"},
		{"text/html", "application/json"},
		{"text/html; encoding=utf8", "application/json"},
		{"application/json", "text/html; encoding=utf8"},
		{"application/json; encoding=utf8", "text/html"},
		{"application/json; encoding=utf8", "text/html; encoding=utf8"},
	}

	for _, tc := range ctCases {
		testcases = append(testcases,
			addBadContentTypeTC(tc, "Content-Type not allowed: text/html", "auth in header", setExpectedSessionCookieParamsInHeader),
			addBadContentTypeTC(tc, "Content-Type not allowed: text/html", "auth in params", setExpectedSessionCookieParamsInQuery),
		)
	}

	invalidCTCases := [][]string{
		{"123=;;;444-22"}, // invalid
		// two headers, one invalid. Invalid, but we must ensure our code does not get confused somehow
		{"123=;;;444-22", "application/json"},
	}

	for _, tc := range invalidCTCases {
		testcases = append(testcases,
			addBadContentTypeTC(tc, "check Content-Type: mime: expected slash after first token", "auth in header", setExpectedSessionCookieParamsInHeader),
			addBadContentTypeTC(tc, "check Content-Type: mime: expected slash after first token", "auth in params", setExpectedSessionCookieParamsInQuery),
		)
	}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			api, mrClient, client, req, requestCount, _, _, _, mockUserAccessRequestCounter, mockUserAccessAgentsCounter, mockUserAccessEventTracker, _, _, _, mockRequestsTotalCounter, _, mockRequestsErrorCounter := setupProxyWithHandler(t, "/", tc.handler(t))
			mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockRequestsErrorCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockUserAccessRequestCounter.EXPECT().Inc()
			mockUserAccessAgentsCounter.EXPECT().Add(testhelpers.AgentkKey1.ID)
			mockUserAccessEventTracker.EXPECT().EmitEvent(K8sAPIProxyRequestsEvent{UserID: testhelpers.UserID, ProjectID: testhelpers.ConfigProjectID})
			requestCount.EXPECT().Inc()
			api.EXPECT().
				HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), matcher.ErrorEq(tc.expectedErr), logz.AgentKey(testhelpers.AgentkKey1))
			contentLength := int64(len(requestPayload))
			send := mockSendStream(t, mrClient, false,
				&grpctool.HttpRequest{
					Message: &grpctool.HttpRequest_Header_{
						Header: &grpctool.HttpRequest_Header{
							Request:       tc.request,
							ContentLength: &contentLength,
						},
					},
				},
				&grpctool.HttpRequest{
					Message: &grpctool.HttpRequest_Data_{
						Data: &grpctool.HttpRequest_Data{
							Data: []byte(requestPayload),
						},
					},
				},
				&grpctool.HttpRequest{
					Message: &grpctool.HttpRequest_Trailer_{
						Trailer: &grpctool.HttpRequest_Trailer{},
					},
				},
			)
			recv := mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Do(testhelpers.RecvMsg(&grpctool.HttpResponse{
					Message: &grpctool.HttpResponse_Header_{
						Header: &grpctool.HttpResponse_Header{
							Response: tc.response,
						},
					},
				}))
			var calls []any //nolint:prealloc
			calls = append(calls, send...)
			calls = append(calls, recv)
			gomock.InOrder(calls...)
			tc.authn(req)
			req.Header.Set(httpz.OriginHeader, "kas.gitlab.example.com")
			req.Header.Set(httpz.UserAgentHeader, "test-agent") // added manually to override what is added by the Go client
			resp, err := client.Do(req)
			require.NoError(t, err)
			defer func() {
				assert.NoError(t, resp.Body.Close())
			}()
			assert.Equal(t, tc.expectedStatusCode, resp.StatusCode)
			actualStatus := readStatus(t, resp)
			assert.True(t, strings.HasPrefix(actualStatus.Message, tc.expectedMessagePrefix), actualStatus.Message)
			assert.Empty(t, cmp.Diff(tc.expectedStatus, actualStatus, cmpopts.IgnoreFields(metav1.Status{}, "Message")))
		})
	}
}

func configUserAccessGitLabHandler(t *testing.T, auth *gapi.AuthorizeProxyUserAPIResponse) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		if !assertUserAccessCredentials(t, req) {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		testhelpers.RespondWithJSON(t, w, auth)
	}
}

func assertUserAccessCredentials(t *testing.T, req *http.Request) bool {
	data, err := io.ReadAll(req.Body)
	assert.NoError(t, err)
	auth := &gapi.AuthorizeProxyUserAPIRequest{}
	err = protojson.Unmarshal(data, auth)
	return assert.NoError(t, err) &&
		assert.Equal(t, testhelpers.AgentkKey1.ID, auth.AgentId) &&
		assert.Equal(t, "session_cookie", auth.AccessType) &&
		assert.Equal(t, "encrypted-session-cookie", auth.AccessKey) &&
		assert.Equal(t, "the-csrf-token", auth.CsrfToken)
}

func setExpectedSessionCookieParamsInHeader(req *http.Request) {
	req.AddCookie(
		&http.Cookie{
			Name:  gitLabKASCookieName,
			Value: "encrypted-session-cookie",
		},
	)
	req.Header[httpz.GitlabAgentIDHeader] = []string{strconv.FormatInt(testhelpers.AgentkKey1.ID, 10)}
	req.Header[httpz.CSRFTokenHeader] = []string{"the-csrf-token"}
}

func setExpectedSessionCookieParamsInQuery(req *http.Request) {
	req.AddCookie(
		&http.Cookie{
			Name:  gitLabKASCookieName,
			Value: "encrypted-session-cookie",
		},
	)
	query := req.URL.Query()
	query[httpz.GitlabAgentIDQueryParam] = []string{strconv.FormatInt(testhelpers.AgentkKey1.ID, 10)}
	query[httpz.CSRFTokenQueryParam] = []string{"the-csrf-token"}
	req.URL.RawQuery = query.Encode()
}

func TestProxy_PersonalAccessTokenAccessHappyPath(t *testing.T) {
	testcases := []struct {
		name          string
		urlPathPrefix string
		auth          *gapi.AuthorizeProxyUserAPIResponse
		expectedExtra *rpc.HeaderExtra
	}{
		{
			name:          "impersonate agent",
			urlPathPrefix: "/",
			auth: &gapi.AuthorizeProxyUserAPIResponse{
				AccessAs: &gapi.AccessAsProxyAuthorization{
					AccessAs: &gapi.AccessAsProxyAuthorization_Agent{
						Agent: &gapi.AccessAsAgentAuthorization{},
					},
				},
			},
		},
		{
			name:          "impersonate user",
			urlPathPrefix: "/",
			auth: &gapi.AuthorizeProxyUserAPIResponse{
				AccessAs: &gapi.AccessAsProxyAuthorization{
					AccessAs: &gapi.AccessAsProxyAuthorization_User{
						User: &gapi.AccessAsUserAuthorization{
							Projects: []*gapi.ProjectAccessCF{
								{
									Id:    1234,
									Roles: []string{"guest", "developer"},
								},
							},
						},
					},
				},
			},
			expectedExtra: &rpc.HeaderExtra{
				ImpConfig: &kubernetes_api.ImpersonationConfig{
					Username: "gitlab:user:testuser",
					Groups: []string{
						"gitlab:user",
						"gitlab:project_role:1234:guest",
						"gitlab:project_role:1234:developer",
					},
					Extra: []*kubernetes_api.ExtraKeyVal{
						{
							Key: "agent.gitlab.com/id",
							Val: []string{"123"},
						},
						{
							Key: "agent.gitlab.com/username",
							Val: []string{"testuser"},
						},
						{
							Key: "agent.gitlab.com/access_type",
							Val: []string{"personal_access_token"},
						},
						{
							Key: "agent.gitlab.com/config_project_id",
							Val: []string{"5"},
						},
					},
				},
			},
		},
	}
	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			tc.auth.Agent = &gapi.AuthorizedAgentForUser{
				Id:            testhelpers.AgentkKey1.ID,
				ConfigProject: &gapi.ConfigProject{Id: 5},
			}
			tc.auth.User = &gapi.User{
				Id:       testhelpers.UserID,
				Username: "testuser",
			}
			_, mrClient, client, req, requestCount, _, _, _, _, _, _, mockPatAccessRequestCounter, mockPatAccessAgentsCounter, mockPatAccessEventTracker, mockRequestsTotalCounter, mockRequestsSuccessCounter, _ := setupProxyWithHandler(t, tc.urlPathPrefix, configPersonalAccessTokenAccessGitLabHandler(t, tc.auth))
			mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockRequestsSuccessCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockPatAccessRequestCounter.EXPECT().Inc()
			mockPatAccessAgentsCounter.EXPECT().Add(testhelpers.AgentkKey1.ID)
			mockPatAccessEventTracker.EXPECT().EmitEvent(K8sAPIProxyRequestsEvent{UserID: testhelpers.UserID, ProjectID: testhelpers.ConfigProjectID})
			testProxyHappyPath(t, setExpectedPersonalAccessToken, tc.expectedExtra, mrClient, client, req, requestCount, false)
		})
	}
}

func TestProxy_WebSocketTokenAuthHappyPath(t *testing.T) {
	testcases := []struct {
		name      string
		impConfig *kubernetes_api.ImpersonationConfig
	}{
		{
			name: "impersonate agent",
		},
		{
			name: "impersonate user",
			impConfig: &kubernetes_api.ImpersonationConfig{
				Username: "gitlab:user:testuser",
				Groups: []string{
					"gitlab:user",
					"gitlab:project_role:1234:guest",
					"gitlab:project_role:1234:developer",
				},
				Extra: []*kubernetes_api.ExtraKeyVal{
					{
						Key: "agent.gitlab.com/id",
						Val: []string{"123"},
					},
					{
						Key: "agent.gitlab.com/username",
						Val: []string{"testuser"},
					},
					{
						Key: "agent.gitlab.com/access_type",
						Val: []string{"personal_access_token"},
					},
					{
						Key: "agent.gitlab.com/config_project_id",
						Val: []string{"5"},
					},
				},
			},
		},
	}
	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			_, mrClient, client, req, requestCount, _, _, _, _, _, _, _, _, _, mockRequestsTotalCounter, mockRequestsSuccessCounter, _ := setupProxyWithHandler(t, "/", func(w http.ResponseWriter, r *http.Request) {
				t.Fatal("not expecting any GitLab API calls")
			})
			mockRequestsTotalCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			mockRequestsSuccessCounter.EXPECT().Add(gomock.Any(), int64(1), gomock.Any())
			var expectedExtra *rpc.HeaderExtra
			if tc.impConfig != nil {
				expectedExtra = &rpc.HeaderExtra{ImpConfig: tc.impConfig}
			}
			testProxyHappyPath(t, func(r *http.Request) {
				et := encodeTokenWithPrefix(generateValidTestWebSocketToken(t, testHost+requestPath, tc.impConfig))
				r.Header[httpz.ConnectionHeader] = []string{"Upgrade"}
				r.Header[httpz.UpgradeHeader] = []string{"WebSocket"}
				r.Header[httpz.SecWebSocketProtocolHeader] = []string{et}
			}, expectedExtra, mrClient, client, req, requestCount, true)
		})
	}
}

func setExpectedPersonalAccessToken(req *http.Request) {
	req.Header.Set(httpz.AuthorizationHeader, fmt.Sprintf("Bearer %s:%d:%s", tokenTypePat, testhelpers.AgentkKey1.ID, personalAccessToken))
}

func configPersonalAccessTokenAccessGitLabHandler(t *testing.T, auth *gapi.AuthorizeProxyUserAPIResponse) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, req *http.Request) {
		if !assertPersonalAccessTokenAccessCredentials(t, req) {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		testhelpers.RespondWithJSON(t, w, auth)
	}
}

func assertPersonalAccessTokenAccessCredentials(t *testing.T, req *http.Request) bool {
	data, err := io.ReadAll(req.Body)
	assert.NoError(t, err)
	auth := &gapi.AuthorizeProxyUserAPIRequest{}
	err = protojson.Unmarshal(data, auth)
	return assert.NoError(t, err) &&
		assert.Equal(t, testhelpers.AgentkKey1.ID, auth.AgentId) &&
		assert.Equal(t, "personal_access_token", auth.AccessType) &&
		assert.Equal(t, personalAccessToken, auth.AccessKey) &&
		assert.Empty(t, auth.CsrfToken)
}

func testProxyHappyPath(t *testing.T, prepareRequest func(*http.Request), expectedExtra *rpc.HeaderExtra,
	mrClient *mock_kubernetes_api.MockKubernetesApi_MakeRequestClient, client *http.Client, req *http.Request,
	requestCount *mock_usage_metrics.MockCounter, expectWebSocketHeaders bool) {
	prepareRequest(req)
	requestCount.EXPECT().Inc()
	var extra *anypb.Any
	if expectedExtra != nil {
		var err error
		extra, err = anypb.New(expectedExtra)
		require.NoError(t, err)
	}
	contentLength := int64(len(requestPayload))
	sendMockHeaders := []*prototool.HeaderKV{
		mock_prototool.NewHeaderKV("Req-Header", "x1", "x2"),
		mock_prototool.NewHeaderKV("Accept-Encoding", "gzip"), // added by the Go client
		mock_prototool.NewHeaderKV(httpz.UserAgentHeader, "test-agent"),
		mock_prototool.NewHeaderKV("Content-Length", strconv.Itoa(len(requestPayload))), // added by the Go client
		mock_prototool.NewHeaderKV(httpz.ViaHeader, "gRPC/1.0 sv1"),
		mock_prototool.NewHeaderKV(httpz.OriginHeader, "kas.gitlab.example.com"),
		mock_prototool.NewHeaderKV(mock_prototool.ValidHeaderName(), mock_prototool.ValidHeaderValue()),
	}
	if expectWebSocketHeaders {
		sendMockHeaders = append(sendMockHeaders,
			mock_prototool.NewHeaderKV(httpz.ConnectionHeader, "upgrade"),
			mock_prototool.NewHeaderKV(httpz.UpgradeHeader, "WebSocket"),
		)
	}
	reqHeader := &grpctool.HttpRequest{
		Message: &grpctool.HttpRequest_Header_{
			Header: &grpctool.HttpRequest_Header{
				Request: &prototool.HttpRequest{
					Method:  http.MethodPost,
					Header:  sendMockHeaders,
					UrlPath: requestPath,
					Query: []*prototool.QueryKV{
						mock_prototool.NewQueryKV(queryParamName, queryParamValue),
					},
				},
				Extra:         extra,
				ContentLength: &contentLength,
			},
		},
	}
	reqData := &grpctool.HttpRequest{
		Message: &grpctool.HttpRequest_Data_{
			Data: &grpctool.HttpRequest_Data{
				Data: []byte(requestPayload),
			},
		},
	}
	reqTrailer := &grpctool.HttpRequest{
		Message: &grpctool.HttpRequest_Trailer_{
			Trailer: &grpctool.HttpRequest_Trailer{},
		},
	}
	calls := []any{
		mrClient.EXPECT().
			Send(mock_prototool.ProtoEq(t, reqHeader)),
		mrClient.EXPECT().
			Send(mock_prototool.ProtoEq(t, reqData)),
		mrClient.EXPECT().
			Send(mock_prototool.ProtoEq(t, reqTrailer)),
	}
	respHeader := &grpctool.HttpResponse{
		Message: &grpctool.HttpResponse_Header_{
			Header: &grpctool.HttpResponse_Header{
				Response: &prototool.HttpResponse{
					StatusCode: http.StatusOK, // respond with 200 even for an upgrade request
					Status:     http.StatusText(http.StatusOK),
					Header: []*prototool.HeaderKV{
						mock_prototool.NewHeaderKV("Content-Type", "application/json; encoding=utf-8"),
						mock_prototool.NewHeaderKV("Date", "NOW!"),
						// These headers are CORS headers and are being removed by the proxy
						mock_prototool.NewHeaderKV(httpz.AccessControlMaxAgeHeader, "42"),
					},
				},
			},
		},
	}
	respData := &grpctool.HttpResponse{
		Message: &grpctool.HttpResponse_Data_{
			Data: &grpctool.HttpResponse_Data{
				Data: []byte(responsePayload),
			},
		},
	}
	respTrailer := &grpctool.HttpResponse{
		Message: &grpctool.HttpResponse_Trailer_{
			Trailer: &grpctool.HttpResponse_Trailer{},
		},
	}
	if expectWebSocketHeaders {
		calls = append(calls,
			mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Do(testhelpers.RecvMsg(respHeader)),
			mrClient.EXPECT().
				CloseSend(),
			mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Do(testhelpers.RecvMsg(respData)),
			mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Do(testhelpers.RecvMsg(respTrailer)),
			mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Return(io.EOF),
		)
	} else {
		calls = append(calls,
			mrClient.EXPECT().
				CloseSend(),
			mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Do(testhelpers.RecvMsg(respHeader)),
			mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Do(testhelpers.RecvMsg(respData)),
			mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Do(testhelpers.RecvMsg(respTrailer)),
			mrClient.EXPECT().
				RecvMsg(gomock.Any()).
				Return(io.EOF),
		)
	}
	gomock.InOrder(calls...)

	req.Header.Set("Req-Header", "x1")
	req.Header.Add("Req-Header", "x2")
	req.Header.Set(httpz.UserAgentHeader, "test-agent") // added manually to override what is added by the Go client
	resp, err := client.Do(req)
	require.NoError(t, err)
	defer func() {
		assert.NoError(t, resp.Body.Close())
	}()
	assert.Equal(t, http.StatusOK, resp.StatusCode)
	assert.JSONEq(t, responsePayload, string(readAll(t, resp.Body)))
	delete(resp.Header, "Date")
	expectedHeader := http.Header{
		"Content-Type":                     {"application/json; encoding=utf-8"},
		"Via":                              {"gRPC/1.0 sv1"},
		"Access-Control-Allow-Credentials": {"true"},
		"Access-Control-Allow-Origin":      {"kas.gitlab.example.com"},
		"Vary":                             {"Origin"},
		"Nel":                              {`{"max_age": 0}`}, // see https://gitlab.com/gitlab-org/gitlab/-/issues/504707
		"X-Content-Type-Options":           {"nosniff"},        // https://fetch.spec.whatwg.org/#x-content-type-options-header
	}
	assert.Empty(t, cmp.Diff(expectedHeader, resp.Header))
}

func TestFormatStatusMessage(t *testing.T) {
	ctx, traceID := testhelpers.CtxWithSpanContext(t)
	tests := []struct {
		name            string
		ctx             context.Context //nolint:containedctx
		err             error
		expectedMessage string
	}{
		{
			name:            "no err, no trace",
			ctx:             context.Background(),
			expectedMessage: "GitLab Agent Server: msg",
		},
		{
			name:            "no err, trace",
			ctx:             ctx,
			expectedMessage: "GitLab Agent Server: msg. Trace ID: " + traceID.String(),
		},
		{
			name:            "err, no trace",
			ctx:             context.Background(),
			err:             errors.New("boom"),
			expectedMessage: "GitLab Agent Server: msg: boom",
		},
		{
			name:            "err, trace",
			ctx:             ctx,
			err:             errors.New("boom"),
			expectedMessage: "GitLab Agent Server: msg: boom. Trace ID: " + traceID.String(),
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			msg := formatStatusMessage(test.ctx, "msg", test.err)
			assert.Equal(t, test.expectedMessage, msg)
		})
	}
}

func assertToken(t *testing.T, r *http.Request) bool {
	return assert.Equal(t, jobToken, r.Header.Get("Job-Token"))
}

func setupProxy(t *testing.T) (
	*mock_modserver.MockAPI, *mock_kubernetes_api.MockKubernetesApi_MakeRequestClient, *http.Client, *http.Request,
	*mock_usage_metrics.MockCounter, *mock_usage_metrics.MockCounter, *mock_usage_metrics.MockUniqueCounter,
	*mock_event_tracker.MockEventsInterface, *mock_usage_metrics.MockCounter, *mock_usage_metrics.MockUniqueCounter,
	*mock_event_tracker.MockEventsInterface, *mock_usage_metrics.MockCounter, *mock_usage_metrics.MockUniqueCounter,
	*mock_event_tracker.MockEventsInterface,
	*mock_otel.MockInt64Counter, *mock_otel.MockInt64Counter, *mock_otel.MockInt64Counter) {
	return setupProxyWithHandler(t, "/", configCIAccessGitLabHandler(t, nil, nil))
}

func configCIAccessGitLabHandler(t *testing.T, config *gapi.Configuration, env *gapi.Environment) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		if !assertToken(t, r) {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		testhelpers.RespondWithJSON(t, w, &gapi.AllowedAgentsForJobAPIResponse{
			AllowedAgents: []*gapi.AllowedAgent{
				{
					Id: testhelpers.AgentkKey1.ID,
					ConfigProject: &gapi.ConfigProject{
						Id: 5,
					},
					Configuration: config,
				},
			},
			Job: &gapi.Job{
				Id: 1,
			},
			Pipeline: &gapi.Pipeline{
				Id: 2,
			},
			Project: &gapi.Project{
				Id: 3,
				Groups: []*gapi.Group{
					{
						Id: 6,
					},
				},
			},
			User: &gapi.User{
				Id:       testhelpers.UserID,
				Username: "testuser",
			},
			Environment: env,
		})
	}
}

func setupProxyWithHandler(t *testing.T, urlPathPrefix string, handler func(http.ResponseWriter, *http.Request)) (
	*mock_modserver.MockAPI, *mock_kubernetes_api.MockKubernetesApi_MakeRequestClient, *http.Client, *http.Request,
	*mock_usage_metrics.MockCounter, *mock_usage_metrics.MockCounter, *mock_usage_metrics.MockUniqueCounter,
	*mock_event_tracker.MockEventsInterface, *mock_usage_metrics.MockCounter, *mock_usage_metrics.MockUniqueCounter,
	*mock_event_tracker.MockEventsInterface, *mock_usage_metrics.MockCounter, *mock_usage_metrics.MockUniqueCounter,
	*mock_event_tracker.MockEventsInterface,
	*mock_otel.MockInt64Counter, *mock_otel.MockInt64Counter, *mock_otel.MockInt64Counter) {
	ctrl := gomock.NewController(t)
	mockAPI := mock_modserver.NewMockAPI(ctrl)
	mrClient := mock_kubernetes_api.NewMockKubernetesApi_MakeRequestClient(gomock.NewController(t))
	requestCount := mock_usage_metrics.NewMockCounter(ctrl)
	mockCIAccessRequestCounter := mock_usage_metrics.NewMockCounter(ctrl)
	mockCIAccessAgentsCounter := mock_usage_metrics.NewMockUniqueCounter(ctrl)
	mockCIAccessEventTracker := mock_event_tracker.NewMockEventsInterface(ctrl)
	mockUserAccessRequestCounter := mock_usage_metrics.NewMockCounter(ctrl)
	mockUserAccessAgentsCounter := mock_usage_metrics.NewMockUniqueCounter(ctrl)
	mockUserAccessEventTracker := mock_event_tracker.NewMockEventsInterface(ctrl)
	mockPatAccessRequestCounter := mock_usage_metrics.NewMockCounter(ctrl)
	mockPatAccessAgentsCounter := mock_usage_metrics.NewMockUniqueCounter(ctrl)
	mockPatAccessEventTracker := mock_event_tracker.NewMockEventsInterface(ctrl)
	mockRequestsTotalCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockRequestsSuccessCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockRequestsErrorCounter := mock_otel.NewMockInt64Counter(ctrl)
	errCache := mock_cache.NewMockErrCacher[string](ctrl)
	proxyErrCache := mock_cache.NewMockErrCacher[proxyUserCacheKey](ctrl)

	tp := tracesdk.NewTracerProvider(tracesdk.WithSpanProcessor(tracetest.NewSpanRecorder()))
	tracer := tp.Tracer(kubernetes_api.ModuleName)

	v, err := protovalidate.New()
	require.NoError(t, err)

	// We don't really care about the feature flag counter here, so we just allow as many calls to it as it does
	mockFeatureFlagCheckCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockFeatureFlagCheckCounter.EXPECT().Add(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()

	p := kubernetesAPIProxy{
		log: testlogger.New(t),
		api: mockAPI,
		makeRequest: func(ctx context.Context, agentKey api.AgentKey) (grpctool.HTTPRequestClient, error) {
			return mrClient, nil
		},
		gitLabClient:      mock_gitlab.SetupClient(t, "", "/", handler, gitlab.WithFeatureFlagCheckCounter(mockFeatureFlagCheckCounter)),
		allowedOriginURLs: []string{"kas.gitlab.example.com"},
		allowedAgentsCache: cache.NewWithError[string, *gapi.AllowedAgentsForJob](0, 0, errCache, tracer,
			func(err error) bool { return false }),
		authorizeProxyUserCache: cache.NewWithError[proxyUserCacheKey, *gapi.AuthorizeProxyUserResponse](0, 0, proxyErrCache, tracer,
			func(err error) bool { return false }),
		utRequestCounter:           requestCount,
		utCIAccessRequestCounter:   mockCIAccessRequestCounter,
		utCIAccessAgentsCounter:    mockCIAccessAgentsCounter,
		utCIAccessEventTracker:     mockCIAccessEventTracker,
		utUserAccessRequestCounter: mockUserAccessRequestCounter,
		utUserAccessAgentsCounter:  mockUserAccessAgentsCounter,
		utUserAccessEventTracker:   mockUserAccessEventTracker,
		utPatAccessRequestCounter:  mockPatAccessRequestCounter,
		utPatAccessAgentsCounter:   mockPatAccessAgentsCounter,
		utPatAccessEventTracker:    mockPatAccessEventTracker,
		requestsTotalCounter:       mockRequestsTotalCounter,
		requestsSuccessCounter:     mockRequestsSuccessCounter,
		requestsErrorCounter:       mockRequestsErrorCounter,
		responseSerializer:         serializer.NewCodecFactory(runtime.NewScheme()),
		traceProvider:              tp,
		tracePropagator:            propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}, propagation.Baggage{}),
		meterProvider:              metricnoop.NewMeterProvider(),
		validator:                  v,
		serverName:                 "sv1",
		serverVia:                  "gRPC/1.0 sv1",
		urlPathPrefix:              urlPathPrefix,
		httpServerTracing:          true,
		webSocketToken: &webSocketToken{
			validator:        v,
			jwtSigningMethod: jwttool.SigningMethodHS3_512,
			jwtSecret:        validWebSocketTokenSecret,
			ffCheckCounter:   mockFeatureFlagCheckCounter,
		},
		allowedResponseHeaders: defaultAllowedResponseHeaders,
	}
	listener := grpcz.NewPipeDialListener()
	var wg wait.Group
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(func() {
		cancel()
		wg.Wait()
		listener.Close()
	})
	wg.Start(func() {
		assert.NoError(t, p.Run(ctx, listener))
	})
	client := &http.Client{
		Transport: &http.Transport{
			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
				return listener.DialContext(ctx, addr)
			},
		},
	}
	req, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		"http://"+testHost+path.Join(urlPathPrefix, requestPath)+"?"+url.QueryEscape(queryParamName)+"="+url.QueryEscape(queryParamValue),
		strings.NewReader(requestPayload),
	)
	require.NoError(t, err)
	req.Header.Set(httpz.OriginHeader, "kas.gitlab.example.com")
	req.Header.Set(mock_prototool.ValidHeaderName(), mock_prototool.ValidHeaderValue())
	return mockAPI, mrClient, client, req, requestCount,
		mockCIAccessRequestCounter, mockCIAccessAgentsCounter, mockCIAccessEventTracker,
		mockUserAccessRequestCounter, mockUserAccessAgentsCounter, mockUserAccessEventTracker,
		mockPatAccessRequestCounter, mockPatAccessAgentsCounter, mockPatAccessEventTracker,
		mockRequestsTotalCounter, mockRequestsSuccessCounter, mockRequestsErrorCounter
}

func mockSendStream(t *testing.T, client *mock_kubernetes_api.MockKubernetesApi_MakeRequestClient, isUpgrade bool, msgs ...*grpctool.HttpRequest) []any {
	res := make([]any, 0, len(msgs)+1)
	for _, msg := range msgs {
		call := client.EXPECT().
			Send(mock_prototool.ProtoEq(t, msg))
		res = append(res, call)
	}
	if !isUpgrade {
		res = append(res, client.EXPECT().CloseSend())
	}
	return res
}

func readAll(t *testing.T, r io.Reader) []byte {
	data, err := io.ReadAll(r)
	require.NoError(t, err)
	return data
}

func readStatus(t *testing.T, resp *http.Response) metav1.Status {
	data := readAll(t, resp.Body)
	var s metav1.Status
	negotiator := runtime.NewClientNegotiator(serializer.NewCodecFactory(runtime.NewScheme()), schema.GroupVersion{})
	decoder, err := negotiator.Decoder(resp.Header.Get(httpz.ContentTypeHeader), nil)
	require.NoError(t, err)
	obj, _, err := decoder.Decode(data, nil, &s)
	require.NoError(t, err)
	return *obj.(*metav1.Status)
}

func Test_MergeProxiedResponseHeaders(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name                   string
		outboundHeaders        http.Header
		inboundHeaders         http.Header
		expectedInboundHeaders http.Header
	}{
		{
			name:            "no outbound, no inbound headers -> expect default headers",
			outboundHeaders: http.Header{},
			inboundHeaders:  http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`}, // see https://gitlab.com/gitlab-org/gitlab/-/issues/504707
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "remove CORS headers from outbound headers",
			outboundHeaders: http.Header{
				"Access-Control-Allow-Origin":      []string{"any"},
				"Access-Control-Allow-Methods":     []string{"any"},
				"Access-Control-Allow-Headers":     []string{"any"},
				"Access-Control-Allow-Credentials": []string{"any"},
				"Access-Control-Max-Age":           []string{"any"},
			},
			inboundHeaders: http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`}, // see https://gitlab.com/gitlab-org/gitlab/-/issues/504707
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "don't overwrite inbound with outbound headers",
			outboundHeaders: http.Header{
				"Date": []string{"overwrite"},
			},
			inboundHeaders: http.Header{
				"Date": []string{"expected-to-see-this"},
			},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`}, // see https://gitlab.com/gitlab-org/gitlab/-/issues/504707
				"Via":                    []string{"test-via"},
				"Date":                   []string{"expected-to-see-this"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "proxy Via header is appended to the outbound Via headers",
			outboundHeaders: http.Header{
				"Via": []string{"any-via"},
			},
			inboundHeaders: http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`}, // see https://gitlab.com/gitlab-org/gitlab/-/issues/504707
				"Via":                    []string{"any-via", "test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "append outbound Vary headers from inbound",
			outboundHeaders: http.Header{
				"Vary": []string{"any-vary"},
			},
			inboundHeaders: http.Header{
				"Vary": []string{"Origin"},
			},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`}, // see https://gitlab.com/gitlab-org/gitlab/-/issues/504707
				"Via":                    []string{"test-via"},
				"Vary":                   []string{"Origin", "any-vary"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name:            "always send NEL policy remove payload",
			outboundHeaders: http.Header{},
			inboundHeaders:  http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`},
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "clear existing NEL headers from outbound requests and send NEL policy remove payload",
			outboundHeaders: http.Header{
				"Nel": []string{`{"max_age": 42}`},
			},
			inboundHeaders: http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`},
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "clear existing Report-To headers from outbound requests and send NEL policy remove payload",
			outboundHeaders: http.Header{
				"Nel":       []string{`{"max_age": 42}`},
				"Report-To": []string{`{}`},
			},
			inboundHeaders: http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`},
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "clear existing caching related headers from outbound requests and send 'disabling' caching control headers instead",
			outboundHeaders: http.Header{
				"Cache-Control":                []string{"max-age=9999", "private=x-content-type-options"},
				"Surrogate-Control":            []string{"max-age=9999", "private=x-content-type-options"},
				"Cdn-Cache-Control":            []string{"max-age=9999", "private=x-content-type-options"},
				"Cloudflare-Cdn-Cache-Control": []string{"max-age=9999", "private=x-content-type-options"},
				"Pragma":                       []string{""},
				"Expires":                      []string{"999"},
			},
			inboundHeaders: http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`},
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "remove nginx internal redirect header",
			outboundHeaders: http.Header{
				"X-Accel-Redirect": []string{"some-value"},
			},
			inboundHeaders: http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`},
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name:            "remove own Server header",
			outboundHeaders: http.Header{},
			inboundHeaders: http.Header{
				httpz.ServerHeader: []string{"kas"},
			},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`},
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "keep origin Server header",
			outboundHeaders: http.Header{
				httpz.ServerHeader: []string{"kube"},
			},
			inboundHeaders: http.Header{
				httpz.ServerHeader: []string{"kas"},
			},
			expectedInboundHeaders: http.Header{
				httpz.ServerHeader:       []string{"kube"},
				"Nel":                    []string{`{"max_age": 0}`},
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
		{
			name: "remove Envoy headers",
			outboundHeaders: http.Header{
				"X-Envoy-Upstream-Service-Time": []string{"123"},
				"X-Envoy-Decorator-Operation":   []string{"bla"},
			},
			inboundHeaders: http.Header{},
			expectedInboundHeaders: http.Header{
				"Nel":                    []string{`{"max_age": 0}`},
				"Via":                    []string{"test-via"},
				"X-Content-Type-Options": []string{"nosniff"}, // https://fetch.spec.whatwg.org/#x-content-type-options-header
			},
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			ctrl := gomock.NewController(t)
			mockAPI := mock_modserver.NewMockAPI(ctrl)
			mockAPI.EXPECT().
				HandleProcessingError(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
				AnyTimes()

			log := testlogger.New(t)
			p := kubernetesAPIProxy{
				log: log,
				api: mockAPI,
				// NOTE: p.serverVia is the only field accessed by the function under test
				serverVia:              "test-via",
				allowedResponseHeaders: defaultAllowedResponseHeaders,
			}

			i, o := tc.inboundHeaders.Clone(), tc.outboundHeaders.Clone()

			p.mergeProxiedResponseHeaders(log, testhelpers.AgentkKey1, o, i)

			assert.Equal(t, tc.expectedInboundHeaders, i)
		})
	}
}

func Test_MergeProxiedResponseHeaders_ReportBlockedHeaders(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modserver.NewMockAPI(ctrl)
	log := testlogger.New(t)
	p := kubernetesAPIProxy{
		log:                    log,
		api:                    mockAPI,
		allowedResponseHeaders: sets.New[string](), // for testing purposes we don't allow any headers, which means that every header will be blocked
	}
	outbound := http.Header{
		httpz.SetCookieHeader: []string{"any"},
	}
	inbound := http.Header{}

	expectedFields := append(
		[]slog.Attr{logz.AgentKey(testhelpers.AgentkKey1)},
		logz.ResponseHeaders(http.Header{"Set-Cookie": []string{"any"}}),
	)
	mockAPI.EXPECT().HandleProcessingError(
		gomock.Any(),
		gomock.Any(),
		"",
		errBlockedResponseHeader,
		gomock.InAnyOrder(expectedFields),
	)

	// WHEN
	p.mergeProxiedResponseHeaders(log, testhelpers.AgentkKey1, outbound, inbound)

	// THEN
	assert.NotContains(t, inbound, httpz.SetCookieHeader)
}

func Test_MergeProxiedResponseHeaders_ReportMultipleBlockedHeadersAllInOne(t *testing.T) {
	// GIVEN
	ctrl := gomock.NewController(t)
	mockAPI := mock_modserver.NewMockAPI(ctrl)
	log := testlogger.New(t)
	p := kubernetesAPIProxy{
		log:                    log,
		api:                    mockAPI,
		allowedResponseHeaders: sets.New[string](), // for testing purposes we don't allow any headers, which means that every header will be blocked
	}
	outbound := http.Header{
		httpz.SetCookieHeader: []string{"any-1"},
		httpz.TrailerHeader:   []string{"any-2"},
	}
	inbound := http.Header{}

	expectedFields := append(
		[]slog.Attr{logz.AgentKey(testhelpers.AgentkKey1)},
		logz.ResponseHeaders(http.Header{
			httpz.SetCookieHeader: []string{"any-1"},
			httpz.TrailerHeader:   []string{"any-2"},
		}),
	)
	mockAPI.EXPECT().HandleProcessingError(
		gomock.Any(),
		gomock.Any(),
		"",
		errBlockedResponseHeader,
		gomock.InAnyOrder(expectedFields),
	)

	// WHEN
	p.mergeProxiedResponseHeaders(log, testhelpers.AgentkKey1, outbound, inbound)

	// THEN
	assert.NotContains(t, inbound, httpz.SetCookieHeader)
}

func Test_MergeProxiedResponseHeaders_NotReportKnownBlockedHeaders(t *testing.T) {
	someKnownBlockedHeaders := maps.Clone(knownBlockedHeaders)
	someKnownBlockedHeaders.Delete(httpz.XContentTypeOptionsHeader) // the proxy sets it itself.

	for exampleKnownBlockedHeader := range someKnownBlockedHeaders {
		t.Run(exampleKnownBlockedHeader, func(t *testing.T) {
			// GIVEN
			ctrl := gomock.NewController(t)
			mockAPI := mock_modserver.NewMockAPI(ctrl)
			log := testlogger.New(t)
			p := kubernetesAPIProxy{
				log:                    log,
				api:                    mockAPI,
				allowedResponseHeaders: sets.New[string](), // for testing purposes we don't allow any headers, which means that every header will be blocked
			}
			outbound := http.Header{
				exampleKnownBlockedHeader: []string{"any"},
			}
			inbound := http.Header{}

			mockAPI.EXPECT().HandleProcessingError(
				gomock.Any(),
				gomock.Any(),
				"",
				errBlockedResponseHeader,
				gomock.Any(),
			).Times(0)

			// WHEN
			p.mergeProxiedResponseHeaders(log, testhelpers.AgentkKey1, outbound, inbound)

			// THEN
			assert.NotContains(t, inbound, exampleKnownBlockedHeader)
		})
	}
}

func TestGetAuthorizationInfoFromRequest(t *testing.T) {
	testcases := []struct {
		name             string
		givenRequest     *http.Request
		expectedAgentKey api.AgentKey
		expectedErrorMsg string
	}{
		{
			name: "missing agent id",
			givenRequest: &http.Request{
				URL: &url.URL{},
			},
			expectedAgentKey: api.AgentKey{ID: 0, Type: api.AgentTypeKubernetes},
			expectedErrorMsg: "exactly one agent id must be provided either in the \"Gitlab-Agent-Id\" header or \"gitlab-agent-id\" query parameter",
		},
		{
			name: "invalid agent id in header",
			givenRequest: &http.Request{
				Header: http.Header{
					httpz.GitlabAgentIDHeader: {"foobar"},
				},
			},
			expectedAgentKey: api.AgentKey{ID: 0, Type: api.AgentTypeKubernetes},
			expectedErrorMsg: "agent id in request: invalid value: \"foobar\"",
		},
		{
			name: "invalid agent id in query param",
			givenRequest: &http.Request{
				URL: &url.URL{
					RawQuery: fmt.Sprintf("%s=foobar", httpz.GitlabAgentIDQueryParam),
				},
			},
			expectedAgentKey: api.AgentKey{ID: 0, Type: api.AgentTypeKubernetes},
			expectedErrorMsg: "agent id in request: invalid value: \"foobar\"",
		},
		{
			name: "multiple agent ids in headers",
			givenRequest: &http.Request{
				Header: http.Header{
					httpz.GitlabAgentIDHeader: {"1", "2"},
				},
			},
			expectedAgentKey: api.AgentKey{ID: 0, Type: api.AgentTypeKubernetes},
			expectedErrorMsg: fmt.Sprintf("%s: header must have exactly one value", httpz.GitlabAgentIDHeader),
		},
		{
			name: "multiple agent ids in query params",
			givenRequest: &http.Request{
				URL: &url.URL{
					RawQuery: fmt.Sprintf("%[1]s=1&%[1]s=2", httpz.GitlabAgentIDQueryParam),
				},
			},
			expectedAgentKey: api.AgentKey{ID: 0, Type: api.AgentTypeKubernetes},
			expectedErrorMsg: fmt.Sprintf("%s: header must have exactly one value", httpz.GitlabAgentIDHeader),
		},
		{
			name: "valid agent id in header",
			givenRequest: &http.Request{
				Header: http.Header{
					httpz.GitlabAgentIDHeader: {"1"},
				},
			},
			expectedAgentKey: api.AgentKey{ID: 1, Type: api.AgentTypeKubernetes},
		},
		{
			name: "valid agent id in query param",
			givenRequest: &http.Request{
				URL: &url.URL{
					RawQuery: fmt.Sprintf("%s=1", httpz.GitlabAgentIDQueryParam),
				},
			},
			expectedAgentKey: api.AgentKey{ID: 1, Type: api.AgentTypeKubernetes},
		},
	}
	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			// WHEN
			actualAgentKey, actualErr := getAgentKeyForSessionCookieRequest(tc.givenRequest)

			// THEN
			require.Equal(t, tc.expectedAgentKey, actualAgentKey)
			if tc.expectedAgentKey == modshared.NoAgentKey {
				require.Error(t, actualErr, tc.expectedErrorMsg)
			}
		})
	}
}

func TestGetCSRFTokenForSessionCookieRequest(t *testing.T) {
	testcases := []struct {
		name              string
		givenRequest      *http.Request
		expectedCSRFToken string
		expectedErrorMsg  string
	}{
		{
			name: "missing CSRF token",
			givenRequest: &http.Request{
				URL: &url.URL{},
			},
			expectedCSRFToken: "",
			expectedErrorMsg:  "exactly one CSRF token must be provided either in the \"X-Csrf-Token\" header or \"gitlab-csrf-token\" query parameter",
		},
		{
			name: "multiple CSRF tokens in headers",
			givenRequest: &http.Request{
				Header: http.Header{
					httpz.CSRFTokenHeader: {"csrf-token-1", "csrf-token-2"},
				},
			},
			expectedCSRFToken: "",
			expectedErrorMsg:  fmt.Sprintf("%s: header must have exactly one value", httpz.CSRFTokenHeader),
		},
		{
			name: "multiple CSRF tokens in query params",
			givenRequest: &http.Request{
				URL: &url.URL{
					RawQuery: fmt.Sprintf("%[1]s=csrf-token-1&%[1]s=csrf-token-2", httpz.CSRFTokenQueryParam),
				},
			},
			expectedCSRFToken: "",
			expectedErrorMsg:  fmt.Sprintf("%s: header must have exactly one value", httpz.CSRFTokenHeader),
		},
		{
			name: "CSRF token in header",
			givenRequest: &http.Request{
				Header: http.Header{
					httpz.CSRFTokenHeader: {"csrf-token"},
				},
			},
			expectedCSRFToken: "csrf-token",
		},
		{
			name: "CSRF token in query param",
			givenRequest: &http.Request{
				URL: &url.URL{
					RawQuery: fmt.Sprintf("%s=csrf-token", httpz.CSRFTokenQueryParam),
				},
			},
			expectedCSRFToken: "csrf-token",
		},
	}
	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			// WHEN
			actualCsrfToken, actualErr := getCSRFTokenForSessionCookieRequest(tc.givenRequest)

			// THEN
			require.Equal(t, tc.expectedCSRFToken, actualCsrfToken)
			if tc.expectedCSRFToken == "" {
				require.Error(t, actualErr, tc.expectedErrorMsg)
			}
		})
	}
}

func TestProxy_WarningHeaderForOutdatedAgents(t *testing.T) {
	tests := []struct {
		name               string
		serverVersion      version.Version
		agentVersion       version.Version
		warningHeaderValue string
	}{
		{
			name:               "agent version header is not set",
			serverVersion:      mustParseVersion(t, "v17.1.0"),
			agentVersion:       version.Version{},
			warningHeaderValue: "",
		},
		{
			name:               "agent version header is set and equal to server",
			serverVersion:      mustParseVersion(t, "v17.1.0"),
			agentVersion:       mustParseVersion(t, "v17.1.0"),
			warningHeaderValue: "",
		},
		{
			name:               "agent version header is set and not compatible",
			serverVersion:      mustParseVersion(t, "v17.1.0"),
			agentVersion:       mustParseVersion(t, "v16.1.0"),
			warningHeaderValue: `299 - "The agent for Kubernetes (agentk) is outdated by more than four minor releases, including a major release, and might not be compatible with the agent server for Kubernetes (KAS). You should upgrade agentk."`,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			// GIVEN
			proxy := &kubernetesAPIProxy{
				serverVersion:          tc.serverVersion,
				gitlabReleasesList:     fakeGitLabReleasesList,
				allowedResponseHeaders: defaultAllowedResponseHeaders,
			}
			outbound := http.Header{}
			if (version.Version{}) != tc.agentVersion {
				outbound.Set(httpz.GitlabAgentVersionHeader, tc.agentVersion.String())
			}
			inbound := http.Header{}

			// WHEN
			proxy.addWarningHeaderIfOutdatedAgent(outbound, inbound)

			// THEN
			assert.Equal(t, tc.warningHeaderValue, inbound.Get(httpz.WarningHeader))
		})
	}
}

func TestProxy_WarningHeaderForOutdatedAgents_InvalidAgentVersion(t *testing.T) {
	// GIVEN
	proxy := &kubernetesAPIProxy{
		log:                    testlogger.New(t),
		allowedResponseHeaders: defaultAllowedResponseHeaders,
	}
	outbound := http.Header{httpz.GitlabAgentVersionHeader: {"invalid.version"}}
	inbound := http.Header{}

	// WHEN
	proxy.addWarningHeaderIfOutdatedAgent(outbound, inbound)

	// THEN
	assert.Empty(t, inbound.Get(httpz.WarningHeader))
}

func TestProxy_resolveStatusCodeClassMetricOptions(t *testing.T) {
	// GIVEN
	tests := []struct {
		statusCode            int
		expectedMetricOptions []otelmetric.AddOption
	}{
		{
			statusCode:            -1,
			expectedMetricOptions: statusCodeClassInvalidMetricOptions,
		},
		{
			statusCode:            0,
			expectedMetricOptions: statusCodeClassInvalidMetricOptions,
		},
		{
			statusCode:            99,
			expectedMetricOptions: statusCodeClassInvalidMetricOptions,
		},
		{
			statusCode:            600,
			expectedMetricOptions: statusCodeClassInvalidMetricOptions,
		},
		{
			statusCode:            4242,
			expectedMetricOptions: statusCodeClassInvalidMetricOptions,
		},
		{
			statusCode:            100,
			expectedMetricOptions: statusCodeClassMetricOptions[0],
		},
		{
			statusCode:            199,
			expectedMetricOptions: statusCodeClassMetricOptions[0],
		},
		{
			statusCode:            200,
			expectedMetricOptions: statusCodeClassMetricOptions[1],
		},
		{
			statusCode:            299,
			expectedMetricOptions: statusCodeClassMetricOptions[1],
		},
		{
			statusCode:            300,
			expectedMetricOptions: statusCodeClassMetricOptions[2],
		},
		{
			statusCode:            399,
			expectedMetricOptions: statusCodeClassMetricOptions[2],
		},
		{
			statusCode:            400,
			expectedMetricOptions: statusCodeClassMetricOptions[3],
		},
		{
			statusCode:            499,
			expectedMetricOptions: statusCodeClassMetricOptions[3],
		},
		{
			statusCode:            500,
			expectedMetricOptions: statusCodeClassMetricOptions[4],
		},
		{
			statusCode:            599,
			expectedMetricOptions: statusCodeClassMetricOptions[4],
		},
	}

	for _, tc := range tests {
		t.Run(strconv.Itoa(tc.statusCode), func(t *testing.T) {
			// WHEN
			class := resolveStatusCodeClassMetricOptions(tc.statusCode)

			// THEN
			assert.Equal(t, tc.expectedMetricOptions, class)
		})
	}
}

func TestProxy_isHeaderAllowed(t *testing.T) {
	// GIVEN
	tests := []struct {
		header                  string
		expectedIsHeaderAllowed bool
	}{
		{
			header:                  "Content-Type",
			expectedIsHeaderAllowed: true,
		},
		{
			header:                  "Cache-Control",
			expectedIsHeaderAllowed: false,
		},
		{
			header:                  "Unknown",
			expectedIsHeaderAllowed: false,
		},
	}

	p := &kubernetesAPIProxy{
		log:                    testlogger.New(t),
		allowedResponseHeaders: defaultAllowedResponseHeaders,
	}

	for _, tc := range tests {
		t.Run(tc.header, func(t *testing.T) {
			// WHEN
			actualIsHeaderAllowed := p.isHeaderAllowed(tc.header)

			// THEN
			assert.Equal(t, tc.expectedIsHeaderAllowed, actualIsHeaderAllowed)
		})
	}
}

func mustParseVersion(t *testing.T, v string) version.Version {
	t.Helper()

	version, err := version.NewVersion(v)
	require.NoError(t, err)
	return version
}

func generateValidTestWebSocketToken(t *testing.T, path string, impConfig *kubernetes_api.ImpersonationConfig) string {
	wt := &webSocketToken{
		jwtSigningMethod: jwttool.SigningMethodHS3_512,
		jwtSecret:        validWebSocketTokenSecret,
	}
	agentKey := api.AgentKey{ID: 42, Type: api.AgentTypeKubernetes}
	token, err := wt.generate(path, agentKey, 84, impConfig, featureflag.Set{})
	require.NoError(t, err)
	return token
}

func generateInvalidTestWebSocketToken(t *testing.T, path string) string {
	wt := &webSocketToken{
		jwtSigningMethod: jwttool.SigningMethodHS3_512,
		jwtSecret:        invalidWebSocketTokenSecret,
	}
	agentKey := api.AgentKey{ID: 42, Type: api.AgentTypeKubernetes}
	token, err := wt.generate(path, agentKey, 84, nil, featureflag.Set{})
	require.NoError(t, err)
	return token
}

func encodeTokenWithPrefix(token string) string {
	return webSocketTokenSubProtocolPrefix + token
}
