From 0499bdaa043734fd036074b72a8f7552734fafca Mon Sep 17 00:00:00 2001
From: Aaron Lehmann <alehmann@netflix.com>
Date: Thu, 9 Jan 2025 10:54:54 -0800
Subject: [PATCH] feat: --ssh-host-prefix flag for "coder ssh"

This adds a flag matching `--ssh-host-prefix` from `coder config-ssh` to
`coder ssh`. By trimming a custom prefix from the argument, we can set
up wildcard-based `Host` entries in SSH config for the IDE plugins (and
eventually `coder config-ssh`).

We also replace `--` in the argument with `/`, so ownership can be
specified in wildcard-based SSH hosts like `<owner>--<workspace>`.
---
 cli/ssh.go                           | 13 +++++-
 cli/ssh_test.go                      | 63 ++++++++++++++++++++++++++++
 cli/testdata/coder_ssh_--help.golden |  5 +++
 docs/reference/cli/ssh.md            |  9 ++++
 4 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/cli/ssh.go b/cli/ssh.go
index 4fa836e44e389..884c5500d703c 100644
--- a/cli/ssh.go
+++ b/cli/ssh.go
@@ -61,6 +61,7 @@ var (
 func (r *RootCmd) ssh() *serpent.Command {
 	var (
 		stdio               bool
+		hostPrefix          string
 		forwardAgent        bool
 		forwardGPG          bool
 		identityAgent       string
@@ -195,7 +196,11 @@ func (r *RootCmd) ssh() *serpent.Command {
 				parsedEnv = append(parsedEnv, [2]string{k, v})
 			}
 
-			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, !disableAutostart, inv.Args[0])
+			namedWorkspace := strings.TrimPrefix(inv.Args[0], hostPrefix)
+			// Support "--" as a delimiter between owner and workspace name
+			namedWorkspace = strings.Replace(namedWorkspace, "--", "/", 1)
+
+			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, !disableAutostart, namedWorkspace)
 			if err != nil {
 				return err
 			}
@@ -509,6 +514,12 @@ func (r *RootCmd) ssh() *serpent.Command {
 			Description: "Specifies whether to emit SSH output over stdin/stdout.",
 			Value:       serpent.BoolOf(&stdio),
 		},
+		{
+			Flag:        "ssh-host-prefix",
+			Env:         "CODER_SSH_SSH_HOST_PREFIX",
+			Description: "Strip this prefix from the provided hostname to determine the workspace name. This is useful when used as part of an OpenSSH proxy command.",
+			Value:       serpent.StringOf(&hostPrefix),
+		},
 		{
 			Flag:          "forward-agent",
 			FlagShorthand: "A",
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
index fa6ab32b59035..23c7a01898cd1 100644
--- a/cli/ssh_test.go
+++ b/cli/ssh_test.go
@@ -1568,6 +1568,69 @@ func TestSSH(t *testing.T) {
 			})
 		}
 	})
