Merge branch '181-migrate-gitlab-shell-checks-fallback' into 'master'

Support falling back to ruby version of checkers

See merge request gitlab-org/gitlab-shell!318

Makefile

@@ -16,8 +16,14 @@ test_ruby:
 	# bin/gitlab-shell must exist and needs to be the Ruby version for
 	# rspec to be able to test.
 	cp bin/gitlab-shell-ruby bin/gitlab-shell
+	# bin/gitlab-shell-authorized-keys-check and bin/gitlab-shell-authorized-principals-check
+	# should link to ruby scripts for rspec to be able to test.
+	ln -sf ./gitlab-shell-authorized-keys-check-ruby bin/gitlab-shell-authorized-keys-check
+	ln -sf ./gitlab-shell-authorized-principals-check-ruby bin/gitlab-shell-authorized-principals-check
 	bundle exec rspec --color --tag '~go' --format d spec
 	rm -f bin/gitlab-shell
+	ln -sf ./gitlab-shell bin/gitlab-shell-authorized-keys-check
+	ln -sf ./gitlab-shell bin/gitlab-shell-authorized-principals-check
 
 test_golang:
 	support/go-test

bin/gitlab-shell-authorized-keys-check

-#!/usr/bin/env ruby
-
-#
-# GitLab shell authorized_keys helper. Query GitLab API to get the authorized
-# command for a given ssh key fingerprint
-#
-# Ex.
-#   bin/gitlab-shell-authorized-keys-check <username> <public-key>
-#
-# Returns
-#   command="/bin/gitlab-shell key-#",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty ssh-rsa AAAA...
-#
-# Expects to be called by the SSH daemon, via configuration like:
-#     AuthorizedKeysCommandUser git
-#     AuthorizedKeysCommand /bin/gitlab-shell-authorized-keys-check git %u %k
-
-abort "# Wrong number of arguments. #{ARGV.size}. Usage:
-#     gitlab-shell-authorized-keys-check <expected-username> <actual-username> <key>" unless ARGV.size == 3
-
-expected_username = ARGV[0]
-abort '# No username provided' if expected_username.nil? || expected_username == ''
-
-actual_username = ARGV[1]
-abort '# No username provided' if actual_username.nil? || actual_username == ''
-
-# Only check access if the requested username matches the configured username.
-# Normally, these would both be 'git', but it can be configured by the user
-exit 0 unless expected_username == actual_username
-
-key = ARGV[2]
-abort "# No key provided" if key.nil? || key == ''
-
-require_relative '../lib/gitlab_init'
-require_relative '../lib/gitlab_net'
-require_relative '../lib/gitlab_keys'
-
-authorized_key = GitlabNet.new.authorized_key(key)
-if authorized_key.nil?
-  puts "# No key was found for #{key}"
-else
-  puts GitlabKeys.key_line("key-#{authorized_key['id']}", authorized_key['key'])
-end

bin/gitlab-shell-authorized-keys-check

+./gitlab-shell
\ No newline at end of file

bin/gitlab-shell-authorized-keys-check-ruby

+#!/usr/bin/env ruby
+
+#
+# GitLab shell authorized_keys helper. Query GitLab API to get the authorized
+# command for a given ssh key fingerprint
+#
+# Ex.
+#   bin/gitlab-shell-authorized-keys-check <username> <public-key>
+#
+# Returns
+#   command="/bin/gitlab-shell key-#",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty ssh-rsa AAAA...
+#
+# Expects to be called by the SSH daemon, via configuration like:
+#     AuthorizedKeysCommandUser git
+#     AuthorizedKeysCommand /bin/gitlab-shell-authorized-keys-check git %u %k
+
+abort "# Wrong number of arguments. #{ARGV.size}. Usage:
+#     gitlab-shell-authorized-keys-check <expected-username> <actual-username> <key>" unless ARGV.size == 3
+
+expected_username = ARGV[0]
+abort '# No username provided' if expected_username.nil? || expected_username == ''
+
+actual_username = ARGV[1]
+abort '# No username provided' if actual_username.nil? || actual_username == ''
+
+# Only check access if the requested username matches the configured username.
+# Normally, these would both be 'git', but it can be configured by the user
+exit 0 unless expected_username == actual_username
+
+key = ARGV[2]
+abort "# No key provided" if key.nil? || key == ''
+
+require_relative '../lib/gitlab_init'
+require_relative '../lib/gitlab_net'
+require_relative '../lib/gitlab_keys'
+
+authorized_key = GitlabNet.new.authorized_key(key)
+if authorized_key.nil?
+  puts "# No key was found for #{key}"
+else
+  puts GitlabKeys.key_line("key-#{authorized_key['id']}", authorized_key['key'])
+end

