From 8a7b5dcd519b444ec7af884bd562faf8a07f44d1 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: parse structured hostname arguments to "coder ssh"

If the argument to `coder ssh` fits the pattern
`<prefix>--<owner>--<workspace>--<agent?>`, then parse this instead of
treating the argument as a literal workspace name or
`<owner>/<workspace>`.

This will enable the use of `coder ssh` for the VS Code and JetBrains
plugins in combination with wildcard `Host` entries in the SSH config.

Part of #14986.
---
 cli/ssh.go      | 12 +++++++++-
 cli/ssh_test.go | 63 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/cli/ssh.go b/cli/ssh.go
index 7a1d5940bfd01..01265780b1e10 100644
--- a/cli/ssh.go
+++ b/cli/ssh.go
@@ -180,7 +180,17 @@ 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 := inv.Args[0]
+			if parts := strings.Split(namedWorkspace, "--"); len(parts) >= 3 {
+				owner := parts[1]
+				name := parts[2]
+				namedWorkspace = owner + "/" + name
+				if len(parts) > 3 {
+					namedWorkspace += "." + parts[3]
+				}
+			}
+
+			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, !disableAutostart, namedWorkspace)
 			if err != nil {
 				return err
 			}
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
index bd107852251f7..ce4173e90fc35 100644
--- a/cli/ssh_test.go
+++ b/cli/ssh_test.go
@@ -1482,6 +1482,69 @@ func TestSSH(t *testing.T) {
 			})
 		}
 	})
+
+	t.Run("ParseHostname", 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", fmt.Sprintf("coder-vscode.coder.prod.netflix.net--%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.