Merge branch 'sh-improve-key-matching-sshd' into 'main'

Relax key and username matching for sshd

See merge request gitlab-org/gitlab-shell!540

2 files changed, 36 insertions, 11 deletions

diff --git a/cmd/gitlab-shell/command/command_test.go b/cmd/gitlab-shell/command/command_test.go
index 2aeee59..ba0db7d 100644
--- a/cmd/gitlab-shell/command/command_test.go
+++ b/cmd/gitlab-shell/command/command_test.go
@@ -170,6 +170,27 @@ func TestParseSuccess(t *testing.T) {
 			expectedArgs: &commandargs.Shell{Arguments: []string{"hello", "username-key-123"}, SshArgs: []string{}, CommandType: commandargs.Discover, GitlabUsername: "key-123", Env: sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"}},
 		},
 		{
+			desc:         "It finds the key id if the key is listed as the last argument",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"},
+			arguments:    []string{"hello", "gitlab-shell -c key-123"},
+			expectedArgs: &commandargs.Shell{Arguments: []string{"hello", "gitlab-shell -c key-123"}, SshArgs: []string{}, CommandType: commandargs.Discover, GitlabKeyId: "123", Env: sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"}},
+		},
+		{
+			desc:         "It finds the username if the username is listed as the last argument",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"},
+			arguments:    []string{"hello", "gitlab-shell -c username-jane-doe"},
+			expectedArgs: &commandargs.Shell{Arguments: []string{"hello", "gitlab-shell -c username-jane-doe"}, SshArgs: []string{}, CommandType: commandargs.Discover, GitlabUsername: "jane-doe", Env: sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"}},
+		},
+		{
+			desc:         "It finds the key id only if the last argument is of <key-id> format",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"},
+			arguments:    []string{"hello", "gitlab-shell -c username-key-123"},
+			expectedArgs: &commandargs.Shell{Arguments: []string{"hello", "gitlab-shell -c username-key-123"}, SshArgs: []string{}, CommandType: commandargs.Discover, GitlabUsername: "key-123", Env: sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"}},
+		},
+		{
 			desc:         "It finds the username in any passed arguments",
 			executable:   &executable.Executable{Name: executable.GitlabShell},
 			env:          sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"},
diff --git a/internal/command/commandargs/shell.go b/internal/command/commandargs/shell.go
index 7a76be5..a783f93 100644
--- a/internal/command/commandargs/shell.go
+++ b/internal/command/commandargs/shell.go
@@ -3,6 +3,7 @@ package commandargs
 import (
 	"fmt"
 	"regexp"
+	"strings"
 
 	"github.com/mattn/go-shellwords"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
@@ -73,26 +74,29 @@ func (s *Shell) parseWho() {
 	}
 }
 
-func tryParseKeyId(argument string) string {
-	matchInfo := whoKeyRegex.FindStringSubmatch(argument)
+func tryParse(r *regexp.Regexp, argument string) string {
+	// sshd may execute the session for AuthorizedKeysCommand in multiple ways:
+	// 1. key-id
+	// 2. /path/to/shell -c key-id
+	args := strings.Split(argument, " ")
+	lastArg := args[len(args)-1]
+
+	matchInfo := r.FindStringSubmatch(lastArg)
 	if len(matchInfo) == 2 {
 		// The first element is the full matched string
-		// The second element is the named `keyid`
+		// The second element is the named `keyid` or `username`
 		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]
-	}
+func tryParseKeyId(argument string) string {
+	return tryParse(whoKeyRegex, argument)
+}
 
-	return ""
+func tryParseUsername(argument string) string {
+	return tryParse(whoUsernameRegex, argument)
 }
 
 func (s *Shell) ParseCommand(commandString string) error {