bin/gitlab-shell-authorized-principals-check

-#!/usr/bin/env ruby
-
-#
-# GitLab shell authorized principals helper. Emits the same sort of
-# command="..." line as gitlab-shell-authorized-principals-check, with
-# the right options.
-#
-# Ex.
-#   bin/gitlab-shell-authorized-keys-check <key-id> <principal1> [<principal2>...]
-#
-# Returns one line per principal passed in, e.g.:
-#   command="/bin/gitlab-shell username-{KEY_ID}",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty {PRINCIPAL}
-#   [command="/bin/gitlab-shell username-{KEY_ID}",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty {PRINCIPAL2}]
-#
-# Expects to be called by the SSH daemon, via configuration like:
-#     AuthorizedPrincipalsCommandUser root
-#     AuthorizedPrincipalsCommand /bin/gitlab-shell-authorized-principals-check git %i sshUsers
-
-abort "# Wrong number of arguments. #{ARGV.size}. Usage:
-#     gitlab-shell-authorized-principals-check <key-id> <principal1> [<principal2>...]" unless ARGV.size >= 2
-
-key_id = ARGV[0]
-abort '# No key_id provided' if key_id.nil? || key_id == ''
-
-principals = ARGV[1..-1]
-principals.each { |principal|
-  abort '# An invalid principal was provided' if principal.nil? || principal == ''  
-}
-
-require_relative '../lib/gitlab_init'
-require_relative '../lib/gitlab_net'
-require_relative '../lib/gitlab_keys'
-
-principals.each { |principal|
-  puts GitlabKeys.principal_line("username-#{key_id}", principal.dup)
-}

bin/gitlab-shell-authorized-principals-check

+./gitlab-shell
\ No newline at end of file

bin/gitlab-shell-authorized-principals-check-ruby

+#!/usr/bin/env ruby
+
+#
+# GitLab shell authorized principals helper. Emits the same sort of
+# command="..." line as gitlab-shell-authorized-principals-check, with
+# the right options.
+#
+# Ex.
+#   bin/gitlab-shell-authorized-keys-check <key-id> <principal1> [<principal2>...]
+#
+# Returns one line per principal passed in, e.g.:
+#   command="/bin/gitlab-shell username-{KEY_ID}",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty {PRINCIPAL}
+#   [command="/bin/gitlab-shell username-{KEY_ID}",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty {PRINCIPAL2}]
+#
+# Expects to be called by the SSH daemon, via configuration like:
+#     AuthorizedPrincipalsCommandUser root
+#     AuthorizedPrincipalsCommand /bin/gitlab-shell-authorized-principals-check git %i sshUsers
+
+abort "# Wrong number of arguments. #{ARGV.size}. Usage:
+#     gitlab-shell-authorized-principals-check <key-id> <principal1> [<principal2>...]" unless ARGV.size >= 2
+
+key_id = ARGV[0]
+abort '# No key_id provided' if key_id.nil? || key_id == ''
+
+principals = ARGV[1..-1]
+principals.each { |principal|
+  abort '# An invalid principal was provided' if principal.nil? || principal == ''
+}
+
+require_relative '../lib/gitlab_init'
+require_relative '../lib/gitlab_net'
+require_relative '../lib/gitlab_keys'
+
+principals.each { |principal|
+  puts GitlabKeys.principal_line("username-#{key_id}", principal.dup)
+}

go/cmd/gitlab-shell/main.go

@@ -3,42 +3,13 @@ package main
 import (
 	"fmt"
 	"os"
-	"path/filepath"
 
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command"
-	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/fallback"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/readwriter"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 )
 
