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

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

+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

+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

+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

+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

+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

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

go/internal/config/config.go

@@ -2,8 +2,10 @@ package config
 
 import (
 	"io/ioutil"
+	"net/url"
 	"os"
 	"path"
+	"strings"
 
 	yaml "gopkg.in/yaml.v2"
 )
@@ -23,6 +25,7 @@ type Config struct {
 	LogFile   string          `yaml:"log_file"`
 	LogFormat string          `yaml:"log_format"`
 	Migration MigrationConfig `yaml:"migration"`
+	GitlabUrl string          `yaml:"gitlab_url"`
 }
 
 func New() (*Config, error) {
@@ -38,6 +41,24 @@ func NewFromDir(dir string) (*Config, error) {
 	return newFromFile(path.Join(dir, configFile))
 }
 
+func (c *Config) FeatureEnabled(featureName string) bool {
+	if !c.Migration.Enabled {
+		return false
+	}
+
+	if !strings.HasPrefix(c.GitlabUrl, "http+unix://") {
+		return false
+	}
+
+	for _, enabledFeature := range c.Migration.Features {
+		if enabledFeature == featureName {
+			return true
+		}
+	}
+
+	return false
+}
+
 func newFromFile(filename string) (*Config, error) {
 	cfg := &Config{RootDir: path.Dir(filename)}
 
@@ -71,5 +92,14 @@ func parseConfig(configBytes []byte, cfg *Config) error {
 		cfg.LogFormat = "text"
 	}
 
+	if cfg.GitlabUrl != "" {
+		unescapedUrl, err := url.PathUnescape(cfg.GitlabUrl)
+		if err != nil {
+			return err
+		}
+
+		cfg.GitlabUrl = unescapedUrl
+	}
+
 	return nil
 }

go/internal/config/config_test.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestParseConfig(t *testing.T) {
@@ -12,6 +14,7 @@ func TestParseConfig(t *testing.T) {
 		yaml      string
 		path      string
 		format    string
+		gitlabUrl string
 		migration MigrationConfig
 	}{
 		{path: "/foo/bar/gitlab-shell.log", format: "text"},
@@ -24,6 +27,12 @@ func TestParseConfig(t *testing.T) {
 			format:    "text",
 			migration: MigrationConfig{Enabled: true, Features: []string{"foo", "bar"}},
 		},
+		{
+			yaml:      "gitlab_url: http+unix://%2Fpath%2Fto%2Fgitlab%2Fgitlab.socket",
+			path:      "/foo/bar/gitlab-shell.log",
+			format:    "text",
+			gitlabUrl: "http+unix:///path/to/gitlab/gitlab.socket",
+		},
 	}
 
 	for _, tc := range testCases {
@@ -48,6 +57,60 @@ func TestParseConfig(t *testing.T) {
 			if cfg.LogFormat != tc.format {
 				t.Fatalf("expected %q, got %q", tc.format, cfg.LogFormat)
 			}
+
+			assert.Equal(t, tc.gitlabUrl, cfg.GitlabUrl)
+		})
+	}
+}
+
+func TestFeatureEnabled(t *testing.T) {
+	testCases := []struct {
+		desc          string
+		config        *Config
+		feature       string
+		expectEnabled bool
+	}{
+		{
+			desc: "When the protocol is supported and the feature enabled",
+			config: &Config{
+				GitlabUrl: "http+unix://gitlab.socket",
+				Migration: MigrationConfig{Enabled: true, Features: []string{"discover"}},
+			},
+			feature:       "discover",
+			expectEnabled: true,
+		},
+		{
+			desc: "When the protocol is supported and the feature is not enabled",
+			config: &Config{
+				GitlabUrl: "http+unix://gitlab.socket",
+				Migration: MigrationConfig{Enabled: true, Features: []string{}},
+			},
+			feature:       "discover",
+			expectEnabled: false,
+		},
+		{
+			desc: "When the protocol is supported and all features are disabled",
+			config: &Config{
+				GitlabUrl: "http+unix://gitlab.socket",
+				Migration: MigrationConfig{Enabled: false, Features: []string{"discover"}},
+			},
+			feature:       "discover",
+			expectEnabled: false,
+		},
+		{
+			desc: "When the protocol is not supported",
+			config: &Config{
+				GitlabUrl: "https://localhost:3000",
+				Migration: MigrationConfig{Enabled: true, Features: []string{"discover"}},
+			},
+			feature:       "discover",
+			expectEnabled: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			assert.Equal(t, tc.expectEnabled, tc.config.FeatureEnabled(string(tc.feature)))
 		})
 	}
 }

