Support falling back to ruby version of checkers

Rename the ruby scripts to have `-ruby` suffix and add a symlink
for both to `./gitlab-shell`. The executable name will be used to
determine how args will be parsed.

For now, we only parse the arguments for gitlab-shell commands. If
the executable is `gitlab-shell-authorized-keys-check` or
`gitlab-shell-authorized-principals-check`, it'll always fallback
to the ruby version.

Ruby specs test the ruby script, the fallback from go to ruby and
go implementation of both (still pending).

14 files changed, 331 insertions, 167 deletions

diff --git a/Makefile b/Makefile
index 84cba6b..aaa3098 100644
--- a/Makefile
+++ b/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
diff --git a/bin/gitlab-shell-authorized-keys-check b/bin/gitlab-shell-authorized-keys-check
index 2ea1a74..3dc14d1 100755..120000
--- a/bin/gitlab-shell-authorized-keys-check
+++ b/bin/gitlab-shell-authorized-keys-check
@@ -1,42 +1 @@
-#!/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
+./gitlab-shell
\ No newline at end of file
diff --git a/bin/gitlab-shell-authorized-keys-check-ruby b/bin/gitlab-shell-authorized-keys-check-ruby
new file mode 100755
index 0000000..2ea1a74
--- /dev/null
+++ b/bin/gitlab-shell-authorized-keys-check-ruby
@@ -0,0 +1,42 @@
+#!/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
diff --git a/bin/gitlab-shell-authorized-principals-check b/bin/gitlab-shell-authorized-principals-check
index aa6d427..3dc14d1 100755..120000
--- a/bin/gitlab-shell-authorized-principals-check
+++ b/bin/gitlab-shell-authorized-principals-check
@@ -1,36 +1 @@
-#!/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)
-}
+./gitlab-shell
\ No newline at end of file
diff --git a/bin/gitlab-shell-authorized-principals-check-ruby b/bin/gitlab-shell-authorized-principals-check-ruby
new file mode 100755
index 0000000..25ee612
--- /dev/null
+++ b/bin/gitlab-shell-authorized-principals-check-ruby
@@ -0,0 +1,36 @@
+#!/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)
+}
diff --git a/go/cmd/gitlab-shell/main.go b/go/cmd/gitlab-shell/main.go
index f4d519f..be388d3 100644
--- a/go/cmd/gitlab-shell/main.go
+++ b/go/cmd/gitlab-shell/main.go
@@ -6,7 +6,6 @@ import (
 	"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"
 )
@@ -28,17 +27,6 @@ func findRootDir() (string, error) {
 	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,
@@ -52,12 +40,10 @@ func main() {
 		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)
 	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)
diff --git a/go/internal/command/command.go b/go/internal/command/command.go
index a1dde42..77feda8 100644
--- a/go/internal/command/command.go
+++ b/go/internal/command/command.go
@@ -19,7 +19,6 @@ type Command interface {
 
 func New(arguments []string, config *config.Config, readWriter *readwriter.ReadWriter) (Command, error) {
 	args, err := commandargs.Parse(arguments)
-
 	if err != nil {
 		return nil, err
 	}
@@ -30,7 +29,7 @@ func New(arguments []string, config *config.Config, readWriter *readwriter.ReadW
 		}
 	}
 
-	return &fallback.Command{RootDir: config.RootDir, Args: arguments}, nil
+	return &fallback.Command{RootDir: config.RootDir, Args: args}, nil
 }
 
 func buildCommand(args *commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command {
diff --git a/go/internal/command/command_test.go b/go/internal/command/command_test.go
index 07260dd..2341dbb 100644
--- a/go/internal/command/command_test.go
+++ b/go/internal/command/command_test.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/discover"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/fallback"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/lfsauthenticate"
@@ -21,6 +22,7 @@ func TestNew(t *testing.T) {
 		desc         string
 		config       *config.Config
 		environment  map[string]string
+		arguments    []string
 		expectedType interface{}
 	}{
 		{
@@ -33,6 +35,7 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
+			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &discover.Command{},
 		},
 		{
@@ -45,6 +48,7 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
+			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &fallback.Command{},
 		},
 		{
@@ -57,6 +61,7 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
 			},
+			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &twofactorrecover.Command{},
 		},
 		{
@@ -69,6 +74,7 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-lfs-authenticate",
 			},
+			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &lfsauthenticate.Command{},
 		},
 		{
@@ -81,6 +87,7 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-receive-pack",
 			},
+			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &receivepack.Command{},
 		},
 		{
@@ -93,6 +100,7 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-upload-pack",
 			},
