Merge branch 'vg/fix-topology-classify-claims-chain' into 'main'

Send Classify claims via the claims chain

See merge request https://gitlab.com/gitlab-org/gitlab-shell/-/merge_requests/1541

Merged-by: Igor Drozdov <idrozdov@gitlab.com>
Approved-by: Igor Drozdov <idrozdov@gitlab.com>
Co-authored-by: Vasilii Iakliushin <viakliushin@gitlab.com>
Co-authored-by: Vladimir Glafirov <vglafirov@gitlab.com>

Author: 3 people

go.mod

@@ -17,7 +17,7 @@ require (
 	github.com/pires/go-proxyproto v0.14.0
 	github.com/prometheus/client_golang v1.23.2
 	github.com/stretchr/testify v1.11.1
-	gitlab.com/gitlab-org/cells/topology-service v0.0.0-20260522095121-2c761c2e9850
+	gitlab.com/gitlab-org/cells/topology-service v0.0.0-20260827174901-e70472732254
 	gitlab.com/gitlab-org/gitaly/v18 v18.9.0-rc4
 	gitlab.com/gitlab-org/labkit v1.64.11
 	gitlab.com/gitlab-org/labkit/v2 v2.35.0

go.sum

@@ -516,8 +516,8 @@ github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9dec
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
 github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
-gitlab.com/gitlab-org/cells/topology-service v0.0.0-20260522095121-2c761c2e9850 h1:ooD76K9Xy1IFB6840t808xPzT4sBx3OGoH6Uz37wCYw=
-gitlab.com/gitlab-org/cells/topology-service v0.0.0-20260522095121-2c761c2e9850/go.mod h1:3cPQOfQvgI26owOpGzADzB3IJJQ3pTYuLQaurpp1k/A=
+gitlab.com/gitlab-org/cells/topology-service v0.0.0-20260827174901-e70472732254 h1:omotGPkyvMHt473hG3xtaZehAeHRm3W1U5UJxg5OXMk=
+gitlab.com/gitlab-org/cells/topology-service v0.0.0-20260827174901-e70472732254/go.mod h1:xOxWBmF9zM1HaoAEfjoSb54FpXqIzG0qyD1oBTOQSQw=
 gitlab.com/gitlab-org/gitaly/v18 v18.9.0-rc4 h1:2PQlCQF5ioVdmcjY6Ut8zxwFY2TIgHWnl5ekT4IJEKQ=
 gitlab.com/gitlab-org/gitaly/v18 v18.9.0-rc4/go.mod h1:ijDPC9GyV3UdmrZG4LiBQJyMjOuzLcbH5AuQ8GpSSIk=
 gitlab.com/gitlab-org/go/reopen v1.0.0 h1:6BujZ0lkkjGIejTUJdNO1w56mN1SI10qcVQyQlOPM+8=

internal/gitlabnet/authorizedcerts/client_test.go

@@ -163,7 +163,7 @@ func TestGetByKeyWithTopologyService(t *testing.T) {
 		require.True(t, cellReceived, "request should have been sent to the cell server")
 		require.False(t, defaultReceived, "request should NOT have been sent to the default server")
 
-		require.Equal(t, "fingerprint-value", mock.LastRequest.GetClaim().GetSshKeyFingerprint())
+		require.Equal(t, "fingerprint-value", mock.LastClaim(t).GetSshKeyFingerprint())
 	})
 
 	t.Run("falls back to default", func(t *testing.T) {

internal/gitlabnet/authorizedkeys/client_test.go

@@ -169,7 +169,7 @@ func TestGetByKeyWithTopologyService(t *testing.T) {
 		require.True(t, cellReceived, "request should have been sent to the cell server")
 		require.False(t, defaultReceived, "request should NOT have been sent to the default server")
 
-		require.Equal(t, expectedFingerprint, mock.LastRequest.GetClaim().GetSshKeyFingerprint())
+		require.Equal(t, expectedFingerprint, mock.LastClaim(t).GetSshKeyFingerprint())
 	})
 
 	t.Run("falls back to default", func(t *testing.T) {

internal/topology/client.go

@@ -98,10 +98,9 @@ func (c *Client) Classify(ctx context.Context, claim *types_proto.Claim) (resp *
 	ctx, cancel := context.WithTimeout(ctx, c.config.Timeout)
 	defer cancel()
 
+	// GitLab Shell resolves one claim per lookup.
 	req := &pb.ClassifyRequest{
-		ClassificationKey: &pb.ClassifyRequest_Claim{
-			Claim: claim,
-		},
+		Claims: []*types_proto.Claim{claim},
 	}
 
 	resp, err = client.Classify(ctx, req)

internal/topology/client_test.go

@@ -14,11 +14,13 @@ import (
 	"github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/require"
 	pb "gitlab.com/gitlab-org/cells/topology-service/clients/go/proto"
+	types_proto "gitlab.com/gitlab-org/cells/topology-service/clients/go/proto/types/v1"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/metrics"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/testhelper"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/topology/topologytest"
 	"gitlab.com/gitlab-org/labkit/correlation"
 	"google.golang.org/grpc"
+	"google.golang.org/protobuf/proto"
 )
 
 const localhostAddr = "localhost:9090"
@@ -128,11 +130,7 @@ func TestClient_Classify(t *testing.T) {
 		require.Equal(t, pb.ClassifyAction_PROXY, result.GetAction())
 		require.Equal(t, "cell-1.gitlab.com:443", result.GetProxy().GetAddress())
 
-		// Verify the request was constructed correctly
-		require.NotNil(t, mock.LastRequest.GetClassificationKey())
-		require.Equal(t, "my-group/my-project", mock.LastRequest.GetClaim().GetRoute())
-		require.Equal(t, pb.ClassifyType_UNSPECIFIED, mock.LastRequest.GetType())
-		require.Empty(t, mock.LastRequest.GetValue())
+		require.Equal(t, "my-group/my-project", mock.LastClaim(t).GetRoute())
 	})
 
 	t.Run("successful SSH key claim returns proxy info", func(t *testing.T) {
@@ -152,7 +150,7 @@ func TestClient_Classify(t *testing.T) {
 
 		require.NoError(t, err)
 		require.Equal(t, pb.ClassifyAction_PROXY, result.GetAction())
-		require.Equal(t, "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ", mock.LastRequest.GetClaim().GetSshKey())
+		require.Equal(t, "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ", mock.LastClaim(t).GetSshKey())
 	})
 
 	t.Run("successful project ID claim returns proxy info", func(t *testing.T) {
@@ -172,7 +170,46 @@ func TestClient_Classify(t *testing.T) {
 
 		require.NoError(t, err)
 		require.Equal(t, pb.ClassifyAction_PROXY, result.GetAction())
-		require.Equal(t, int64(42), mock.LastRequest.GetClaim().GetProjectId())
+		require.Equal(t, int64(42), mock.LastClaim(t).GetProjectId())
+	})
+
+	t.Run("sends the claim via the claims chain, not type/value", func(t *testing.T) {
+		mock := &topologytest.MockClassifyServer{}
+		addr, stop := topologytest.StartMockServer(t, mock)
+		defer stop()
+
+		client := NewClient(&Config{
+			Enabled: true,
+			Address: addr,
+			Timeout: 5 * time.Second,
+		})
+		defer client.Close()
+
+		tests := []struct {
+			name  string
+			claim *types_proto.Claim
+		}{
+			{name: "route", claim: RouteClaim("my-group")},
+			{name: "ssh_key", claim: SSHKeyClaim("ssh-rsa AAAAB3")},
+			{name: "ssh_fingerprint", claim: SSHFingerprintClaim("W3THTJOKxMaZp0VIOrjVSBVDnFjyzVSMFGMLmSPcaGo")},
+			{name: "project_id", claim: ProjectIDClaim(42)},
+			{name: "username", claim: UsernameClaim("jane.doe")},
+		}
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				_, err := client.Classify(ctx, tt.claim)
+				require.NoError(t, err)
+
+				claims := mock.LastRequest.GetClaims()
+				require.Len(t, claims, 1)
+				require.True(t, proto.Equal(tt.claim, claims[0]),
+					"claim must be sent unmodified in the claims chain")
+
+				require.Equal(t, pb.ClassifyType_UNSPECIFIED, mock.LastRequest.GetType())
+				require.Empty(t, mock.LastRequest.GetValue())
+			})
+		}
 	})
 
 	t.Run("server error is propagated", func(t *testing.T) {

internal/topology/resolver_test.go

@@ -320,7 +320,7 @@ func TestResolveByRoute(t *testing.T) {
 		require.Equal(t, "http://cell-1:8080", result)
 
 		// Verify the claim sent to the server used just the top-level namespace
-		require.Equal(t, groupNamespace, mock.LastRequest.GetClaim().GetRoute())
+		require.Equal(t, groupNamespace, mock.LastClaim(t).GetRoute())
 	})
 
 	t.Run("empty repo path returns empty string", func(t *testing.T) {
@@ -497,7 +497,7 @@ func TestResolveBySSHKey(t *testing.T) {
 		require.Equal(t, "http://cell-2:8080", result)
 
 		// Verify the claim sent to the server used the SSH key
-		require.Equal(t, "ssh-rsa AAAAB3...", mock.LastRequest.GetClaim().GetSshKey())
+		require.Equal(t, "ssh-rsa AAAAB3...", mock.LastClaim(t).GetSshKey())
 	})
 
 	t.Run("empty key returns empty string", func(t *testing.T) {
@@ -557,7 +557,7 @@ func TestResolveBySSHFingerprint(t *testing.T) {
 		result := resolver.resolveBySSHFingerprint(context.Background(), "W3THTJOKxMaZp0VIOrjVSBVDnFjyzVSMFGMLmSPcaGo")
 		require.Equal(t, "http://cell-2:8080", result)
 
-		require.Equal(t, "W3THTJOKxMaZp0VIOrjVSBVDnFjyzVSMFGMLmSPcaGo", mock.LastRequest.GetClaim().GetSshKeyFingerprint())
+		require.Equal(t, "W3THTJOKxMaZp0VIOrjVSBVDnFjyzVSMFGMLmSPcaGo", mock.LastClaim(t).GetSshKeyFingerprint())
 	})
 
 	t.Run("empty fingerprint returns empty string", func(t *testing.T) {
@@ -618,7 +618,7 @@ func TestResolveByUserArgs(t *testing.T) {
 		require.Equal(t, "http://cell-2:8080", result)
 
 		// Verify the claim sent to the server used the username
-		require.Equal(t, janeDoe, mock.LastRequest.GetClaim().GetUsername())
+		require.Equal(t, janeDoe, mock.LastClaim(t).GetUsername())
 	})
 
 	fallbackTests := []struct {

internal/topology/topologytest/mock_server.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 	pb "gitlab.com/gitlab-org/cells/topology-service/clients/go/proto"
+	types_proto "gitlab.com/gitlab-org/cells/topology-service/clients/go/proto/types/v1"
 	"google.golang.org/grpc"
 )
 
@@ -34,6 +35,15 @@ type MockClassifyServer struct {
 	ErrUntilAttempt int
 }
 
+// LastClaim returns the first claim from the most recent request.
+func (m *MockClassifyServer) LastClaim(t *testing.T) *types_proto.Claim {
+	t.Helper()
+	require.NotNil(t, m.LastRequest)
+	require.NotEmpty(t, m.LastRequest.GetClaims())
+
+	return m.LastRequest.GetClaims()[0]
+}
+
 // Classify implements the ClassifyService RPC.
 func (m *MockClassifyServer) Classify(_ context.Context, req *pb.ClassifyRequest) (*pb.ClassifyResponse, error) {
 	m.CallCount++