Add Executable struct

This struct is responsible for determining the name and
root dir of the executable.

The `RootDir` property will be used to find the config.

The `Name` property will be used to determine what `Command`
and `CommandArgs` to be built.

go/cmd/gitlab-shell/main.go

@@ -3,30 +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/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
-}
-
 func main() {
 	readWriter := &readwriter.ReadWriter{
 		Out:    os.Stdout,
@@ -34,19 +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)
 	}
 
-	config, err := config.NewFromDir(rootDir)
+	config, err := config.NewFromDir(executable.RootDir)
 	if err != nil {
 		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,28 +11,29 @@ 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 cmd := buildCommand(args, config, readWriter); cmd != nil {
+	if cmd := buildCommand(e, args, config, readWriter); cmd != nil {
 		return cmd, nil
 	}
 
-	return &fallback.Command{RootDir: config.RootDir, Args: args}, nil
+	return &fallback.Command{Executable: e, RootDir: config.RootDir, Args: args}, nil
 }
 
-func buildCommand(args commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command {
-	switch args.Executable() {
-	case commandargs.GitlabShell:
+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)
 	}
 

go/internal/command/command_test.go

@@ -5,7 +5,6 @@ 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"
@@ -14,19 +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"}},
@@ -35,11 +37,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
-			arguments:    []string{string(commandargs.GitlabShell)},
+			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},
@@ -48,11 +51,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "",
 			},
-			arguments:    []string{string(commandargs.GitlabShell)},
+			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"}},
@@ -61,11 +65,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
 			},
-			arguments:    []string{string(commandargs.GitlabShell)},
+			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"}},
@@ -74,11 +79,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-lfs-authenticate",
 			},
-			arguments:    []string{string(commandargs.GitlabShell)},
+			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"}},
@@ -87,11 +93,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-receive-pack",
 			},
-			arguments:    []string{string(commandargs.GitlabShell)},
+			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"}},
@@ -100,11 +107,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-upload-pack",
 			},
-			arguments:    []string{string(commandargs.GitlabShell)},
+			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"}},
@@ -113,11 +121,12 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-upload-archive",
 			},
-			arguments:    []string{string(commandargs.GitlabShell)},
+			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"}},
@@ -126,13 +135,14 @@ func TestNew(t *testing.T) {
 				"SSH_CONNECTION":       "1",
 				"SSH_ORIGINAL_COMMAND": "git-unimplemented-feature",
 			},
-			arguments:    []string{string(commandargs.GitlabShell)},
+			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{"unknown"},
+			arguments:    []string{},
 			expectedType: &fallback.Command{},
 		},
 	}