-// findRootDir determines the root directory (and so, the location of the config
-// file) from os.Executable()
-func findRootDir() (string, error) {
-	if path := os.Getenv("GITLAB_SHELL_DIR"); path != "" {
-		return path, nil
-	}
-
-	path, err := os.Executable()
-	if err != nil {
-		return "", err
-	}
-
-	// Start: /opt/.../gitlab-shell/bin/gitlab-shell
-	// Ends:  /opt/.../gitlab-shell
-	return filepath.Dir(filepath.Dir(path)), nil
-}
-
-// rubyExec will never return. It either replaces the current process with a
-// Ruby interpreter, or outputs an error and kills the process.
-func execRuby(rootDir string, readWriter *readwriter.ReadWriter) {
-	cmd := &fallback.Command{RootDir: rootDir, Args: os.Args}
-
-	if err := cmd.Execute(); err != nil {
-		fmt.Fprintf(readWriter.ErrOut, "Failed to exec: %v\n", err)
-		os.Exit(1)
-	}
-}
-
 func main() {
 	readWriter := &readwriter.ReadWriter{
 		Out:    os.Stdout,
@@ -46,21 +17,19 @@ func main() {
 		ErrOut: os.Stderr,
 	}
 
-	rootDir, err := findRootDir()
+	executable, err := executable.New()
 	if err != nil {
-		fmt.Fprintln(readWriter.ErrOut, "Failed to determine root directory, exiting")
+		fmt.Fprintln(readWriter.ErrOut, "Failed to determine executable, exiting")
 		os.Exit(1)
 	}
 
-	// Fall back to Ruby in case of problems reading the config, but issue a
-	// warning as this isn't something we can sustain indefinitely
-	config, err := config.NewFromDir(rootDir)
+	config, err := config.NewFromDir(executable.RootDir)
 	if err != nil {
-		fmt.Fprintln(readWriter.ErrOut, "Failed to read config, falling back to gitlab-shell-ruby")
-		execRuby(rootDir, readWriter)
+		fmt.Fprintln(readWriter.ErrOut, "Failed to read config, exiting")
+		os.Exit(1)
 	}
 
-	cmd, err := command.New(os.Args, config, readWriter)
+	cmd, err := command.New(executable, os.Args[1:], config, readWriter)
 	if err != nil {
 		// For now this could happen if `SSH_CONNECTION` is not set on
 		// the environment

go/internal/command/command.go

@@ -11,29 +11,40 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/uploadarchive"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/uploadpack"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 )
 
 type Command interface {
 	Execute() error
 }
 
-func New(arguments []string, config *config.Config, readWriter *readwriter.ReadWriter) (Command, error) {
-	args, err := commandargs.Parse(arguments)
-
+func New(e *executable.Executable, arguments []string, config *config.Config, readWriter *readwriter.ReadWriter) (Command, error) {
+	args, err := commandargs.Parse(e, arguments)
 	if err != nil {
 		return nil, err
 	}
 
-	if config.FeatureEnabled(string(args.CommandType)) {
-		if cmd := buildCommand(args, config, readWriter); cmd != nil {
-			return cmd, nil
-		}
+	if cmd := buildCommand(e, args, config, readWriter); cmd != nil {
+		return cmd, nil
+	}
+
+	return &fallback.Command{Executable: e, RootDir: config.RootDir, Args: args}, nil
+}
+
+func buildCommand(e *executable.Executable, args commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command {
+	switch e.Name {
+	case executable.GitlabShell:
+		return buildShellCommand(args.(*commandargs.Shell), config, readWriter)
 	}
 
-	return &fallback.Command{RootDir: config.RootDir, Args: arguments}, nil
+	return nil
 }
 
-func buildCommand(args *commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command {
+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}

go/internal/command/command_test.go

@@ -13,18 +13,22 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/uploadarchive"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/uploadpack"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
 )
 
 func TestNew(t *testing.T) {
 	testCases := []struct {
 		desc         string
+		executable   *executable.Executable
 		config       *config.Config
 		environment  map[string]string
+		arguments    []string
 		expectedType interface{}
 	}{
 		{
-			desc: "it returns a Discover command if the feature is enabled",
+			desc:       "it returns a Discover command if the feature is enabled",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
 				GitlabUrl: "http+unix://gitlab.socket",
 				Migration: config.MigrationConfig{Enabled: true, Features: []string{"discover"}},
@@ -33,10 +37,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
+			arguments:    []string{},
 			expectedType: &discover.Command{},
 		},
 		{
-			desc: "it returns a Fallback command no feature is enabled",
+			desc:       "it returns a Fallback command no feature is enabled",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
 				GitlabUrl: "http+unix://gitlab.socket",
 				Migration: config.MigrationConfig{Enabled: false},
@@ -45,10 +51,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
+			arguments:    []string{},
 			expectedType: &fallback.Command{},
 		},
 		{
-			desc: "it returns a TwoFactorRecover command if the feature is enabled",
+			desc:       "it returns a TwoFactorRecover command if the feature is enabled",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
 				GitlabUrl: "http+unix://gitlab.socket",
 				Migration: config.MigrationConfig{Enabled: true, Features: []string{"2fa_recovery_codes"}},
@@ -57,10 +65,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
 			},
+			arguments:    []string{},
 			expectedType: &twofactorrecover.Command{},
 		},
 		{
-			desc: "it returns an LfsAuthenticate command if the feature is enabled",
+			desc:       "it returns an LfsAuthenticate command if the feature is enabled",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
 				GitlabUrl: "http+unix://gitlab.socket",
 				Migration: config.MigrationConfig{Enabled: true, Features: []string{"git-lfs-authenticate"}},
@@ -69,10 +79,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-lfs-authenticate",
 			},
+			arguments:    []string{},
 			expectedType: &lfsauthenticate.Command{},
 		},
 		{
-			desc: "it returns a ReceivePack command if the feature is enabled",
+			desc:       "it returns a ReceivePack command if the feature is enabled",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
 				GitlabUrl: "http+unix://gitlab.socket",
 				Migration: config.MigrationConfig{Enabled: true, Features: []string{"git-receive-pack"}},
@@ -81,10 +93,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-receive-pack",
 			},
+			arguments:    []string{},
 			expectedType: &receivepack.Command{},
 		},
 		{
-			desc: "it returns a UploadPack command if the feature is enabled",
+			desc:       "it returns an UploadPack command if the feature is enabled",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
 				GitlabUrl: "http+unix://gitlab.socket",
 				Migration: config.MigrationConfig{Enabled: true, Features: []string{"git-upload-pack"}},
@@ -93,10 +107,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-upload-pack",
 			},
+			arguments:    []string{},
 			expectedType: &uploadpack.Command{},
 		},
 		{
-			desc: "it returns a UploadArchive command if the feature is enabled",
+			desc:       "it returns an UploadArchive command if the feature is enabled",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
 				GitlabUrl: "http+unix://gitlab.socket",
 				Migration: config.MigrationConfig{Enabled: true, Features: []string{"git-upload-archive"}},
@@ -105,10 +121,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-upload-archive",
 			},
+			arguments:    []string{},
 			expectedType: &uploadarchive.Command{},
 		},
 		{
-			desc: "it returns a Fallback command if the feature is unimplemented",
+			desc:       "it returns a Fallback command if the feature is unimplemented",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
 				GitlabUrl: "http+unix://gitlab.socket",
 				Migration: config.MigrationConfig{Enabled: true, Features: []string{"git-unimplemented-feature"}},
@@ -117,6 +135,14 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-unimplemented-feature",
 			},
+			arguments:    []string{},
+			expectedType: &fallback.Command{},
+		},
+		{
+			desc:         "it returns a Fallback command if executable is unknown",
+			executable:   &executable.Executable{Name: "unknown"},
+			config:       &config.Config{},
+			arguments:    []string{},
 			expectedType: &fallback.Command{},
 		},
 	}
