From 8de769b103d3316cd473feaddb20255be2296e58 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Tue, 9 Jan 2024 19:27:01 -0700 Subject: [PATCH 01/15] Add new `tctl bots update` command This adds a new command to update roles and logins for existing bots, with flags to both set and append as desired. changelog: Add new `tctl bots update` command to update bot roles and logins --- tool/tctl/common/bots_command.go | 198 +++++++++++++++++++++++++- tool/tctl/common/bots_command_test.go | 94 ++++++++++++ 2 files changed, 289 insertions(+), 3 deletions(-) create mode 100644 tool/tctl/common/bots_command_test.go diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index ff3802a6abf6d..f887bcf6b22a6 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -56,13 +56,17 @@ type BotsCommand struct { botRoles string tokenID string tokenTTL time.Duration + addRole string allowedLogins []string + addLogins string + setLogins string botsList *kingpin.CmdClause botsAdd *kingpin.CmdClause botsRemove *kingpin.CmdClause botsLock *kingpin.CmdClause + botsUpdate *kingpin.CmdClause } // Initialize sets up the "tctl bots" command. @@ -88,6 +92,13 @@ func (c *BotsCommand) Initialize(app *kingpin.Application, config *servicecfg.Co c.botsLock.Flag("expires", "Time point (RFC3339) when the lock expires.").StringVar(&c.lockExpires) c.botsLock.Flag("ttl", "Time duration after which the lock expires.").DurationVar(&c.lockTTL) c.botsLock.Hidden() + + c.botsUpdate = bots.Command("update", "Update an existing bot.") + c.botsUpdate.Arg("name", "Name of an existing bot to update.").Required().StringVar(&c.botName) + c.botsUpdate.Flag("set-roles", "A comma-separated list of roles to replace.").StringVar(&c.botRoles) + c.botsUpdate.Flag("add-role", "A comma-separated list of roles to add to an existing bot.").StringVar(&c.addRole) + c.botsUpdate.Flag("set-logins", "A comma-separated list of allowed logins to replace").StringVar(&c.setLogins) + c.botsUpdate.Flag("add-login", "A comma-separated list of logins to add to an existing bot.").StringVar(&c.addLogins) } // TryRun attempts to run subcommands. @@ -101,6 +112,8 @@ func (c *BotsCommand) TryRun(ctx context.Context, cmd string, client auth.Client err = c.RemoveBot(ctx, client) case c.botsLock.FullCommand(): err = c.LockBot(ctx, client) + case c.botsUpdate.FullCommand(): + err = c.UpdateBot(ctx, client) default: return false, nil } @@ -228,7 +241,7 @@ Please note: // TODO(noah): DELETE IN 16.0.0 func (c *BotsCommand) addBotLegacy(ctx context.Context, client auth.ClientI) error { - roles := splitRoles(c.botRoles) + roles := splitEntries(c.botRoles) if len(roles) == 0 { log.Warning("No roles specified. The bot will not be able to produce outputs until a role is added to the bot.") } @@ -298,7 +311,7 @@ func (c *BotsCommand) AddBot(ctx context.Context, client auth.ClientI) error { } } - roles := splitRoles(c.botRoles) + roles := splitEntries(c.botRoles) if len(roles) == 0 { log.Warning("No roles specified. The bot will not be able to produce outputs until a role is added to the bot.") } @@ -483,7 +496,137 @@ func (c *BotsCommand) LockBot(ctx context.Context, client auth.ClientI) error { return nil } -func splitRoles(flag string) []string { +// updateBotUser applies updates to a bot user that were specified in the CLI +// arguments. +func (c *BotsCommand) updateBotUser(ctx context.Context, client auth.ClientI) error { + resource := machineidv1.BotResourceName(c.botName) + user, err := client.GetUser(ctx, resource, false) + if err != nil { + return trace.Wrap(err) + } + + traits := user.GetTraits() + if traits == nil { + traits = map[string][]string{} + } + + var currentLogins map[string]struct{} + if logins, exists := traits[constants.TraitLogins]; exists { + currentLogins = arrayToSet(logins) + } else { + currentLogins = map[string]struct{}{} + } + + var desiredLogins map[string]struct{} + if c.setLogins != "" { + desiredLogins = arrayToSet(splitEntries(c.setLogins)) + } else { + desiredLogins = setUnion(currentLogins) + } + + addLogins := splitEntries(c.addLogins) + if len(addLogins) > 0 { + desiredLogins = setUnion(desiredLogins, arrayToSet(addLogins)) + } + + log.Infof("Desired logins for bot %q: %+v", c.botName, setToArray(desiredLogins)) + if setsEqual(currentLogins, desiredLogins) { + log.Infof("Desired logins match existing set, will not update bot user") + return nil + } + + if len(desiredLogins) == 0 { + delete(traits, constants.TraitLogins) + log.Infof("Removing logins trait from bot user") + } else { + traits[constants.TraitLogins] = setToArray(desiredLogins) + } + user.SetTraits(traits) + + if _, err := client.UpdateUser(ctx, user); err != nil { + return trace.Wrap(err) + } + + return nil +} + +// updateBotRole performs any updates to a bot's role that were specified in the +// CLI arguments. +func (c *BotsCommand) updateBotRole(ctx context.Context, client auth.ClientI) error { + resource := machineidv1.BotResourceName(c.botName) + role, err := client.GetRole(ctx, resource) + if err != nil { + return trace.Wrap(err) + } + + allowCond := role.GetImpersonateConditions(types.Allow) + currentRoles := arrayToSet(allowCond.Roles) + + log.Infof("Existing roles for bot %q: %+v", c.botName, setToArray(currentRoles)) + + var desiredRoles map[string]struct{} + if c.botRoles != "" { + desiredRoles = arrayToSet(splitEntries(c.botRoles)) + } else { + desiredRoles = setUnion(currentRoles) + } + + if c.addRole != "" { + desiredRoles = setUnion(desiredRoles, arrayToSet(splitEntries(c.addRole))) + } + + log.Infof("Desired roles for bot %q: %+v", c.botName, setToArray(desiredRoles)) + if setsEqual(currentRoles, desiredRoles) { + log.Infof("Desired roles match existing roles, will not update bot role") + return nil + } + + // Validate roles (server does not do this yet). + for roleName := range desiredRoles { + if _, err := client.GetRole(ctx, roleName); err != nil { + return trace.Wrap(err) + } + } + + allowCond.Roles = setToArray(desiredRoles) + role.SetImpersonateConditions(types.Allow, allowCond) + + if _, err := client.UpdateRole(ctx, role); err != nil { + return trace.Wrap(err) + } + return nil +} + +// UpdateBot performs various updates to existing bot users and roles. +func (c *BotsCommand) UpdateBot(ctx context.Context, client auth.ClientI) error { + changes := false + + if len(c.allowedLogins) > 0 || len(c.addLogins) > 0 { + if err := c.updateBotUser(ctx, client); err != nil { + return trace.Wrap(err) + } + + changes = true + } + + if c.botRoles != "" || c.addRole != "" { + if err := c.updateBotRole(ctx, client); err != nil { + return trace.Wrap(err) + } + + changes = true + } + + if changes { + log.Infof("Bot %q has been updated. Roles will take affect on its next renewal.", c.botName) + } + + return nil +} + +// splitEntries splits a comma separated string into an array of entries, +// ignoring empty or whitespace-only elements. +func splitEntries(flag string) []string { var roles []string for _, s := range strings.Split(flag, ",") { s = strings.TrimSpace(s) @@ -494,3 +637,52 @@ func splitRoles(flag string) []string { } return roles } + +// setsEqual determines if two sets contain the same keys. +func setsEqual[T comparable](a map[T]struct{}, b map[T]struct{}) bool { + if len(a) != len(b) { + return false + } + + for k := range a { + if _, ok := b[k]; !ok { + return false + } + } + + return true +} + +// setToArray converts a set to an array +func setToArray[T comparable](m map[T]struct{}) []T { + var ret []T + for entry := range m { + ret = append(ret, entry) + } + + return ret +} + +// arrayToSet converts an array to a set, removing duplicates +func arrayToSet[T comparable](arr []T) map[T]struct{} { + ret := make(map[T]struct{}) + + for _, entry := range arr { + ret[entry] = struct{}{} + } + + return ret +} + +// setUnion returns a new set containing a union of all provided sets. If only +// one set is given, it is effectively shallow copied. +func setUnion[T comparable](sets ...map[T]struct{}) map[T]struct{} { + ret := make(map[T]struct{}) + for _, set := range sets { + for key := range set { + ret[key] = struct{}{} + } + } + + return ret +} diff --git a/tool/tctl/common/bots_command_test.go b/tool/tctl/common/bots_command_test.go new file mode 100644 index 0000000000000..bad3d516ac59b --- /dev/null +++ b/tool/tctl/common/bots_command_test.go @@ -0,0 +1,94 @@ +/* + * Teleport + * Copyright (C) 2024 Gravitational, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +package common + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSetsEqual(t *testing.T) { + tests := []struct { + name string + a map[string]struct{} + b map[string]struct{} + assert require.BoolAssertionFunc + }{ + { + name: "out of order true", + a: arrayToSet([]string{"a", "b", "c", "d"}), + b: arrayToSet([]string{"d", "b", "c", "a"}), + assert: require.True, + }, + { + name: "length mismatch", + a: arrayToSet([]string{"a", "b", "c"}), + b: arrayToSet([]string{"d", "b", "c", "a"}), + assert: require.False, + }, + { + name: "simple false", + a: arrayToSet([]string{"a", "b", "c", "d"}), + b: arrayToSet([]string{"d", "b", "c", "e"}), + assert: require.False, + }, + { + name: "duplicates ignored", + a: arrayToSet([]string{"a", "b", "c", "d"}), + b: arrayToSet([]string{"d", "b", "c", "a", "a", "b", "c"}), + assert: require.True, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tt.assert(t, setsEqual(tt.a, tt.b)) + }) + } +} + +func TestArrayToSetRoundtrip(t *testing.T) { + a := []string{"a", "b", "c", "d"} + require.ElementsMatch(t, a, setToArray(arrayToSet(a))) + + // It should also remove duplicates + require.ElementsMatch(t, setToArray(arrayToSet(append(a, a...))), a) +} + +func TestSetUnion(t *testing.T) { + a := []string{"a", "b", "c", "d"} + b := []string{"c", "d", "e", "f"} + + // Self union clones + require.ElementsMatch( + t, + setToArray(setUnion(arrayToSet(a), arrayToSet(a), arrayToSet(a))), + a, + ) + + require.ElementsMatch( + t, + setToArray(setUnion(arrayToSet(a), arrayToSet(b))), + []string{"a", "b", "c", "d", "e", "f"}, + ) +} From 8566578bab14531cc2cde204e4d1ef085166f9f7 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Wed, 10 Jan 2024 16:41:27 -0700 Subject: [PATCH 02/15] Pluralize "add-" flags for consistency --- tool/tctl/common/bots_command.go | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index f887bcf6b22a6..cc680b2c47e66 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -56,7 +56,7 @@ type BotsCommand struct { botRoles string tokenID string tokenTTL time.Duration - addRole string + addRoles string allowedLogins []string addLogins string @@ -96,9 +96,9 @@ func (c *BotsCommand) Initialize(app *kingpin.Application, config *servicecfg.Co c.botsUpdate = bots.Command("update", "Update an existing bot.") c.botsUpdate.Arg("name", "Name of an existing bot to update.").Required().StringVar(&c.botName) c.botsUpdate.Flag("set-roles", "A comma-separated list of roles to replace.").StringVar(&c.botRoles) - c.botsUpdate.Flag("add-role", "A comma-separated list of roles to add to an existing bot.").StringVar(&c.addRole) + c.botsUpdate.Flag("add-roles", "A comma-separated list of roles to add to an existing bot.").StringVar(&c.addRoles) c.botsUpdate.Flag("set-logins", "A comma-separated list of allowed logins to replace").StringVar(&c.setLogins) - c.botsUpdate.Flag("add-login", "A comma-separated list of logins to add to an existing bot.").StringVar(&c.addLogins) + c.botsUpdate.Flag("add-logins", "A comma-separated list of logins to add to an existing bot.").StringVar(&c.addLogins) } // TryRun attempts to run subcommands. @@ -571,8 +571,8 @@ func (c *BotsCommand) updateBotRole(ctx context.Context, client auth.ClientI) er desiredRoles = setUnion(currentRoles) } - if c.addRole != "" { - desiredRoles = setUnion(desiredRoles, arrayToSet(splitEntries(c.addRole))) + if c.addRoles != "" { + desiredRoles = setUnion(desiredRoles, arrayToSet(splitEntries(c.addRoles))) } log.Infof("Desired roles for bot %q: %+v", c.botName, setToArray(desiredRoles)) @@ -601,7 +601,7 @@ func (c *BotsCommand) updateBotRole(ctx context.Context, client auth.ClientI) er func (c *BotsCommand) UpdateBot(ctx context.Context, client auth.ClientI) error { changes := false - if len(c.allowedLogins) > 0 || len(c.addLogins) > 0 { + if c.setLogins != "" || c.addLogins != "" { if err := c.updateBotUser(ctx, client); err != nil { return trace.Wrap(err) } @@ -609,7 +609,7 @@ func (c *BotsCommand) UpdateBot(ctx context.Context, client auth.ClientI) error changes = true } - if c.botRoles != "" || c.addRole != "" { + if c.botRoles != "" || c.addRoles != "" { if err := c.updateBotRole(ctx, client); err != nil { return trace.Wrap(err) } @@ -619,6 +619,8 @@ func (c *BotsCommand) UpdateBot(ctx context.Context, client auth.ClientI) error if changes { log.Infof("Bot %q has been updated. Roles will take affect on its next renewal.", c.botName) + } else { + log.Infof("No changes requested. Specify one or more flags.") } return nil From 91ca4db0afc1f72b7d93065642b746b12cbe52b7 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Wed, 10 Jan 2024 16:48:49 -0700 Subject: [PATCH 03/15] Add docs for `tctl bots update` Also merges in some other misc docs changes for `tctl bots ...` --- docs/pages/reference/cli/tctl.mdx | 54 +++++++++++++++++++++++++++++-- 1 file changed, 52 insertions(+), 2 deletions(-) diff --git a/docs/pages/reference/cli/tctl.mdx b/docs/pages/reference/cli/tctl.mdx index 6ff1a1a04c6db..7db3990b5b7aa 100644 --- a/docs/pages/reference/cli/tctl.mdx +++ b/docs/pages/reference/cli/tctl.mdx @@ -405,6 +405,15 @@ $ tctl bots add [] | `--logins` | none | Comma-separated list of allowed SSH logins to set Bot's login trait to. | Optional. Specifies the values that should be configured as the Bot's logins trait for the purpose of role templating. | | `--token` | none | String name of an existing join token. | Optional. Specifies an existing join token to be used rather than creating one as part of the Bot creation. | | `--ttl` | 15m | Duration. | Optional. Overrides the TTL of the token that will be created if `--token` is not specified. | +| `--format` | `text` | `text`, `json` | If set to `json`, return new bot information as a machine readable JSON string. | + +### Examples + +Create a new bot named `example` that may assume the `access` role and login as `root`: + +```code +$ tctl bots add example --roles=access --logins=root +``` ### Global flags @@ -441,6 +450,49 @@ $ tctl bots rm [] These flags are available for all commands: `--debug, --config`. Run `tctl help ` or see the [Global Flags section](#tctl-global-flags). +## tctl bots update + +Update a Machine ID bot: + +```code +$ tctl bots update [] +``` + +### Arguments + +* `` - The name of the bot to update + +### Flags + +| Name | Default Value(s) | Allowed Value(s) | Description | +| - | - | - | - | +| `--add-roles` | none | Comma-separated list of roles the bot may assume | Appends the given roles to the existing roles the bot may assume. | +| `--set-roles` | none | Comma-separated list of roles the bot may assume | Replaces the bots's current roles with the list provided. | +| `--add-logins` | none | Comma-separated list of allowed Unix logins | Appends the given logins to the bot's current allowed logins. | +| `--set-logins` | none | Comma-separated list of allowed Unix logins | Replaces the bot's current allowed logins with the given list. | + +### Examples + +Replace the bot `example` roles and add a new allowed Unix login: + +```code +$ tctl bots update example --add-login root --set-roles=access +``` + +Remove all implicitly allowed Unix logins from a bot named `example` by passing +an empty list: + +```code +$ tctl bots update example --set-logins=, +``` + +(Note that the bot may still be granted additional logins via roles.) + +### Global flags + +These flags are available for all commands: `--debug, --config`. Run +`tctl help ` or see the [Global Flags section](#tctl-global-flags). + ## tctl create Create or update a Teleport resource from a YAML file. @@ -1492,5 +1544,3 @@ Print the version of your `tctl` binary: ```code tctl version ``` - - From 3bbcb5b4cba747f7a60c650e39e02420bb46c5a6 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Wed, 10 Jan 2024 16:56:29 -0700 Subject: [PATCH 04/15] Fix docs lint --- docs/pages/reference/cli/tctl.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/pages/reference/cli/tctl.mdx b/docs/pages/reference/cli/tctl.mdx index 7db3990b5b7aa..1f1ce7670fc48 100644 --- a/docs/pages/reference/cli/tctl.mdx +++ b/docs/pages/reference/cli/tctl.mdx @@ -460,7 +460,7 @@ $ tctl bots update [] ### Arguments -* `` - The name of the bot to update +- `` - The name of the bot to update ### Flags From 89e901ee04b4eaf23e50df187f1bc8c831481103 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Wed, 10 Jan 2024 16:57:18 -0700 Subject: [PATCH 05/15] Fix incorrect example --- docs/pages/reference/cli/tctl.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/pages/reference/cli/tctl.mdx b/docs/pages/reference/cli/tctl.mdx index 1f1ce7670fc48..094fe79a4c8f9 100644 --- a/docs/pages/reference/cli/tctl.mdx +++ b/docs/pages/reference/cli/tctl.mdx @@ -476,7 +476,7 @@ $ tctl bots update [] Replace the bot `example` roles and add a new allowed Unix login: ```code -$ tctl bots update example --add-login root --set-roles=access +$ tctl bots update example --add-logins root --set-roles=access ``` Remove all implicitly allowed Unix logins from a bot named `example` by passing From b6e9078dd9c6abeae9fd72e91257b463e3da9fa4 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Thu, 11 Jan 2024 17:06:40 -0700 Subject: [PATCH 06/15] Apply suggestions from code review Co-authored-by: Paul Gottschling --- docs/pages/reference/cli/tctl.mdx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/pages/reference/cli/tctl.mdx b/docs/pages/reference/cli/tctl.mdx index 094fe79a4c8f9..edbebd34d9ec2 100644 --- a/docs/pages/reference/cli/tctl.mdx +++ b/docs/pages/reference/cli/tctl.mdx @@ -405,11 +405,11 @@ $ tctl bots add [] | `--logins` | none | Comma-separated list of allowed SSH logins to set Bot's login trait to. | Optional. Specifies the values that should be configured as the Bot's logins trait for the purpose of role templating. | | `--token` | none | String name of an existing join token. | Optional. Specifies an existing join token to be used rather than creating one as part of the Bot creation. | | `--ttl` | 15m | Duration. | Optional. Overrides the TTL of the token that will be created if `--token` is not specified. | -| `--format` | `text` | `text`, `json` | If set to `json`, return new bot information as a machine readable JSON string. | +| `--format` | `text` | `text`, `json` | If set to `json`, return new bot information as a machine-readable JSON string. | ### Examples -Create a new bot named `example` that may assume the `access` role and login as `root`: +Create a new bot named `example` that may assume the `access` role and log in as `root`: ```code $ tctl bots add example --roles=access --logins=root @@ -486,7 +486,7 @@ an empty list: $ tctl bots update example --set-logins=, ``` -(Note that the bot may still be granted additional logins via roles.) +Note that the bot may still be granted additional logins via roles. ### Global flags From 257a2eba444e2da4fae377ed8a3a7ab6d18d4809 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Thu, 11 Jan 2024 18:15:44 -0700 Subject: [PATCH 07/15] Use the bot service's `UpdateBot()` RPC instead of modifying directly This switches to use the new `UpdateBot()` RPC instead of modifying the bot user and role directly. Also, tweaks CLI messaging to be a bit less chatty. --- tool/tctl/common/bots_command.go | 102 ++++++++++++++++--------------- 1 file changed, 53 insertions(+), 49 deletions(-) diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index cc680b2c47e66..6f8e9c1114fce 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -31,6 +31,7 @@ import ( "github.com/google/uuid" "github.com/gravitational/trace" log "github.com/sirupsen/logrus" + "google.golang.org/protobuf/types/known/fieldmaskpb" "github.com/gravitational/teleport" "github.com/gravitational/teleport/api/client/proto" @@ -496,18 +497,12 @@ func (c *BotsCommand) LockBot(ctx context.Context, client auth.ClientI) error { return nil } -// updateBotUser applies updates to a bot user that were specified in the CLI -// arguments. -func (c *BotsCommand) updateBotUser(ctx context.Context, client auth.ClientI) error { - resource := machineidv1.BotResourceName(c.botName) - user, err := client.GetUser(ctx, resource, false) - if err != nil { - return trace.Wrap(err) - } - - traits := user.GetTraits() - if traits == nil { - traits = map[string][]string{} +// updateBotLogins applies updates from CLI arguments to a bot's logins trait, +// updating the field mask if any updates were made. +func (c *BotsCommand) updateBotLogins(ctx context.Context, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error { + traits := map[string][]string{} + for _, t := range bot.Spec.GetTraits() { + traits[t.Name] = t.Values } var currentLogins map[string]struct{} @@ -529,40 +524,37 @@ func (c *BotsCommand) updateBotUser(ctx context.Context, client auth.ClientI) er desiredLogins = setUnion(desiredLogins, arrayToSet(addLogins)) } - log.Infof("Desired logins for bot %q: %+v", c.botName, setToArray(desiredLogins)) if setsEqual(currentLogins, desiredLogins) { - log.Infof("Desired logins match existing set, will not update bot user") + log.Infof("Requested logins match existing, nothing to do: %+v", setToArray(desiredLogins)) return nil } + log.Infof("Desired logins for bot %q: %+v", c.botName, setToArray(desiredLogins)) + if len(desiredLogins) == 0 { delete(traits, constants.TraitLogins) log.Infof("Removing logins trait from bot user") } else { traits[constants.TraitLogins] = setToArray(desiredLogins) } - user.SetTraits(traits) - if _, err := client.UpdateUser(ctx, user); err != nil { - return trace.Wrap(err) + traitsArray := []*machineidv1pb.Trait{} + for k, v := range traits { + traitsArray = append(traitsArray, &machineidv1pb.Trait{ + Name: k, + Values: v, + }) } - return nil -} - -// updateBotRole performs any updates to a bot's role that were specified in the -// CLI arguments. -func (c *BotsCommand) updateBotRole(ctx context.Context, client auth.ClientI) error { - resource := machineidv1.BotResourceName(c.botName) - role, err := client.GetRole(ctx, resource) - if err != nil { - return trace.Wrap(err) - } + bot.Spec.Traits = traitsArray - allowCond := role.GetImpersonateConditions(types.Allow) - currentRoles := arrayToSet(allowCond.Roles) + return trace.Wrap(mask.Append(&machineidv1pb.Bot{}, "spec.traits")) +} - log.Infof("Existing roles for bot %q: %+v", c.botName, setToArray(currentRoles)) +// updateBotRoles applies updates from CLI arguments to a bot's roles, updating +// the field mask as necessary if any updates were made. +func (c *BotsCommand) updateBotRoles(ctx context.Context, client auth.ClientI, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error { + currentRoles := arrayToSet(bot.Spec.Roles) var desiredRoles map[string]struct{} if c.botRoles != "" { @@ -575,12 +567,13 @@ func (c *BotsCommand) updateBotRole(ctx context.Context, client auth.ClientI) er desiredRoles = setUnion(desiredRoles, arrayToSet(splitEntries(c.addRoles))) } - log.Infof("Desired roles for bot %q: %+v", c.botName, setToArray(desiredRoles)) if setsEqual(currentRoles, desiredRoles) { - log.Infof("Desired roles match existing roles, will not update bot role") + log.Infof("Requested roles match existing, nothing to do: %+v", setToArray(desiredRoles)) return nil } + log.Infof("Desired roles for bot %q: %+v", c.botName, setToArray(desiredRoles)) + // Validate roles (server does not do this yet). for roleName := range desiredRoles { if _, err := client.GetRole(ctx, roleName); err != nil { @@ -588,41 +581,52 @@ func (c *BotsCommand) updateBotRole(ctx context.Context, client auth.ClientI) er } } - allowCond.Roles = setToArray(desiredRoles) - role.SetImpersonateConditions(types.Allow, allowCond) + bot.Spec.Roles = setToArray(desiredRoles) - if _, err := client.UpdateRole(ctx, role); err != nil { - return trace.Wrap(err) - } - return nil + return trace.Wrap(mask.Append(&machineidv1pb.Bot{}, "spec.roles")) } // UpdateBot performs various updates to existing bot users and roles. func (c *BotsCommand) UpdateBot(ctx context.Context, client auth.ClientI) error { - changes := false + bot, err := client.BotServiceClient().GetBot(ctx, &machineidv1pb.GetBotRequest{ + BotName: c.botName, + }) + if err != nil { + return trace.Wrap(err) + } + + fieldMask, err := fieldmaskpb.New(&machineidv1pb.Bot{}) + if err != nil { + return trace.Wrap(err) + } if c.setLogins != "" || c.addLogins != "" { - if err := c.updateBotUser(ctx, client); err != nil { + if err := c.updateBotLogins(ctx, bot, fieldMask); err != nil { return trace.Wrap(err) } - - changes = true } if c.botRoles != "" || c.addRoles != "" { - if err := c.updateBotRole(ctx, client); err != nil { + if err := c.updateBotRoles(ctx, client, bot, fieldMask); err != nil { return trace.Wrap(err) } + } - changes = true + if len(fieldMask.Paths) == 0 { + log.Infof("No changes requested, nothing to do.") + return nil } - if changes { - log.Infof("Bot %q has been updated. Roles will take affect on its next renewal.", c.botName) - } else { - log.Infof("No changes requested. Specify one or more flags.") + _, err = client.BotServiceClient().UpdateBot(ctx, &machineidv1pb.UpdateBotRequest{ + Bot: bot, + UpdateMask: fieldMask, + }) + if err != nil { + return trace.Wrap(err) } + log.Infof("Bot %q has been updated. Roles will take affect on its next renewal.", c.botName) + return nil } From a9ae2213a4e2b16497817f206db8f571fcec212d Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Tue, 16 Jan 2024 16:18:35 -0700 Subject: [PATCH 08/15] Update tool/tctl/common/bots_command.go Co-authored-by: Hugo Shaka --- tool/tctl/common/bots_command.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index 6f8e9c1114fce..074f6b8dd981f 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -625,7 +625,7 @@ func (c *BotsCommand) UpdateBot(ctx context.Context, client auth.ClientI) error return trace.Wrap(err) } - log.Infof("Bot %q has been updated. Roles will take affect on its next renewal.", c.botName) + log.Infof("Bot %q has been updated. Roles will take effect on its next renewal.", c.botName) return nil } From e670493fc09141b0b10abfb8e90125bc5aac338b Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Tue, 16 Jan 2024 17:25:31 -0700 Subject: [PATCH 09/15] Move set functions into their own reusable package The set package adds a little extra functionality but is still missing various important functions, like subset/superset, difference, intersection, xor/symmetric difference, etc. This changes the API slightly, so the tctl command is updated to use the new package. --- lib/utils/set/set.go | 76 +++++++++++++++++++++++ lib/utils/set/set_test.go | 101 +++++++++++++++++++++++++++++++ tool/tctl/common/bots_command.go | 90 +++++++-------------------- 3 files changed, 198 insertions(+), 69 deletions(-) create mode 100644 lib/utils/set/set.go create mode 100644 lib/utils/set/set_test.go diff --git a/lib/utils/set/set.go b/lib/utils/set/set.go new file mode 100644 index 0000000000000..d1b0513f248a1 --- /dev/null +++ b/lib/utils/set/set.go @@ -0,0 +1,76 @@ +/* + * Teleport + * Copyright (C) 2024 Gravitational, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +package set + +import "golang.org/x/exp/maps" + +// Set represents a set data structure, using the map type with an empty value. +type Set[T comparable] map[T]struct{} + +// Empty returns a new, empty set. +func Empty[T comparable]() Set[T] { + return make(map[T]struct{}) +} + +// Union returns a new Set that is the union of one or more other sets. A single +// argument effectively clones the set. +func Union[T comparable](sets ...Set[T]) Set[T] { + ret := make(map[T]struct{}) + for _, set := range sets { + for key := range set { + ret[key] = struct{}{} + } + } + + return ret +} + +// New returns a new Set from the given array of entries. Duplicate entries will +// be removed. +func New[T comparable](entries ...T) Set[T] { + ret := make(map[T]struct{}) + + for _, entry := range entries { + ret[entry] = struct{}{} + } + + return ret +} + +// Union returns a new Set that is the union of this set and one or more other +// sets. +func (s Set[T]) Union(sets ...Set[T]) Set[T] { + args := append([]Set[T]{s}, sets...) + return Union(args...) +} + +// Clone returns a clone of this set. +func (s Set[T]) Clone() Set[T] { + return Union(s) +} + +// ToArray converts a set to an array. Note that the output order is undefined. +func (s Set[T]) ToArray() []T { + return maps.Keys(s) +} + +// Equal determines if this set contains the same keys as another set. +func (s Set[T]) Equals(other Set[T]) bool { + return maps.Equal(s, other) +} diff --git a/lib/utils/set/set_test.go b/lib/utils/set/set_test.go new file mode 100644 index 0000000000000..468817887fae8 --- /dev/null +++ b/lib/utils/set/set_test.go @@ -0,0 +1,101 @@ +/* + * Teleport + * Copyright (C) 2024 Gravitational, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +package set + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSetsEqual(t *testing.T) { + tests := []struct { + name string + a Set[string] + b Set[string] + assert require.BoolAssertionFunc + }{ + { + name: "out of order true", + a: New("a", "b", "c", "d"), + b: New("d", "b", "c", "a"), + assert: require.True, + }, + { + name: "length mismatch", + a: New("a", "b", "c"), + b: New("d", "b", "c", "a"), + assert: require.False, + }, + { + name: "simple false", + a: New("a", "b", "c", "d"), + b: New("d", "b", "c", "e"), + assert: require.False, + }, + { + name: "duplicates ignored", + a: New("a", "b", "c", "d"), + b: New("d", "b", "c", "a", "a", "b", "c"), + assert: require.True, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + tt.assert(t, tt.a.Equals(tt.b)) + }) + } +} + +func TestArrayToSetRoundtrip(t *testing.T) { + a := []string{"a", "b", "c", "d"} + require.ElementsMatch(t, a, New(a...).ToArray()) + + // It should also remove duplicates + require.ElementsMatch(t, New(append(a, a...)...).ToArray(), a) +} + +func TestSetUnion(t *testing.T) { + a := []string{"a", "b", "c", "d"} + b := []string{"c", "d", "e", "f"} + + // Self union clones + require.ElementsMatch( + t, + Union(New(a...), New(a...), New(a...)).ToArray(), + a, + ) + + // Clone also clones (and is a union) + require.ElementsMatch( + t, + New(a...).Clone().ToArray(), + a, + ) + + require.ElementsMatch( + t, + New(a...).Union(New(b...)).ToArray(), + []string{"a", "b", "c", "d", "e", "f"}, + ) +} diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index 6f8e9c1114fce..efec188e7e68a 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -45,6 +45,7 @@ import ( "github.com/gravitational/teleport/lib/defaults" "github.com/gravitational/teleport/lib/service/servicecfg" "github.com/gravitational/teleport/lib/utils" + "github.com/gravitational/teleport/lib/utils/set" ) type BotsCommand struct { @@ -505,37 +506,37 @@ func (c *BotsCommand) updateBotLogins(ctx context.Context, bot *machineidv1pb.Bo traits[t.Name] = t.Values } - var currentLogins map[string]struct{} + var currentLogins set.Set[string] if logins, exists := traits[constants.TraitLogins]; exists { - currentLogins = arrayToSet(logins) + currentLogins = set.New(logins...) } else { - currentLogins = map[string]struct{}{} + currentLogins = set.Empty[string]() } - var desiredLogins map[string]struct{} + var desiredLogins set.Set[string] if c.setLogins != "" { - desiredLogins = arrayToSet(splitEntries(c.setLogins)) + desiredLogins = set.New(splitEntries(c.setLogins)...) } else { - desiredLogins = setUnion(currentLogins) + desiredLogins = currentLogins.Clone() } addLogins := splitEntries(c.addLogins) if len(addLogins) > 0 { - desiredLogins = setUnion(desiredLogins, arrayToSet(addLogins)) + desiredLogins = desiredLogins.Union(set.New(addLogins...)) } - if setsEqual(currentLogins, desiredLogins) { - log.Infof("Requested logins match existing, nothing to do: %+v", setToArray(desiredLogins)) + if currentLogins.Equals(desiredLogins) { + log.Infof("Requested logins match existing, nothing to do: %+v", desiredLogins.ToArray()) return nil } - log.Infof("Desired logins for bot %q: %+v", c.botName, setToArray(desiredLogins)) + log.Infof("Desired logins for bot %q: %+v", c.botName, desiredLogins.ToArray()) if len(desiredLogins) == 0 { delete(traits, constants.TraitLogins) log.Infof("Removing logins trait from bot user") } else { - traits[constants.TraitLogins] = setToArray(desiredLogins) + traits[constants.TraitLogins] = desiredLogins.ToArray() } traitsArray := []*machineidv1pb.Trait{} @@ -554,25 +555,25 @@ func (c *BotsCommand) updateBotLogins(ctx context.Context, bot *machineidv1pb.Bo // updateBotRoles applies updates from CLI arguments to a bot's roles, updating // the field mask as necessary if any updates were made. func (c *BotsCommand) updateBotRoles(ctx context.Context, client auth.ClientI, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error { - currentRoles := arrayToSet(bot.Spec.Roles) + currentRoles := set.New(bot.Spec.Roles...) - var desiredRoles map[string]struct{} + var desiredRoles set.Set[string] if c.botRoles != "" { - desiredRoles = arrayToSet(splitEntries(c.botRoles)) + desiredRoles = set.New(splitEntries(c.botRoles)...) } else { - desiredRoles = setUnion(currentRoles) + desiredRoles = currentRoles.Clone() } if c.addRoles != "" { - desiredRoles = setUnion(desiredRoles, arrayToSet(splitEntries(c.addRoles))) + desiredRoles = set.New(splitEntries(c.addRoles)...).Union(desiredRoles) } - if setsEqual(currentRoles, desiredRoles) { - log.Infof("Requested roles match existing, nothing to do: %+v", setToArray(desiredRoles)) + if currentRoles.Equals(desiredRoles) { + log.Infof("Requested roles match existing, nothing to do: %+v", desiredRoles.ToArray()) return nil } - log.Infof("Desired roles for bot %q: %+v", c.botName, setToArray(desiredRoles)) + log.Infof("Desired roles for bot %q: %+v", c.botName, desiredRoles.ToArray()) // Validate roles (server does not do this yet). for roleName := range desiredRoles { @@ -581,7 +582,7 @@ func (c *BotsCommand) updateBotRoles(ctx context.Context, client auth.ClientI, b } } - bot.Spec.Roles = setToArray(desiredRoles) + bot.Spec.Roles = desiredRoles.ToArray() return trace.Wrap(mask.Append(&machineidv1pb.Bot{}, "spec.roles")) } @@ -643,52 +644,3 @@ func splitEntries(flag string) []string { } return roles } - -// setsEqual determines if two sets contain the same keys. -func setsEqual[T comparable](a map[T]struct{}, b map[T]struct{}) bool { - if len(a) != len(b) { - return false - } - - for k := range a { - if _, ok := b[k]; !ok { - return false - } - } - - return true -} - -// setToArray converts a set to an array -func setToArray[T comparable](m map[T]struct{}) []T { - var ret []T - for entry := range m { - ret = append(ret, entry) - } - - return ret -} - -// arrayToSet converts an array to a set, removing duplicates -func arrayToSet[T comparable](arr []T) map[T]struct{} { - ret := make(map[T]struct{}) - - for _, entry := range arr { - ret[entry] = struct{}{} - } - - return ret -} - -// setUnion returns a new set containing a union of all provided sets. If only -// one set is given, it is effectively shallow copied. -func setUnion[T comparable](sets ...map[T]struct{}) map[T]struct{} { - ret := make(map[T]struct{}) - for _, set := range sets { - for key := range set { - ret[key] = struct{}{} - } - } - - return ret -} From 23738de2531cb69d9868fdc51001aa497854f788 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Tue, 16 Jan 2024 17:37:06 -0700 Subject: [PATCH 10/15] Remove duplicate tests --- tool/tctl/common/bots_command_test.go | 94 --------------------------- 1 file changed, 94 deletions(-) delete mode 100644 tool/tctl/common/bots_command_test.go diff --git a/tool/tctl/common/bots_command_test.go b/tool/tctl/common/bots_command_test.go deleted file mode 100644 index bad3d516ac59b..0000000000000 --- a/tool/tctl/common/bots_command_test.go +++ /dev/null @@ -1,94 +0,0 @@ -/* - * Teleport - * Copyright (C) 2024 Gravitational, Inc. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as published by - * the Free Software Foundation, either version 3 of the License, or - * (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - */ - -package common - -import ( - "testing" - - "github.com/stretchr/testify/require" -) - -func TestSetsEqual(t *testing.T) { - tests := []struct { - name string - a map[string]struct{} - b map[string]struct{} - assert require.BoolAssertionFunc - }{ - { - name: "out of order true", - a: arrayToSet([]string{"a", "b", "c", "d"}), - b: arrayToSet([]string{"d", "b", "c", "a"}), - assert: require.True, - }, - { - name: "length mismatch", - a: arrayToSet([]string{"a", "b", "c"}), - b: arrayToSet([]string{"d", "b", "c", "a"}), - assert: require.False, - }, - { - name: "simple false", - a: arrayToSet([]string{"a", "b", "c", "d"}), - b: arrayToSet([]string{"d", "b", "c", "e"}), - assert: require.False, - }, - { - name: "duplicates ignored", - a: arrayToSet([]string{"a", "b", "c", "d"}), - b: arrayToSet([]string{"d", "b", "c", "a", "a", "b", "c"}), - assert: require.True, - }, - } - - for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - tt.assert(t, setsEqual(tt.a, tt.b)) - }) - } -} - -func TestArrayToSetRoundtrip(t *testing.T) { - a := []string{"a", "b", "c", "d"} - require.ElementsMatch(t, a, setToArray(arrayToSet(a))) - - // It should also remove duplicates - require.ElementsMatch(t, setToArray(arrayToSet(append(a, a...))), a) -} - -func TestSetUnion(t *testing.T) { - a := []string{"a", "b", "c", "d"} - b := []string{"c", "d", "e", "f"} - - // Self union clones - require.ElementsMatch( - t, - setToArray(setUnion(arrayToSet(a), arrayToSet(a), arrayToSet(a))), - a, - ) - - require.ElementsMatch( - t, - setToArray(setUnion(arrayToSet(a), arrayToSet(b))), - []string{"a", "b", "c", "d", "e", "f"}, - ) -} From e1f7684e954078d1c099ceba86edc20c47a2dd40 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Tue, 16 Jan 2024 17:37:25 -0700 Subject: [PATCH 11/15] Reword flag help strings for clarity --- tool/tctl/common/bots_command.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index c80cfcccf21c1..c462fcb67e3e6 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -97,10 +97,10 @@ func (c *BotsCommand) Initialize(app *kingpin.Application, config *servicecfg.Co c.botsUpdate = bots.Command("update", "Update an existing bot.") c.botsUpdate.Arg("name", "Name of an existing bot to update.").Required().StringVar(&c.botName) - c.botsUpdate.Flag("set-roles", "A comma-separated list of roles to replace.").StringVar(&c.botRoles) - c.botsUpdate.Flag("add-roles", "A comma-separated list of roles to add to an existing bot.").StringVar(&c.addRoles) - c.botsUpdate.Flag("set-logins", "A comma-separated list of allowed logins to replace").StringVar(&c.setLogins) - c.botsUpdate.Flag("add-logins", "A comma-separated list of logins to add to an existing bot.").StringVar(&c.addLogins) + c.botsUpdate.Flag("set-roles", "Sets the bot's roles to the given comma-separated list, replacing any existing roles.").StringVar(&c.botRoles) + c.botsUpdate.Flag("add-roles", "Adds a comma-separated list of roles to an existing bot.").StringVar(&c.addRoles) + c.botsUpdate.Flag("set-logins", "Sets the bot's logins to the given comma-separated list, replacing any existing logins.").StringVar(&c.setLogins) + c.botsUpdate.Flag("add-logins", "Adds a comma-separated list of logins to an existing bot.").StringVar(&c.addLogins) } // TryRun attempts to run subcommands. From 9da8659994dd992a4230206d1325cc0b68497050 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Wed, 17 Jan 2024 16:20:48 -0700 Subject: [PATCH 12/15] Remove set abstraction Removes the Set package and unrolls all related function calls in bots_command. --- lib/utils/set/set.go | 76 ----------------------- lib/utils/set/set_test.go | 101 ------------------------------- tool/tctl/common/bots_command.go | 67 +++++++++++++------- 3 files changed, 45 insertions(+), 199 deletions(-) delete mode 100644 lib/utils/set/set.go delete mode 100644 lib/utils/set/set_test.go diff --git a/lib/utils/set/set.go b/lib/utils/set/set.go deleted file mode 100644 index d1b0513f248a1..0000000000000 --- a/lib/utils/set/set.go +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Teleport - * Copyright (C) 2024 Gravitational, Inc. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as published by - * the Free Software Foundation, either version 3 of the License, or - * (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - */ - -package set - -import "golang.org/x/exp/maps" - -// Set represents a set data structure, using the map type with an empty value. -type Set[T comparable] map[T]struct{} - -// Empty returns a new, empty set. -func Empty[T comparable]() Set[T] { - return make(map[T]struct{}) -} - -// Union returns a new Set that is the union of one or more other sets. A single -// argument effectively clones the set. -func Union[T comparable](sets ...Set[T]) Set[T] { - ret := make(map[T]struct{}) - for _, set := range sets { - for key := range set { - ret[key] = struct{}{} - } - } - - return ret -} - -// New returns a new Set from the given array of entries. Duplicate entries will -// be removed. -func New[T comparable](entries ...T) Set[T] { - ret := make(map[T]struct{}) - - for _, entry := range entries { - ret[entry] = struct{}{} - } - - return ret -} - -// Union returns a new Set that is the union of this set and one or more other -// sets. -func (s Set[T]) Union(sets ...Set[T]) Set[T] { - args := append([]Set[T]{s}, sets...) - return Union(args...) -} - -// Clone returns a clone of this set. -func (s Set[T]) Clone() Set[T] { - return Union(s) -} - -// ToArray converts a set to an array. Note that the output order is undefined. -func (s Set[T]) ToArray() []T { - return maps.Keys(s) -} - -// Equal determines if this set contains the same keys as another set. -func (s Set[T]) Equals(other Set[T]) bool { - return maps.Equal(s, other) -} diff --git a/lib/utils/set/set_test.go b/lib/utils/set/set_test.go deleted file mode 100644 index 468817887fae8..0000000000000 --- a/lib/utils/set/set_test.go +++ /dev/null @@ -1,101 +0,0 @@ -/* - * Teleport - * Copyright (C) 2024 Gravitational, Inc. - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as published by - * the Free Software Foundation, either version 3 of the License, or - * (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. - * - * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . - */ - -package set - -import ( - "testing" - - "github.com/stretchr/testify/require" -) - -func TestSetsEqual(t *testing.T) { - tests := []struct { - name string - a Set[string] - b Set[string] - assert require.BoolAssertionFunc - }{ - { - name: "out of order true", - a: New("a", "b", "c", "d"), - b: New("d", "b", "c", "a"), - assert: require.True, - }, - { - name: "length mismatch", - a: New("a", "b", "c"), - b: New("d", "b", "c", "a"), - assert: require.False, - }, - { - name: "simple false", - a: New("a", "b", "c", "d"), - b: New("d", "b", "c", "e"), - assert: require.False, - }, - { - name: "duplicates ignored", - a: New("a", "b", "c", "d"), - b: New("d", "b", "c", "a", "a", "b", "c"), - assert: require.True, - }, - } - - for _, tt := range tests { - tt := tt - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - tt.assert(t, tt.a.Equals(tt.b)) - }) - } -} - -func TestArrayToSetRoundtrip(t *testing.T) { - a := []string{"a", "b", "c", "d"} - require.ElementsMatch(t, a, New(a...).ToArray()) - - // It should also remove duplicates - require.ElementsMatch(t, New(append(a, a...)...).ToArray(), a) -} - -func TestSetUnion(t *testing.T) { - a := []string{"a", "b", "c", "d"} - b := []string{"c", "d", "e", "f"} - - // Self union clones - require.ElementsMatch( - t, - Union(New(a...), New(a...), New(a...)).ToArray(), - a, - ) - - // Clone also clones (and is a union) - require.ElementsMatch( - t, - New(a...).Clone().ToArray(), - a, - ) - - require.ElementsMatch( - t, - New(a...).Union(New(b...)).ToArray(), - []string{"a", "b", "c", "d", "e", "f"}, - ) -} diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index c462fcb67e3e6..a1aa7d1791ad7 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -22,6 +22,7 @@ import ( "context" "encoding/json" "fmt" + "maps" "os" "strings" "text/template" @@ -45,7 +46,6 @@ import ( "github.com/gravitational/teleport/lib/defaults" "github.com/gravitational/teleport/lib/service/servicecfg" "github.com/gravitational/teleport/lib/utils" - "github.com/gravitational/teleport/lib/utils/set" ) type BotsCommand struct { @@ -506,37 +506,47 @@ func (c *BotsCommand) updateBotLogins(ctx context.Context, bot *machineidv1pb.Bo traits[t.Name] = t.Values } - var currentLogins set.Set[string] + currentLogins := make(map[string]struct{}) if logins, exists := traits[constants.TraitLogins]; exists { - currentLogins = set.New(logins...) - } else { - currentLogins = set.Empty[string]() + for _, login := range logins { + currentLogins[login] = struct{}{} + } } - var desiredLogins set.Set[string] + var desiredLogins map[string]struct{} if c.setLogins != "" { - desiredLogins = set.New(splitEntries(c.setLogins)...) + desiredLogins = make(map[string]struct{}) + for _, login := range splitEntries(c.setLogins) { + desiredLogins[login] = struct{}{} + } } else { - desiredLogins = currentLogins.Clone() + desiredLogins = maps.Clone(currentLogins) } addLogins := splitEntries(c.addLogins) if len(addLogins) > 0 { - desiredLogins = desiredLogins.Union(set.New(addLogins...)) + for _, login := range addLogins { + desiredLogins[login] = struct{}{} + } + } + + var desiredLoginsArray []string + for login := range desiredLogins { + desiredLoginsArray = append(desiredLoginsArray, login) } - if currentLogins.Equals(desiredLogins) { - log.Infof("Requested logins match existing, nothing to do: %+v", desiredLogins.ToArray()) + if maps.Equal(currentLogins, desiredLogins) { + log.Infof("Logins will be left unchanged: %+v", desiredLoginsArray) return nil } - log.Infof("Desired logins for bot %q: %+v", c.botName, desiredLogins.ToArray()) + log.Infof("Desired logins for bot %q: %+v", c.botName, desiredLoginsArray) if len(desiredLogins) == 0 { delete(traits, constants.TraitLogins) log.Infof("Removing logins trait from bot user") } else { - traits[constants.TraitLogins] = desiredLogins.ToArray() + traits[constants.TraitLogins] = desiredLoginsArray } traitsArray := []*machineidv1pb.Trait{} @@ -555,25 +565,38 @@ func (c *BotsCommand) updateBotLogins(ctx context.Context, bot *machineidv1pb.Bo // updateBotRoles applies updates from CLI arguments to a bot's roles, updating // the field mask as necessary if any updates were made. func (c *BotsCommand) updateBotRoles(ctx context.Context, client auth.ClientI, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error { - currentRoles := set.New(bot.Spec.Roles...) + currentRoles := make(map[string]struct{}) + for _, role := range bot.Spec.Roles { + currentRoles[role] = struct{}{} + } - var desiredRoles set.Set[string] + var desiredRoles map[string]struct{} if c.botRoles != "" { - desiredRoles = set.New(splitEntries(c.botRoles)...) + desiredRoles = make(map[string]struct{}) + for _, role := range splitEntries(c.botRoles) { + desiredRoles[role] = struct{}{} + } } else { - desiredRoles = currentRoles.Clone() + desiredRoles = maps.Clone(currentRoles) } if c.addRoles != "" { - desiredRoles = set.New(splitEntries(c.addRoles)...).Union(desiredRoles) + for _, role := range splitEntries(c.addRoles) { + desiredRoles[role] = struct{}{} + } + } + + var desiredRolesArray []string + for role := range desiredRoles { + desiredRolesArray = append(desiredRolesArray, role) } - if currentRoles.Equals(desiredRoles) { - log.Infof("Requested roles match existing, nothing to do: %+v", desiredRoles.ToArray()) + if maps.Equal(currentRoles, desiredRoles) { + log.Infof("Roles will be left unchanged: %+v", desiredRolesArray) return nil } - log.Infof("Desired roles for bot %q: %+v", c.botName, desiredRoles.ToArray()) + log.Infof("Desired roles for bot %q: %+v", c.botName, desiredRolesArray) // Validate roles (server does not do this yet). for roleName := range desiredRoles { @@ -582,7 +605,7 @@ func (c *BotsCommand) updateBotRoles(ctx context.Context, client auth.ClientI, b } } - bot.Spec.Roles = desiredRoles.ToArray() + bot.Spec.Roles = desiredRolesArray return trace.Wrap(mask.Append(&machineidv1pb.Bot{}, "spec.roles")) } From f6ea1b36e8e23212dd857c71c088e2446be465bf Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Thu, 18 Jan 2024 15:12:10 -0700 Subject: [PATCH 13/15] Add unit tests for `updateBotLogins()` and `updateBotRoles()` This also makes a minimal API wrapper interface for mocking. --- tool/tctl/common/bots_command.go | 11 +- tool/tctl/common/bots_command_test.go | 228 ++++++++++++++++++++++++++ 2 files changed, 236 insertions(+), 3 deletions(-) create mode 100644 tool/tctl/common/bots_command_test.go diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index a1aa7d1791ad7..8809299018f38 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -500,7 +500,7 @@ func (c *BotsCommand) LockBot(ctx context.Context, client auth.ClientI) error { // updateBotLogins applies updates from CLI arguments to a bot's logins trait, // updating the field mask if any updates were made. -func (c *BotsCommand) updateBotLogins(ctx context.Context, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error { +func (c *BotsCommand) updateBotLogins(bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error { traits := map[string][]string{} for _, t := range bot.Spec.GetTraits() { traits[t.Name] = t.Values @@ -562,9 +562,14 @@ func (c *BotsCommand) updateBotLogins(ctx context.Context, bot *machineidv1pb.Bo return trace.Wrap(mask.Append(&machineidv1pb.Bot{}, "spec.traits")) } +// clientRoleGetter is a minimal mockable interface for the client API +type clientRoleGetter interface { + GetRole(context.Context, string) (types.Role, error) +} + // updateBotRoles applies updates from CLI arguments to a bot's roles, updating // the field mask as necessary if any updates were made. -func (c *BotsCommand) updateBotRoles(ctx context.Context, client auth.ClientI, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error { +func (c *BotsCommand) updateBotRoles(ctx context.Context, client clientRoleGetter, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error { currentRoles := make(map[string]struct{}) for _, role := range bot.Spec.Roles { currentRoles[role] = struct{}{} @@ -625,7 +630,7 @@ func (c *BotsCommand) UpdateBot(ctx context.Context, client auth.ClientI) error } if c.setLogins != "" || c.addLogins != "" { - if err := c.updateBotLogins(ctx, bot, fieldMask); err != nil { + if err := c.updateBotLogins(bot, fieldMask); err != nil { return trace.Wrap(err) } } diff --git a/tool/tctl/common/bots_command_test.go b/tool/tctl/common/bots_command_test.go new file mode 100644 index 0000000000000..af32d0c261360 --- /dev/null +++ b/tool/tctl/common/bots_command_test.go @@ -0,0 +1,228 @@ +/* + * Teleport + * Copyright (C) 2024 Gravitational, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published by + * the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +package common + +import ( + "context" + "slices" + "testing" + + "github.com/gravitational/trace" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/types/known/fieldmaskpb" + + "github.com/gravitational/teleport/api/constants" + headerv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/header/v1" + machineidv1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/machineid/v1" + "github.com/gravitational/teleport/api/types" +) + +func TestUpdateBotLogins(t *testing.T) { + tests := []struct { + desc string + add string + set string + initialLogins []string + assert func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) + }{ + { + desc: "should add and set with existing logins", + set: "a,b,c", + add: "d,e,e,e,e", + initialLogins: []string{"a"}, + assert: func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) { + require.NoError(t, err) + require.ElementsMatch(t, mask.Paths, []string{"spec.traits"}) + require.ElementsMatch(t, bot.Spec.Traits[0].Values, splitEntries("a,b,c,d,e")) + }, + }, + { + desc: "should not update with no changes", + set: "a,b,c", + add: "d,e,e,e,e", + initialLogins: splitEntries("a,b,c,d,e"), + assert: func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) { + require.NoError(t, err) + require.Empty(t, mask.Paths) + require.ElementsMatch(t, bot.Spec.Traits[0].Values, splitEntries("a,b,c,d,e")) + }, + }, + { + desc: "should add with empty initial logins trait", + set: "a,b,c", + assert: func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) { + require.NoError(t, err) + require.ElementsMatch(t, mask.Paths, []string{"spec.traits"}) + require.ElementsMatch(t, bot.Spec.Traits[0].Values, splitEntries("a,b,c")) + }, + }, + { + desc: "should remove on set if necessary", + set: "a,b,c", + initialLogins: splitEntries("a,b,c,d,e"), + assert: func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) { + require.NoError(t, err) + require.ElementsMatch(t, mask.Paths, []string{"spec.traits"}) + require.ElementsMatch(t, bot.Spec.Traits[0].Values, splitEntries("a,b,c")) + }, + }, + } + + for _, tt := range tests { + tt := tt + + const botName = "test" + + t.Run(tt.desc, func(t *testing.T) { + traits := []*machineidv1pb.Trait{} + if len(tt.initialLogins) > 0 { + traits = append(traits, &machineidv1pb.Trait{ + Name: constants.TraitLogins, + Values: tt.initialLogins, + }) + } + + bot := &machineidv1pb.Bot{ + Metadata: &headerv1.Metadata{ + Name: botName, + }, + Spec: &machineidv1pb.BotSpec{ + Roles: []string{}, + Traits: traits, + }, + } + + fieldMask, err := fieldmaskpb.New(&machineidv1pb.Bot{}) + require.NoError(t, err) + + cmd := BotsCommand{ + botName: botName, + addLogins: tt.add, + setLogins: tt.set, + } + + err = cmd.updateBotLogins(bot, fieldMask) + tt.assert(t, bot, fieldMask, err) + }) + } +} + +// mockAPIClient is a minimal API client used for testing +type mockRoleGetterClient struct { + roles []string +} + +func (m *mockRoleGetterClient) GetRole(ctx context.Context, name string) (types.Role, error) { + if !slices.Contains(m.roles, name) { + return nil, trace.NotFound("invalid role %s", name) + } + + return types.NewRole(name, types.RoleSpecV6{}) +} + +func TestUpdateBotRoles(t *testing.T) { + tests := []struct { + desc string + add string + set string + initialRoles []string + knownRoles []string + assert func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) + }{ + { + desc: "should add and set without duplicating roles", + set: "a,b,c", + add: "d,e,e,e,e", + knownRoles: splitEntries("a,b,c,d,e"), + initialRoles: []string{"a"}, + assert: func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) { + require.NoError(t, err) + require.ElementsMatch(t, mask.Paths, []string{"spec.roles"}) + require.ElementsMatch(t, bot.Spec.Roles, splitEntries("a,b,c,d,e")) + }, + }, + { + desc: "should not update with no changes", + set: "a,b,c", + add: "d,e,e,e,e", + knownRoles: splitEntries("a,b,c,d,e"), + initialRoles: splitEntries("a,b,c,d,e"), + assert: func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) { + require.NoError(t, err) + require.Empty(t, mask.Paths) + require.ElementsMatch(t, bot.Spec.Roles, splitEntries("a,b,c,d,e")) + }, + }, + { + desc: "should remove on set if necessary", + set: "a,b,c", + knownRoles: splitEntries("a,b,c,d"), + initialRoles: splitEntries("a,b,c,d"), + assert: func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) { + require.NoError(t, err) + require.ElementsMatch(t, mask.Paths, []string{"spec.roles"}) + require.ElementsMatch(t, bot.Spec.Roles, splitEntries("a,b,c")) + }, + }, + { + desc: "should fail if an unknown role is specified and leave bot unmodified", + add: "d", + knownRoles: splitEntries("a,b,c"), + initialRoles: splitEntries("a,b,c"), + assert: func(t *testing.T, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask, err error) { + require.True(t, trace.IsNotFound(err)) + require.Empty(t, mask.Paths) + require.ElementsMatch(t, bot.Spec.Roles, splitEntries("a,b,c")) + }, + }, + } + + for _, tt := range tests { + tt := tt + + const botName = "test" + + t.Run(tt.desc, func(t *testing.T) { + mockClient := mockRoleGetterClient{ + roles: tt.knownRoles, + } + + bot := &machineidv1pb.Bot{ + Metadata: &headerv1.Metadata{ + Name: botName, + }, + Spec: &machineidv1pb.BotSpec{ + Roles: tt.initialRoles, + }, + } + + fieldMask, err := fieldmaskpb.New(&machineidv1pb.Bot{}) + require.NoError(t, err) + + cmd := BotsCommand{ + botName: botName, + addRoles: tt.add, + botRoles: tt.set, + } + + err = cmd.updateBotRoles(context.TODO(), &mockClient, bot, fieldMask) + tt.assert(t, bot, fieldMask, err) + }) + } +} From 37e0df45b8d8688d6b9bf7e6266ad0e952669cf4 Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Thu, 18 Jan 2024 15:56:41 -0700 Subject: [PATCH 14/15] Use utils function to convert sets to a string slice --- tool/tctl/common/bots_command.go | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/tool/tctl/common/bots_command.go b/tool/tctl/common/bots_command.go index 8809299018f38..ce3aa40386e55 100644 --- a/tool/tctl/common/bots_command.go +++ b/tool/tctl/common/bots_command.go @@ -530,10 +530,7 @@ func (c *BotsCommand) updateBotLogins(bot *machineidv1pb.Bot, mask *fieldmaskpb. } } - var desiredLoginsArray []string - for login := range desiredLogins { - desiredLoginsArray = append(desiredLoginsArray, login) - } + desiredLoginsArray := utils.StringsSliceFromSet(desiredLogins) if maps.Equal(currentLogins, desiredLogins) { log.Infof("Logins will be left unchanged: %+v", desiredLoginsArray) @@ -591,10 +588,7 @@ func (c *BotsCommand) updateBotRoles(ctx context.Context, client clientRoleGette } } - var desiredRolesArray []string - for role := range desiredRoles { - desiredRolesArray = append(desiredRolesArray, role) - } + desiredRolesArray := utils.StringsSliceFromSet(desiredRoles) if maps.Equal(currentRoles, desiredRoles) { log.Infof("Roles will be left unchanged: %+v", desiredRolesArray) From 43e64d1c55e8aa60dd7c7f11540d8d5d55235d1b Mon Sep 17 00:00:00 2001 From: Tim Buckley Date: Thu, 18 Jan 2024 16:55:59 -0700 Subject: [PATCH 15/15] Update docs/pages/reference/cli/tctl.mdx Co-authored-by: Zac Bergquist --- docs/pages/reference/cli/tctl.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/pages/reference/cli/tctl.mdx b/docs/pages/reference/cli/tctl.mdx index edbebd34d9ec2..d8eab86f1c9e4 100644 --- a/docs/pages/reference/cli/tctl.mdx +++ b/docs/pages/reference/cli/tctl.mdx @@ -476,7 +476,7 @@ $ tctl bots update [] Replace the bot `example` roles and add a new allowed Unix login: ```code -$ tctl bots update example --add-logins root --set-roles=access +$ tctl bots update example --add-logins=root --set-roles=access ``` Remove all implicitly allowed Unix logins from a bot named `example` by passing