Support different CommandArgs type

`CommandArgs` has been renamed to `Shell`.

An interface has been added that includes `Executable()` and
`Arguments()` method. The `BaseArgs` implement this methods
and should be embeeded in each type.

8 files changed, 238 insertions, 156 deletions

diff --git a/go/internal/command/command.go b/go/internal/command/command.go
index 77feda8..66bed6f 100644
--- a/go/internal/command/command.go
+++ b/go/internal/command/command.go
@@ -23,16 +23,27 @@ func New(arguments []string, config *config.Config, readWriter *readwriter.ReadW
 		return nil, err
 	}
 
-	if config.FeatureEnabled(string(args.CommandType)) {
-		if cmd := buildCommand(args, config, readWriter); cmd != nil {
-			return cmd, nil
-		}
+	if cmd := buildCommand(args, config, readWriter); cmd != nil {
+		return cmd, nil
 	}
 
 	return &fallback.Command{RootDir: config.RootDir, Args: args}, nil
 }
 
-func buildCommand(args *commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command {
+func buildCommand(args commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command {
+	switch args.Executable() {
+	case commandargs.GitlabShell:
+		return buildShellCommand(args.(*commandargs.Shell), config, readWriter)
+	}
+
+	return nil
+}
+
+func buildShellCommand(args *commandargs.Shell, config *config.Config, readWriter *readwriter.ReadWriter) Command {
+	if !config.FeatureEnabled(string(args.CommandType)) {
+		return nil
+	}
+
 	switch args.CommandType {
 	case commandargs.Discover:
 		return &discover.Command{Config: config, Args: args, ReadWriter: readWriter}
diff --git a/go/internal/command/command_test.go b/go/internal/command/command_test.go
index 2341dbb..d8ef295 100644
--- a/go/internal/command/command_test.go
+++ b/go/internal/command/command_test.go
@@ -129,6 +129,12 @@ func TestNew(t *testing.T) {
 			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &fallback.Command{},
 		},
+		{
+			desc:         "it returns a Fallback command if executable is unknown",
+			config:       &config.Config{},
+			arguments:    []string{"unknown"},
+			expectedType: &fallback.Command{},
+		},
 	}
 
 	for _, tc := range testCases {
@@ -145,12 +151,9 @@ func TestNew(t *testing.T) {
 }
 
 func TestFailingNew(t *testing.T) {
-	t.Run("It returns an error when SSH_CONNECTION is not set", func(t *testing.T) {
-		restoreEnv := testhelper.TempEnv(map[string]string{})
-		defer restoreEnv()
-
-		_, err := New([]string{string(commandargs.GitlabShell)}, &config.Config{}, nil)
+	t.Run("It returns an error parsing arguments failed", func(t *testing.T) {
+		_, err := New([]string{}, &config.Config{}, nil)
 
-		require.Error(t, err, "Only ssh allowed")
+		require.Error(t, err)
 	})
 }
diff --git a/go/internal/command/commandargs/base_args.go b/go/internal/command/commandargs/base_args.go
new file mode 100644
index 0000000..f65373e
--- /dev/null
+++ b/go/internal/command/commandargs/base_args.go
@@ -0,0 +1,34 @@
+package commandargs
+
+import (
+	"errors"
+	"path/filepath"
+)
+
+type BaseArgs struct {
+	arguments []string
+}
+
+func (b *BaseArgs) Parse() error {
+	if b.hasEmptyArguments() {
+		return errors.New("arguments should include the executable")
+	}
+
+	return nil
+}
+
+func (b *BaseArgs) Executable() Executable {
+	if b.hasEmptyArguments() {
+		return Executable("")
+	}
+
+	return Executable(filepath.Base(b.arguments[0]))
+}
+
+func (b *BaseArgs) Arguments() []string {
+	return b.arguments[1:]
+}
+
+func (b *BaseArgs) hasEmptyArguments() bool {
+	return len(b.arguments) == 0
+}
diff --git a/go/internal/command/commandargs/command_args.go b/go/internal/command/commandargs/command_args.go
index d68f5d2..9f28817 100644
--- a/go/internal/command/commandargs/command_args.go
+++ b/go/internal/command/commandargs/command_args.go
@@ -1,128 +1,29 @@
 package commandargs
 
-import (
-	"errors"
-	"os"
-	"path/filepath"
-	"regexp"
-
-	"github.com/mattn/go-shellwords"
-)
-
 type CommandType string
 type Executable string
 
 const (
-	Discover         CommandType = "discover"
-	TwoFactorRecover CommandType = "2fa_recovery_codes"
-	LfsAuthenticate  CommandType = "git-lfs-authenticate"
-	ReceivePack      CommandType = "git-receive-pack"
-	UploadPack       CommandType = "git-upload-pack"
-	UploadArchive    CommandType = "git-upload-archive"
-	GitlabShell      Executable  = "gitlab-shell"
-)
-
-var (
-	whoKeyRegex      = regexp.MustCompile(`\bkey-(?P<keyid>\d+)\b`)
-	whoUsernameRegex = regexp.MustCompile(`\busername-(?P<username>\S+)\b`)
+	GitlabShell Executable = "gitlab-shell"
 )
 
-type CommandArgs struct {
-	arguments      []string
-	GitlabUsername string
-	GitlabKeyId    string
-	SshArgs        []string
-	CommandType    CommandType
-}
-
-func Parse(arguments []string) (*CommandArgs, error) {
-	args := &CommandArgs{arguments: arguments}
-
-	if args.Executable() == GitlabShell {
-		if sshConnection := os.Getenv("SSH_CONNECTION"); sshConnection == "" {
-			return nil, errors.New("Only ssh allowed")
-		}
-
-		args.parseWho()
-
-		if err := args.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND")); err != nil {
-			return nil, errors.New("Invalid ssh command")
-		}
-		args.defineCommandType()
-
-		return args, nil
-	} else {
-		return args, nil
-	}
-}
-
-func (c *CommandArgs) Executable() Executable {
-	return Executable(filepath.Base(c.arguments[0]))
-}
-
-func (c *CommandArgs) Arguments() []string {
-	return c.arguments[1:]
+type CommandArgs interface {
+	Parse() error
+	Executable() Executable
+	Arguments() []string
 }
 
-func (c *CommandArgs) parseWho() {
-	for _, argument := range c.arguments {
-		if keyId := tryParseKeyId(argument); keyId != "" {
-			c.GitlabKeyId = keyId
-			break
-		}
+func Parse(arguments []string) (CommandArgs, error) {
+	var args CommandArgs = &BaseArgs{arguments: arguments}
 
-		if username := tryParseUsername(argument); username != "" {
-			c.GitlabUsername = username
-			break
-		}
+	switch args.Executable() {
+	case GitlabShell:
+		args = &Shell{BaseArgs: args.(*BaseArgs)}
 	}
-}
-
-func tryParseKeyId(argument string) string {
-	matchInfo := whoKeyRegex.FindStringSubmatch(argument)
-	if len(matchInfo) == 2 {
-		// The first element is the full matched string
-		// The second element is the named `keyid`
-		return matchInfo[1]
-	}
-
-	return ""
-}
-
-func tryParseUsername(argument string) string {
-	matchInfo := whoUsernameRegex.FindStringSubmatch(argument)
-	if len(matchInfo) == 2 {
-		// The first element is the full matched string
-		// The second element is the named `username`
-		return matchInfo[1]
-	}
-
-	return ""
-}
-
-func (c *CommandArgs) parseCommand(commandString string) error {
-	args, err := shellwords.Parse(commandString)
-	if err != nil {
-		return err
-	}
-
-	// Handle Git for Windows 2.14 using "git upload-pack" instead of git-upload-pack
-	if len(args) > 1 && args[0] == "git" {
-		command := args[0] + "-" + args[1]
-		commandArgs := args[2:]
 
-		args = append([]string{command}, commandArgs...)
+	if err := args.Parse(); err != nil {
+		return nil, err
 	}
 
-	c.SshArgs = args
-
-	return nil
-}
-
-func (c *CommandArgs) defineCommandType() {
-	if len(c.SshArgs) == 0 {
-		c.CommandType = Discover
-	} else {
-		c.CommandType = CommandType(c.SshArgs[0])
-	}
+	return args, nil
 }
diff --git a/go/internal/command/commandargs/command_args_test.go b/go/internal/command/commandargs/command_args_test.go
index 21a676c..1fe3ba7 100644
--- a/go/internal/command/commandargs/command_args_test.go
+++ b/go/internal/command/commandargs/command_args_test.go
@@ -13,7 +13,7 @@ func TestParseSuccess(t *testing.T) {
 		desc         string
 		environment  map[string]string
 		arguments    []string
-		expectedArgs *CommandArgs
+		expectedArgs CommandArgs
 	}{
 		// Setting the used env variables for every case to ensure we're
 		// not using anything set in the original env.
@@ -24,7 +24,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "",
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{}, CommandType: Discover},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, SshArgs: []string{}, CommandType: Discover},
 		},
 		{
 			desc: "It finds the key id in any passed arguments",
@@ -33,7 +33,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "",
 			},
 			arguments:    []string{string(GitlabShell), "hello", "key-123"},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell), "hello", "key-123"}, SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell), "hello", "key-123"}}, SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
 		}, {
 			desc: "It finds the username in any passed arguments",
 			environment: map[string]string{
@@ -41,7 +41,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "",
 			},
 			arguments:    []string{string(GitlabShell), "hello", "username-jane-doe"},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell), "hello", "username-jane-doe"}, SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell), "hello", "username-jane-doe"}}, SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
 		}, {
 			desc: "It parses 2fa_recovery_codes command",
 			environment: map[string]string{
@@ -49,7 +49,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"2fa_recovery_codes"}, CommandType: TwoFactorRecover},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, SshArgs: []string{"2fa_recovery_codes"}, CommandType: TwoFactorRecover},
 		}, {
 			desc: "It parses git-receive-pack command",
 			environment: map[string]string{
@@ -57,7 +57,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "git-receive-pack group/repo",
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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{
@@ -65,7 +65,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "git receive-pack 'group/repo'",
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		}, {
 			desc: `It parses "git receive-pack" command`,
 			environment: map[string]string{
@@ -73,7 +73,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": `git receive-pack "group/repo"`,
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		}, {
 			desc: `It parses a command followed by control characters`,
 			environment: map[string]string{
@@ -81,7 +81,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": `git-receive-pack group/repo; any command`,
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		}, {
 			desc: "It parses git-upload-pack command",
 			environment: map[string]string{
@@ -89,7 +89,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": `git upload-pack "group/repo"`,
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-upload-pack", "group/repo"}, CommandType: UploadPack},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, SshArgs: []string{"git-upload-pack", "group/repo"}, CommandType: UploadPack},
 		}, {
 			desc: "It parses git-upload-archive command",
 			environment: map[string]string{
@@ -97,7 +97,7 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "git-upload-archive 'group/repo'",
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-upload-archive", "group/repo"}, CommandType: UploadArchive},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, SshArgs: []string{"git-upload-archive", "group/repo"}, CommandType: UploadArchive},
 		}, {
 			desc: "It parses git-lfs-authenticate command",
 			environment: map[string]string{
@@ -105,7 +105,11 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": "git-lfs-authenticate 'group/repo' download",
 			},
 			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate},
+			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate},
+		}, {
+			desc:         "Unknown executable",
+			arguments:    []string{"unknown"},
+			expectedArgs: &BaseArgs{arguments: []string{"unknown"}},
 		},
 	}
 
@@ -123,22 +127,41 @@ func TestParseSuccess(t *testing.T) {
 }
 
 func TestParseFailure(t *testing.T) {
-	t.Run("It fails if SSH connection is not set", func(t *testing.T) {
-		_, err := Parse([]string{string(GitlabShell)})
-
-		require.Error(t, err, "Only ssh allowed")
-	})
+	testCases := []struct {
+		desc          string
+		environment   map[string]string
+		arguments     []string
+		expectedError string
+	}{
+		{
+			desc:          "It fails if SSH connection is not set",
+			arguments:     []string{string(GitlabShell)},
+			expectedError: "Only ssh allowed",
+		},
+		{
+			desc: "It fails if SSH command is invalid",
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": `git receive-pack "`,
+			},
+			arguments:     []string{string(GitlabShell)},
+			expectedError: "Only ssh allowed",
+		},
+		{
+			desc:          "It fails if arguments is empty",
+			arguments:     []string{},
+			expectedError: "arguments should include the executable",
+		},
+	}
 
-	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()
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			restoreEnv := testhelper.TempEnv(tc.environment)
+			defer restoreEnv()
 
-		_, err := Parse([]string{string(GitlabShell)})
+			_, err := Parse(tc.arguments)
 
-		require.Error(t, err, "Invalid ssh command")
-	})
+			require.Error(t, err, tc.expectedError)
+		})
+	}
 }
diff --git a/go/internal/command/commandargs/shell.go b/go/internal/command/commandargs/shell.go
new file mode 100644
index 0000000..04b1040
--- /dev/null
+++ b/go/internal/command/commandargs/shell.go
@@ -0,0 +1,110 @@
+package commandargs
+
+import (
+	"errors"
+	"os"
+	"regexp"
+
+	"github.com/mattn/go-shellwords"
+)
+
+const (
+	Discover         CommandType = "discover"
+	TwoFactorRecover CommandType = "2fa_recovery_codes"
+	LfsAuthenticate  CommandType = "git-lfs-authenticate"
+	ReceivePack      CommandType = "git-receive-pack"
+	UploadPack       CommandType = "git-upload-pack"
+	UploadArchive    CommandType = "git-upload-archive"
+)
+
+var (
+	whoKeyRegex      = regexp.MustCompile(`\bkey-(?P<keyid>\d+)\b`)
+	whoUsernameRegex = regexp.MustCompile(`\busername-(?P<username>\S+)\b`)
+)
+
+type Shell struct {
+	*BaseArgs
+	GitlabUsername string
+	GitlabKeyId    string
+	SshArgs        []string
+	CommandType    CommandType
+}
+
+func (s *Shell) Parse() error {
+	if sshConnection := os.Getenv("SSH_CONNECTION"); sshConnection == "" {
+		return errors.New("Only ssh allowed")
+	}
+
+	s.parseWho()
+
+	if err := s.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND")); err != nil {
+		return errors.New("Invalid ssh command")
+	}
+
+	s.defineCommandType()
+
+	return nil
+}
+
+func (s *Shell) parseWho() {
+	for _, argument := range s.arguments {
+		if keyId := tryParseKeyId(argument); keyId != "" {
+			s.GitlabKeyId = keyId
+			break
+		}
+
+		if username := tryParseUsername(argument); username != "" {
+			s.GitlabUsername = username
+			break
+		}
+	}
+}
+
+func tryParseKeyId(argument string) string {
+	matchInfo := whoKeyRegex.FindStringSubmatch(argument)
+	if len(matchInfo) == 2 {
+		// The first element is the full matched string
+		// The second element is the named `keyid`
+		return matchInfo[1]
+	}
+
+	return ""
+}
+
+func tryParseUsername(argument string) string {
+	matchInfo := whoUsernameRegex.FindStringSubmatch(argument)
+	if len(matchInfo) == 2 {
+		// The first element is the full matched string
+		// The second element is the named `username`
+		return matchInfo[1]
+	}
+
+	return ""
+}
+
+func (s *Shell) parseCommand(commandString string) error {
+	args, err := shellwords.Parse(commandString)
+	if err != nil {
+		return err
+	}
+
+	// Handle Git for Windows 2.14 using "git upload-pack" instead of git-upload-pack
+	if len(args) > 1 && args[0] == "git" {
+		command := args[0] + "-" + args[1]
+		commandArgs := args[2:]
+
+		args = append([]string{command}, commandArgs...)
+	}
+
+	s.SshArgs = args
+
+	return nil
+}
+
+func (s *Shell) defineCommandType() {
+	if len(s.SshArgs) == 0 {
+		s.CommandType = Discover
+	} else {
+		s.CommandType = CommandType(s.SshArgs[0])
+	}
+}
diff --git a/go/internal/command/fallback/fallback.go b/go/internal/command/fallback/fallback.go
index cec94d5..81baaf5 100644
--- a/go/internal/command/fallback/fallback.go
+++ b/go/internal/command/fallback/fallback.go
@@ -9,14 +9,9 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
 )
 
-type CommandArgs interface {
-	Executable() commandargs.Executable
-	Arguments() []string
-}
-
 type Command struct {
 	RootDir string
-	Args    CommandArgs
+	Args    commandargs.CommandArgs
 }
 
 var (
diff --git a/go/internal/command/fallback/fallback_test.go b/go/internal/command/fallback/fallback_test.go
index 91dcb53..669aad1 100644
--- a/go/internal/command/fallback/fallback_test.go
+++ b/go/internal/command/fallback/fallback_test.go
@@ -25,6 +25,11 @@ type FakeCommandArgs struct {
 	arguments  []string
 }
 
+func (f *FakeCommandArgs) Parse() error {
+	// Do nothing as no need to parse anything
+	return nil
+}
+
 func (f *FakeCommandArgs) Executable() commandargs.Executable {
 	return f.executable
 }