+			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &uploadpack.Command{},
 		},
 		{
@@ -105,6 +113,7 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-upload-archive",
 			},
+			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &uploadarchive.Command{},
 		},
 		{
@@ -117,6 +126,7 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-unimplemented-feature",
 			},
+			arguments:    []string{string(commandargs.GitlabShell)},
 			expectedType: &fallback.Command{},
 		},
 	}
@@ -126,7 +136,7 @@ func TestNew(t *testing.T) {
 			restoreEnv := testhelper.TempEnv(tc.environment)
 			defer restoreEnv()
 
-			command, err := New([]string{}, tc.config, nil)
+			command, err := New(tc.arguments, tc.config, nil)
 
 			require.NoError(t, err)
 			require.IsType(t, tc.expectedType, command)
@@ -139,7 +149,7 @@ func TestFailingNew(t *testing.T) {
 		restoreEnv := testhelper.TempEnv(map[string]string{})
 		defer restoreEnv()
 
-		_, err := New([]string{}, &config.Config{}, nil)
+		_, err := New([]string{string(commandargs.GitlabShell)}, &config.Config{}, nil)
 
 		require.Error(t, err, "Only ssh allowed")
 	})
diff --git a/go/internal/command/commandargs/command_args.go b/go/internal/command/commandargs/command_args.go
index d8fe32d..d68f5d2 100644
--- a/go/internal/command/commandargs/command_args.go
+++ b/go/internal/command/commandargs/command_args.go
@@ -3,12 +3,14 @@ package commandargs
 import (
 	"errors"
 	"os"
+	"path/filepath"
 	"regexp"
 
 	"github.com/mattn/go-shellwords"
 )
 
 type CommandType string
