Allow enabling gitlab-shell "discover"-feature

This adds the possibility to enable features for GitLab shell. The first feature being recognized is "Discover": It's the command that is executed when running `ssh git@gitlab.example.com` and is called without a command. The gitlab key id or username is already parsed from the command line arguments. Currently we only support communicating with GitLab-rails using unix sockets. So features will not be enabled if the GitLab-url is using a different protocol. The url for this read from the config yaml. Pending ruby-specs have been added for the gitlab-shell command. Refactor to have separate command packages

Author: Reprazent

go/cmd/gitlab-shell/main.go

@@ -4,8 +4,9 @@ import (
 "fmt"
 "os"
 "path/filepath"
-"syscall"
 
+"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/config"
 )
 
@@ -19,20 +20,12 @@ func init() {
 rootDir = filepath.Dir(binDir)
 }
 
-func migrate(*config.Config) (int, bool) {
-// TODO: Dispatch appropriate requests to Go handlers and return
-// <exitstatus, true> depending on how they fare
-return 0, false
-}
-
 // rubyExec will never return. It either replaces the current process with a
 // Ruby interpreter, or outputs an error and kills the process.
 func execRuby() {
-rubyCmd := filepath.Join(binDir, "gitlab-shell-ruby")
-
-execErr := syscall.Exec(rubyCmd, os.Args, os.Environ())
-if execErr != nil {
-fmt.Fprintf(os.Stderr, "Failed to exec(%q): %v\n", rubyCmd, execErr)
+cmd := &fallback.Command{}
+if err := cmd.Execute(); err != nil {
+fmt.Fprintf(os.Stderr, "Failed to exec: %v\n", err)
 os.Exit(1)
 }
 }
@@ -46,11 +39,19 @@ func main() {
 execRuby()
 }
 
-// Try to handle the command with the Go implementation
-if exitCode, done := migrate(config); done {
-os.Exit(exitCode)
+cmd, err := command.New(os.Args, config)
+if err != nil {
+// Failed to build the command, fall back to ruby.
+// For now this could happen if `SSH_CONNECTION` is not set on
+// the environment
+fmt.Fprintf(os.Stderr, "Failed to build command: %v\n", err)
+execRuby()
 }
 
-// Since a migration has not handled the command, fall back to Ruby to do so
-execRuby()
+// The command will write to STDOUT on execution or replace the current
+// process in case of the `fallback.Command`
+if err = cmd.Execute(); err != nil {
+fmt.Fprintf(os.Stderr, "%s\n", err)
+os.Exit(1)
+}
 }

go/internal/command/command.go

@@ -0,0 +1,35 @@
+package command
+
+import (
+"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/config"
+)
+
+type Command interface {
+Execute() error
+}
+
+func New(arguments []string, config *config.Config) (Command, error) {
+args, err := commandargs.Parse(arguments)
+
+if err != nil {
+return nil, err
+}
+
+if config.FeatureEnabled(string(args.CommandType)) {
+return buildCommand(args, config), nil
+}
+
+return &fallback.Command{}, nil
+}
+
+func buildCommand(args *commandargs.CommandArgs, config *config.Config) Command {
+switch args.CommandType {
+case commandargs.Discover:
+return &discover.Command{Config: config, Args: args}
+}
+
+return nil
+}

go/internal/command/command_test.go

@@ -0,0 +1,71 @@
+package command
+
+import (
+"testing"
+
+"github.com/stretchr/testify/assert"
+"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/config"
+"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
+)
+
+func TestNew(t *testing.T) {
+testCases := []struct {
+desc string
+arguments []string
+config *config.Config
+environment map[string]string
+expectedType interface{}
+}{
+{
+desc: "it returns a Discover command if the feature is enabled",
+arguments: []string{},
+config: &config.Config{
+GitlabUrl: "http+unix://gitlab.socket",
+Migration: config.MigrationConfig{Enabled: true, Features: []string{"discover"}},
+},
+environment: map[string]string{
+"SSH_CONNECTION": "1",
+"SSH_ORIGINAL_COMMAND": "",
+},
+expectedType: &discover.Command{},
+},
+{
+desc: "it returns a Fallback command no feature is enabled",
+arguments: []string{},
+config: &config.Config{
+GitlabUrl: "http+unix://gitlab.socket",
+Migration: config.MigrationConfig{Enabled: false},
+},
+environment: map[string]string{
+"SSH_CONNECTION": "1",
+"SSH_ORIGINAL_COMMAND": "",
+},
+expectedType: &fallback.Command{},
+},
+}
+
+for _, tc := range testCases {
+t.Run(tc.desc, func(t *testing.T) {
+restoreEnv := testhelper.TempEnv(tc.environment)
+defer restoreEnv()
+
+command, err := New(tc.arguments, tc.config)
+
+assert.NoError(t, err)
+assert.IsType(t, tc.expectedType, command)
+})
+}
+}
+
+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{})
+
+assert.Error(t, err, "Only ssh allowed")
+})
+}

