Improvements to Bot resource (additional validation and label propagation)

lib/auth/machineid/machineidv1/bot_service.go

@@ -21,6 +21,7 @@ package machineidv1
 import (
 	"context"
 	"fmt"
+	"slices"
 	"strings"
 	"time"
 
@@ -440,6 +441,11 @@ func (bs *BotService) UpdateBot(
 	for _, path := range req.UpdateMask.Paths {
 		switch {
 		case path == "spec.roles":
+			if slices.Contains(req.Bot.Spec.Roles, "") {
+				return nil, trace.BadParameter(
+					"spec.roles: must not contain empty strings",
+				)
+			}
 			role.SetImpersonateConditions(types.Allow, types.ImpersonateConditions{
 				Roles: req.Bot.Spec.Roles,
 			})
@@ -582,6 +588,9 @@ func validateBot(b *pb.Bot) error {
 	if b.Spec == nil {
 		return trace.BadParameter("spec: must be non-nil")
 	}
+	if slices.Contains(b.Spec.Roles, "") {
+		return trace.BadParameter("spec.roles: must not contain empty strings")
+	}
 	return nil
 }
 
@@ -612,6 +621,22 @@ func botFromUserAndRole(user types.User, role types.Role) (*pb.Bot, error) {
 		},
 	}
 
+	// Copy in labels from the user
+	b.Metadata.Labels = map[string]string{}
+	for k, v := range user.GetMetadata().Labels {
+		// We exclude the two labels that are implicitly added to the user by
+		// the bot service.
+		specialLabels := []string{
+			types.BotLabel,
+			types.BotGenerationLabel,
+		}
+		if slices.Contains(specialLabels, k) {
+			continue
+		}
+		b.Metadata.Labels[k] = v
+	}
+
+	// Copy in traits
 	for k, v := range user.GetTraits() {
 		if len(v) == 0 {
 			continue
@@ -661,12 +686,19 @@ func botToUserAndRole(bot *pb.Bot, now time.Time, createdBy string) (types.User,
 	}
 	user.SetRoles([]string{resourceName})
 	userMeta := user.GetMetadata()
-	userMeta.Labels = map[string]string{
-		types.BotLabel: bot.Metadata.Name,
-		// We always set this to zero here - but in Upsert, we copy from the
-		// previous user before writing if necessary.
-		types.BotGenerationLabel: "0",
+
+	// First copy in the labels from the Bot resource
+	userMeta.Labels = map[string]string{}
+	for k, v := range bot.Metadata.Labels {
+		userMeta.Labels[k] = v
 	}
+	// Then set these labels over the top - we exclude these when converting
+	// back.
+	userMeta.Labels[types.BotLabel] = bot.Metadata.Name
+	// We always set this to zero here - but in Upsert, we copy from the
+	// previous user before writing if necessary
+	userMeta.Labels[types.BotGenerationLabel] = "0"
+
 	user.SetMetadata(userMeta)
 
 	traits := map[string][]string{}

lib/auth/machineid/machineidv1/machineidv1_test.go

@@ -115,6 +115,10 @@ func TestCreateBot(t *testing.T) {
 				Bot: &machineidv1pb.Bot{
 					Metadata: &headerv1.Metadata{
 						Name: "success",
+						Labels: map[string]string{
+							"my-label":       "my-value",
+							"my-other-label": "my-other-value",
+						},
 					},
 					Spec: &machineidv1pb.BotSpec{
 						Roles: []string{testRole.GetName()},
@@ -138,6 +142,10 @@ func TestCreateBot(t *testing.T) {
 				Version: types.V1,
 				Metadata: &headerv1.Metadata{
 					Name: "success",
+					Labels: map[string]string{
+						"my-label":       "my-value",
+						"my-other-label": "my-other-value",
+					},
 				},
 				Spec: &machineidv1pb.BotSpec{
 					Roles: []string{testRole.GetName()},
@@ -162,6 +170,8 @@ func TestCreateBot(t *testing.T) {
 					Labels: map[string]string{
 						types.BotLabel:           "success",
 						types.BotGenerationLabel: "0",
+						"my-label":               "my-value",
+						"my-other-label":         "my-other-value",
 					},
 				},
 				Spec: types.UserSpecV2{
@@ -282,6 +292,25 @@ func TestCreateBot(t *testing.T) {
 				require.True(t, trace.IsBadParameter(err), "error should be bad parameter")
 			},
 		},
+		{
+			name: "validation - empty role",
+			user: botCreator.GetName(),
+			req: &machineidv1pb.CreateBotRequest{
+				Bot: &machineidv1pb.Bot{
+					Metadata: &headerv1.Metadata{
+						Name: "empty-string-role",
+					},
+					Spec: &machineidv1pb.BotSpec{
+						Roles:  []string{""},
+						Traits: []*machineidv1pb.Trait{},
+					},
+				},
+			},
+			assertError: func(t require.TestingT, err error, i ...interface{}) {
+				require.ErrorContains(t, err, "spec.roles: must not contain empty strings")
+				require.True(t, trace.IsBadParameter(err), "error should be bad parameter")
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -610,6 +639,28 @@ func TestUpdateBot(t *testing.T) {
 				require.True(t, trace.IsBadParameter(err), "error should be bad parameter")
 			},
 		},
+		{
+			name: "validation - empty string role",
+			user: botUpdaterUser.GetName(),
+			req: &machineidv1pb.UpdateBotRequest{
+				Bot: &machineidv1pb.Bot{
+					Metadata: &headerv1.Metadata{
+						Name: preExistingBot.Metadata.Name,
+					},
+					Spec: &machineidv1pb.BotSpec{
+						Roles:  []string{""},
+						Traits: []*machineidv1pb.Trait{},
+					},
+				},
+				UpdateMask: &fieldmaskpb.FieldMask{
+					Paths: []string{"spec.roles"},
+				},
+			},
+			assertError: func(t require.TestingT, err error, i ...interface{}) {
+				require.ErrorContains(t, err, "spec.roles: must not contain empty strings")
+				require.True(t, trace.IsBadParameter(err), "error should be bad parameter")
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -684,6 +735,10 @@ func TestUpsertBot(t *testing.T) {
 		Bot: &machineidv1pb.Bot{
 			Metadata: &headerv1.Metadata{
 				Name: "pre-existing",
+				Labels: map[string]string{
+					"my-label":       "my-value",
+					"my-other-label": "my-other-value",
+				},
 			},
 			Spec: &machineidv1pb.BotSpec{
 				Roles: []string{testRole.GetName()},
@@ -721,6 +776,10 @@ func TestUpsertBot(t *testing.T) {
 				Bot: &machineidv1pb.Bot{
 					Metadata: &headerv1.Metadata{
 						Name: "new",
+						Labels: map[string]string{
+							"my-label":       "my-value",
+							"my-other-label": "my-other-value",
+						},
 					},
 					Spec: &machineidv1pb.BotSpec{
 						Roles: []string{testRole.GetName()},
@@ -740,6 +799,10 @@ func TestUpsertBot(t *testing.T) {
 				Version: types.V1,
 				Metadata: &headerv1.Metadata{
 					Name: "new",
+					Labels: map[string]string{
+						"my-label":       "my-value",
+						"my-other-label": "my-other-value",
+					},
 				},
 				Spec: &machineidv1pb.BotSpec{
 					Roles: []string{testRole.GetName()},
@@ -764,6 +827,8 @@ func TestUpsertBot(t *testing.T) {
 					Labels: map[string]string{
 						types.BotLabel:           "new",
 						types.BotGenerationLabel: "0",
+						"my-label":               "my-value",
+						"my-other-label":         "my-other-value",
 					},
 				},
 				Spec: types.UserSpecV2{
@@ -820,6 +885,8 @@ func TestUpsertBot(t *testing.T) {
 					Labels: map[string]string{
 						types.BotLabel:           "pre-existing",
 						types.BotGenerationLabel: "1337",
+						"my-label":               "my-value",
+						"my-other-label":         "my-other-value",
 					},
 				},
 				Spec: types.UserSpecV2{
@@ -911,6 +978,25 @@ func TestUpsertBot(t *testing.T) {
 				require.True(t, trace.IsBadParameter(err), "error should be bad parameter")
 			},
 		},
+		{
+			name: "validation - empty role",
+			user: botCreator.GetName(),
+			req: &machineidv1pb.UpsertBotRequest{
+				Bot: &machineidv1pb.Bot{
+					Metadata: &headerv1.Metadata{
+						Name: "empty-string-role",
+					},
+					Spec: &machineidv1pb.BotSpec{
+						Roles:  []string{""},
+						Traits: []*machineidv1pb.Trait{},
+					},
+				},
+			},
+			assertError: func(t require.TestingT, err error, i ...interface{}) {
+				require.ErrorContains(t, err, "spec.roles: must not contain empty strings")
+				require.True(t, trace.IsBadParameter(err), "error should be bad parameter")
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -983,6 +1069,10 @@ func TestGetBot(t *testing.T) {
 			Bot: &machineidv1pb.Bot{
 				Metadata: &headerv1.Metadata{
 					Name: "pre-existing",
+					Labels: map[string]string{
+						"my-label":       "my-value",
+						"my-other-label": "my-other-value",
+					},
 				},
 				Spec: &machineidv1pb.BotSpec{
 					Roles: []string{testRole.GetName()},
@@ -1090,6 +1180,10 @@ func TestListBots(t *testing.T) {
 			Bot: &machineidv1pb.Bot{
 				Metadata: &headerv1.Metadata{
 					Name: "pre-existing",
+					Labels: map[string]string{
+						"my-label":       "my-value",
+						"my-other-label": "my-other-value",
+					},
 				},
 				Spec: &machineidv1pb.BotSpec{
 					Roles: []string{testRole.GetName()},