@@ -126,7 +152,7 @@ func TestNew(t *testing.T) {
 			restoreEnv := testhelper.TempEnv(tc.environment)
 			defer restoreEnv()
 
-			command, err := New([]string{}, tc.config, nil)
+			command, err := New(tc.executable, tc.arguments, tc.config, nil)
 
 			require.NoError(t, err)
 			require.IsType(t, tc.expectedType, command)
@@ -135,12 +161,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{}, &config.Config{}, nil)
+	t.Run("It returns an error parsing arguments failed", func(t *testing.T) {
+		_, err := New(&executable.Executable{Name: executable.GitlabShell}, []string{}, &config.Config{}, nil)
 
-		require.Error(t, err, "Only ssh allowed")
+		require.Error(t, err)
 	})
 }

go/internal/command/commandargs/command_args.go

 package commandargs
 
 import (
-	"errors"
-	"os"
-	"regexp"
-
-	"github.com/mattn/go-shellwords"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 )
 
 type CommandType 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"
-)
-
-var (
-	whoKeyRegex      = regexp.MustCompile(`\bkey-(?P<keyid>\d+)\b`)
-	whoUsernameRegex = regexp.MustCompile(`\busername-(?P<username>\S+)\b`)
-)
-
-type CommandArgs struct {
-	GitlabUsername string
-	GitlabKeyId    string
-	SshArgs        []string
-	CommandType    CommandType
-}
-
-func Parse(arguments []string) (*CommandArgs, error) {
-	if sshConnection := os.Getenv("SSH_CONNECTION"); sshConnection == "" {
-		return nil, errors.New("Only ssh allowed")
-	}
-
-	args := &CommandArgs{}
-	args.parseWho(arguments)
-
-	if err := args.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND")); err != nil {
-		return nil, errors.New("Invalid ssh command")
-	}
-	args.defineCommandType()
-
-	return args, nil
-}
-
-func (c *CommandArgs) parseWho(arguments []string) {
-	for _, argument := range arguments {
-		if keyId := tryParseKeyId(argument); keyId != "" {
-			c.GitlabKeyId = keyId
-			break
-		}
-
-		if username := tryParseUsername(argument); username != "" {
-			c.GitlabUsername = username
-			break
-		}
-	}
+type CommandArgs interface {
+	Parse() error
+	GetArguments() []string
 }
 
-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 Parse(e *executable.Executable, arguments []string) (CommandArgs, error) {
+	var args CommandArgs = &GenericArgs{Arguments: arguments}
 
-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]
+	switch e.Name {
+	case executable.GitlabShell:
+		args = &Shell{Arguments: arguments}
 	}
 
-	return ""
-}
-
-func (c *CommandArgs) parseCommand(commandString string) error {
-	args, err := shellwords.Parse(commandString)
-	if err != nil {
-		return err
+	if err := args.Parse(); err != nil {
+		return nil, 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...)
-	}
-
-	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
 }

