Skip to content

Commit 2cd3f99

Browse files
brettkolodnybrettkolodny
authored
feat: add one shot commands to the coder ssh command (#17779)
Closes #2154 > [!WARNING] > The tests in this PR were co-authored by AI
1 parent cb0f778 commit 2cd3f99

File tree

7 files changed

+193
-39
lines changed

7 files changed

+193
-39
lines changed

‎cli/ssh.go

+56-32
Original file line numberDiff line numberDiff line change
@@ -90,15 +90,33 @@ func (r *RootCmd) ssh() *serpent.Command {
9090
wsClient := workspacesdk.New(client)
9191
cmd := &serpent.Command{
9292
Annotations: workspaceCommand,
93-
Use: "ssh <workspace>",
94-
Short: "Start a shell into a workspace",
95-
Long: "This command does not have full parity with the standard SSH command. For users who need the full functionality of SSH, create an ssh configuration with `coder config-ssh`.",
93+
Use: "ssh <workspace> [command]",
94+
Short: "Start a shell into a workspace or run a command",
95+
Long: "This command does not have full parity with the standard SSH command. For users who need the full functionality of SSH, create an ssh configuration with `coder config-ssh`.\n\n" +
96+
FormatExamples(
97+
Example{
98+
Description: "Use `--` to separate and pass flags directly to the command executed via SSH.",
99+
Command: "coder ssh <workspace> -- ls -la",
100+
},
101+
),
96102
Middleware: serpent.Chain(
97-
serpent.RequireNArgs(1),
103+
// Require at least one arg for the workspace name
104+
func(next serpent.HandlerFunc) serpent.HandlerFunc {
105+
return func(i *serpent.Invocation) error {
106+
got := len(i.Args)
107+
if got < 1 {
108+
return xerrors.New("expected the name of a workspace")
109+
}
110+
111+
return next(i)
112+
}
113+
},
98114
r.InitClient(client),
99115
initAppearance(client, &appearanceConfig),
100116
),
101117
Handler: func(inv *serpent.Invocation) (retErr error) {
118+
command := strings.Join(inv.Args[1:], " ")
119+
102120
// Before dialing the SSH server over TCP, capture Interrupt signals
103121
// so that if we are interrupted, we have a chance to tear down the
104122
// TCP session cleanly before exiting. If we don't, then the TCP
@@ -548,40 +566,46 @@ func (r *RootCmd) ssh() *serpent.Command {
548566
sshSession.Stdout = inv.Stdout
549567
sshSession.Stderr = inv.Stderr
550568

551-
err = sshSession.Shell()
552-
if err != nil {
553-
return xerrors.Errorf("start shell: %w", err)
554-
}
569+
if command != "" {
570+
err := sshSession.Run(command)
571+
if err != nil {
572+
return xerrors.Errorf("run command: %w", err)
573+
}
574+
} else {
575+
err = sshSession.Shell()
576+
if err != nil {
577+
return xerrors.Errorf("start shell: %w", err)
578+
}
555579

556-
// Put cancel at the top of the defer stack to initiate
557-
// shutdown of services.
558-
defer cancel()
580+
// Put cancel at the top of the defer stack to initiate
581+
// shutdown of services.
582+
defer cancel()
559583

560-
if validOut {
561-
// Set initial window size.
562-
width, height, err := term.GetSize(int(stdoutFile.Fd()))
563-
if err == nil {
564-
_ = sshSession.WindowChange(height, width)
584+
if validOut {
585+
// Set initial window size.
586+
width, height, err := term.GetSize(int(stdoutFile.Fd()))
587+
if err == nil {
588+
_ = sshSession.WindowChange(height, width)
589+
}
565590
}
566-
}
567591

568-
err = sshSession.Wait()
569-
conn.SendDisconnectedTelemetry()
570-
if err != nil {
571-
if exitErr := (&gossh.ExitError{}); errors.As(err, &exitErr) {
572-
// Clear the error since it's not useful beyond
573-
// reporting status.
574-
return ExitError(exitErr.ExitStatus(), nil)
575-
}
576-
// If the connection drops unexpectedly, we get an
577-
// ExitMissingError but no other error details, so try to at
578-
// least give the user a better message
579-
if errors.Is(err, &gossh.ExitMissingError{}) {
580-
return ExitError(255, xerrors.New("SSH connection ended unexpectedly"))
592+
err = sshSession.Wait()
593+
conn.SendDisconnectedTelemetry()
594+
if err != nil {
595+
if exitErr := (&gossh.ExitError{}); errors.As(err, &exitErr) {
596+
// Clear the error since it's not useful beyond
597+
// reporting status.
598+
return ExitError(exitErr.ExitStatus(), nil)
599+
}
600+
// If the connection drops unexpectedly, we get an
601+
// ExitMissingError but no other error details, so try to at
602+
// least give the user a better message
603+
if errors.Is(err, &gossh.ExitMissingError{}) {
604+
return ExitError(255, xerrors.New("SSH connection ended unexpectedly"))
605+
}
606+
return xerrors.Errorf("session ended: %w", err)
581607
}
582-
return xerrors.Errorf("session ended: %w", err)
583608
}
584-
585609
return nil
586610
},
587611
}

‎cli/ssh_test.go

+121
Original file line numberDiff line numberDiff line change
@@ -2200,6 +2200,127 @@ func TestSSH_CoderConnect(t *testing.T) {
22002200

22012201
<-cmdDone
22022202
})
2203+
2204+
t.Run("OneShot", func(t *testing.T) {
2205+
t.Parallel()
2206+
2207+
client, workspace, agentToken := setupWorkspaceForAgent(t)
2208+
inv, root := clitest.New(t, "ssh", workspace.Name, "echo 'hello world'")
2209+
clitest.SetupConfig(t, client, root)
2210+
2211+
// Capture command output
2212+
output := new(bytes.Buffer)
2213+
inv.Stdout = output
2214+
2215+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
2216+
defer cancel()
2217+
2218+
cmdDone := tGo(t, func() {
2219+
err := inv.WithContext(ctx).Run()
2220+
assert.NoError(t, err)
2221+
})
2222+
2223+
_ = agenttest.New(t, client.URL, agentToken)
2224+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
2225+
2226+
<-cmdDone
2227+
2228+
// Verify command output
2229+
assert.Contains(t, output.String(), "hello world")
2230+
})
2231+
2232+
t.Run("OneShotExitCode", func(t *testing.T) {
2233+
t.Parallel()
2234+
2235+
client, workspace, agentToken := setupWorkspaceForAgent(t)
2236+
2237+
// Setup agent first to avoid race conditions
2238+
_ = agenttest.New(t, client.URL, agentToken)
2239+
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
2240+
2241+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
2242+
defer cancel()
2243+
2244+
// Test successful exit code
2245+
t.Run("Success", func(t *testing.T) {
2246+
inv, root := clitest.New(t, "ssh", workspace.Name, "exit 0")
2247+
clitest.SetupConfig(t, client, root)
2248+
2249+
err := inv.WithContext(ctx).Run()
2250+
assert.NoError(t, err)
2251+
})
2252+
2253+
// Test error exit code
2254+
t.Run("Error", func(t *testing.T) {
2255+
inv, root := clitest.New(t, "ssh", workspace.Name, "exit 1")
2256+
clitest.SetupConfig(t, client, root)
2257+
2258+
err := inv.WithContext(ctx).Run()
2259+
assert.Error(t, err)
2260+
var exitErr *ssh.ExitError
2261+
assert.True(t, errors.As(err, &exitErr))
2262+
assert.Equal(t, 1, exitErr.ExitStatus())
2263+
})
2264+
})
2265+
2266+
t.Run("OneShotStdio", func(t *testing.T) {
2267+
t.Parallel()
2268+
client, workspace, agentToken := setupWorkspaceForAgent(t)
2269+
_, _ = tGoContext(t, func(ctx context.Context) {
2270+
// Run this async so the SSH command has to wait for
2271+
// the build and agent to connect!
2272+
_ = agenttest.New(t, client.URL, agentToken)
2273+
<-ctx.Done()
2274+
})
2275+
2276+
clientOutput, clientInput := io.Pipe()
2277+
serverOutput, serverInput := io.Pipe()
2278+
defer func() {
2279+
for _, c := range []io.Closer{clientOutput, clientInput, serverOutput, serverInput} {
2280+
_ = c.Close()
2281+
}
2282+
}()
2283+
2284+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
2285+
defer cancel()
2286+
2287+
inv, root := clitest.New(t, "ssh", "--stdio", workspace.Name, "echo 'hello stdio'")
2288+
clitest.SetupConfig(t, client, root)
2289+
inv.Stdin = clientOutput
2290+
inv.Stdout = serverInput
2291+
inv.Stderr = io.Discard
2292+
2293+
cmdDone := tGo(t, func() {
2294+
err := inv.WithContext(ctx).Run()
2295+
assert.NoError(t, err)
2296+
})
2297+
2298+
conn, channels, requests, err := ssh.NewClientConn(&testutil.ReaderWriterConn{
2299+
Reader: serverOutput,
2300+
Writer: clientInput,
2301+
}, "", &ssh.ClientConfig{
2302+
// #nosec
2303+
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
2304+
})
2305+
require.NoError(t, err)
2306+
defer conn.Close()
2307+
2308+
sshClient := ssh.NewClient(conn, channels, requests)
2309+
session, err := sshClient.NewSession()
2310+
require.NoError(t, err)
2311+
defer session.Close()
2312+
2313+
// Capture and verify command output
2314+
output, err := session.Output("echo 'hello back'")
2315+
require.NoError(t, err)
2316+
assert.Contains(t, string(output), "hello back")
2317+
2318+
err = sshClient.Close()
2319+
require.NoError(t, err)
2320+
_ = clientOutput.Close()
2321+
2322+
<-cmdDone
2323+
})
22032324
}
22042325

22052326
type fakeCoderConnectDialer struct{}

‎cli/testdata/coder_--help.golden

+1-1
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ SUBCOMMANDS:
4646
show Display details of a workspace's resources and agents
4747
speedtest Run upload and download tests from your machine to a
4848
workspace
49-
ssh Start a shell into a workspace
49+
ssh Start a shell into a workspace or run a command
5050
start Start a workspace
5151
stat Show resource usage for the current workspace.
5252
state Manually manage Terraform state to fix broken workspaces

‎cli/testdata/coder_ssh_--help.golden

+7-2
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,18 @@
11
coder v0.0.0-devel
22

33
USAGE:
4-
coder ssh [flags] <workspace>
4+
coder ssh [flags] <workspace> [command]
55

6-
Start a shell into a workspace
6+
Start a shell into a workspace or run a command
77

88
This command does not have full parity with the standard SSH command. For
99
users who need the full functionality of SSH, create an ssh configuration with
1010
`coder config-ssh`.
11+
12+
- Use `--` to separate and pass flags directly to the command executed via
13+
SSH.:
14+
15+
$ coder ssh <workspace> -- ls -la
1116

1217
OPTIONS:
1318
--disable-autostart bool, $CODER_SSH_DISABLE_AUTOSTART (default: false)

‎docs/manifest.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -1460,7 +1460,7 @@
14601460
},
14611461
{
14621462
"title": "ssh",
1463-
"description": "Start a shell into a workspace",
1463+
"description": "Start a shell into a workspace or run a command",
14641464
"path": "reference/cli/ssh.md"
14651465
},
14661466
{

‎docs/reference/cli/index.md

+1-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/reference/cli/ssh.md

+6-2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)