+type Executable string
 
 const (
 	Discover         CommandType = "discover"
@@ -17,6 +19,7 @@ const (
 	ReceivePack      CommandType = "git-receive-pack"
 	UploadPack       CommandType = "git-upload-pack"
 	UploadArchive    CommandType = "git-upload-archive"
+	GitlabShell      Executable  = "gitlab-shell"
 )
 
 var (
@@ -25,6 +28,7 @@ var (
 )
 
 type CommandArgs struct {
+	arguments      []string
 	GitlabUsername string
 	GitlabKeyId    string
 	SshArgs        []string
@@ -32,23 +36,36 @@ type CommandArgs struct {
 }
 
 func Parse(arguments []string) (*CommandArgs, error) {
-	if sshConnection := os.Getenv("SSH_CONNECTION"); sshConnection == "" {
-		return nil, errors.New("Only ssh allowed")
-	}
+	args := &CommandArgs{arguments: arguments}
+
+	if args.Executable() == GitlabShell {
+		if sshConnection := os.Getenv("SSH_CONNECTION"); sshConnection == "" {
+			return nil, errors.New("Only ssh allowed")
+		}
 
-	args := &CommandArgs{}
-	args.parseWho(arguments)
+		args.parseWho()
+
+		if err := args.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND")); err != nil {
+			return nil, errors.New("Invalid ssh command")
+		}
+		args.defineCommandType()
 
-	if err := args.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND")); err != nil {
-		return nil, errors.New("Invalid ssh command")
+		return args, nil
+	} else {
+		return args, nil
 	}
-	args.defineCommandType()
+}
+
+func (c *CommandArgs) Executable() Executable {
+	return Executable(filepath.Base(c.arguments[0]))
+}
 
-	return args, nil
+func (c *CommandArgs) Arguments() []string {
+	return c.arguments[1:]
 }
 
-func (c *CommandArgs) parseWho(arguments []string) {
-	for _, argument := range arguments {
+func (c *CommandArgs) parseWho() {
+	for _, argument := range c.arguments {
 		if keyId := tryParseKeyId(argument); keyId != "" {
 			c.GitlabKeyId = keyId
 			break
diff --git a/go/internal/command/commandargs/command_args_test.go b/go/internal/command/commandargs/command_args_test.go
index e60bb92..21a676c 100644
--- a/go/internal/command/commandargs/command_args_test.go
+++ b/go/internal/command/commandargs/command_args_test.go
@@ -3,16 +3,16 @@ package commandargs
 import (
 	"testing"
 
-	"github.com/stretchr/testify/require"
-
 	"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
 		environment  map[string]string
+		arguments    []string
 		expectedArgs *CommandArgs
 	}{
 		// Setting the used env variables for every case to ensure we're
@@ -23,7 +23,8 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
-			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover},
+			arguments:    []string{string(GitlabShell)},
+			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{}, CommandType: Discover},
 		},
 		{
 			desc: "It finds the key id in any passed arguments",
@@ -31,72 +32,80 @@ func TestParseSuccess(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
-			arguments:    []string{"hello", "key-123"},
-			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
+			arguments:    []string{string(GitlabShell), "hello", "key-123"},
+			expectedArgs: &CommandArgs{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{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
-			arguments:    []string{"hello", "username-jane-doe"},
-			expectedArgs: &CommandArgs{SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
+			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"},
 		}, {
 			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},
+			arguments:    []string{string(GitlabShell)},
+			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, 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},
+			arguments:    []string{string(GitlabShell)},
+			expectedArgs: &CommandArgs{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{
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git receive-pack 'group/repo'",
 			},
-			expectedArgs: &CommandArgs{SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			arguments:    []string{string(GitlabShell)},
+			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, 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},
+			arguments:    []string{string(GitlabShell)},
+			expectedArgs: &CommandArgs{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{
 				"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{string(GitlabShell)},
+			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
 		}, {
 			desc: "It parses git-upload-pack command",
 			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{string(GitlabShell)},
+			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-upload-pack", "group/repo"}, CommandType: UploadPack},
 		}, {
 			desc: "It parses git-upload-archive command",
 			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{string(GitlabShell)},
+			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-upload-archive", "group/repo"}, CommandType: UploadArchive},
 		}, {
 			desc: "It parses git-lfs-authenticate command",
 			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{string(GitlabShell)},
+			expectedArgs: &CommandArgs{arguments: []string{string(GitlabShell)}, SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate},
 		},
 	}
 
@@ -115,7 +124,7 @@ 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{})
+		_, err := Parse([]string{string(GitlabShell)})
 
 		require.Error(t, err, "Only ssh allowed")
 	})
@@ -128,7 +137,7 @@ func TestParseFailure(t *testing.T) {
 		restoreEnv := testhelper.TempEnv(environment)
 		defer restoreEnv()
 
-		_, err := Parse([]string{})
+		_, err := Parse([]string{string(GitlabShell)})
 
 		require.Error(t, err, "Invalid ssh command")
 	})
diff --git a/go/internal/command/fallback/fallback.go b/go/internal/command/fallback/fallback.go
index f525a57..cec94d5 100644
--- a/go/internal/command/fallback/fallback.go
+++ b/go/internal/command/fallback/fallback.go
@@ -1,14 +1,22 @@
 package fallback
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"syscall"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
 )
 
+type CommandArgs interface {
+	Executable() commandargs.Executable
+	Arguments() []string
+}
+
 type Command struct {
 	RootDir string
-	Args    []string
+	Args    CommandArgs
 }
 
 var (
@@ -16,15 +24,15 @@ var (
 	execFunc = syscall.Exec
 )
 
-const (
-	RubyProgram = "gitlab-shell-ruby"
-)
-
 func (c *Command) Execute() error {
-	rubyCmd := filepath.Join(c.RootDir, "bin", RubyProgram)
+	rubyCmd := filepath.Join(c.RootDir, "bin", 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.Arguments()...)
 
 	return execFunc(rubyCmd, rubyArgs, os.Environ())
 }
+
+func (c *Command) fallbackProgram() string {
+	return fmt.Sprintf("%s-ruby", c.Args.Executable())
+}
diff --git a/go/internal/command/fallback/fallback_test.go b/go/internal/command/fallback/fallback_test.go
index afd752b..91dcb53 100644
--- a/go/internal/command/fallback/fallback_test.go
+++ b/go/internal/command/fallback/fallback_test.go
@@ -6,6 +6,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
 )
 
 type fakeExec struct {
@@ -18,8 +20,21 @@ type fakeExec struct {
 	Env      []string
 }
 
+type FakeCommandArgs struct {
+	executable commandargs.Executable
+	arguments  []string
+}
+
+func (f *FakeCommandArgs) Executable() commandargs.Executable {
+	return f.executable
+}
+
+func (f *FakeCommandArgs) Arguments() []string {
+	return f.arguments
+}
+
 var (
-	fakeArgs = []string{"./test", "foo", "bar"}
+	fakeArgs = &FakeCommandArgs{executable: commandargs.GitlabShell, arguments: []string{"foo", "bar"}}
 )
 
 func (f *fakeExec) Exec(filename string, args []string, env []string) error {
diff --git a/spec/gitlab_shell_authorized_keys_check_spec.rb b/spec/gitlab_shell_authorized_keys_check_spec.rb
index 7050604..eb0cbca 100644
--- a/spec/gitlab_shell_authorized_keys_check_spec.rb
+++ b/spec/gitlab_shell_authorized_keys_check_spec.rb
@@ -21,54 +21,84 @@ describe 'bin/gitlab-shell-authorized-keys-check' do
     end
   end
 
-  before(:all) do
-    write_config(
-      "gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}",
-    )
-  end
-
   let(:authorized_keys_check_path) { File.join(tmp_root_path, 'bin', 'gitlab-shell-authorized-keys-check') }
 
-  it 'succeeds when a valid key is given' do
-    output, status = run!
+  shared_examples 'authorized keys check' do
+    it 'succeeds when a valid key is given' do
+      output, status = run!
 
-    expect(output).to eq("command=\"#{gitlab_shell_path} key-1\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty known-rsa-key\n")
-    expect(status).to be_success
-  end
+      expect(output).to eq("command=\"#{gitlab_shell_path} key-1\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty known-rsa-key\n")
+      expect(status).to be_success
+    end
 
-  it 'returns nothing when an unknown key is given' do
-    output, status = run!(key: 'unknown-key')
+    it 'returns nothing when an unknown key is given' do
+      output, status = run!(key: 'unknown-key')
 
-    expect(output).to eq("# No key was found for unknown-key\n")
-    expect(status).to be_success
-  end
+      expect(output).to eq("# No key was found for unknown-key\n")
+      expect(status).to be_success
+    end
+
+    it' fails when not enough arguments are given' do
+      output, status = run!(key: nil)
 
-  it' fails when not enough arguments are given' do
-    output, status = run!(key: nil)
+      expect(output).to eq('')
+      expect(status).not_to be_success
+    end
 
-    expect(output).to eq('')
-    expect(status).not_to be_success
+    it' fails when too many arguments are given' do
+      output, status = run!(key: ['a', 'b'])
+
+      expect(output).to eq('')
+      expect(status).not_to be_success
+    end
+
+    it 'skips when run as the wrong user' do
+      output, status = run!(expected_user: 'unknown-user')
+
+      expect(output).to eq('')
+      expect(status).to be_success
+    end
+
+    it 'skips when the wrong users connects' do
+      output, status = run!(actual_user: 'unknown-user')
+
+      expect(output).to eq('')
+      expect(status).to be_success
+    end
   end
 
-  it' fails when too many arguments are given' do
-    output, status = run!(key: ['a', 'b'])
+  describe 'without go features' do
+    before(:all) do
+      write_config(
+        "gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}",
+      )
+    end
 
-    expect(output).to eq('')
-    expect(status).not_to be_success
+    it_behaves_like 'authorized keys check'
   end
 
-  it 'skips when run as the wrong user' do
-    output, status = run!(expected_user: 'unknown-user')
+  describe 'without go features (via go)', :go do
+    before(:all) do
+      write_config(
+        "gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}",
+      )
+    end
 
-    expect(output).to eq('')
-    expect(status).to be_success
+    it_behaves_like 'authorized keys check'
   end
 
-  it 'skips when the wrong users connects' do
-    output, status = run!(actual_user: 'unknown-user')
+  pending 'with the go authorized-keys-check feature', :go do
+    before(:all) do
+      write_config(
+        'gitlab_url' => "http+unix://#{CGI.escape(tmp_socket_path)}",
+        'migration' => {
+          'enabled' => true,
+          'features' => ['authorized-keys-check']
+        }
+      )
+    end
 
-    expect(output).to eq('')
-    expect(status).to be_success
+    it_behaves_like 'authorized keys check'
   end
 
   def run!(expected_user: 'git', actual_user: 'git', key: 'known-rsa-key')
diff --git a/spec/gitlab_shell_authorized_principals_check_spec.rb b/spec/gitlab_shell_authorized_principals_check_spec.rb
new file mode 100644
index 0000000..8b946bc
--- /dev/null
+++ b/spec/gitlab_shell_authorized_principals_check_spec.rb
@@ -0,0 +1,82 @@
+require_relative 'spec_helper'
+
+describe 'bin/gitlab-shell-authorized-principals-check' do
+  include_context 'gitlab shell'
+
+  def mock_server(server)
+    # Do nothing as we're not connecting to a server in this check.
+  end
+
+  let(:authorized_principals_check_path) { File.join(tmp_root_path, 'bin', 'gitlab-shell-authorized-principals-check') }
+
+  shared_examples 'authorized principals check' do
+    it 'succeeds when a valid principal is given' do
+      output, status = run!
+
+      expect(output).to eq("command=\"#{gitlab_shell_path} username-key\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty principal\n")
+      expect(status).to be_success
+    end
+
+    it 'fails when not enough arguments are given' do
+      output, status = run!(key_id: nil, principals: [])
+
+      expect(output).to eq('')
+      expect(status).not_to be_success
+    end
+
+    it 'fails when key_id is blank' do
+      output, status = run!(key_id: '')
+
+      expect(output).to eq('')
+      expect(status).not_to be_success
+    end
+
+    it 'fails when principals include an empty item' do
+      output, status = run!(principals: ['principal', ''])
+
+      expect(output).to eq('')
+      expect(status).not_to be_success
+    end
+  end
+
+  describe 'without go features' do
+    before(:all) do
+      write_config({})
+    end
+
+    it_behaves_like 'authorized principals check'
+  end
+
+  describe 'without go features (via go)', :go do
+    before(:all) do
+      write_config({})
+    end
+
+    it_behaves_like 'authorized principals check'
+  end
+
+  pending 'with the go authorized-principals-check feature', :go do
+    before(:all) do
+      write_config(
+        'migration' => {
+          'enabled' => true,
+          'features' => ['authorized-principals-check']
+        }
+      )
+    end
+
+    it_behaves_like 'authorized principals check'
+  end
+
+  def run!(key_id: 'key', principals: ['principal'])
+    cmd = [
+      authorized_principals_check_path,
+      key_id,
+      principals,
+    ].flatten.compact
+
+    output = IO.popen(cmd, &:read)
+
+    [output, $?]
+  end
+end