go/internal/command/commandargs/command_args_test.go

@@ -3,100 +3,127 @@ package commandargs
 import (
 	"testing"
 
-	"github.com/stretchr/testify/require"
-
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestParseSuccess(t *testing.T) {
 	testCases := []struct {
 		desc         string
-		arguments    []string
+		executable   *executable.Executable
 		environment  map[string]string
-		expectedArgs *CommandArgs
+		arguments    []string
+		expectedArgs CommandArgs
 	}{
 		// Setting the used env variables for every case to ensure we're
 		// not using anything set in the original env.
 		{
-			desc: "It sets discover as the command when the command string was empty",
+			desc:       "It sets discover as the command when the command string was empty",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			environment: map[string]string{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
-			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{}, CommandType: Discover},
 		},
 		{
-			desc: "It finds the key id in any passed arguments",
+			desc:       "It finds the key id in any passed arguments",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			environment: map[string]string{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
 			arguments:    []string{"hello", "key-123"},
-			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
+			expectedArgs: &Shell{Arguments: []string{"hello", "key-123"}, SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
 		}, {
-			desc: "It finds the username in any passed arguments",
+			desc:       "It finds the username in any passed arguments",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			environment: map[string]string{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
 			arguments:    []string{"hello", "username-jane-doe"},
-			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
+			expectedArgs: &Shell{Arguments: []string{"hello", "username-jane-doe"}, SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
 		}, {
-			desc: "It parses 2fa_recovery_codes command",
+			desc:       "It parses 2fa_recovery_codes command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			environment: map[string]string{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
 			},
-			expectedArgs: &CommandArgs{SshArgs: []string{"2fa_recovery_codes"}, CommandType: TwoFactorRecover},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"2fa_recovery_codes"}, CommandType: TwoFactorRecover},
 		}, {
-			desc: "It parses git-receive-pack command",
+			desc:       "It parses git-receive-pack command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			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},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		}, {
-			desc: "It parses git-receive-pack command and a project with single quotes",
+			desc:       "It parses git-receive-pack command and a project with single quotes",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			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},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		}, {
-			desc: `It parses "git receive-pack" command`,
+			desc:       `It parses "git receive-pack" command`,
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			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},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		}, {
-			desc: `It parses a command followed by control characters`,
+			desc:       `It parses a command followed by control characters`,
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			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},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		}, {
-			desc: "It parses git-upload-pack command",
+			desc:       "It parses git-upload-pack command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			environment: map[string]string{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": `git upload-pack "group/repo"`,
 			},
-			expectedArgs: &CommandArgs{SshArgs: []string{"git-upload-pack", "group/repo"}, CommandType: UploadPack},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-upload-pack", "group/repo"}, CommandType: UploadPack},
 		}, {
-			desc: "It parses git-upload-archive command",
+			desc:       "It parses git-upload-archive command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			environment: map[string]string{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-upload-archive 'group/repo'",
 			},
