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, 19 insertions, 12 deletions

diff --git a/internal/handler/exec.go b/internal/handler/exec.go
index 8f4a7df..9bd8018 100644
--- a/internal/handler/exec.go
+++ b/internal/handler/exec.go
@@ -29,21 +29,23 @@ import (
 // GitalyHandlerFunc implementations are responsible for making
 // an appropriate Gitaly call using the provided client and context
 // and returning an error from the Gitaly call.
-type GitalyHandlerFunc func(ctx context.Context, client *grpc.ClientConn) (int32, error)
+type GitalyHandlerFunc func(ctx context.Context, client *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error)
 
 type GitalyCommand struct {
-	Config      *config.Config
-	ServiceName string
-	Address     string
-	Token       string
-	Features    map[string]string
+	Config          *config.Config
+	ServiceName     string
+	Address         string
+	Token           string
+	Features        map[string]string
+	DialSidechannel bool
 }
 
 // RunGitalyCommand provides a bootstrap for Gitaly commands executed
 // through GitLab-Shell. It ensures that logging, tracing and other
 // common concerns are configured before executing the `handler`.
 func (gc *GitalyCommand) RunGitalyCommand(ctx context.Context, handler GitalyHandlerFunc) error {
-	conn, err := getConn(ctx, gc)
+	registry := client.NewSidechannelRegistry(log.ContextLogger(ctx))
+	conn, err := getConn(ctx, gc, registry)
 	if err != nil {
 		log.ContextLogger(ctx).WithError(fmt.Errorf("RunGitalyCommand: %v", err)).Error("Failed to get connection to execute Git command")
 
@@ -53,7 +55,7 @@ func (gc *GitalyCommand) RunGitalyCommand(ctx context.Context, handler GitalyHan
 
 	childCtx := withOutgoingMetadata(ctx, gc.Features)
 	ctxlog := log.ContextLogger(childCtx)
-	exitStatus, err := handler(childCtx, conn)
+	exitStatus, err := handler(childCtx, conn, registry)
 
 	if err != nil {
 		if grpcstatus.Convert(err).Code() == grpccodes.Unavailable {
@@ -116,7 +118,7 @@ func withOutgoingMetadata(ctx context.Context, features map[string]string) conte
 	return metadata.NewOutgoingContext(ctx, md)
 }
 
-func getConn(ctx context.Context, gc *GitalyCommand) (*grpc.ClientConn, error) {
+func getConn(ctx context.Context, gc *GitalyCommand, registry *client.SidechannelRegistry) (*grpc.ClientConn, error) {
 	if gc.Address == "" {
 		return nil, fmt.Errorf("no gitaly_address given")
 	}
@@ -160,5 +162,9 @@ func getConn(ctx context.Context, gc *GitalyCommand) (*grpc.ClientConn, error) {
 		)
 	}
 
+	if gc.DialSidechannel {
+		return client.DialSidechannel(ctx, gc.Address, registry, connOpts)
+	}
+
 	return client.DialContext(ctx, gc.Address, connOpts)
 }
diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go
index ba6bd6a..784cf19 100644
--- a/internal/handler/exec_test.go
+++ b/internal/handler/exec_test.go
@@ -11,14 +11,15 @@ import (
 	"google.golang.org/grpc/metadata"
 	grpcstatus "google.golang.org/grpc/status"
 
+	"gitlab.com/gitlab-org/gitaly/v14/client"
 	pb "gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 )
 
-func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn) (int32, error) {
-	return func(ctx context.Context, client *grpc.ClientConn) (int32, error) {
+func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn, *client.SidechannelRegistry) (int32, error) {
+	return func(ctx context.Context, client *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error) {
 		require.NotNil(t, ctx)
 		require.NotNil(t, client)
 
@@ -86,7 +87,7 @@ func TestRunGitalyCommandMetadata(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			cmd := tt.gc
 
-			err := cmd.RunGitalyCommand(context.Background(), func(ctx context.Context, _ *grpc.ClientConn) (int32, error) {
+			err := cmd.RunGitalyCommand(context.Background(), func(ctx context.Context, _ *grpc.ClientConn, _ *client.SidechannelRegistry) (int32, error) {
 				md, exists := metadata.FromOutgoingContext(ctx)
 				require.True(t, exists)
 				require.Equal(t, len(tt.want), md.Len())