@@ -142,7 +152,7 @@ func TestNew(t *testing.T) {
 			restoreEnv := testhelper.TempEnv(tc.environment)
 			defer restoreEnv()
 
-			command, err := New(tc.arguments, tc.config, nil)
+			command, err := New(tc.executable, tc.arguments, tc.config, nil)
 
 			require.NoError(t, err)
 			require.IsType(t, tc.expectedType, command)
@@ -152,7 +162,7 @@ func TestNew(t *testing.T) {
 
 func TestFailingNew(t *testing.T) {
 	t.Run("It returns an error parsing arguments failed", func(t *testing.T) {
-		_, err := New([]string{}, &config.Config{}, nil)
+		_, err := New(&executable.Executable{Name: executable.GitlabShell}, []string{}, &config.Config{}, nil)
 
 		require.Error(t, err)
 	})

go/internal/command/commandargs/base_args.go

-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
-}

go/internal/command/commandargs/command_args.go

 package commandargs
 
-type CommandType string
-type Executable string
-
-const (
-	GitlabShell Executable = "gitlab-shell"
+import (
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 )
 
+type CommandType string
+
 type CommandArgs interface {
 	Parse() error
-	Executable() Executable
-	Arguments() []string
+	GetArguments() []string
 }
 
-func Parse(arguments []string) (CommandArgs, error) {
-	var args CommandArgs = &BaseArgs{arguments: arguments}
+func Parse(e *executable.Executable, arguments []string) (CommandArgs, error) {
+	var args CommandArgs = &GenericArgs{Arguments: arguments}
 
-	switch args.Executable() {
-	case GitlabShell:
-		args = &Shell{BaseArgs: args.(*BaseArgs)}
+	switch e.Name {
+	case executable.GitlabShell:
+		args = &Shell{Arguments: arguments}
 	}
 
 	if err := args.Parse(); err != nil {

go/internal/command/commandargs/command_args_test.go

@@ -3,6 +3,7 @@ package commandargs
 import (
 	"testing"
 
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/executable"
 	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
 
 	"github.com/stretchr/testify/require"
@@ -11,6 +12,7 @@ import (
 func TestParseSuccess(t *testing.T) {
 	testCases := []struct {
 		desc         string
+		executable   *executable.Executable
 		environment  map[string]string
 		arguments    []string
 		expectedArgs CommandArgs
@@ -18,98 +20,110 @@ func TestParseSuccess(t *testing.T) {
 		// 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": "",
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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{string(GitlabShell), "hello", "key-123"},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell), "hello", "key-123"}}, SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
+			arguments:    []string{"hello", "key-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{string(GitlabShell), "hello", "username-jane-doe"},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell), "hello", "username-jane-doe"}}, SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
+			arguments:    []string{"hello", "username-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",
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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",
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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'",
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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"`,
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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`,
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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"`,
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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'",
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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",
 			},
-			arguments:    []string{string(GitlabShell)},
-			expectedArgs: &Shell{BaseArgs: &BaseArgs{arguments: []string{string(GitlabShell)}}, 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",
-			arguments:    []string{"unknown"},
-			expectedArgs: &BaseArgs{arguments: []string{"unknown"}},
+			executable:   &executable.Executable{Name: "unknown"},
+			arguments:    []string{},
+			expectedArgs: &GenericArgs{Arguments: []string{}},
 		},
 	}
 
@@ -118,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)
@@ -129,28 +143,26 @@ func TestParseSuccess(t *testing.T) {
 func TestParseFailure(t *testing.T) {
 	testCases := []struct {
 		desc          string
+		executable    *executable.Executable
 		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",
+			executable:    &executable.Executable{Name: executable.GitlabShell},
+			arguments:     []string{},
+			expectedError: "Only SSH allowed",
 		},
 		{
-			desc: "It fails if SSH command is invalid",
+			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{string(GitlabShell)},
-			expectedError: "Only ssh allowed",
-		},
-		{
-			desc:          "It fails if arguments is empty",
 			arguments:     []string{},
-			expectedError: "arguments should include the executable",
+			expectedError: "Invalid SSH allowed",
 		},
 	}
 
@@ -159,7 +171,7 @@ func TestParseFailure(t *testing.T) {
 			restoreEnv := testhelper.TempEnv(tc.environment)
 			defer restoreEnv()
 
-			_, err := Parse(tc.arguments)
+			_, err := Parse(tc.executable, tc.arguments)
 
 			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

@@ -23,7 +23,7 @@ var (
 )
 
 type Shell struct {
-	*BaseArgs
+	Arguments      []string
 	GitlabUsername string
 	GitlabKeyId    string
 	SshArgs        []string
@@ -31,23 +31,44 @@ type Shell struct {
 }
 
 func (s *Shell) Parse() error {
-	if sshConnection := os.Getenv("SSH_CONNECTION"); sshConnection == "" {
-		return errors.New("Only ssh allowed")
+	if err := s.validate(); err != nil {
+		return err
 	}
 
 	s.parseWho()
+	s.defineCommandType()
+
+	return nil
+}
+
+func (s *Shell) GetArguments() []string {
+	return s.Arguments
+}
 
-	if err := s.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND")); err != nil {
-		return errors.New("Invalid ssh command")
+func (s *Shell) validate() error {
+	if !s.isSshConnection() {
+		return errors.New("Only SSH allowed")
 	}
 
-	s.defineCommandType()
+	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 {
+	for _, argument := range s.Arguments {
 		if keyId := tryParseKeyId(argument); keyId != "" {
 			s.GitlabKeyId = keyId
 			break

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    commandargs.CommandArgs
+	Executable *executable.Executable
+	RootDir    string
+	Args       commandargs.CommandArgs
 }
 
 var (
 	// execFunc is overridden in tests
-	execFunc = syscall.Exec
+	execFunc  = syscall.Exec
+	whitelist = []string{
+		executable.GitlabShell,
+		executable.AuthorizedKeysCheck,
+		executable.AuthorizedPrincipalsCheck,
+	}
 )
 
 func (c *Command) Execute() error {
-	rubyCmd := filepath.Join(c.RootDir, "bin", c.fallbackProgram())
+	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.Arguments()...)
+	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 {
-	return fmt.Sprintf("%s-ruby", c.Args.Executable())
+	fileName := fmt.Sprintf("%s-ruby", c.Executable.Name)
+
+	return filepath.Join(c.RootDir, "bin", fileName)
 }

go/internal/command/fallback/fallback_test.go

@@ -8,6 +8,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/executable"
 )
 
 type fakeExec struct {
@@ -20,26 +21,8 @@ type fakeExec struct {
 	Env      []string
 }
 
-type FakeCommandArgs struct {
-	executable commandargs.Executable
-	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
-}
-
-func (f *FakeCommandArgs) Arguments() []string {
-	return f.arguments
-}
-
 var (
-	fakeArgs = &FakeCommandArgs{executable: commandargs.GitlabShell, arguments: []string{"foo", "bar"}}
+	fakeArgs = &commandargs.GenericArgs{Arguments: []string{"foo", "bar"}}
 )
 
 func (f *fakeExec) Exec(filename string, args []string, env []string) error {
@@ -62,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{}
@@ -76,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")}
@@ -89,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/executable/executable.go

+package executable
+
+import (
+	"os"
+	"path/filepath"
+)
+
+const (
+	GitlabShell               = "gitlab-shell"
+	AuthorizedKeysCheck       = "gitlab-shell-authorized-keys-check"
+	AuthorizedPrincipalsCheck = "gitlab-shell-authorized-principals-check"
+)
+
+type Executable struct {
+	Name    string
+	RootDir string
+}
+
+var (
+	// osExecutable is overridden in tests
+	osExecutable = os.Executable
+)
+
+func New() (*Executable, error) {
+	path, err := osExecutable()
+	if err != nil {
+		return nil, err
+	}
+
+	rootDir, err := findRootDir(path)
+	if err != nil {
+		return nil, err
+	}
+
+	executable := &Executable{
+		Name:    filepath.Base(path),
+		RootDir: rootDir,
+	}
+
+	return executable, nil
+}
+
+func findRootDir(path string) (string, error) {
+	// Start: /opt/.../gitlab-shell/bin/gitlab-shell
+	// Ends:  /opt/.../gitlab-shell
+	rootDir := filepath.Dir(filepath.Dir(path))
+	pathFromEnv := os.Getenv("GITLAB_SHELL_DIR")
+
+	if pathFromEnv != "" {
+		if _, err := os.Stat(pathFromEnv); os.IsNotExist(err) {
+			return "", err
+		}
+
+		rootDir = pathFromEnv
+	}
+
+	return rootDir, nil
+}

go/internal/executable/executable_test.go

+package executable
+
+import (
+	"errors"
+	"testing"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
+
+	"github.com/stretchr/testify/require"
+)
+
+type fakeOs struct {
+	OldExecutable func() (string, error)
+	Path          string
+	Error         error
+}
+
+func (f *fakeOs) Executable() (string, error) {
+	return f.Path, f.Error
+}
+
+func (f *fakeOs) Setup() {
+	f.OldExecutable = osExecutable
+	osExecutable = f.Executable
+}
+
+func (f *fakeOs) Cleanup() {
+	osExecutable = f.OldExecutable
+}
+
+func TestNewSuccess(t *testing.T) {
+	testCases := []struct {
+		desc            string
+		fakeOs          *fakeOs
+		environment     map[string]string
+		expectedRootDir string
+	}{
+		{
+			desc:            "GITLAB_SHELL_DIR env var is not defined",
+			fakeOs:          &fakeOs{Path: "/tmp/bin/gitlab-shell"},
+			expectedRootDir: "/tmp",
+		},
+		{
+			desc:   "GITLAB_SHELL_DIR env var is defined",
+			fakeOs: &fakeOs{Path: "/opt/bin/gitlab-shell"},
+			environment: map[string]string{
+				"GITLAB_SHELL_DIR": "/tmp",
+			},
+			expectedRootDir: "/tmp",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			restoreEnv := testhelper.TempEnv(tc.environment)
+			defer restoreEnv()
+
+			fake := tc.fakeOs
+			fake.Setup()
+			defer fake.Cleanup()
+
+			result, err := New()
+
+			require.NoError(t, err)
+			require.Equal(t, result.Name, "gitlab-shell")
+			require.Equal(t, result.RootDir, tc.expectedRootDir)
+		})
+	}
+}
+
+func TestNewFailure(t *testing.T) {
+	testCases := []struct {
+		desc        string
+		fakeOs      *fakeOs
+		environment map[string]string
+	}{
+		{
+			desc:   "failed to determine executable",
+			fakeOs: &fakeOs{Path: "", Error: errors.New("error")},
+		},
+		{
+			desc:   "GITLAB_SHELL_DIR doesn't exist",
+			fakeOs: &fakeOs{Path: "/tmp/bin/gitlab-shell"},
+			environment: map[string]string{
+				"GITLAB_SHELL_DIR": "/tmp/non/existing/directory",
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			restoreEnv := testhelper.TempEnv(tc.environment)
+			defer restoreEnv()
+
+			fake := tc.fakeOs
+			fake.Setup()
+			defer fake.Cleanup()
+
+			_, err := New()
+
+			require.Error(t, err)
+		})
+	}
+}

spec/gitlab_shell_gitlab_shell_spec.rb

@@ -123,10 +123,10 @@ describe 'bin/gitlab-shell' do
 
     it_behaves_like 'results with keys'
 
-    it 'outputs "Only ssh allowed"' do
+    it 'outputs "Only SSH allowed"' do
       _, stderr, status = run!(["-c/usr/share/webapps/gitlab-shell/bin/gitlab-shell", "username-someuser"], env: {})
 
-      expect(stderr).to eq("Only ssh allowed\n")
+      expect(stderr).to eq("Only SSH allowed\n")
       expect(status).not_to be_success
     end