go/internal/command/commandargs/command_args.go

@@ -0,0 +1,82 @@
+package commandargs
+
+import (
+"errors"
+"os"
+"regexp"
+)
+
+type CommandType string
+
+const (
+Discover CommandType = "discover"
+)
+
+var (
+whoKeyRegex = regexp.MustCompile(`\bkey-(?P<keyid>\d+)\b`)
+whoUsernameRegex = regexp.MustCompile(`\busername-(?P<username>\S+)\b`)
+)
+
+type CommandArgs struct {
+GitlabUsername string
+GitlabKeyId string
+SshCommand string
+CommandType CommandType
+}
+
+func Parse(arguments []string) (*CommandArgs, error) {
+if sshConnection := os.Getenv("SSH_CONNECTION"); sshConnection == "" {
+return nil, errors.New("Only ssh allowed")
+}
+
+info := &CommandArgs{}
+
+info.parseWho(arguments)
+info.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND"))
+
+return info, nil
+}
+
+func (info *CommandArgs) parseWho(arguments []string) {
+for _, argument := range arguments {
+if keyId := tryParseKeyId(argument); keyId != "" {
+info.GitlabKeyId = keyId
+break
+}
+
+if username := tryParseUsername(argument); username != "" {
+info.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 (c *CommandArgs) parseCommand(commandString string) {
+c.SshCommand = commandString
+
+if commandString == "" {
+c.CommandType = Discover
+}
+}

go/internal/command/commandargs/command_args_test.go

@@ -0,0 +1,73 @@
+package commandargs
+
+import (
+"testing"
+
+"github.com/stretchr/testify/assert"
+"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
+)
+
+func TestParseSuccess(t *testing.T) {
+testCases := []struct {
+desc string
+arguments []string
+environment map[string]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",
+environment: map[string]string{
+"SSH_CONNECTION": "1",
+"SSH_ORIGINAL_COMMAND": "",
+},
+expectedArgs: &CommandArgs{CommandType: Discover},
+},
+{
+desc: "It passes on the original ssh command from the environment",
+environment: map[string]string{
+"SSH_CONNECTION": "1",
+"SSH_ORIGINAL_COMMAND": "hello world",
+},
+expectedArgs: &CommandArgs{SshCommand: "hello world"},
+}, {
+desc: "It finds the key id in any passed arguments",
+environment: map[string]string{
+"SSH_CONNECTION": "1",
+"SSH_ORIGINAL_COMMAND": "",
+},
+arguments: []string{"hello", "key-123"},
+expectedArgs: &CommandArgs{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{CommandType: Discover, GitlabUsername: "jane-doe"},
+},
+}
+
+for _, tc := range testCases {
+t.Run(tc.desc, func(t *testing.T) {
+restoreEnv := testhelper.TempEnv(tc.environment)
+defer restoreEnv()
+
+result, err := Parse(tc.arguments)
+
+assert.NoError(t, err)
+assert.Equal(t, tc.expectedArgs, result)
+})
+}
+}
+
+func TestParseFailure(t *testing.T) {
+t.Run("It fails if SSH connection is not set", func(t *testing.T) {
+_, err := Parse([]string{})
+
+assert.Error(t, err, "Only ssh allowed")
+})
+
+}

go/internal/command/discover/discover.go

@@ -0,0 +1,17 @@
+package discover
+
+import (
+"fmt"
+
+"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
+"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+)
+
+type Command struct {
+Config *config.Config
+Args *commandargs.CommandArgs
+}
+
+func (c *Command) Execute() error {
+return fmt.Errorf("No feature is implemented yet")
+}

go/internal/command/fallback/fallback.go

@@ -0,0 +1,19 @@
+package fallback
+
+import (
+"os"
+"path/filepath"
+"syscall"
+)
+
+type Command struct{}
+
+var (
+binDir = filepath.Dir(os.Args[0])
+)
+
+func (c *Command) Execute() error {
+rubyCmd := filepath.Join(binDir, "gitlab-shell-ruby")
+execErr := syscall.Exec(rubyCmd, os.Args, os.Environ())
+return execErr
+}