[v16] Support traits in workload_identity_labels (#56249)

docs/pages/includes/role-spec.mdx

@@ -293,6 +293,14 @@ spec:
     cluster_labels:
       'env': 'prod'
 
+    # workload_identity_labels: a user/bot with this role will be allowed to
+    # issue Workload Identities with labels matching below.
+    #
+    # Supports role templating with traits.
+    workload_identity_labels:
+      'env': 'prod'
+      'team': '{{external.team}}'
+
     # node_labels_expression has the same purpose as node_labels but
     # supports predicate expressions to configure custom logic.
     # A user with this role will be allowed to access nodes if they are in the

lib/auth/machineid/workloadidentityv1/workloadidentityv1_test.go

@@ -395,6 +395,34 @@ func TestIssueWorkloadIdentity(t *testing.T) {
 	specificAccessClient, err := tp.srv.NewClient(auth.TestUser(specificAccess.GetName()))
 	require.NoError(t, err)
 
+	traitAccess, _, err := auth.CreateUserAndRole(
+		tp.srv.Auth(),
+		"traity",
+		[]string{},
+		[]types.Rule{
+			types.NewRule(
+				types.KindWorkloadIdentity,
+				[]string{types.VerbRead, types.VerbList},
+			),
+		},
+		auth.WithUserMutator(func(user types.User) {
+			tr := user.GetTraits()
+			if tr == nil {
+				tr = map[string][]string{}
+			}
+			tr["custom"] = []string{"trait-value-a", "trait-value-b"}
+			user.SetTraits(tr)
+		}),
+		auth.WithRoleMutator(func(role types.Role) {
+			role.SetWorkloadIdentityLabels(types.Allow, types.Labels{
+				"trait-label": []string{"{{external.custom}}"},
+			})
+		}),
+	)
+	require.NoError(t, err)
+	traitAccessClient, err := tp.srv.NewClient(auth.TestUser(traitAccess.GetName()))
+	require.NoError(t, err)
+
 	// Generate a keypair to generate x509 SVIDs for.
 	workloadKey, err := native.GenerateRSAPrivateKey()
 	require.NoError(t, err)
@@ -493,6 +521,23 @@ func TestIssueWorkloadIdentity(t *testing.T) {
 	})
 	require.NoError(t, err)
 
+	traitsRequired, err := tp.srv.Auth().CreateWorkloadIdentity(ctx, &workloadidentityv1pb.WorkloadIdentity{
+		Kind:    types.KindWorkloadIdentity,
+		Version: types.V1,
+		Metadata: &headerv1.Metadata{
+			Name: "traits-required",
+			Labels: map[string]string{
+				"trait-label": "trait-value-b",
+			},
+		},
+		Spec: &workloadidentityv1pb.WorkloadIdentitySpec{
+			Spiffe: &workloadidentityv1pb.WorkloadIdentitySPIFFE{
+				Id: "/foo",
+			},
+		},
+	})
+	require.NoError(t, err)
+
 	workloadAttrs := func(f func(attrs *workloadidentityv1pb.WorkloadAttrs)) *workloadidentityv1pb.WorkloadAttrs {
 		attrs := &workloadidentityv1pb.WorkloadAttrs{
 			Kubernetes: &workloadidentityv1pb.WorkloadAttrsKubernetes{
@@ -864,6 +909,23 @@ func TestIssueWorkloadIdentity(t *testing.T) {
 				))
 			},
 		},
+		{
+			name:   "x509 svid - access via traits in labels",
+			client: traitAccessClient,
+			req: &workloadidentityv1pb.IssueWorkloadIdentityRequest{
+				Name: traitsRequired.GetMetadata().GetName(),
+				Credential: &workloadidentityv1pb.IssueWorkloadIdentityRequest_X509SvidParams{
+					X509SvidParams: &workloadidentityv1pb.X509SVIDParams{
+						PublicKey: workloadKeyPubBytes,
+					},
+				},
+				WorkloadAttrs: workloadAttrs(nil),
+			},
+			requireErr: require.NoError,
+			assert: func(t *testing.T, res *workloadidentityv1pb.IssueWorkloadIdentityResponse) {
+				require.NotNil(t, res.Credential)
+			},
+		},
 		{
 			name:   "unauthorized by rules",
 			client: wilcardAccessClient,

lib/services/role.go

@@ -497,6 +497,7 @@ func ApplyTraits(r types.Role, traits map[string][]string) (types.Role, error) {
 			types.KindDatabaseService,
 			types.KindWindowsDesktop,
 			types.KindUserGroup,
+			types.KindWorkloadIdentity,
 		} {
 			labelMatchers, err := r.GetLabelMatchers(condition, kind)
 			if err != nil {