feat: Add --env flag for "coder ssh"

aaronlehmann · aaronlehmann · commit e99d1846d0a5 · 2024-04-18T08:57:24.000-07:00
This allows environment variables to be set on the SSH session.

Example:

   coder ssh myworkspace --env VAR1=val1,VAR2=val2
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -56,6 +56,7 @@ func (r *RootCmd) ssh() *serpent.Command {
 		noWait           bool
 		logDirPath       string
 		remoteForwards   []string
+		env              []string
 		disableAutostart bool
 	)
 	client := new(codersdk.Client)
@@ -145,16 +146,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])
@@ -369,6 +377,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)
@@ -477,6 +491,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
@@ -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
@@ -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.
