From 8a1216254eac8ab95967d16e1c5041f32afaedc5 Mon Sep 17 00:00:00 2001
From: Aaron Lehmann <aaronl@vitelus.com>
Date: Mon, 22 Apr 2024 03:13:48 -0700
Subject: [PATCH] feat(cli): add `--env` flag for `coder ssh` (#12991)

This allows environment variables to be set on the SSH session.

Example:

   coder ssh myworkspace --env VAR1=val1,VAR2=val2
---
 cli/ssh.go                           | 39 +++++++++++++++++++------
 cli/ssh_test.go                      | 43 ++++++++++++++++++++++++++++
 cli/testdata/coder_ssh_--help.golden |  3 ++
 docs/cli/ssh.md                      |  9 ++++++
 4 files changed, 85 insertions(+), 9 deletions(-)

diff --git a/cli/ssh.go b/cli/ssh.go
index a291b3764b..1aa832fcda 100644
--- a/cli/ssh.go
+++ b/cli/ssh.go
@@ -55,6 +55,7 @@ func (r *RootCmd) ssh() *serpent.Command {
 		noWait           bool
 		logDirPath       string
 		remoteForwards   []string
+		env              []string
 		disableAutostart bool
 	)
 	client := new(codersdk.Client)
@@ -144,16 +145,23 @@ func (r *RootCmd) ssh() *serpent.Command {
 			stack := newCloserStack(ctx, logger)
 			defer stack.close(nil)
 
-			if len(remoteForwards) > 0 {
-				for _, remoteForward := range remoteForwards {
-					isValid := validateRemoteForward(remoteForward)
-					if !isValid {
-						return xerrors.Errorf(`invalid format of remote-forward, expected: remote_port:local_address:local_port`)
-					}
-					if isValid && stdio {
-						return xerrors.Errorf(`remote-forward can't be enabled in the stdio mode`)
-					}
+			for _, remoteForward := range remoteForwards {
+				isValid := validateRemoteForward(remoteForward)
+				if !isValid {
+					return xerrors.Errorf(`invalid format of remote-forward, expected: remote_port:local_address:local_port`)
 				}
+				if isValid && stdio {
+					return xerrors.Errorf(`remote-forward can't be enabled in the stdio mode`)
+				}
+			}
+
+			var parsedEnv [][2]string
+			for _, e := range env {
+				k, v, ok := strings.Cut(e, "=")
+				if !ok {
+					return xerrors.Errorf("invalid environment variable setting %q", e)
+				}
+				parsedEnv = append(parsedEnv, [2]string{k, v})
 			}
 
 			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, !disableAutostart, inv.Args[0])
@@ -375,6 +383,12 @@ func (r *RootCmd) ssh() *serpent.Command {
 				}()
 			}
 
+			for _, kv := range parsedEnv {
+				if err := sshSession.Setenv(kv[0], kv[1]); err != nil {
+					return xerrors.Errorf("setenv: %w", err)
+				}
+			}
+
 			err = sshSession.RequestPty("xterm-256color", 128, 128, gossh.TerminalModes{})
 			if err != nil {
 				return xerrors.Errorf("request pty: %w", err)
@@ -483,6 +497,13 @@ func (r *RootCmd) ssh() *serpent.Command {
 			FlagShorthand: "R",
 			Value:         serpent.StringArrayOf(&remoteForwards),
 		},
+		{
+			Flag:          "env",
+			Description:   "Set environment variable(s) for session (key1=value1,key2=value2,...).",
+			Env:           "CODER_SSH_ENV",
+			FlagShorthand: "e",
+			Value:         serpent.StringArrayOf(&env),
+		},
 		sshDisableAutostartOption(serpent.BoolOf(&disableAutostart)),
 	}
 	return cmd
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
index f1e52b1a0c..8c3c1a4e40 100644
--- a/cli/ssh_test.go
+++ b/cli/ssh_test.go
@@ -968,6 +968,49 @@ func TestSSH(t *testing.T) {
 		<-cmdDone
 	})
 
+	t.Run("Env", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			t.Skip("Test not supported on windows")
+		}
+
+		t.Parallel()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+		_ = agenttest.New(t, client.URL, agentToken)
+		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+		inv, root := clitest.New(t,
+			"ssh",
+			workspace.Name,
+			"--env",
+			"foo=bar,baz=qux",
+		)
+		clitest.SetupConfig(t, client, root)
+
+		pty := ptytest.New(t).Attach(inv)
+		inv.Stderr = pty.Output()
+
+		// Wait super long so this doesn't flake on -race test.
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
+		defer cancel()
+
+		w := clitest.StartWithWaiter(t, inv.WithContext(ctx))
+		defer w.Wait() // We don't care about any exit error (exit code 255: SSH connection ended unexpectedly).
+
+		// Since something was output, it should be safe to write input.
+		// This could show a prompt or "running startup scripts", so it's
+		// not indicative of the SSH connection being ready.
+		_ = pty.Peek(ctx, 1)
+
+		// Ensure the SSH connection is ready by testing the shell
+		// input/output.
+		pty.WriteLine("echo $foo $baz")
+		pty.ExpectMatchContext(ctx, "bar qux")
+
+		// And we're done.
+		pty.WriteLine("exit")
+	})
+
 	t.Run("RemoteForwardUnixSocket", func(t *testing.T) {
 		if runtime.GOOS == "windows" {
 			t.Skip("Test not supported on windows")
diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden
index ce53948c70..80aaa3c204 100644
--- a/cli/testdata/coder_ssh_--help.golden
+++ b/cli/testdata/coder_ssh_--help.golden
@@ -9,6 +9,9 @@ OPTIONS:
       --disable-autostart bool, $CODER_SSH_DISABLE_AUTOSTART (default: false)
           Disable starting the workspace automatically when connecting via SSH.
 
+  -e, --env string-array, $CODER_SSH_ENV
+          Set environment variable(s) for session (key1=value1,key2=value2,...).
+
   -A, --forward-agent bool, $CODER_SSH_FORWARD_AGENT
           Specifies whether to forward the SSH agent specified in
           $SSH_AUTH_SOCK.
diff --git a/docs/cli/ssh.md b/docs/cli/ssh.md
index c2945a0be2..d2110628fe 100644
--- a/docs/cli/ssh.md
+++ b/docs/cli/ssh.md
@@ -95,6 +95,15 @@ Specify the directory containing SSH diagnostic log files.
 
 Enable remote port forwarding (remote_port:local_address:local_port).
 
+### -e, --env
+
+|             |                             |
+| ----------- | --------------------------- |
+| Type        | <code>string-array</code>   |
+| Environment | <code>$CODER_SSH_ENV</code> |
+
+Set environment variable(s) for session (key1=value1,key2=value2,...).
+
 ### --disable-autostart
 
 |             |                                           |