+
+	t.Run("SSHHostPrefix", func(t *testing.T) {
+		t.Parallel()
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+		_, _ = tGoContext(t, func(ctx context.Context) {
+			// Run this async so the SSH command has to wait for
+			// the build and agent to connect!
+			_ = agenttest.New(t, client.URL, agentToken)
+			<-ctx.Done()
+		})
+
+		clientOutput, clientInput := io.Pipe()
+		serverOutput, serverInput := io.Pipe()
+		defer func() {
+			for _, c := range []io.Closer{clientOutput, clientInput, serverOutput, serverInput} {
+				_ = c.Close()
+			}
+		}()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		user, err := client.User(ctx, codersdk.Me)
+		require.NoError(t, err)
+
+		inv, root := clitest.New(t, "ssh", "--stdio", "--ssh-host-prefix", "coder.dummy.com--", fmt.Sprintf("coder.dummy.com--%s--%s", user.Username, workspace.Name))
+		clitest.SetupConfig(t, client, root)
+		inv.Stdin = clientOutput
+		inv.Stdout = serverInput
+		inv.Stderr = io.Discard
+
+		cmdDone := tGo(t, func() {
+			err := inv.WithContext(ctx).Run()
+			assert.NoError(t, err)
+		})
+
+		conn, channels, requests, err := ssh.NewClientConn(&stdioConn{
+			Reader: serverOutput,
+			Writer: clientInput,
+		}, "", &ssh.ClientConfig{
+			// #nosec
+			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		})
+		require.NoError(t, err)
+		defer conn.Close()
+
+		sshClient := ssh.NewClient(conn, channels, requests)
+		session, err := sshClient.NewSession()
+		require.NoError(t, err)
+		defer session.Close()
+
+		command := "sh -c exit"
+		if runtime.GOOS == "windows" {
+			command = "cmd.exe /c exit"
+		}
+		err = session.Run(command)
+		require.NoError(t, err)
+		err = sshClient.Close()
+		require.NoError(t, err)
+		_ = clientOutput.Close()
+
+		<-cmdDone
+	})
 }
 
 //nolint:paralleltest // This test uses t.Setenv, parent test MUST NOT be parallel.
diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden
index d847e9d7abb03..3d2f584727cd9 100644
--- a/cli/testdata/coder_ssh_--help.golden
+++ b/cli/testdata/coder_ssh_--help.golden
@@ -45,6 +45,11 @@ OPTIONS:
   -R, --remote-forward string-array, $CODER_SSH_REMOTE_FORWARD
           Enable remote port forwarding (remote_port:local_address:local_port).
 
+      --ssh-host-prefix string, $CODER_SSH_SSH_HOST_PREFIX
+          Strip this prefix from the provided hostname to determine the
+          workspace name. This is useful when used as part of an OpenSSH proxy
+          command.
+
       --stdio bool, $CODER_SSH_STDIO
           Specifies whether to emit SSH output over stdin/stdout.
 
diff --git a/docs/reference/cli/ssh.md b/docs/reference/cli/ssh.md
index 74e28837ad7e4..72d63a1f003af 100644
--- a/docs/reference/cli/ssh.md
+++ b/docs/reference/cli/ssh.md
@@ -20,6 +20,15 @@ coder ssh [flags] <workspace>
 
 Specifies whether to emit SSH output over stdin/stdout.
 
+### --ssh-host-prefix
+
+|             |                                         |
+|-------------|-----------------------------------------|
+| Type        | <code>string</code>                     |
+| Environment | <code>$CODER_SSH_SSH_HOST_PREFIX</code> |
+
+Strip this prefix from the provided hostname to determine the workspace name. This is useful when used as part of an OpenSSH proxy command.
+
 ### -A, --forward-agent
 
 |             |                                       |

<!DOCTYPE html PUBLIC '-//W3C//DTD XHTML 1.0 Transitional//EN' 'http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd'>
<html xmlns='http://www.w3.org/1999/xhtml'>
<head>
<title>pFad - Phonifier reborn</title>
<meta http-equiv='Content-Type' content='text/html; charset=utf-8' />
</head>
<body>
<h1>Pfad - The Proxy pFad of &#169; 2024 Garber Painting. All rights reserved.</h1>


<!-- Disclaimer -->
<p>Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.</p>
<br>
<p>Alternative Proxies:</p><p><a href="http://rainy.clevelandohioweatherforecast.com/php-proxy/index.php?q=https://patch-diff.githubusercontent.com/raw/coder/coder/pull/16088.patch" target="_blank">Alternative Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/index.php?u=https://patch-diff.githubusercontent.com/raw/coder/coder/pull/16088.patch" target="_blank">pFad Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v3index.php?u=https://patch-diff.githubusercontent.com/raw/coder/coder/pull/16088.patch" target="_blank">pFad v3 Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v4index.php?u=https://patch-diff.githubusercontent.com/raw/coder/coder/pull/16088.patch" target="_blank">pFad v4 Proxy</a></p></body>
</html>