Optionally use SSHUploadPackWithSidechanneljv-ssh-sidechannel

If the GitLab API returns an allowed response with use_sidechannel set
to true, gitlab-shell will establish a sidechannel connection and use
SSHUploadPackWithSidechannel instead of SSHUploadPack. This is an
efficiency improvement.

2 files changed, 74 insertions, 1 deletions

diff --git a/internal/command/uploadpack/gitalycall.go b/internal/command/uploadpack/gitalycall.go
index 0e6a6d0..43ae4f8 100644
--- a/internal/command/uploadpack/gitalycall.go
+++ b/internal/command/uploadpack/gitalycall.go
@@ -21,13 +21,30 @@ func (c *Command) performGitalyCall(ctx context.Context, response *accessverifie
 		Features:    response.Gitaly.Features,
 	}
 
+	if response.Gitaly.UseSidechannel {
+		gc.DialSidechannel = true
+		request := &pb.SSHUploadPackWithSidechannelRequest{
+			Repository:       &response.Gitaly.Repo,
+			GitProtocol:      c.Args.Env.GitProtocolVersion,
+			GitConfigOptions: response.GitConfigOptions,
+		}
+
+		return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error) {
+			ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
+			defer cancel()
+
+			rw := c.ReadWriter
+			return client.UploadPackWithSidechannel(ctx, conn, registry, rw.In, rw.Out, rw.ErrOut, request)
+		})
+	}
+
 	request := &pb.SSHUploadPackRequest{
 		Repository:       &response.Gitaly.Repo,
 		GitProtocol:      c.Args.Env.GitProtocolVersion,
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
-	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
+	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error) {
 		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
 		defer cancel()
 
diff --git a/internal/command/uploadpack/gitalycall_test.go b/internal/command/uploadpack/gitalycall_test.go
index 926d2c9..213b6f9 100644
--- a/internal/command/uploadpack/gitalycall_test.go
+++ b/internal/command/uploadpack/gitalycall_test.go
@@ -71,3 +71,59 @@ func TestUploadPack(t *testing.T) {
 	require.Empty(t, testServer.ReceivedMD["some-other-ff"])
 	require.Equal(t, testServer.ReceivedMD["x-gitlab-correlation-id"][0], "a-correlation-id")
 }
+
+func TestUploadPack_withSidechannel(t *testing.T) {
+	gitalyAddress, testServer := testserver.StartGitalyServer(t)
+
+	requests := requesthandlers.BuildAllowedWithGitalyHandlersWithSidechannel(t, gitalyAddress)
+	url := testserver.StartHttpServer(t, requests)
+
+	output := &bytes.Buffer{}
+	input := &bytes.Buffer{}
+
+	userId := "1"
+	repo := "group/repo"
+
+	env := sshenv.Env{
+		IsSSHConnection: true,
+		OriginalCommand: "git-upload-pack " + repo,
+		RemoteAddr:      "127.0.0.1",
+	}
+
+	args := &commandargs.Shell{
+		GitlabKeyId: userId,
+		CommandType: commandargs.UploadPack,
+		SshArgs:     []string{"git-upload-pack", repo},
+		Env:         env,
+	}
+
+	cmd := &Command{
+		Config:     &config.Config{GitlabUrl: url},
+		Args:       args,
+		ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
+	}
+
+	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
+	ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
+
+	err := cmd.Execute(ctx)
+	require.NoError(t, err)
+
+	require.Equal(t, "SSHUploadPackWithSidechannel: "+repo, output.String())
+
+	for k, v := range map[string]string{
+		"gitaly-feature-cache_invalidator":        "true",
+		"gitaly-feature-inforef_uploadpack_cache": "false",
+		"x-gitlab-client-name":                    "gitlab-shell-tests-git-upload-pack",
+		"key_id":                                  "123",
+		"user_id":                                 "1",
+		"remote_ip":                               "127.0.0.1",
+		"key_type":                                "key",
+	} {
+		actual := testServer.ReceivedMD[k]
+		require.Len(t, actual, 1)
+		require.Equal(t, v, actual[0])
+	}
+	require.Empty(t, testServer.ReceivedMD["some-other-ff"])
+	require.Equal(t, testServer.ReceivedMD["x-gitlab-correlation-id"][0], "a-correlation-id")
+}