chore: merge master

coder · deansheather · May 18, 2022 · May 9, 2022 · May 9, 2022 · May 17, 2022
commit ee256233b3d24b3b54017e351517beff893466de
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -76,7 +76,7 @@ func portForward() *cobra.Command {
 				return err
 			}
 
-			workspace, agent, err := getWorkspaceAndAgent(cmd, client, organization.ID, codersdk.Me, args[0])
+			workspace, agent, err := getWorkspaceAndAgent(cmd, client, organization.ID, codersdk.Me, args[0], false)
 			if err != nil {
 				return err
 			}

diff --git a/cli/ssh.go b/cli/ssh.go
@@ -50,7 +50,19 @@ func ssh() *cobra.Command {
 				return err
 			}
 
-			workspace, agent, err := getWorkspaceAndAgent(cmd, client, organization.ID, codersdk.Me, args[0])
+			if shuffle {
+				err := cobra.ExactArgs(0)(cmd, args)
+				if err != nil {
+					return err
+				}
+			} else {
+				err := cobra.MinimumNArgs(1)(cmd, args)
+				if err != nil {
+					return err
+				}
+			}
+
+			workspace, agent, err := getWorkspaceAndAgent(cmd, client, organization.ID, codersdk.Me, args[0], shuffle)
 			if err != nil {
 				return err
 			}
@@ -150,20 +162,42 @@ func ssh() *cobra.Command {
 	return cmd
 }
 
-func getWorkspaceAndAgent(cmd *cobra.Command, client *codersdk.Client, orgID uuid.UUID, userID string, in string) (codersdk.Workspace, codersdk.WorkspaceAgent, error) {
+// getWorkspaceAgent returns the workspace and agent selected using either the
+// `<workspace>[.<agent>]` syntax via `in` or picks a random workspace and agent
+// if `shuffle` is true.
+func getWorkspaceAndAgent(cmd *cobra.Command, client *codersdk.Client, orgID uuid.UUID, userID string, in string, shuffle bool) (codersdk.Workspace, codersdk.WorkspaceAgent, error) { //nolint:revive
 	ctx := cmd.Context()
 
-	workspaceParts := strings.Split(in, ".")
-	workspace, err := client.WorkspaceByOwnerAndName(ctx, orgID, userID, workspaceParts[0])
-	if err != nil {
-		return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("get workspace %q: %w", workspaceParts[0], err)
+	var (
+		workspace      codersdk.Workspace
+		workspaceParts = strings.Split(in, ".")
+		err            error
+	)
+	if shuffle {
+		workspaces, err := client.WorkspacesByOwner(cmd.Context(), orgID, userID)
+		if err != nil {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
+		}
+		if len(workspaces) == 0 {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.New("no workspaces to shuffle")
+		}
+
+		workspace, err = cryptorand.Element(workspaces)
+		if err != nil {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
+		}
+	} else {
+		workspace, err = client.WorkspaceByOwnerAndName(cmd.Context(), orgID, userID, workspaceParts[0])
+		if err != nil {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
+		}
 	}
 
 	if workspace.LatestBuild.Transition != database.WorkspaceTransitionStart {
 		return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.New("workspace must be in start transition to ssh")
 	}
 	if workspace.LatestBuild.Job.CompletedAt == nil {
-		err = cliui.WorkspaceBuild(ctx, cmd.ErrOrStderr(), client, workspace.LatestBuild.ID, workspace.CreatedAt)
+		err := cliui.WorkspaceBuild(ctx, cmd.ErrOrStderr(), client, workspace.LatestBuild.ID, workspace.CreatedAt)
 		if err != nil {
 			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
 		}
@@ -199,9 +233,16 @@ func getWorkspaceAndAgent(cmd *cobra.Command, client *codersdk.Client, orgID uui
 	}
 	if agent.ID == uuid.Nil {
 		if len(agents) > 1 {
-			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.New("you must specify the name of an agent")
+			if !shuffle {
+				return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.New("you must specify the name of an agent")
+			}
+			agent, err = cryptorand.Element(agents)
+			if err != nil {
+				return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
+			}
+		} else {
+			agent = agents[0]
 		}
-		agent = agents[0]
 	}
 
 	return workspace, agent, nil

diff --git a/cryptorand/slices.go b/cryptorand/slices.go
@@ -0,0 +1,20 @@
+package cryptorand
+
+import (
+	"golang.org/x/xerrors"
+)
+
+// Element returns a random element of the slice. An error will be returned if
+// the slice has no elements in it.
+func Element[T any](s []T) (out T, err error) {
+	if len(s) == 0 {
+		return out, xerrors.New("slice must have at least one element")
+	}
+
+	i, err := Intn(len(s))
+	if err != nil {
+		return out, xerrors.Errorf("generate random integer from 0-%v: %w", len(s), err)
+	}
+
+	return s[i], nil
+}
diff --git a/cryptorand/slices_test.go b/cryptorand/slices_test.go
@@ -0,0 +1,56 @@
+package cryptorand_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/cryptorand"
+)
+
+func TestRandomElement(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Empty", func(t *testing.T) {
+		t.Parallel()
+
+		s := []string{}
+		v, err := cryptorand.Element(s)
+		require.Error(t, err)
+		require.ErrorContains(t, err, "slice must have at least one element")
+		require.Empty(t, v)
+	})
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		// Generate random slices of ints and strings
+		var (
+			ints    = make([]int, 20)
+			strings = make([]string, 20)
+		)
+		for i := range ints {
+			v, err := cryptorand.Intn(1024)
+			require.NoError(t, err, "generate random int for test slice")
+			ints[i] = v
+		}
+		for i := range strings {
+			v, err := cryptorand.String(10)
+			require.NoError(t, err, "generate random string for test slice")
+			strings[i] = v
+		}
+
+		// Get a random value from each 20 times.
+		for i := 0; i < 20; i++ {
+			iv, err := cryptorand.Element(ints)
+			require.NoError(t, err, "unexpected error from Element(ints)")
+			t.Logf("random int slice element: %v", iv)
+			require.Contains(t, ints, iv)
+
+			sv, err := cryptorand.Element(strings)
+			require.NoError(t, err, "unexpected error from Element(strings)")
+			t.Logf("random string slice element: %v", sv)
+			require.Contains(t, strings, sv)
+		}
+	})
+}