-			expectedArgs: &CommandArgs{SshArgs: []string{"git-upload-archive", "group/repo"}, CommandType: UploadArchive},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-upload-archive", "group/repo"}, CommandType: UploadArchive},
 		}, {
-			desc: "It parses git-lfs-authenticate command",
+			desc:       "It parses git-lfs-authenticate command",
+			executable: &executable.Executable{Name: executable.GitlabShell},
 			environment: map[string]string{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-lfs-authenticate 'group/repo' download",
 			},
-			expectedArgs: &CommandArgs{SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate},
+			arguments:    []string{},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate},
+		}, {
+			desc:         "Unknown executable",
+			executable:   &executable.Executable{Name: "unknown"},
+			arguments:    []string{},
+			expectedArgs: &GenericArgs{Arguments: []string{}},
 		},
 	}
 
@@ -105,7 +132,7 @@ func TestParseSuccess(t *testing.T) {
 			restoreEnv := testhelper.TempEnv(tc.environment)
 			defer restoreEnv()
 
-			result, err := Parse(tc.arguments)
+			result, err := Parse(tc.executable, tc.arguments)
 
 			require.NoError(t, err)
 			require.Equal(t, tc.expectedArgs, result)
@@ -114,22 +141,39 @@ 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{})
-
-		require.Error(t, err, "Only ssh allowed")
-	})
+	testCases := []struct {
+		desc          string
+		executable    *executable.Executable
+		environment   map[string]string
+		arguments     []string
+		expectedError string
+	}{
+		{
+			desc:          "It fails if SSH connection is not set",
+			executable:    &executable.Executable{Name: executable.GitlabShell},
+			arguments:     []string{},
+			expectedError: "Only SSH allowed",
+		},
+		{
+			desc:       "It fails if SSH command is invalid",
+			executable: &executable.Executable{Name: executable.GitlabShell},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": `git receive-pack "`,
+			},
+			arguments:     []string{},
+			expectedError: "Invalid 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()
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			restoreEnv := testhelper.TempEnv(tc.environment)
+			defer restoreEnv()
 
-		_, err := Parse([]string{})
+			_, err := Parse(tc.executable, tc.arguments)
 
-		require.Error(t, err, "Invalid ssh command")
-	})
+			require.Error(t, err, tc.expectedError)
+		})
+	}
 }

go/internal/command/commandargs/generic_args.go

+package commandargs
+
+type GenericArgs struct {
+	Arguments []string
+}
+
+func (b *GenericArgs) Parse() error {
+	// Do nothing
+	return nil
+}
+
+func (b *GenericArgs) GetArguments() []string {
+	return b.Arguments
+}

go/internal/command/commandargs/shell.go

+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 {
+	Arguments      []string
+	GitlabUsername string
+	GitlabKeyId    string
+	SshArgs        []string
+	CommandType    CommandType
+}
+
+func (s *Shell) Parse() error {
+	if err := s.validate(); err != nil {
+		return err
+	}
+
+	s.parseWho()
+	s.defineCommandType()
+
+	return nil
+}
+
+func (s *Shell) GetArguments() []string {
+	return s.Arguments
+}
+
+func (s *Shell) validate() error {
+	if !s.isSshConnection() {
+		return errors.New("Only SSH allowed")
+	}
+
+	if !s.isValidSshCommand() {
+		return errors.New("Invalid SSH command")
+	}
+
+	return nil
+}
+
+func (s *Shell) isSshConnection() bool {
+	ok := os.Getenv("SSH_CONNECTION")
+	return ok != ""
+}
+
+func (s *Shell) isValidSshCommand() bool {
+	err := s.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND"))
+	return err == 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])
+	}
+}

go/internal/command/discover/discover.go

@@ -11,7 +11,7 @@ import (
 
 type Command struct {
 	Config     *config.Config
-	Args       *commandargs.CommandArgs
+	Args       *commandargs.Shell
 	ReadWriter *readwriter.ReadWriter
 }
 

go/internal/command/discover/discover_test.go

@@ -49,27 +49,27 @@ func TestExecute(t *testing.T) {
 
 	testCases := []struct {
 		desc           string
-		arguments      *commandargs.CommandArgs
+		arguments      *commandargs.Shell
 		expectedOutput string
 	}{
 		{
 			desc:           "With a known username",
-			arguments:      &commandargs.CommandArgs{GitlabUsername: "alex-doe"},
+			arguments:      &commandargs.Shell{GitlabUsername: "alex-doe"},
 			expectedOutput: "Welcome to GitLab, @alex-doe!\n",
 		},
 		{
 			desc:           "With a known key id",
-			arguments:      &commandargs.CommandArgs{GitlabKeyId: "1"},
+			arguments:      &commandargs.Shell{GitlabKeyId: "1"},
 			expectedOutput: "Welcome to GitLab, @alex-doe!\n",
 		},
 		{
 			desc:           "With an unknown key",
-			arguments:      &commandargs.CommandArgs{GitlabKeyId: "-1"},
+			arguments:      &commandargs.Shell{GitlabKeyId: "-1"},
 			expectedOutput: "Welcome to GitLab, Anonymous!\n",
 		},
 		{
 			desc:           "With an unknown username",
-			arguments:      &commandargs.CommandArgs{GitlabUsername: "unknown"},
+			arguments:      &commandargs.Shell{GitlabUsername: "unknown"},
 			expectedOutput: "Welcome to GitLab, Anonymous!\n",
 		},
 	}
@@ -97,22 +97,22 @@ func TestFailingExecute(t *testing.T) {
 
 	testCases := []struct {
 		desc          string
-		arguments     *commandargs.CommandArgs
+		arguments     *commandargs.Shell
 		expectedError string
 	}{
 		{
 			desc:          "With missing arguments",
-			arguments:     &commandargs.CommandArgs{},
+			arguments:     &commandargs.Shell{},
 			expectedError: "Failed to get username: who='' is invalid",
 		},
 		{
 			desc:          "When the API returns an error",
-			arguments:     &commandargs.CommandArgs{GitlabUsername: "broken_message"},
+			arguments:     &commandargs.Shell{GitlabUsername: "broken_message"},
 			expectedError: "Failed to get username: Forbidden!",
 		},
 		{
 			desc:          "When the API fails",
-			arguments:     &commandargs.CommandArgs{GitlabUsername: "broken"},
+			arguments:     &commandargs.Shell{GitlabUsername: "broken"},
 			expectedError: "Failed to get username: Internal API error (500)",
 		},
 	}

go/internal/command/fallback/fallback.go

 package fallback
 
 import (
+	"errors"
+	"fmt"
 	"os"
 	"path/filepath"
 	"syscall"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 )
 
 type Command struct {
-	RootDir string
-	Args    []string
+	Executable *executable.Executable
+	RootDir    string
+	Args       commandargs.CommandArgs
 }
 
 var (
 	// execFunc is overridden in tests
-	execFunc = syscall.Exec
-)
-
-const (
-	RubyProgram = "gitlab-shell-ruby"
+	execFunc  = syscall.Exec
+	whitelist = []string{
+		executable.GitlabShell,
+		executable.AuthorizedKeysCheck,
+		executable.AuthorizedPrincipalsCheck,
+	}
 )
 
 func (c *Command) Execute() error {
-	rubyCmd := filepath.Join(c.RootDir, "bin", RubyProgram)
+	if !c.isWhitelisted() {
+		return errors.New("Failed to execute unknown executable")
+	}
+
+	rubyCmd := c.fallbackProgram()
 
 	// Ensure rubyArgs[0] is the full path to gitlab-shell-ruby
-	rubyArgs := append([]string{rubyCmd}, c.Args[1:]...)
+	rubyArgs := append([]string{rubyCmd}, c.Args.GetArguments()...)
 
 	return execFunc(rubyCmd, rubyArgs, os.Environ())
 }
+
+func (c *Command) isWhitelisted() bool {
+	for _, item := range whitelist {
+		if c.Executable.Name == item {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (c *Command) fallbackProgram() string {
+	fileName := fmt.Sprintf("%s-ruby", c.Executable.Name)
+
+	return filepath.Join(c.RootDir, "bin", fileName)
+}

go/internal/command/fallback/fallback_test.go

@@ -6,6 +6,9 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 )
 
 type fakeExec struct {
@@ -19,7 +22,7 @@ type fakeExec struct {
 }
 
 var (
-	fakeArgs = []string{"./test", "foo", "bar"}
+	fakeArgs = &commandargs.GenericArgs{Arguments: []string{"foo", "bar"}}
 )
 
 func (f *fakeExec) Exec(filename string, args []string, env []string) error {
@@ -42,7 +45,7 @@ func (f *fakeExec) Cleanup() {
 }
 
 func TestExecuteExecsCommandSuccesfully(t *testing.T) {
-	cmd := &Command{RootDir: "/tmp", Args: fakeArgs}
+	cmd := &Command{Executable: &executable.Executable{Name: executable.GitlabShell}, RootDir: "/tmp", Args: fakeArgs}
 
 	// Override the exec func
 	fake := &fakeExec{}
@@ -56,8 +59,14 @@ func TestExecuteExecsCommandSuccesfully(t *testing.T) {
 	require.Equal(t, fake.Env, os.Environ())
 }
 
+func TestExecuteExecsUnknownExecutable(t *testing.T) {
+	cmd := &Command{Executable: &executable.Executable{Name: "unknown"}, RootDir: "/test"}
+
+	require.Error(t, cmd.Execute())
+}
+
 func TestExecuteExecsCommandOnError(t *testing.T) {
-	cmd := &Command{RootDir: "/test", Args: fakeArgs}
+	cmd := &Command{Executable: &executable.Executable{Name: executable.GitlabShell}, RootDir: "/test", Args: fakeArgs}
 
 	// Override the exec func
 	fake := &fakeExec{Error: errors.New("Test error")}
@@ -69,7 +78,7 @@ func TestExecuteExecsCommandOnError(t *testing.T) {
 }
 
 func TestExecuteGivenNonexistentCommand(t *testing.T) {
-	cmd := &Command{RootDir: "/tmp/does/not/exist", Args: fakeArgs}
+	cmd := &Command{Executable: &executable.Executable{Name: executable.GitlabShell}, RootDir: "/tmp/does/not/exist", Args: fakeArgs}
 
 	require.Error(t, cmd.Execute())
 }

go/internal/command/lfsauthenticate/lfsauthenticate.go

@@ -20,7 +20,7 @@ const (
 
 type Command struct {
 	Config     *config.Config
-	Args       *commandargs.CommandArgs
+	Args       *commandargs.Shell
 	ReadWriter *readwriter.ReadWriter
 }
 

go/internal/command/lfsauthenticate/lfsauthenticate_test.go

@@ -25,22 +25,22 @@ func TestFailedRequests(t *testing.T) {
 
 	testCases := []struct {
 		desc           string
-		arguments      *commandargs.CommandArgs
+		arguments      *commandargs.Shell
 		expectedOutput string
 	}{
 		{
 			desc:           "With missing arguments",
-			arguments:      &commandargs.CommandArgs{},
+			arguments:      &commandargs.Shell{},
 			expectedOutput: "> GitLab: Disallowed command",
 		},
 		{
 			desc:           "With disallowed command",
-			arguments:      &commandargs.CommandArgs{GitlabKeyId: "1", SshArgs: []string{"git-lfs-authenticate", "group/repo", "unknown"}},
+			arguments:      &commandargs.Shell{GitlabKeyId: "1", SshArgs: []string{"git-lfs-authenticate", "group/repo", "unknown"}},
 			expectedOutput: "> GitLab: Disallowed command",
 		},
 		{
 			desc:           "With disallowed user",
-			arguments:      &commandargs.CommandArgs{GitlabKeyId: "disallowed", SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}},
+			arguments:      &commandargs.Shell{GitlabKeyId: "disallowed", SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}},
 			expectedOutput: "Disallowed by API call",
 		},
 	}
@@ -140,7 +140,7 @@ func TestLfsAuthenticateRequests(t *testing.T) {
 			output := &bytes.Buffer{}
 			cmd := &Command{
 				Config:     &config.Config{GitlabUrl: url},
-				Args:       &commandargs.CommandArgs{GitlabUsername: tc.username, SshArgs: []string{"git-lfs-authenticate", "group/repo", "upload"}},
+				Args:       &commandargs.Shell{GitlabUsername: tc.username, SshArgs: []string{"git-lfs-authenticate", "group/repo", "upload"}},
 				ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output},
 			}