go/internal/testhelper/testhelper.go

+package testhelper
+
+import "os"
+
+func TempEnv(env map[string]string) func() {
+	var original = make(map[string]string)
+	for key, value := range env {
+		original[key] = os.Getenv(key)
+		os.Setenv(key, value)
+	}
+
+	return func() {
+		for key, originalValue := range original {
+			os.Setenv(key, originalValue)
+		}
+	}
+}

spec/gitlab_shell_gitlab_shell_spec.rb

@@ -43,11 +43,7 @@ describe 'bin/gitlab-shell' do
 
     sleep(0.1) while @webrick_thread.alive? && @server.status != :Running
     raise "Couldn't start stub GitlabNet server" unless @server.status == :Running
-
-    File.open(config_path, 'w') do |f|
-      f.write("---\ngitlab_url: http+unix://#{CGI.escape(tmp_socket_path)}\n")
-    end
-
+    system(original_root_path, 'bin/compile')
     copy_dirs = ['bin', 'lib']
     FileUtils.rm_rf(copy_dirs.map { |d| File.join(tmp_root_path, d) })
     FileUtils.cp_r(copy_dirs, tmp_root_path)
@@ -61,72 +57,100 @@ describe 'bin/gitlab-shell' do
 
   let(:gitlab_shell_path) { File.join(tmp_root_path, 'bin', 'gitlab-shell') }
 
-  # Basic valid input
-  it 'succeeds and prints username when a valid known key id is given' do
-    output, status = run!(["key-100"])
+  shared_examples 'results with keys' do
+    # Basic valid input
+    it 'succeeds and prints username when a valid known key id is given' do
+      output, status = run!(["key-100"])
 
-    expect(output).to eq("Welcome to GitLab, @someuser!\n")
-    expect(status).to be_success
-  end
+      expect(output).to eq("Welcome to GitLab, @someuser!\n")
+      expect(status).to be_success
+    end
 
-  it 'succeeds and prints username when a valid known username is given' do
-    output, status = run!(["username-someuser"])
+    it 'succeeds and prints username when a valid known username is given' do
+      output, status = run!(["username-someuser"])
 
-    expect(output).to eq("Welcome to GitLab, @someuser!\n")
-    expect(status).to be_success
-  end
+      expect(output).to eq("Welcome to GitLab, @someuser!\n")
+      expect(status).to be_success
+    end
 
-  # Valid but unknown input
-  it 'succeeds and prints Anonymous when a valid unknown key id is given' do
-    output, status = run!(["key-12345"])
+    # Valid but unknown input
+    it 'succeeds and prints Anonymous when a valid unknown key id is given' do
+      output, status = run!(["key-12345"])
 
-    expect(output).to eq("Welcome to GitLab, Anonymous!\n")
-    expect(status).to be_success
-  end
+      expect(output).to eq("Welcome to GitLab, Anonymous!\n")
+      expect(status).to be_success
+    end
 
-  it 'succeeds and prints Anonymous when a valid unknown username is given' do
-    output, status = run!(["username-unknown"])
+    it 'succeeds and prints Anonymous when a valid unknown username is given' do
+      output, status = run!(["username-unknown"])
 
