6 files changed, 33 insertions, 13 deletions
diff --git a/internal/command/command.go b/internal/command/command.go
index 5062d15..c0f9090 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -70,7 +70,7 @@ func ContextWithCorrelationID() (context.Context, func()) {
 func buildCommand(e *executable.Executable, args commandargs.CommandArgs, config *config.Config, readWriter *readwriter.ReadWriter) Command {
 	switch e.Name {
 	case executable.GitlabShell:
-		return buildShellCommand(args.(*commandargs.Shell), config, readWriter)
+		return BuildShellCommand(args.(*commandargs.Shell), config, readWriter)
 	case executable.AuthorizedKeysCheck:
 		return buildAuthorizedKeysCommand(args.(*commandargs.AuthorizedKeys), config, readWriter)
 	case executable.AuthorizedPrincipalsCheck:
@@ -82,7 +82,7 @@ func buildCommand(e *executable.Executable, args commandargs.CommandArgs, config
 	return nil
 }
 
-func buildShellCommand(args *commandargs.Shell, config *config.Config, readWriter *readwriter.ReadWriter) Command {
+func BuildShellCommand(args *commandargs.Shell, config *config.Config, readWriter *readwriter.ReadWriter) Command {
 	switch args.CommandType {
 	case commandargs.Discover:
 		return &discover.Command{Config: config, Args: args, ReadWriter: readWriter}
diff --git a/internal/command/commandargs/shell.go b/internal/command/commandargs/shell.go
index 1535ccb..62fc8fa 100644
--- a/internal/command/commandargs/shell.go
+++ b/internal/command/commandargs/shell.go
@@ -2,6 +2,7 @@ package commandargs
 
 import (
 	"errors"
+	"net"
 	"os"
 	"regexp"
 
@@ -32,6 +33,10 @@ type Shell struct {
 	GitlabKeyId    string
 	SshArgs        []string
 	CommandType    CommandType
+
+	// Only set when running standalone
+	RemoteAddr         *net.TCPAddr
+	GitProtocolVersion string
 }
 
 func (s *Shell) Parse() error {
@@ -40,7 +45,6 @@ func (s *Shell) Parse() error {
 	}
 
 	s.parseWho()
-	s.defineCommandType()
 
 	return nil
 }
@@ -67,7 +71,7 @@ func (s *Shell) isSshConnection() bool {
 }
 
 func (s *Shell) isValidSshCommand() bool {
-	err := s.parseCommand(os.Getenv("SSH_ORIGINAL_COMMAND"))
+	err := s.ParseCommand(os.Getenv("SSH_ORIGINAL_COMMAND"))
 	return err == nil
 }
 
@@ -107,7 +111,7 @@ func tryParseUsername(argument string) string {
 	return ""
 }
 
-func (s *Shell) parseCommand(commandString string) error {
+func (s *Shell) ParseCommand(commandString string) error {
 	args, err := shellwords.Parse(commandString)
 	if err != nil {
 		return err
@@ -123,6 +127,8 @@ func (s *Shell) parseCommand(commandString string) error {
 
 	s.SshArgs = args
 
+	s.defineCommandType()
+
 	return nil
 }
 
diff --git a/internal/command/receivepack/gitalycall.go b/internal/command/receivepack/gitalycall.go
index a983c1a..b27b75a 100644
--- a/internal/command/receivepack/gitalycall.go
+++ b/internal/command/receivepack/gitalycall.go
@@ -2,7 +2,6 @@ package receivepack
 
 import (
 	"context"
-	"os"
 
 	"google.golang.org/grpc"
 
@@ -13,7 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/handler"
 )
 
-func (c *Command) performGitalyCall(response *accessverifier.Response) error {
+func (c *Command) performGitalyCall(response *accessverifier.Response, gitProtocolVersion string) error {
 	gc := &handler.GitalyCommand{
 		Config:      c.Config,
 		ServiceName: string(commandargs.ReceivePack),
@@ -27,7 +26,7 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 		GlId:             response.Who,
 		GlRepository:     response.Repo,
 		GlUsername:       response.Username,
-		GitProtocol:      os.Getenv(commandargs.GitProtocolEnv),
+		GitProtocol:      gitProtocolVersion,
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
diff --git a/internal/command/receivepack/receivepack.go b/internal/command/receivepack/receivepack.go
index 4d5c686..5a67c5a 100644
--- a/internal/command/receivepack/receivepack.go
+++ b/internal/command/receivepack/receivepack.go
@@ -2,6 +2,7 @@ package receivepack
 
 import (
 	"context"
+	"os"
 
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
@@ -38,7 +39,14 @@ func (c *Command) Execute(ctx context.Context) error {
 		return customAction.Execute(ctx, response)
 	}
 
-	return c.performGitalyCall(response)
+	var gitProtocolVersion string
+	if c.Args.RemoteAddr != nil {
+		gitProtocolVersion = c.Args.GitProtocolVersion
+	} else {
+		gitProtocolVersion = os.Getenv(commandargs.GitProtocolEnv)
+	}
+
+	return c.performGitalyCall(response, gitProtocolVersion)
 }
 
 func (c *Command) verifyAccess(ctx context.Context, repo string) (*accessverifier.Response, error) {
diff --git a/internal/command/uploadpack/gitalycall.go b/internal/command/uploadpack/gitalycall.go
index ba0fef2..3ebc8b3 100644
--- a/internal/command/uploadpack/gitalycall.go
+++ b/internal/command/uploadpack/gitalycall.go
@@ -2,7 +2,6 @@ package uploadpack
 
 import (
 	"context"
-	"os"
 
 	"google.golang.org/grpc"
 
@@ -13,7 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/handler"
 )
 
-func (c *Command) performGitalyCall(response *accessverifier.Response) error {
+func (c *Command) performGitalyCall(response *accessverifier.Response, gitProtocolVersion string) error {
 	gc := &handler.GitalyCommand{
 		Config:      c.Config,
 		ServiceName: string(commandargs.UploadPack),
@@ -24,7 +23,7 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 
 	request := &pb.SSHUploadPackRequest{
 		Repository:       &response.Gitaly.Repo,
-		GitProtocol:      os.Getenv(commandargs.GitProtocolEnv),
+		GitProtocol:      gitProtocolVersion,
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
diff --git a/internal/command/uploadpack/uploadpack.go b/internal/command/uploadpack/uploadpack.go
index fca3823..bf5db2c 100644
--- a/internal/command/uploadpack/uploadpack.go
+++ b/internal/command/uploadpack/uploadpack.go
@@ -2,6 +2,7 @@ package uploadpack
 
 import (
 	"context"
+	"os"
 
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
@@ -38,7 +39,14 @@ func (c *Command) Execute(ctx context.Context) error {
 		return customAction.Execute(ctx, response)
 	}
 
-	return c.performGitalyCall(response)
+	var gitProtocolVersion string
+	if c.Args.RemoteAddr != nil {
+		gitProtocolVersion = c.Args.GitProtocolVersion
+	} else {
+		gitProtocolVersion = os.Getenv(commandargs.GitProtocolEnv)
+	}
+
+	return c.performGitalyCall(response, gitProtocolVersion)
 }
 
 func (c *Command) verifyAccess(ctx context.Context, repo string) (*accessverifier.Response, error) {