chore: Refactor env introspection to rely on command initialization496-move-env-introspection-to-sshenv

Refactors introspection of execution environment to rely on
per-connection state (`gitlab-shell`) or per request (`gitlab-sshd`)

Relates to https://gitlab.com/gitlab-org/gitlab-shell/-/issues/496

11 files changed, 100 insertions, 173 deletions

diff --git a/internal/command/command.go b/internal/command/command.go
index c0f9090..bb430b7 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -20,6 +20,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/uploadpack"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/executable"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 	"gitlab.com/gitlab-org/labkit/correlation"
 	"gitlab.com/gitlab-org/labkit/log"
 	"gitlab.com/gitlab-org/labkit/tracing"
@@ -29,8 +30,8 @@ type Command interface {
 	Execute(ctx context.Context) error
 }
 
-func New(e *executable.Executable, arguments []string, config *config.Config, readWriter *readwriter.ReadWriter) (Command, error) {
-	args, err := commandargs.Parse(e, arguments)
+func New(e *executable.Executable, arguments []string, env sshenv.Env, config *config.Config, readWriter *readwriter.ReadWriter) (Command, error) {
+	args, err := commandargs.Parse(e, arguments, env)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/command/command_test.go b/internal/command/command_test.go
index d134e61..4b315aa 100644
--- a/internal/command/command_test.go
+++ b/internal/command/command_test.go
@@ -21,7 +21,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/uploadpack"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/executable"
-	"gitlab.com/gitlab-org/gitlab-shell/internal/testhelper"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 	"gitlab.com/gitlab-org/labkit/correlation"
 )
 
@@ -35,10 +35,10 @@ var (
 	advancedConfig = &config.Config{GitlabUrl: "http+unix://gitlab.socket", SslCertDir: "/tmp/certs"}
 )
 
-func buildEnv(command string) map[string]string {
-	return map[string]string{
-		"SSH_CONNECTION":       "1",
-		"SSH_ORIGINAL_COMMAND": command,
+func buildEnv(command string) sshenv.Env {
+	return sshenv.Env{
+		IsSSHConnection: true,
+		OriginalCommand: command,
 	}
 }
 
@@ -46,7 +46,7 @@ func TestNew(t *testing.T) {
 	testCases := []struct {
 		desc               string
 		executable         *executable.Executable
-		environment        map[string]string
+		env                sshenv.Env
 		arguments          []string
 		config             *config.Config
 		expectedType       interface{}
@@ -55,7 +55,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns a Discover command",
 			executable:         gitlabShellExec,
-			environment:        buildEnv(""),
+			env:                buildEnv(""),
 			config:             basicConfig,
 			expectedType:       &discover.Command{},
 			expectedSslCertDir: "",
@@ -63,7 +63,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns a Discover command with SSL_CERT_DIR env var set",
 			executable:         gitlabShellExec,
-			environment:        buildEnv(""),
+			env:                buildEnv(""),
 			config:             advancedConfig,
 			expectedType:       &discover.Command{},
 			expectedSslCertDir: "/tmp/certs",
@@ -71,7 +71,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns a TwoFactorRecover command",
 			executable:         gitlabShellExec,
-			environment:        buildEnv("2fa_recovery_codes"),
+			env:                buildEnv("2fa_recovery_codes"),
 			config:             basicConfig,
 			expectedType:       &twofactorrecover.Command{},
 			expectedSslCertDir: "",
@@ -79,7 +79,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns a TwoFactorVerify command",
 			executable:         gitlabShellExec,
-			environment:        buildEnv("2fa_verify"),
+			env:                buildEnv("2fa_verify"),
 			config:             basicConfig,
 			expectedType:       &twofactorverify.Command{},
 			expectedSslCertDir: "",
@@ -87,7 +87,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns an LfsAuthenticate command",
 			executable:         gitlabShellExec,
-			environment:        buildEnv("git-lfs-authenticate"),
+			env:                buildEnv("git-lfs-authenticate"),
 			config:             basicConfig,
 			expectedType:       &lfsauthenticate.Command{},
 			expectedSslCertDir: "",
@@ -95,7 +95,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns a ReceivePack command",
 			executable:         gitlabShellExec,
-			environment:        buildEnv("git-receive-pack"),
+			env:                buildEnv("git-receive-pack"),
 			config:             basicConfig,
 			expectedType:       &receivepack.Command{},
 			expectedSslCertDir: "",
@@ -103,7 +103,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns an UploadPack command",
 			executable:         gitlabShellExec,
-			environment:        buildEnv("git-upload-pack"),
+			env:                buildEnv("git-upload-pack"),
 			config:             basicConfig,
 			expectedType:       &uploadpack.Command{},
 			expectedSslCertDir: "",
@@ -111,7 +111,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns an UploadArchive command",
 			executable:         gitlabShellExec,
-			environment:        buildEnv("git-upload-archive"),
+			env:                buildEnv("git-upload-archive"),
 			config:             basicConfig,
 			expectedType:       &uploadarchive.Command{},
 			expectedSslCertDir: "",
@@ -142,7 +142,7 @@ func TestNew(t *testing.T) {
 		{
 			desc:               "it returns a PersonalAccessToken command",
 			executable:         gitlabShellExec,
-			environment:        buildEnv("personal_access_token"),
+			env:                buildEnv("personal_access_token"),
 			config:             basicConfig,
 			expectedType:       &personalaccesstoken.Command{},
 			expectedSslCertDir: "",
@@ -151,11 +151,8 @@ func TestNew(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			restoreEnv := testhelper.TempEnv(tc.environment)
-			defer restoreEnv()
-
 			os.Unsetenv("SSL_CERT_DIR")
-			command, err := New(tc.executable, tc.arguments, tc.config, nil)
+			command, err := New(tc.executable, tc.arguments, tc.env, tc.config, nil)
 
 			require.NoError(t, err)
 			require.IsType(t, tc.expectedType, command)
@@ -168,7 +165,7 @@ func TestFailingNew(t *testing.T) {
 	testCases := []struct {
 		desc          string
 		executable    *executable.Executable
-		environment   map[string]string
+		env           sshenv.Env
 		expectedError error
 	}{
 		{
@@ -179,17 +176,14 @@ func TestFailingNew(t *testing.T) {
 		{
 			desc:          "Unknown command given",
 			executable:    gitlabShellExec,
-			environment:   buildEnv("unknown"),
+			env:           buildEnv("unknown"),
 			expectedError: disallowedcommand.Error,
 		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			restoreEnv := testhelper.TempEnv(tc.environment)
-			defer restoreEnv()
-
-			command, err := New(tc.executable, []string{}, basicConfig, nil)
+			command, err := New(tc.executable, []string{}, tc.env, basicConfig, nil)
 			require.Nil(t, command)
 			require.Equal(t, tc.expectedError, err)
 		})
diff --git a/internal/command/commandargs/command_args.go b/internal/command/commandargs/command_args.go
index b4bf334..b7d04a8 100644
--- a/internal/command/commandargs/command_args.go
+++ b/internal/command/commandargs/command_args.go
@@ -2,6 +2,7 @@ package commandargs
 
 import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/executable"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 )
 
 type CommandType string
@@ -11,12 +12,12 @@ type CommandArgs interface {
 	GetArguments() []string
 }
 
-func Parse(e *executable.Executable, arguments []string) (CommandArgs, error) {
+func Parse(e *executable.Executable, arguments []string, env sshenv.Env) (CommandArgs, error) {
 	var args CommandArgs = &GenericArgs{Arguments: arguments}
 
 	switch e.Name {
 	case executable.GitlabShell:
-		args = &Shell{Arguments: arguments}
+		args = &Shell{Arguments: arguments, Env: env}
 	case executable.AuthorizedKeysCheck:
 		args = &AuthorizedKeys{Arguments: arguments}
 	case executable.AuthorizedPrincipalsCheck:
diff --git a/internal/command/commandargs/command_args_test.go b/internal/command/commandargs/command_args_test.go
index aa74237..0329c82 100644
--- a/internal/command/commandargs/command_args_test.go
+++ b/internal/command/commandargs/command_args_test.go
@@ -4,7 +4,7 @@ import (
 	"testing"
 
 	"gitlab.com/gitlab-org/gitlab-shell/internal/executable"
-	"gitlab.com/gitlab-org/gitlab-shell/internal/testhelper"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 
 	"github.com/stretchr/testify/require"
 )
@@ -13,112 +13,77 @@ func TestParseSuccess(t *testing.T) {
 	testCases := []struct {
 		desc         string
 		executable   *executable.Executable
-		environment  map[string]string
+		env          sshenv.Env
 		arguments    []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",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "",
-			},
+			desc:         "It sets discover as the command when the command string was empty",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{}, CommandType: Discover},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{}, CommandType: Discover, Env: sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"}},
 		},
 		{
-			desc:       "It finds the key id in any passed arguments",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "",
-			},
+			desc:         "It finds the key id in any passed arguments",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"},
 			arguments:    []string{"hello", "key-123"},
-			expectedArgs: &Shell{Arguments: []string{"hello", "key-123"}, SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123"},
+			expectedArgs: &Shell{Arguments: []string{"hello", "key-123"}, SshArgs: []string{}, CommandType: Discover, GitlabKeyId: "123", Env: sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"}},
 		}, {
-			desc:       "It finds the username in any passed arguments",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "",
-			},
+			desc:         "It finds the username in any passed arguments",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"},
 			arguments:    []string{"hello", "username-jane-doe"},
-			expectedArgs: &Shell{Arguments: []string{"hello", "username-jane-doe"}, SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe"},
+			expectedArgs: &Shell{Arguments: []string{"hello", "username-jane-doe"}, SshArgs: []string{}, CommandType: Discover, GitlabUsername: "jane-doe", Env: sshenv.Env{IsSSHConnection: true, RemoteAddr: "1"}},
 		}, {
-			desc:       "It parses 2fa_recovery_codes command",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "2fa_recovery_codes",
-			},
+			desc:         "It parses 2fa_recovery_codes command",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, OriginalCommand: "2fa_recovery_codes"},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"2fa_recovery_codes"}, CommandType: TwoFactorRecover},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"2fa_recovery_codes"}, CommandType: TwoFactorRecover, Env: sshenv.Env{IsSSHConnection: true, OriginalCommand: "2fa_recovery_codes"}},
 		}, {
-			desc:       "It parses git-receive-pack command",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "git-receive-pack group/repo",
-			},
+			desc:         "It parses git-receive-pack command",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, OriginalCommand: "git-receive-pack group/repo"},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack, Env: sshenv.Env{IsSSHConnection: true, OriginalCommand: "git-receive-pack group/repo"}},
 		}, {
-			desc:       "It parses git-receive-pack command and a project with single quotes",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "git receive-pack 'group/repo'",
-			},
+			desc:         "It parses git-receive-pack command and a project with single quotes",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, OriginalCommand: "git-receive-pack 'group/repo'"},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack, Env: sshenv.Env{IsSSHConnection: true, OriginalCommand: "git-receive-pack 'group/repo'"}},
 		}, {
-			desc:       `It parses "git receive-pack" command`,
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": `git receive-pack "group/repo"`,
-			},
+			desc:         `It parses "git receive-pack" command`,
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, OriginalCommand: `git-receive-pack "group/repo"`},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack, Env: sshenv.Env{IsSSHConnection: true, OriginalCommand: `git-receive-pack "group/repo"`}},
 		}, {
-			desc:       `It parses a command followed by control characters`,
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": `git-receive-pack group/repo; any command`,
-			},
+			desc:         `It parses a command followed by control characters`,
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, OriginalCommand: `git-receive-pack group/repo; any command`},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-receive-pack", "group/repo"}, CommandType: ReceivePack, Env: sshenv.Env{IsSSHConnection: true, OriginalCommand: `git-receive-pack group/repo; any command`}},
 		}, {
-			desc:       "It parses git-upload-pack command",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": `git upload-pack "group/repo"`,
-			},
+			desc:         "It parses git-upload-pack command",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, OriginalCommand: `git upload-pack "group/repo"`},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-upload-pack", "group/repo"}, CommandType: UploadPack},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-upload-pack", "group/repo"}, CommandType: UploadPack, Env: sshenv.Env{IsSSHConnection: true, OriginalCommand: `git upload-pack "group/repo"`}},
 		}, {
-			desc:       "It parses git-upload-archive command",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "git-upload-archive 'group/repo'",
-			},
+			desc:         "It parses git-upload-archive command",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, OriginalCommand: "git-upload-archive 'group/repo'"},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-upload-archive", "group/repo"}, CommandType: UploadArchive},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-upload-archive", "group/repo"}, CommandType: UploadArchive, Env: sshenv.Env{IsSSHConnection: true, OriginalCommand: "git-upload-archive 'group/repo'"}},
 		}, {
-			desc:       "It parses git-lfs-authenticate command",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": "git-lfs-authenticate 'group/repo' download",
-			},
+			desc:         "It parses git-lfs-authenticate command",
+			executable:   &executable.Executable{Name: executable.GitlabShell},
+			env:          sshenv.Env{IsSSHConnection: true, OriginalCommand: "git-lfs-authenticate 'group/repo' download"},
 			arguments:    []string{},
