1 files changed, 36 insertions, 9 deletions
diff --git a/go/internal/command/fallback/fallback.go b/go/internal/command/fallback/fallback.go
index f525a57..781eda1 100644
--- a/go/internal/command/fallback/fallback.go
+++ b/go/internal/command/fallback/fallback.go
@@ -1,30 +1,57 @@
 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)
+}