Implement AuthorizedKeys command181-authorized-keys-check-go

Build this command when `Executable` name is
`gitlab-shell-authorized-keys-check`. Feature flag is the same
name.

7 files changed, 257 insertions, 2 deletions

diff --git a/go/internal/command/authorizedkeys/authorized_keys.go b/go/internal/command/authorizedkeys/authorized_keys.go
new file mode 100644
index 0000000..d5837b0
--- /dev/null
+++ b/go/internal/command/authorizedkeys/authorized_keys.go
@@ -0,0 +1,61 @@
+package authorizedkeys
+
+import (
+	"fmt"
+	"strconv"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
+	"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/gitlabnet/authorizedkeys"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/keyline"
+)
+
+type Command struct {
+	Config     *config.Config
+	Args       *commandargs.AuthorizedKeys
+	ReadWriter *readwriter.ReadWriter
+}
+
+func (c *Command) Execute() error {
+	// Do and return nothing when the expected and actual user don't match.
+	// This can happen when the user in sshd_config doesn't match the user
+	// trying to login. When nothing is printed, the user will be denied access.
+	if c.Args.ExpectedUser != c.Args.ActualUser {
+		// TODO: Log this event once we have a consistent way to log in Go.
+		// See https://gitlab.com/gitlab-org/gitlab-shell/issues/192 for more info.
+		return nil
+	}
+
+	if err := c.printKeyLine(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (c *Command) printKeyLine() error {
+	response, err := c.getAuthorizedKey()
+	if err != nil {
+		fmt.Fprintln(c.ReadWriter.Out, fmt.Sprintf("# No key was found for %s", c.Args.Key))
+		return nil
+	}
+
+	keyLine, err := keyline.NewPublicKeyLine(strconv.FormatInt(response.Id, 10), response.Key, c.Config.RootDir)
+	if err != nil {
+		return err
+	}
+
+	fmt.Fprintln(c.ReadWriter.Out, keyLine.ToString())
+
+	return nil
+}
+
+func (c *Command) getAuthorizedKey() (*authorizedkeys.Response, error) {
+	client, err := authorizedkeys.NewClient(c.Config)
+	if err != nil {
+		return nil, err
+	}
+
+	return client.GetByKey(c.Args.Key)
+}
diff --git a/go/internal/command/authorizedkeys/authorized_keys_test.go b/go/internal/command/authorizedkeys/authorized_keys_test.go
new file mode 100644
index 0000000..5cde366
--- /dev/null
+++ b/go/internal/command/authorizedkeys/authorized_keys_test.go
@@ -0,0 +1,90 @@
+package authorizedkeys
+
+import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/commandargs"
+	"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/gitlabnet/testserver"
+)
+
+var (
+	requests = []testserver.TestRequestHandler{
+		{
+			Path: "/api/v4/internal/authorized_keys",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				if r.URL.Query().Get("key") == "key" {
+					body := map[string]interface{}{
+						"id":  1,
+						"key": "public-key",
+					}
+					json.NewEncoder(w).Encode(body)
+				} else if r.URL.Query().Get("key") == "broken-message" {
+					body := map[string]string{
+						"message": "Forbidden!",
+					}
+					w.WriteHeader(http.StatusForbidden)
+					json.NewEncoder(w).Encode(body)
+				} else if r.URL.Query().Get("key") == "broken" {
+					w.WriteHeader(http.StatusInternalServerError)
+				} else {
+					w.WriteHeader(http.StatusNotFound)
+				}
+			},
+		},
+	}
+)
+
+func TestExecute(t *testing.T) {
+	url, cleanup := testserver.StartSocketHttpServer(t, requests)
+	defer cleanup()
+
+	testCases := []struct {
+		desc           string
+		arguments      *commandargs.AuthorizedKeys
+		expectedOutput string
+	}{
+		{
+			desc:           "With matching username and key",
+			arguments:      &commandargs.AuthorizedKeys{ExpectedUser: "user", ActualUser: "user", Key: "key"},
+			expectedOutput: "command=\"/tmp/bin/gitlab-shell key-1\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty public-key\n",
+		},
+		{
+			desc:           "When key doesn't match any existing key",
+			arguments:      &commandargs.AuthorizedKeys{ExpectedUser: "user", ActualUser: "user", Key: "not-found"},
+			expectedOutput: "# No key was found for not-found\n",
+		},
+		{
+			desc:           "When the API returns an error",
+			arguments:      &commandargs.AuthorizedKeys{ExpectedUser: "user", ActualUser: "user", Key: "broken-message"},
+			expectedOutput: "# No key was found for broken-message\n",
+		},
+		{
+			desc:           "When the API fails",
+			arguments:      &commandargs.AuthorizedKeys{ExpectedUser: "user", ActualUser: "user", Key: "broken"},
+			expectedOutput: "# No key was found for broken\n",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			buffer := &bytes.Buffer{}
+			cmd := &Command{
+				Config:     &config.Config{RootDir: "/tmp", GitlabUrl: url},
+				Args:       tc.arguments,
+				ReadWriter: &readwriter.ReadWriter{Out: buffer},
+			}
+
+			err := cmd.Execute()
+
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedOutput, buffer.String())
+		})
+	}
+}
diff --git a/go/internal/command/command.go b/go/internal/command/command.go
index 27378aa..d6b96f1 100644
--- a/go/internal/command/command.go
+++ b/go/internal/command/command.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/authorizedkeys"
 	"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"
