diff --git a/cli/ssh.go b/cli/ssh.go index 4ccaa241d5902..b7f82f00dd7c2 100644 --- a/cli/ssh.go +++ b/cli/ssh.go @@ -24,6 +24,7 @@ import ( "github.com/coder/coder/coderd/autobuild/schedule" "github.com/coder/coder/coderd/database" "github.com/coder/coder/codersdk" + "github.com/coder/coder/cryptorand" ) var autostopPollInterval = 30 * time.Second @@ -31,13 +32,14 @@ var autostopNotifyCountdown = []time.Duration{30 * time.Minute} func ssh() *cobra.Command { var ( - stdio bool + stdio bool + shuffle bool ) cmd := &cobra.Command{ Annotations: workspaceCommand, Use: "ssh ", Short: "SSH into a workspace", - Args: cobra.MinimumNArgs(1), + Args: cobra.ArbitraryArgs, RunE: func(cmd *cobra.Command, args []string) error { client, err := createClient(cmd) if err != nil { @@ -48,10 +50,38 @@ func ssh() *cobra.Command { return err } - workspaceParts := strings.Split(args[0], ".") - workspace, err := client.WorkspaceByOwnerAndName(cmd.Context(), organization.ID, codersdk.Me, workspaceParts[0]) - if err != nil { - return err + var workspace codersdk.Workspace + var workspaceParts []string + if shuffle { + err := cobra.ExactArgs(0)(cmd, args) + if err != nil { + return err + } + + workspaces, err := client.WorkspacesByOwner(cmd.Context(), organization.ID, codersdk.Me) + if err != nil { + return err + } + if len(workspaces) == 0 { + return xerrors.New("no workspaces to shuffle") + } + + idx, err := cryptorand.Intn(len(workspaces)) + if err != nil { + return err + } + workspace = workspaces[idx] + } else { + err := cobra.MinimumNArgs(1)(cmd, args) + if err != nil { + return err + } + + workspaceParts = strings.Split(args[0], ".") + workspace, err = client.WorkspaceByOwnerAndName(cmd.Context(), organization.ID, codersdk.Me, workspaceParts[0]) + if err != nil { + return err + } } if workspace.LatestBuild.Transition != database.WorkspaceTransitionStart { @@ -96,9 +126,17 @@ func ssh() *cobra.Command { } if agent.ID == uuid.Nil { if len(agents) > 1 { - return xerrors.New("you must specify the name of an agent") + if !shuffle { + return xerrors.New("you must specify the name of an agent") + } + idx, err := cryptorand.Intn(len(agents)) + if err != nil { + return err + } + agent = agents[idx] + } else { + agent = agents[0] } - agent = agents[0] } // OpenSSH passes stderr directly to the calling TTY. // This is required in "stdio" mode so a connecting indicator can be displayed. @@ -189,6 +227,8 @@ func ssh() *cobra.Command { }, } cliflag.BoolVarP(cmd.Flags(), &stdio, "stdio", "", "CODER_SSH_STDIO", false, "Specifies whether to emit SSH output over stdin/stdout.") + cliflag.BoolVarP(cmd.Flags(), &shuffle, "shuffle", "", "CODER_SSH_SHUFFLE", false, "Specifies whether to choose a random workspace") + _ = cmd.Flags().MarkHidden("shuffle") return cmd }