-			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate},
+			expectedArgs: &Shell{Arguments: []string{}, SshArgs: []string{"git-lfs-authenticate", "group/repo", "download"}, CommandType: LfsAuthenticate, Env: sshenv.Env{IsSSHConnection: true, OriginalCommand: "git-lfs-authenticate 'group/repo' download"}},
 		}, {
 			desc:         "It parses authorized-keys command",
 			executable:   &executable.Executable{Name: executable.AuthorizedKeysCheck},
@@ -139,10 +104,7 @@ func TestParseSuccess(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			restoreEnv := testhelper.TempEnv(tc.environment)
-			defer restoreEnv()
-
-			result, err := Parse(tc.executable, tc.arguments)
+			result, err := Parse(tc.executable, tc.arguments, tc.env)
 
 			require.NoError(t, err)
 			require.Equal(t, tc.expectedArgs, result)
@@ -154,7 +116,7 @@ func TestParseFailure(t *testing.T) {
 	testCases := []struct {
 		desc          string
 		executable    *executable.Executable
-		environment   map[string]string
+		env           sshenv.Env
 		arguments     []string
 		expectedError string
 	}{
@@ -165,12 +127,9 @@ func TestParseFailure(t *testing.T) {
 			expectedError: "Only SSH allowed",
 		},
 		{
-			desc:       "It fails if SSH command is invalid",
-			executable: &executable.Executable{Name: executable.GitlabShell},
-			environment: map[string]string{
-				"SSH_CONNECTION":       "1",
-				"SSH_ORIGINAL_COMMAND": `git receive-pack "`,
-			},
+			desc:          "It fails if SSH command is invalid",
+			executable:    &executable.Executable{Name: executable.GitlabShell},
+			env:           sshenv.Env{IsSSHConnection: true, OriginalCommand: `git receive-pack "`},
 			arguments:     []string{},
 			expectedError: "Invalid SSH command",
 		},
@@ -220,10 +179,7 @@ func TestParseFailure(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			restoreEnv := testhelper.TempEnv(tc.environment)
-			defer restoreEnv()
-
-			_, err := Parse(tc.executable, tc.arguments)
+			_, err := Parse(tc.executable, tc.arguments, tc.env)
 
 			require.EqualError(t, err, tc.expectedError)
 		})
diff --git a/internal/command/commandargs/shell.go b/internal/command/commandargs/shell.go
index 62fc8fa..9cf6720 100644
--- a/internal/command/commandargs/shell.go
+++ b/internal/command/commandargs/shell.go
@@ -2,11 +2,10 @@ package commandargs
 
 import (
 	"errors"
-	"net"
-	"os"
 	"regexp"
 
 	"github.com/mattn/go-shellwords"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 )
 
 const (
@@ -18,8 +17,6 @@ const (
 	UploadPack          CommandType = "git-upload-pack"
 	UploadArchive       CommandType = "git-upload-archive"
 	PersonalAccessToken CommandType = "personal_access_token"
-
-	GitProtocolEnv = "GIT_PROTOCOL"
 )
 
 var (
@@ -33,10 +30,7 @@ type Shell struct {
 	GitlabKeyId    string
 	SshArgs        []string
 	CommandType    CommandType
-
-	// Only set when running standalone
-	RemoteAddr         *net.TCPAddr
-	GitProtocolVersion string
+	Env            sshenv.Env
 }
 
 func (s *Shell) Parse() error {
@@ -54,24 +48,19 @@ func (s *Shell) GetArguments() []string {
 }
 
 func (s *Shell) validate() error {
-	if !s.isSshConnection() {
+	if !s.Env.IsSSHConnection {
 		return errors.New("Only SSH allowed")
 	}
 
-	if !s.isValidSshCommand() {
+	if !s.isValidSSHCommand() {
 		return errors.New("Invalid SSH command")
 	}
 
 	return nil
 }
 
-func (s *Shell) isSshConnection() bool {
-	ok := os.Getenv("SSH_CONNECTION")
-	return ok != ""
-}
-
-func (s *Shell) isValidSshCommand() bool {
-	err := s.ParseCommand(os.Getenv("SSH_ORIGINAL_COMMAND"))
+func (s *Shell) isValidSSHCommand() bool {
+	err := s.ParseCommand(s.Env.OriginalCommand)
 	return err == nil
 }
 
diff --git a/internal/command/receivepack/gitalycall.go b/internal/command/receivepack/gitalycall.go
index b27b75a..713f323 100644
--- a/internal/command/receivepack/gitalycall.go
+++ b/internal/command/receivepack/gitalycall.go
@@ -12,7 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/handler"
 )
 
-func (c *Command) performGitalyCall(response *accessverifier.Response, gitProtocolVersion string) error {
+func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 	gc := &handler.GitalyCommand{
 		Config:      c.Config,
 		ServiceName: string(commandargs.ReceivePack),
@@ -26,12 +26,12 @@ func (c *Command) performGitalyCall(response *accessverifier.Response, gitProtoc
 		GlId:             response.Who,
 		GlRepository:     response.Repo,
 		GlUsername:       response.Username,
-		GitProtocol:      gitProtocolVersion,
+		GitProtocol:      c.Args.Env.GitProtocolVersion,
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
 	return gc.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
-		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, request.GitProtocol)
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
 		defer cancel()
 
 		rw := c.ReadWriter
diff --git a/internal/command/receivepack/gitalycall_test.go b/internal/command/receivepack/gitalycall_test.go
index ea07477..d756248 100644
--- a/internal/command/receivepack/gitalycall_test.go
+++ b/internal/command/receivepack/gitalycall_test.go
@@ -13,6 +13,7 @@ import (
 	"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/config"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/testhelper"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/testhelper/requesthandlers"
 )
@@ -25,10 +26,6 @@ func TestReceivePack(t *testing.T) {
 	url, cleanup := testserver.StartHttpServer(t, requests)
 	defer cleanup()
 
-	envCleanup, err := testhelper.Setenv("SSH_CONNECTION", "127.0.0.1 0")
-	require.NoError(t, err)
-	defer envCleanup()
-
 	testCases := []struct {
 		username string
 		keyId    string
@@ -46,7 +43,12 @@ func TestReceivePack(t *testing.T) {
 		input := &bytes.Buffer{}
 		repo := "group/repo"
 
-		args := &commandargs.Shell{CommandType: commandargs.ReceivePack, SshArgs: []string{"git-receive-pack", repo}}
+		env := sshenv.Env{
+			IsSSHConnection: true,
+			OriginalCommand: "git-receive-pack group/repo",
+			RemoteAddr:      "127.0.0.1",
+		}
+		args := &commandargs.Shell{CommandType: commandargs.ReceivePack, SshArgs: []string{"git-receive-pack", repo}, Env: env}
 
 		if tc.username != "" {
 			args.GitlabUsername = tc.username
@@ -62,7 +64,7 @@ func TestReceivePack(t *testing.T) {
 
 		hook := testhelper.SetupLogger()
 
-		err = cmd.Execute(context.Background())
+		err := cmd.Execute(context.Background())
 		require.NoError(t, err)
 
 		if tc.username != "" {
diff --git a/internal/command/receivepack/receivepack.go b/internal/command/receivepack/receivepack.go
index 5a67c5a..4d5c686 100644
--- a/internal/command/receivepack/receivepack.go
+++ b/internal/command/receivepack/receivepack.go
@@ -2,7 +2,6 @@ package receivepack
 
 import (
 	"context"
-	"os"
 
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
@@ -39,14 +38,7 @@ func (c *Command) Execute(ctx context.Context) error {
 		return customAction.Execute(ctx, response)
 	}
 
-	var gitProtocolVersion string
-	if c.Args.RemoteAddr != nil {
-		gitProtocolVersion = c.Args.GitProtocolVersion
-	} else {
-		gitProtocolVersion = os.Getenv(commandargs.GitProtocolEnv)
-	}
-
-	return c.performGitalyCall(response, gitProtocolVersion)
+	return c.performGitalyCall(response)
 }
 
 func (c *Command) verifyAccess(ctx context.Context, repo string) (*accessverifier.Response, error) {
diff --git a/internal/command/uploadarchive/gitalycall.go b/internal/command/uploadarchive/gitalycall.go
index f17ee50..05bb794 100644
--- a/internal/command/uploadarchive/gitalycall.go
+++ b/internal/command/uploadarchive/gitalycall.go
@@ -24,7 +24,7 @@ func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 	request := &pb.SSHUploadArchiveRequest{Repository: &response.Gitaly.Repo}
 
 	return gc.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
-		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, "")
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
 		defer cancel()
 
 		rw := c.ReadWriter
diff --git a/internal/command/uploadpack/gitalycall.go b/internal/command/uploadpack/gitalycall.go
index 3ebc8b3..5a8a605 100644
--- a/internal/command/uploadpack/gitalycall.go
+++ b/internal/command/uploadpack/gitalycall.go
@@ -12,7 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/internal/handler"
 )
 
-func (c *Command) performGitalyCall(response *accessverifier.Response, gitProtocolVersion string) error {
+func (c *Command) performGitalyCall(response *accessverifier.Response) error {
 	gc := &handler.GitalyCommand{
 		Config:      c.Config,
 		ServiceName: string(commandargs.UploadPack),
@@ -23,12 +23,12 @@ func (c *Command) performGitalyCall(response *accessverifier.Response, gitProtoc
 
 	request := &pb.SSHUploadPackRequest{
 		Repository:       &response.Gitaly.Repo,
-		GitProtocol:      gitProtocolVersion,
+		GitProtocol:      c.Args.Env.GitProtocolVersion,
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
 	return gc.RunGitalyCommand(func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
-		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, request.GitProtocol)
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
 		defer cancel()
 
 		rw := c.ReadWriter
diff --git a/internal/command/uploadpack/uploadpack.go b/internal/command/uploadpack/uploadpack.go
index bf5db2c..fca3823 100644
--- a/internal/command/uploadpack/uploadpack.go
+++ b/internal/command/uploadpack/uploadpack.go
@@ -2,7 +2,6 @@ package uploadpack
 
 import (
 	"context"
-	"os"
 
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/command/readwriter"
@@ -39,14 +38,7 @@ func (c *Command) Execute(ctx context.Context) error {
 		return customAction.Execute(ctx, response)
 	}
 
-	var gitProtocolVersion string
-	if c.Args.RemoteAddr != nil {
-		gitProtocolVersion = c.Args.GitProtocolVersion
-	} else {
-		gitProtocolVersion = os.Getenv(commandargs.GitProtocolEnv)
-	}
-
-	return c.performGitalyCall(response, gitProtocolVersion)
+	return c.performGitalyCall(response)
 }
 
 func (c *Command) verifyAccess(ctx context.Context, repo string) (*accessverifier.Response, error) {