-    expect(output).to eq("Welcome to GitLab, Anonymous!\n")
-    expect(status).to be_success
-  end
+      expect(output).to eq("Welcome to GitLab, Anonymous!\n")
+      expect(status).to be_success
+    end
 
-  # Invalid input. TODO: capture stderr & compare
-  it 'gets an ArgumentError on invalid input (empty)' do
-    output, status = run!([])
+    # Invalid input. TODO: capture stderr & compare
+    it 'gets an ArgumentError on invalid input (empty)' do
+      output, status = run!([])
 
-    expect(output).to eq("")
-    expect(status).not_to be_success
-  end
+      expect(output).to eq("")
+      expect(status).not_to be_success
+    end
 
-  it 'gets an ArgumentError on invalid input (unknown)' do
-    output, status = run!(["whatever"])
+    it 'gets an ArgumentError on invalid input (unknown)' do
+      output, status = run!(["whatever"])
 
-    expect(output).to eq("")
-    expect(status).not_to be_success
-  end
+      expect(output).to eq("")
+      expect(status).not_to be_success
+    end
 
-  it 'gets an ArgumentError on invalid input (multiple unknown)' do
-    output, status = run!(["this", "is", "all", "invalid"])
+    it 'gets an ArgumentError on invalid input (multiple unknown)' do
+      output, status = run!(["this", "is", "all", "invalid"])
 
-    expect(output).to eq("")
-    expect(status).not_to be_success
+      expect(output).to eq("")
+      expect(status).not_to be_success
+    end
+
+    # Not so basic valid input
+    # (https://gitlab.com/gitlab-org/gitlab-shell/issues/145)
+    it 'succeeds and prints username when a valid known key id is given in the middle of other input' do
+      output, status = run!(["-c/usr/share/webapps/gitlab-shell/bin/gitlab-shell", "key-100", "2foo"])
+
+      expect(output).to eq("Welcome to GitLab, @someuser!\n")
+      expect(status).to be_success
+    end
+
+    it 'succeeds and prints username when a valid known username is given in the middle of other input' do
+      output, status = run!(["-c/usr/share/webapps/gitlab-shell/bin/gitlab-shell", "username-someuser" ,"foo"])
+
+      expect(output).to eq("Welcome to GitLab, @someuser!\n")
+      expect(status).to be_success
+    end
   end
 
-  # Not so basic valid input
-  # (https://gitlab.com/gitlab-org/gitlab-shell/issues/145)
-  it 'succeeds and prints username when a valid known key id is given in the middle of other input' do
-    output, status = run!(["-c/usr/share/webapps/gitlab-shell/bin/gitlab-shell", "key-100", "2foo"])
+  describe 'without go features' do
+    before(:context) do
+      write_config("gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}")
+    end
 
-    expect(output).to eq("Welcome to GitLab, @someuser!\n")
-    expect(status).to be_success
+    it_behaves_like 'results with keys'
   end
 
-  it 'succeeds and prints username when a valid known username is given in the middle of other input' do
-    output, status = run!(["-c/usr/share/webapps/gitlab-shell/bin/gitlab-shell", "username-someuser" ,"foo"])
+  describe 'with the go discover feature', :go do
+    before(:context) do
+      write_config(
+        "gitlab_url" => "http+unix://#{CGI.escape(tmp_socket_path)}",
+        "migration" => { "enabled" => true,
+                        "features" => ["discover"] }
+      )
+    end
+
+
 
-    expect(output).to eq("Welcome to GitLab, @someuser!\n")
-    expect(status).to be_success
+    it_behaves_like 'results with keys' do
+      before do
+        pending
+      end
+    end
   end
 
   def run!(args)
@@ -139,4 +163,10 @@ describe 'bin/gitlab-shell' do
 
     [output, $?]
   end
+
+  def write_config(config)
+    File.open(config_path, 'w') do |f|
+      f.write(config.to_yaml)
+    end
+  end
 end