Skip to content

Commit 85f79ed

Browse files
josebaliusjosebalius
authored
Merge pull request cli#159 from github/jg/ssh-cmd-flags
ghcs ssh: ssh flags and command support
2 parents 67a6f0a + 82c1972 commit 85f79ed

File tree

5 files changed

+168
-19
lines changed

5 files changed

+168
-19
lines changed

‎cmd/ghcs/logs.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,9 +82,12 @@ func logs(ctx context.Context, log *output.Logger, codespaceName string, follow
8282
}
8383

8484
dst := fmt.Sprintf("%s@localhost", sshUser)
85-
cmd := codespaces.NewRemoteCommand(
85+
cmd, err := codespaces.NewRemoteCommand(
8686
ctx, localPort, dst, fmt.Sprintf("%s /workspaces/.codespaces/.persistedshare/creation.log", cmdType),
8787
)
88+
if err != nil {
89+
return fmt.Errorf("remote command: %w", err)
90+
}
8891

8992
tunnelClosed := make(chan error, 1)
9093
go func() {

‎cmd/ghcs/ssh.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,10 @@ func newSSHCmd() *cobra.Command {
1818
var sshServerPort int
1919

2020
sshCmd := &cobra.Command{
21-
Use: "ssh",
21+
Use: "ssh [flags] [--] [ssh-flags] [command]",
2222
Short: "SSH into a codespace",
2323
RunE: func(cmd *cobra.Command, args []string) error {
24-
return ssh(context.Background(), sshProfile, codespaceName, sshServerPort)
24+
return ssh(context.Background(), args, sshProfile, codespaceName, sshServerPort)
2525
},
2626
}
2727

@@ -36,7 +36,7 @@ func init() {
3636
rootCmd.AddCommand(newSSHCmd())
3737
}
3838

39-
func ssh(ctx context.Context, sshProfile, codespaceName string, localSSHServerPort int) error {
39+
func ssh(ctx context.Context, sshArgs []string, sshProfile, codespaceName string, localSSHServerPort int) error {
4040
// Ensure all child tasks (e.g. port forwarding) terminate before return.
4141
ctx, cancel := context.WithCancel(ctx)
4242
defer cancel()
@@ -89,7 +89,7 @@ func ssh(ctx context.Context, sshProfile, codespaceName string, localSSHServerPo
8989

9090
shellClosed := make(chan error, 1)
9191
go func() {
92-
shellClosed <- codespaces.Shell(ctx, log, localSSHServerPort, connectDestination, usingCustomPort)
92+
shellClosed <- codespaces.Shell(ctx, log, sshArgs, localSSHServerPort, connectDestination, usingCustomPort)
9393
}()
9494

9595
select {

‎internal/codespaces/ssh.go‎

Lines changed: 50 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package codespaces
22

33
import (
44
"context"
5+
"fmt"
56
"os"
67
"os/exec"
78
"strconv"
@@ -11,8 +12,11 @@ import (
1112
// Shell runs an interactive secure shell over an existing
1213
// port-forwarding session. It runs until the shell is terminated
1314
// (including by cancellation of the context).
14-
func Shell(ctx context.Context, log logger, port int, destination string, usingCustomPort bool) error {
15-
cmd, connArgs := newSSHCommand(ctx, port, destination, "")
15+
func Shell(ctx context.Context, log logger, sshArgs []string, port int, destination string, usingCustomPort bool) error {
16+
cmd, connArgs, err := newSSHCommand(ctx, port, destination, sshArgs)
17+
if err != nil {
18+
return fmt.Errorf("failed to create ssh command: %w", err)
19+
}
1620

1721
if usingCustomPort {
1822
log.Println("Connection Details: ssh " + destination + " " + strings.Join(connArgs, " "))
@@ -23,31 +27,64 @@ func Shell(ctx context.Context, log logger, port int, destination string, usingC
2327

2428
// NewRemoteCommand returns an exec.Cmd that will securely run a shell
2529
// command on the remote machine.
26-
func NewRemoteCommand(ctx context.Context, tunnelPort int, destination, command string) *exec.Cmd {
27-
cmd, _ := newSSHCommand(ctx, tunnelPort, destination, command)
28-
return cmd
30+
func NewRemoteCommand(ctx context.Context, tunnelPort int, destination string, sshArgs ...string) (*exec.Cmd, error) {
31+
cmd, _, err := newSSHCommand(ctx, tunnelPort, destination, sshArgs)
32+
return cmd, err
2933
}
3034

3135
// newSSHCommand populates an exec.Cmd to run a command (or if blank,
3236
// an interactive shell) over ssh.
33-
func newSSHCommand(ctx context.Context, port int, dst, command string) (*exec.Cmd, []string) {
37+
func newSSHCommand(ctx context.Context, port int, dst string, cmdArgs []string) (*exec.Cmd, []string, error) {
3438
connArgs := []string{"-p", strconv.Itoa(port), "-o", "NoHostAuthenticationForLocalhost=yes"}
3539

36-
cmdArgs := []string{dst, "-C"} // Always use Compression
37-
if command == "" {
38-
// if we are in a shell send X11 and X11Trust
39-
cmdArgs = append(cmdArgs, "-X", "-Y")
40+
// The ssh command syntax is: ssh [flags] user@host command [args...]
41+
// There is no way to specify the user@host destination as a flag.
42+
// Unfortunately, that means we need to know which user-provided words are
43+
// SSH flags and which are command arguments so that we can place
44+
// them before or after the destination, and that means we need to know all
45+
// the flags and their arities.
46+
cmdArgs, command, err := parseSSHArgs(cmdArgs)
47+
if err != nil {
48+
return nil, nil, err
4049
}
4150

4251
cmdArgs = append(cmdArgs, connArgs...)
43-
if command != "" {
44-
cmdArgs = append(cmdArgs, command)
52+
cmdArgs = append(cmdArgs, "-C") // Compression
53+
cmdArgs = append(cmdArgs, dst) // user@host
54+
55+
if command != nil {
56+
cmdArgs = append(cmdArgs, command...)
4557
}
4658

4759
cmd := exec.CommandContext(ctx, "ssh", cmdArgs...)
4860
cmd.Stdout = os.Stdout
4961
cmd.Stdin = os.Stdin
5062
cmd.Stderr = os.Stderr
5163

52-
return cmd, connArgs
64+
return cmd, connArgs, nil
65+
}
66+
67+
// parseSSHArgs parses SSH arguments into two distinct slices of flags and command.
68+
// It returns an error if a unary flag is provided without an argument.
69+
func parseSSHArgs(args []string) (cmdArgs, command []string, err error) {
70+
for i := 0; i < len(args); i++ {
71+
arg := args[i]
72+
73+
// if we've started parsing the command, set it to the rest of the args
74+
if !strings.HasPrefix(arg, "-") {
75+
command = args[i:]
76+
break
77+
}
78+
79+
cmdArgs = append(cmdArgs, arg)
80+
if len(arg) == 2 && strings.Contains("bcDeFIiLlmOopRSWw", arg[1:2]) {
81+
if i++; i == len(args) {
82+
return nil, nil, fmt.Errorf("ssh flag: %s requires an argument", arg)
83+
}
84+
85+
cmdArgs = append(cmdArgs, args[i])
86+
}
87+
}
88+
89+
return cmdArgs, command, nil
5390
}

‎internal/codespaces/ssh_test.go‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
package codespaces
2+
3+
import (
4+
"fmt"
5+
"testing"
6+
)
7+
8+
func TestParseSSHArgs(t *testing.T) {
9+
type testCase struct {
10+
Args []string
11+
ParsedArgs []string
12+
Command []string
13+
Error string
14+
}
15+
16+
testCases := []testCase{
17+
{}, // empty test case
18+
{
19+
Args: []string{"-X", "-Y"},
20+
ParsedArgs: []string{"-X", "-Y"},
21+
Command: nil,
22+
},
23+
{
24+
Args: []string{"-X", "-Y", "-o", "someoption=test"},
25+
ParsedArgs: []string{"-X", "-Y", "-o", "someoption=test"},
26+
Command: nil,
27+
},
28+
{
29+
Args: []string{"-X", "-Y", "-o", "someoption=test", "somecommand"},
30+
ParsedArgs: []string{"-X", "-Y", "-o", "someoption=test"},
31+
Command: []string{"somecommand"},
32+
},
33+
{
34+
Args: []string{"-X", "-Y", "-o", "someoption=test", "echo", "test"},
35+
ParsedArgs: []string{"-X", "-Y", "-o", "someoption=test"},
36+
Command: []string{"echo", "test"},
37+
},
38+
{
39+
Args: []string{"somecommand"},
40+
ParsedArgs: []string{},
41+
Command: []string{"somecommand"},
42+
},
43+
{
44+
Args: []string{"echo", "test"},
45+
ParsedArgs: []string{},
46+
Command: []string{"echo", "test"},
47+
},
48+
{
49+
Args: []string{"-v", "echo", "hello", "world"},
50+
ParsedArgs: []string{"-v"},
51+
Command: []string{"echo", "hello", "world"},
52+
},
53+
{
54+
Args: []string{"-L", "-l"},
55+
ParsedArgs: []string{"-L", "-l"},
56+
Command: nil,
57+
},
58+
{
59+
Args: []string{"-v", "echo", "-n", "test"},
60+
ParsedArgs: []string{"-v"},
61+
Command: []string{"echo", "-n", "test"},
62+
},
63+
{
64+
Args: []string{"-v", "echo", "-b", "test"},
65+
ParsedArgs: []string{"-v"},
66+
Command: []string{"echo", "-b", "test"},
67+
},
68+
{
69+
Args: []string{"-b"},
70+
ParsedArgs: nil,
71+
Command: nil,
72+
Error: "ssh flag: -b requires an argument",
73+
},
74+
}
75+
76+
for _, tcase := range testCases {
77+
args, command, err := parseSSHArgs(tcase.Args)
78+
if tcase.Error != "" {
79+
if err == nil {
80+
t.Errorf("expected error and got nil: %#v", tcase)
81+
}
82+
83+
if err.Error() != tcase.Error {
84+
t.Errorf("error does not match expected error, got: '%s', expected: '%s'", err.Error(), tcase.Error)
85+
}
86+
87+
continue
88+
}
89+
90+
if err != nil {
91+
t.Errorf("unexpected error: %v on test case: %#v", err, tcase)
92+
continue
93+
}
94+
95+
argsStr, parsedArgsStr := fmt.Sprintf("%s", args), fmt.Sprintf("%s", tcase.ParsedArgs)
96+
if argsStr != parsedArgsStr {
97+
t.Errorf("args do not match parsed args. got: '%s', expected: '%s'", argsStr, parsedArgsStr)
98+
}
99+
100+
commandStr, parsedCommandStr := fmt.Sprintf("%s", command), fmt.Sprintf("%s", tcase.Command)
101+
if commandStr != parsedCommandStr {
102+
t.Errorf("command does not match parsed command. got: '%s', expected: '%s'", commandStr, parsedCommandStr)
103+
}
104+
}
105+
}

‎internal/codespaces/states.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,10 +89,14 @@ func PollPostCreateStates(ctx context.Context, log logger, apiClient *api.API, u
8989
}
9090

9191
func getPostCreateOutput(ctx context.Context, tunnelPort int, codespace *api.Codespace, user string) ([]PostCreateState, error) {
92-
cmd := NewRemoteCommand(
92+
cmd, err := NewRemoteCommand(
9393
ctx, tunnelPort, fmt.Sprintf("%s@localhost", user),
9494
"cat /workspaces/.codespaces/shared/postCreateOutput.json",
9595
)
96+
if err != nil {
97+
return nil, fmt.Errorf("remote command: %w", err)
98+
}
99+
96100
stdout := new(bytes.Buffer)
97101
cmd.Stdout = stdout
98102
if err := cmd.Run(); err != nil {

0 commit comments

Comments
 (0)