@@ -35,6 +36,8 @@ func buildCommand(e *executable.Executable, args commandargs.CommandArgs, config
 	switch e.Name {
 	case executable.GitlabShell:
 		return buildShellCommand(args.(*commandargs.Shell), config, readWriter)
+	case executable.AuthorizedKeysCheck:
+		return buildAuthorizedKeysCommand(args.(*commandargs.AuthorizedKeys), config, readWriter)
 	}
 
 	return nil
@@ -62,3 +65,11 @@ func buildShellCommand(args *commandargs.Shell, config *config.Config, readWrite
 
 	return nil
 }
+
+func buildAuthorizedKeysCommand(args *commandargs.AuthorizedKeys, config *config.Config, readWriter *readwriter.ReadWriter) Command {
+	if !config.FeatureEnabled(executable.AuthorizedKeysCheck) {
+		return nil
+	}
+
+	return &authorizedkeys.Command{Config: config, Args: args, ReadWriter: readWriter}
+}
diff --git a/go/internal/command/command_test.go b/go/internal/command/command_test.go
index ea88a6a..b651c78 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/authorizedkeys"
 	"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"
@@ -125,6 +126,16 @@ func TestNew(t *testing.T) {
 			expectedType: &uploadarchive.Command{},
 		},
 		{
+			desc:       "it returns a AuthorizedKeys command if the feature is enabled",
+			executable: &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			config: &config.Config{
+				Migration: config.MigrationConfig{Enabled: true, Features: []string{"gitlab-shell-authorized-keys-check"}},
+			},
+			environment:  map[string]string{},
+			arguments:    []string{"git", "git", "key"},
+			expectedType: &authorizedkeys.Command{},
+		},
+		{
 			desc:       "it returns a Fallback command if the feature is unimplemented",
 			executable: &executable.Executable{Name: executable.GitlabShell},
 			config: &config.Config{
diff --git a/go/internal/command/commandargs/authorized_keys.go b/go/internal/command/commandargs/authorized_keys.go
new file mode 100644
index 0000000..2733954
--- /dev/null
+++ b/go/internal/command/commandargs/authorized_keys.go
@@ -0,0 +1,51 @@
+package commandargs
+
+import (
+	"errors"
+	"fmt"
+)
+
+type AuthorizedKeys struct {
+	Arguments    []string
+	ExpectedUser string
+	ActualUser   string
+	Key          string
+}
+
+func (ak *AuthorizedKeys) Parse() error {
+	if err := ak.validate(); err != nil {
+		return err
+	}
+
+	ak.ExpectedUser = ak.Arguments[0]
+	ak.ActualUser = ak.Arguments[1]
+	ak.Key = ak.Arguments[2]
+
+	return nil
+}
+
+func (ak *AuthorizedKeys) GetArguments() []string {
+	return ak.Arguments
+}
+
+func (ak *AuthorizedKeys) validate() error {
+	argsSize := len(ak.Arguments)
+
+	if argsSize != 3 {
+		return errors.New(fmt.Sprintf("# Insufficient arguments. %d. Usage\n#\tgitlab-shell-authorized-keys-check <expected-username> <actual-username> <key>", argsSize))
+	}
+
+	expectedUsername := ak.Arguments[0]
+	actualUsername := ak.Arguments[1]
+	key := ak.Arguments[2]
+
+	if expectedUsername == "" || actualUsername == "" {
+		return errors.New("# No username provided")
+	}
+
+	if key == "" {
+		return errors.New("# No key provided")
+	}
+
+	return nil
+}
diff --git a/go/internal/command/commandargs/command_args.go b/go/internal/command/commandargs/command_args.go
index 5338d6b..dfdf8e5 100644
--- a/go/internal/command/commandargs/command_args.go
+++ b/go/internal/command/commandargs/command_args.go
@@ -17,6 +17,8 @@ func Parse(e *executable.Executable, arguments []string) (CommandArgs, error) {
 	switch e.Name {
 	case executable.GitlabShell:
 		args = &Shell{Arguments: arguments}
+	case executable.AuthorizedKeysCheck:
+		args = &AuthorizedKeys{Arguments: arguments}
 	}
 
 	if err := args.Parse(); err != nil {
diff --git a/go/internal/command/commandargs/command_args_test.go b/go/internal/command/commandargs/command_args_test.go
index 148c987..e981c22 100644
--- a/go/internal/command/commandargs/command_args_test.go
+++ b/go/internal/command/commandargs/command_args_test.go
@@ -120,6 +120,11 @@ func TestParseSuccess(t *testing.T) {
 			arguments:    []string{},
 			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate},
 		}, {
+			desc:         "It parses authorized-keys command",
+			executable:   &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:    []string{"git", "git", "key"},
+			expectedArgs: &AuthorizedKeys{Arguments: []string{"git", "git", "key"}, ExpectedUser: "git", ActualUser: "git", Key: "key"},
+		}, {
 			desc:         "Unknown executable",
 			executable:   &executable.Executable{Name: "unknown"},
 			arguments:    []string{},
@@ -162,7 +167,31 @@ func TestParseFailure(t *testing.T) {
 				"SSH_ORIGINAL_COMMAND": `git receive-pack "`,
 			},
 			arguments:     []string{},
-			expectedError: "Invalid SSH allowed",
+			expectedError: "Invalid SSH command",
+		},
+		{
+			desc:          "With not enough arguments for the AuthorizedKeysCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:     []string{"user"},
+			expectedError: "# Insufficient arguments. 1. Usage\n#\tgitlab-shell-authorized-keys-check <expected-username> <actual-username> <key>",
+		},
+		{
+			desc:          "With too many arguments for the AuthorizedKeysCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:     []string{"user", "user", "key", "something-else"},
+			expectedError: "# Insufficient arguments. 4. Usage\n#\tgitlab-shell-authorized-keys-check <expected-username> <actual-username> <key>",
+		},
+		{
+			desc:          "With missing username for the AuthorizedKeysCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:     []string{"user", "", "key"},
+			expectedError: "# No username provided",
+		},
+		{
+			desc:          "With missing key for the AuthorizedKeysCheck",
+			executable:    &executable.Executable{Name: executable.AuthorizedKeysCheck},
+			arguments:     []string{"user", "user", ""},
+			expectedError: "# No key provided",
 		},
 	}
 
@@ -173,7 +202,7 @@ func TestParseFailure(t *testing.T) {
 
 			_, err := Parse(tc.executable, tc.arguments)
 
-			require.Error(t, err, tc.expectedError)
+			require.EqualError(t, err, tc.expectedError)
 		})
 	}
 }