From 76037ee75367125bdea702aa92885947cff3973c Mon Sep 17 00:00:00 2001
From: Jose Garcia <jose.balius@gmail.com>
Date: Fri, 17 Sep 2021 13:54:00 -0400
Subject: [PATCH] Update docs, simplify loop to append to command

---
 internal/codespaces/ssh.go      | 16 +++++++---------
 internal/codespaces/ssh_test.go |  2 +-
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/internal/codespaces/ssh.go b/internal/codespaces/ssh.go
index e99f8971d..4cd0a4c92 100644
--- a/internal/codespaces/ssh.go
+++ b/internal/codespaces/ssh.go
@@ -64,16 +64,11 @@ func newSSHCommand(ctx context.Context, port int, dst string, cmdArgs []string)
 	return cmd, connArgs, nil
 }
 
-// parseSSHArgs parses SSH arguments into two distinct slices of flags
-// and command. It returns an error if flags are found after a command
-// or if a unary flag is provided without an argument.
+// parseSSHArgs parses SSH arguments into two distinct slices of flags and command.
+// It returns an error if a unary flag is provided without an argument.
 func parseSSHArgs(args []string) (cmdArgs []string, command []string, err error) {
 	for i := 0; i < len(args); i++ {
 		arg := args[i]
-		if command != nil {
-			command = append(command, arg)
-			continue
-		}
 
 		if strings.HasPrefix(arg, "-") {
 			cmdArgs = append(cmdArgs, arg)
@@ -84,9 +79,12 @@ func parseSSHArgs(args []string) (cmdArgs []string, command []string, err error)
 
 				cmdArgs = append(cmdArgs, args[i])
 			}
-		} else {
-			command = append(command, arg)
+			continue
 		}
+
+		// if we've started parsing the command, append all further args to it
+		command = append(command, args[i:]...)
+		break
 	}
 
 	return cmdArgs, command, nil
diff --git a/internal/codespaces/ssh_test.go b/internal/codespaces/ssh_test.go
index 5450adf1a..ed6922762 100644
--- a/internal/codespaces/ssh_test.go
+++ b/internal/codespaces/ssh_test.go
@@ -69,7 +69,7 @@ func TestParseSSHArgs(t *testing.T) {
 
 	for _, tcase := range testCases {
 		args, command, err := parseSSHArgs(tcase.Args)
-		if err != nil && !tcase.Error {
+		if !tcase.Error && err != nil {
 			t.Errorf("unexpected error: %v on test case: %#v", err, tcase)
 			continue
 		}