1 files changed, 55 insertions, 14 deletions

diff --git a/go/internal/command/commandargs/command_args_test.go b/go/internal/command/commandargs/command_args_test.go
index 10c46fe..01202c0 100644
--- a/go/internal/command/commandargs/command_args_test.go
+++ b/go/internal/command/commandargs/command_args_test.go
@@ -3,7 +3,8 @@ package commandargs
 import (
 	"testing"
 
-	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
 )
 
@@ -22,23 +23,16 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
-			expectedArgs: &CommandArgs{CommandType: Discover},
+			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover},
 		},
 		{
-			desc: "It passes on the original ssh command from the environment",
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "hello world",
-			},
-			expectedArgs: &CommandArgs{SshCommand: "hello world"},
-		}, {
 			desc: "It finds the key id in any passed arguments",
 			environment: map[string]string{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
 			arguments:    []string{"hello", "key-123"},
-			expectedArgs: &CommandArgs{CommandType: Discover, GitlabKeyId: "123"},
+			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
 		}, {
 			desc: "It finds the username in any passed arguments",
 			environment: map[string]string{
@@ -46,7 +40,42 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "",
 			},
 			arguments:    []string{"hello", "username-jane-doe"},
-			expectedArgs: &CommandArgs{CommandType: Discover, GitlabUsername: "jane-doe"},
+			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
+		}, {
+			desc: "It parses 2fa_recovery_codes command",
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
+			},
+			expectedArgs: &CommandArgs{SshArgs: []string{"2fa_recovery_codes"}, CommandType: TwoFactorRecover},
+		}, {
+			desc: "It parses git-receive-pack command",
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "git-receive-pack group/repo",
+			},
+			expectedArgs: &CommandArgs{SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+		}, {
+			desc: "It parses git-receive-pack command and a project with single quotes",
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "git receive-pack 'group/repo'",
+			},
+			expectedArgs: &CommandArgs{SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+		}, {
+			desc: `It parses "git receive-pack" command`,
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": `git receive-pack "group/repo"`,
+			},
+			expectedArgs: &CommandArgs{SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+		}, {
+			desc: `It parses a command followed by control characters`,
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": `git-receive-pack group/repo; any command`,
+			},
+			expectedArgs: &CommandArgs{SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		},
 	}
 
@@ -57,8 +86,8 @@ func TestParseSuccess(t *testing.T) {
 
 			result, err := Parse(tc.arguments)
 
-			assert.NoError(t, err)
-			assert.Equal(t, tc.expectedArgs, result)
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedArgs, result)
 		})
 	}
 }
@@ -67,7 +96,19 @@ func TestParseFailure(t *testing.T) {
 	t.Run("It fails if SSH connection is not set", func(t *testing.T) {
 		_, err := Parse([]string{})
 
-		assert.Error(t, err, "Only ssh allowed")
+		require.Error(t, err, "Only ssh allowed")
 	})
 
+	t.Run("It fails if SSH command is invalid", func(t *testing.T) {
+		environment := map[string]string{
+			"SSH_CONNECTION":       "1",
+			"SSH_ORIGINAL_COMMAND": `git receive-pack "`,
+		}
+		restoreEnv := testhelper.TempEnv(environment)
+		defer restoreEnv()
+
+		_, err := Parse([]string{})
+
+		require.Error(t, err, "Invalid ssh command")
+	})
 }