refactor: unify instantiation of command.Shell

3 files changed, 29 insertions, 18 deletions

diff --git a/cmd/gitlab-shell/command/command.go b/cmd/gitlab-shell/command/command.go
index be54ac6..5f828cd 100644
--- a/cmd/gitlab-shell/command/command.go
+++ b/cmd/gitlab-shell/command/command.go
@@ -30,6 +30,20 @@ func New(arguments []string, env sshenv.Env, config *config.Config, readWriter *
 	return nil, disallowedcommand.Error
 }
 
+func NewWithKey(gitlabKeyId string, env sshenv.Env, config *config.Config, readWriter *readwriter.ReadWriter) (command.Command, error) {
+	args, err := Parse(nil, env)
+	if err != nil {
+		return nil, err
+	}
+
+	args.GitlabKeyId = gitlabKeyId
+	if cmd := Build(args, config, readWriter); cmd != nil {
+		return cmd, nil
+	}
+
+	return nil, disallowedcommand.Error
+}
+
 func Parse(arguments []string, env sshenv.Env) (*commandargs.Shell, error) {
 	args := &commandargs.Shell{Arguments: arguments, Env: env}
 
diff --git a/internal/sshd/session.go b/internal/sshd/session.go
index b58598e..d5a0174 100644
--- a/internal/sshd/session.go
+++ b/internal/sshd/session.go
@@ -2,13 +2,14 @@ package sshd
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"golang.org/x/crypto/ssh"
 
 	shellCmd "gitlab.com/gitlab-org/gitlab-shell/cmd/gitlab-shell/command"
-	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/command/shared/disallowedcommand"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 )
@@ -104,19 +105,11 @@ func (s *session) handleShell(ctx context.Context, req *ssh.Request) uint32 {
 		req.Reply(true, []byte{})
 	}
 
-	args := &commandargs.Shell{
-		GitlabKeyId: s.gitlabKeyId,
-		Env: sshenv.Env{
-			IsSSHConnection:    true,
-			OriginalCommand:    s.execCmd,
-			GitProtocolVersion: s.gitProtocolVersion,
-			RemoteAddr:         s.remoteAddr,
-		},
-	}
-
-	if err := args.ParseCommand(s.execCmd); err != nil {
-		s.toStderr("Failed to parse command: %v\n", err.Error())
-		return 128
+	env := sshenv.Env{
+		IsSSHConnection:    true,
+		OriginalCommand:    s.execCmd,
+		GitProtocolVersion: s.gitProtocolVersion,
+		RemoteAddr:         s.remoteAddr,
 	}
 
 	rw := &readwriter.ReadWriter{
@@ -125,9 +118,13 @@ func (s *session) handleShell(ctx context.Context, req *ssh.Request) uint32 {
 		ErrOut: s.channel.Stderr(),
 	}
 
-	cmd := shellCmd.Build(args, s.cfg, rw)
-	if cmd == nil {
-		s.toStderr("Unknown command: %v\n", args.CommandType)
+	cmd, err := shellCmd.NewWithKey(s.gitlabKeyId, env, s.cfg, rw)
+
+	if err != nil {
+		if !errors.Is(err, disallowedcommand.Error) {
+			s.toStderr("Failed to parse command: %v\n", err.Error())
+		}
+		s.toStderr("Unknown command: %v\n", s.execCmd)
 		return 128
 	}
 
diff --git a/internal/sshd/session_test.go b/internal/sshd/session_test.go
index e4e312c..b014c0c 100644
--- a/internal/sshd/session_test.go
+++ b/internal/sshd/session_test.go
@@ -145,7 +145,7 @@ func TestHandleShell(t *testing.T) {
 		{
 			desc:             "fails to parse command",
 			cmd:              `\`,
-			errMsg:           "Failed to parse command: invalid command line string\n",
+			errMsg:           "Failed to parse command: Invalid SSH command\nUnknown command: \\\n",
 			gitlabKeyId:      "root",
 			expectedExitCode: 128,
 		}, {