gitaly-ssh: add sidechannel support for upload-pack

If GITALY_USE_SIDECHANNEL=1 is set, gitaly-ssh will use
SSHUploadPackWithSidechannel instead of SSHUploadPack.

5 files changed, 51 insertions, 13 deletions

diff --git a/cmd/gitaly-ssh/main.go b/cmd/gitaly-ssh/main.go
index b7a4f51dd..137ec690b 100644
--- a/cmd/gitaly-ssh/main.go
+++ b/cmd/gitaly-ssh/main.go
@@ -8,6 +8,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/sirupsen/logrus"
 	gitalyauth "gitlab.com/gitlab-org/gitaly/v14/auth"
 	"gitlab.com/gitlab-org/gitaly/v14/client"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/metadata/featureflag"
@@ -15,7 +16,7 @@ import (
 	"google.golang.org/grpc"
 )
 
-type packFn func(_ context.Context, _ *grpc.ClientConn, _ string) (int32, error)
+type packFn func(context.Context, *grpc.ClientConn, *client.SidechannelRegistry, string) (int32, error)
 
 type gitalySSHCommand struct {
 	// The git packer that shall be executed. One of receivePack,
@@ -37,6 +38,7 @@ type gitalySSHCommand struct {
 // GITALY_PAYLOAD="{repo...}"
 // GITALY_WD="/path/to/working-directory"
 // GITALY_FEATUREFLAGS="upload_pack_filter:false,hooks_rpc:true"
+// GITALY_USE_SIDECHANNEL=1 if desired
 // gitaly-ssh upload-pack <git-garbage-x2>
 func main() {
 	// < 4 since git throws on 2x garbage here
@@ -49,7 +51,11 @@ func main() {
 	var packer packFn
 	switch command {
 	case "upload-pack":
-		packer = uploadPack
+		if useSidechannel() {
+			packer = uploadPackWithSidechannel
+		} else {
+			packer = uploadPack
+		}
 	case "receive-pack":
 		packer = receivePack
 	case "upload-archive":
@@ -104,13 +110,14 @@ func (cmd gitalySSHCommand) run() (int, error) {
 		}
 	}
 
-	conn, err := getConnection(cmd.address)
+	registry := client.NewSidechannelRegistry(logrus.NewEntry(logrus.StandardLogger()))
+	conn, err := getConnection(ctx, cmd.address, registry)
 	if err != nil {
 		return 1, err
 	}
 	defer conn.Close()
 
-	code, err := cmd.packer(ctx, conn, cmd.payload)
+	code, err := cmd.packer(ctx, conn, registry, cmd.payload)
 	if err != nil {
 		return 1, err
 	}
@@ -118,12 +125,16 @@ func (cmd gitalySSHCommand) run() (int, error) {
 	return int(code), nil
 }
 
-func getConnection(url string) (*grpc.ClientConn, error) {
+func getConnection(ctx context.Context, url string, registry *client.SidechannelRegistry) (*grpc.ClientConn, error) {
 	if url == "" {
 		return nil, fmt.Errorf("gitaly address can not be empty")
 	}
 
-	return client.Dial(url, dialOpts())
+	if useSidechannel() {
+		return client.DialSidechannel(ctx, url, registry, dialOpts())
+	}
+
+	return client.DialContext(ctx, url, dialOpts())
 }
 
 func dialOpts() []grpc.DialOption {
@@ -134,3 +145,5 @@ func dialOpts() []grpc.DialOption {
 
 	return connOpts
 }
+
+func useSidechannel() bool { return os.Getenv("GITALY_USE_SIDECHANNEL") == "1" }
diff --git a/cmd/gitaly-ssh/main_test.go b/cmd/gitaly-ssh/main_test.go
index 2ec16c736..76327740a 100644
--- a/cmd/gitaly-ssh/main_test.go
+++ b/cmd/gitaly-ssh/main_test.go
@@ -7,14 +7,21 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"gitlab.com/gitlab-org/gitaly/v14/client"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 	"google.golang.org/grpc"
 )
 
 func TestRun(t *testing.T) {
-	var successPacker packFn = func(_ context.Context, _ *grpc.ClientConn, _ string) (int32, error) { return 0, nil }
-	var exitCodePacker packFn = func(_ context.Context, _ *grpc.ClientConn, _ string) (int32, error) { return 123, nil }
-	var errorPacker packFn = func(_ context.Context, _ *grpc.ClientConn, _ string) (int32, error) { return 1, fmt.Errorf("fail") }
+	var successPacker packFn = func(context.Context, *grpc.ClientConn, *client.SidechannelRegistry, string) (int32, error) {
+		return 0, nil
+	}
+	var exitCodePacker packFn = func(context.Context, *grpc.ClientConn, *client.SidechannelRegistry, string) (int32, error) {
+		return 123, nil
+	}
+	var errorPacker packFn = func(context.Context, *grpc.ClientConn, *client.SidechannelRegistry, string) (int32, error) {
+		return 1, fmt.Errorf("fail")
+	}
 
 	gitalyTCPAddress := "tcp://localhost:9999"
 	gitalyUnixAddress := fmt.Sprintf("unix://%s", testhelper.GetTemporaryGitalySocketFileName(t))
diff --git a/cmd/gitaly-ssh/receive_pack.go b/cmd/gitaly-ssh/receive_pack.go
index 57711db5c..62cc84e65 100644
--- a/cmd/gitaly-ssh/receive_pack.go
+++ b/cmd/gitaly-ssh/receive_pack.go
@@ -11,7 +11,7 @@ import (
 	"google.golang.org/protobuf/encoding/protojson"
 )
 
-func receivePack(ctx context.Context, conn *grpc.ClientConn, req string) (int32, error) {
+func receivePack(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry, req string) (int32, error) {
 	var request gitalypb.SSHReceivePackRequest
 
 	if err := protojson.Unmarshal([]byte(req), &request); err != nil {
diff --git a/cmd/gitaly-ssh/upload_archive.go b/cmd/gitaly-ssh/upload_archive.go
index 051ceb311..618bb92d6 100644
--- a/cmd/gitaly-ssh/upload_archive.go
+++ b/cmd/gitaly-ssh/upload_archive.go
@@ -11,7 +11,7 @@ import (
 	"google.golang.org/protobuf/encoding/protojson"
 )
 
-func uploadArchive(ctx context.Context, conn *grpc.ClientConn, req string) (int32, error) {
+func uploadArchive(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry, req string) (int32, error) {
 	var request gitalypb.SSHUploadArchiveRequest
 	if err := protojson.Unmarshal([]byte(req), &request); err != nil {
 		return 0, fmt.Errorf("json unmarshal: %w", err)
diff --git a/cmd/gitaly-ssh/upload_pack.go b/cmd/gitaly-ssh/upload_pack.go
index 64e0f580a..97ba6b7c7 100644
--- a/cmd/gitaly-ssh/upload_pack.go
+++ b/cmd/gitaly-ssh/upload_pack.go
@@ -18,16 +18,34 @@ const (
 	GitConfigShowAllRefs = "transfer.hideRefs=!refs"
 )
 
-func uploadPack(ctx context.Context, conn *grpc.ClientConn, req string) (int32, error) {
+func uploadPackConfig(config []string) []string {
+	return append([]string{GitConfigShowAllRefs}, config...)
+}
+
+func uploadPack(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry, req string) (int32, error) {
 	var request gitalypb.SSHUploadPackRequest
 	if err := protojson.Unmarshal([]byte(req), &request); err != nil {
 		return 0, fmt.Errorf("json unmarshal: %w", err)
 	}
 
-	request.GitConfigOptions = append([]string{GitConfigShowAllRefs}, request.GitConfigOptions...)
+	request.GitConfigOptions = uploadPackConfig(request.GitConfigOptions)
 
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	return client.UploadPack(ctx, conn, os.Stdin, os.Stdout, os.Stderr, &request)
 }
+
+func uploadPackWithSidechannel(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry, req string) (int32, error) {
+	var request gitalypb.SSHUploadPackWithSidechannelRequest
+	if err := protojson.Unmarshal([]byte(req), &request); err != nil {
+		return 0, fmt.Errorf("json unmarshal: %w", err)
+	}
+
+	request.GitConfigOptions = uploadPackConfig(request.GitConfigOptions)
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	return client.UploadPackWithSidechannel(ctx, conn, registry, os.Stdin, os.Stdout, os.Stderr, &request)
+}