Merge branch 'pks-ssh-upload-pack-filter' into 'master'

Support partial clones with SSH transports

See merge request gitlab-org/gitaly!1893

5 files changed, 198 insertions, 70 deletions

diff --git a/changelogs/unreleased/pks-ssh-upload-pack-filter.yml b/changelogs/unreleased/pks-ssh-upload-pack-filter.yml
new file mode 100644
index 000000000..cf47ec014
--- /dev/null
+++ b/changelogs/unreleased/pks-ssh-upload-pack-filter.yml
@@ -0,0 +1,5 @@
+---
+title: Support partial clones with SSH transports
+merge_request: 1893
+author:
+type: added
diff --git a/cmd/gitaly-ssh/main.go b/cmd/gitaly-ssh/main.go
index 4049cc47b..5c70de194 100644
--- a/cmd/gitaly-ssh/main.go
+++ b/cmd/gitaly-ssh/main.go
@@ -5,10 +5,12 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"strings"
 
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	gitalyauth "gitlab.com/gitlab-org/gitaly/auth"
 	"gitlab.com/gitlab-org/gitaly/client"
+	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
 	grpccorrelation "gitlab.com/gitlab-org/labkit/correlation/grpc"
 	"gitlab.com/gitlab-org/labkit/tracing"
 	grpctracing "gitlab.com/gitlab-org/labkit/tracing/grpc"
@@ -17,10 +19,26 @@ import (
 
 type packFn func(_ context.Context, _ *grpc.ClientConn, _ string) (int32, error)
 
+type gitalySSHCommand struct {
+	// The git packer that shall be executed. One of receivePack,
+	// uploadPack or uploadArchive
+	packer packFn
+	// Working directory to execute the packer in
+	workingDir string
+	// Address of the server we want to post the request to
+	address string
+	// Marshalled gRPC payload to pass to the remote server
+	payload string
+	// Comma separated list of feature flags that shall be enabled on the
+	// remote server
+	featureFlags string
+}
+
 // GITALY_ADDRESS="tcp://1.2.3.4:9999" or "unix:/var/run/gitaly.sock"
 // GITALY_TOKEN="foobar1234"
 // GITALY_PAYLOAD="{repo...}"
 // GITALY_WD="/path/to/working-directory"
+// GITALY_FEATUREFLAGS="upload_pack_filter,hooks_rpc"
 // gitaly-ssh upload-pack <git-garbage-x2>
 func main() {
 	// < 4 since git throws on 2x garbage here
@@ -42,11 +60,15 @@ func main() {
 		log.Fatalf("invalid pack command: %q", command)
 	}
 
-	gitalyWorkingDir := os.Getenv("GITALY_WD")
-	gitalyAddress := os.Getenv("GITALY_ADDRESS")
-	gitalyPayload := os.Getenv("GITALY_PAYLOAD")
+	cmd := gitalySSHCommand{
+		packer:       packer,
+		workingDir:   os.Getenv("GITALY_WD"),
+		address:      os.Getenv("GITALY_ADDRESS"),
+		payload:      os.Getenv("GITALY_PAYLOAD"),
+		featureFlags: os.Getenv("GITALY_FEATUREFLAGS"),
+	}
 
-	code, err := run(packer, gitalyWorkingDir, gitalyAddress, gitalyPayload)
+	code, err := cmd.run()
 	if err != nil {
 		log.Printf("%s: %v", command, err)
 	}
@@ -54,7 +76,7 @@ func main() {
 	os.Exit(code)
 }
 
-func run(packer packFn, gitalyWorkingDir string, gitalyAddress string, gitalyPayload string) (int, error) {
+func (cmd gitalySSHCommand) run() (int, error) {
 	// Configure distributed tracing
 	closer := tracing.Initialize(tracing.WithServiceName("gitaly-ssh"))
 	defer closer.Close()
@@ -62,19 +84,25 @@ func run(packer packFn, gitalyWorkingDir string, gitalyAddress string, gitalyPay
 	ctx, finished := tracing.ExtractFromEnv(context.Background())
 	defer finished()
 
-	if gitalyWorkingDir != "" {
-		if err := os.Chdir(gitalyWorkingDir); err != nil {
-			return 1, fmt.Errorf("unable to chdir to %v", gitalyWorkingDir)
+	if cmd.featureFlags != "" {
+		for _, flag := range strings.Split(cmd.featureFlags, ",") {
+			ctx = featureflag.OutgoingCtxWithFeatureFlag(ctx, flag)
+		}
+	}
+
+	if cmd.workingDir != "" {
+		if err := os.Chdir(cmd.workingDir); err != nil {
+			return 1, fmt.Errorf("unable to chdir to %v", cmd.workingDir)
 		}
 	}
 
-	conn, err := getConnection(gitalyAddress)
+	conn, err := getConnection(cmd.address)
 	if err != nil {
 		return 1, err
 	}
 	defer conn.Close()
 
-	code, err := packer(ctx, conn, gitalyPayload)
+	code, err := cmd.packer(ctx, conn, cmd.payload)
 	if err != nil {
 		return 1, err
 	}
diff --git a/cmd/gitaly-ssh/main_test.go b/cmd/gitaly-ssh/main_test.go
index d66a57a0f..7256a7e02 100644
--- a/cmd/gitaly-ssh/main_test.go
+++ b/cmd/gitaly-ssh/main_test.go
@@ -89,7 +89,14 @@ func TestRun(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			gotCode, err := run(tt.packer, tt.workingDir, tt.gitalyAddress, "{}")
+			cmd := gitalySSHCommand{
+				packer:     tt.packer,
+				workingDir: tt.workingDir,
+				address:    tt.gitalyAddress,
+				payload:    "{}",
+			}
+
+			gotCode, err := cmd.run()
 			if tt.wantErr {
 				assert.Error(t, err)
 			} else {
diff --git a/internal/service/ssh/upload_pack.go b/internal/service/ssh/upload_pack.go
index d7ed643bd..94161fead 100644
--- a/internal/service/ssh/upload_pack.go
+++ b/internal/service/ssh/upload_pack.go
@@ -10,6 +10,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/git"
 	"gitlab.com/gitlab-org/gitaly/internal/git/pktline"
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/internal/service/inspect"
 	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitaly/streamio"
@@ -74,6 +75,10 @@ func (s *server) sshUploadPack(stream gitalypb.SSHService_SSHUploadPackServer, r
 	git.WarnIfTooManyBitmaps(ctx, repoPath)
 
 	var globalOpts []git.Option
+	if featureflag.IsEnabled(ctx, featureflag.UploadPackFilter) {
+		globalOpts = append(globalOpts, git.UploadPackFilterConfig()...)
+	}
+
 	for _, o := range req.GitConfigOptions {
 		globalOpts = append(globalOpts, git.ValueFlag{"-c", o})
 	}
diff --git a/internal/service/ssh/upload_pack_test.go b/internal/service/ssh/upload_pack_test.go
index d3936461d..30ef62590 100644
--- a/internal/service/ssh/upload_pack_test.go
+++ b/internal/service/ssh/upload_pack_test.go
@@ -1,7 +1,6 @@
 package ssh
 
 import (
-	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -15,11 +14,72 @@ import (
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/internal/git"
+	"gitlab.com/gitlab-org/gitaly/internal/helper/text"
+	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
 )
 
+type cloneCommand struct {
+	command      *exec.Cmd
+	repository   *gitalypb.Repository
+	server       string
+	featureFlags []string
+	gitConfig    string
+	gitProtocol  string
+}
+
+func (cmd cloneCommand) execute(t *testing.T) error {
+	req := &gitalypb.SSHUploadPackRequest{
+		Repository:  cmd.repository,
+		GitProtocol: cmd.gitProtocol,
+	}
+	if cmd.gitConfig != "" {
+		req.GitConfigOptions = strings.Split(cmd.gitConfig, " ")
+	}
+	pbMarshaler := &jsonpb.Marshaler{}
+	payload, err := pbMarshaler.MarshalToString(req)
+
+	require.NoError(t, err)
+
+	cmd.command.Env = []string{
+		fmt.Sprintf("GITALY_ADDRESS=%s", cmd.server),
+		fmt.Sprintf("GITALY_PAYLOAD=%s", payload),
+		fmt.Sprintf("GITALY_FEATUREFLAGS=%s", strings.Join(cmd.featureFlags, ",")),
+		fmt.Sprintf("PATH=.:%s", os.Getenv("PATH")),
+		fmt.Sprintf(`GIT_SSH_COMMAND=%s upload-pack`, gitalySSHPath),
+	}
+
+	out, err := cmd.command.CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("%v: %q", err, out)
+	}
+	if !cmd.command.ProcessState.Success() {
+		return fmt.Errorf("Failed to run `git clone`: %q", out)
+	}
+
+	return nil
+}
+
+func (cmd cloneCommand) test(t *testing.T, localRepoPath string) (string, string, string, string) {
+	defer os.RemoveAll(localRepoPath)
+
+	err := cmd.execute(t)
+	require.NoError(t, err)
+
+	storagePath := testhelper.GitlabTestStoragePath()
+	testRepoPath := path.Join(storagePath, testRepo.GetRelativePath())
+
+	remoteHead := text.ChompBytes(testhelper.MustRunCommand(t, nil, "git", "-C", testRepoPath, "rev-parse", "master"))
+	localHead := text.ChompBytes(testhelper.MustRunCommand(t, nil, "git", "-C", localRepoPath, "rev-parse", "master"))
+
+	remoteTags := text.ChompBytes(testhelper.MustRunCommand(t, nil, "git", "-C", testRepoPath, "tag"))
+	localTags := text.ChompBytes(testhelper.MustRunCommand(t, nil, "git", "-C", localRepoPath, "tag"))
+
+	return localHead, remoteHead, localTags, remoteTags
+}
+
 func TestFailedUploadPackRequestDueToTimeout(t *testing.T) {
 	server, serverSocketPath := runSSHServer(t, WithUploadPackRequestTimeout(10*time.Microsecond))
 
@@ -147,13 +207,68 @@ func TestUploadPackCloneSuccess(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.desc, func(t *testing.T) {
-			lHead, rHead, _, _, err := testClone(t, serverSocketPath, testRepo.GetStorageName(), testRepo.GetRelativePath(), localRepoPath, "", tc.cmd, "")
-			require.NoError(t, err, "clone failed")
+			cmd := cloneCommand{
+				repository: testRepo,
+				command:    tc.cmd,
+				server:     serverSocketPath,
+			}
+			lHead, rHead, _, _ := cmd.test(t, localRepoPath)
 			require.Equal(t, lHead, rHead, "local and remote head not equal")
 		})
 	}
 }
 
+func TestUploadPackCloneWithPartialCloneFilter(t *testing.T) {
+	server, serverSocketPath := runSSHServer(t)
+	defer server.Stop()
+
+	// Ruby file which is ~1kB in size and not present in HEAD
+	blobLessThanLimit := "6ee41e85cc9bf33c10b690df09ca735b22f3790f"
+	// Image which is ~100kB in size and not present in HEAD
+	blobGreaterThanLimit := "18079e308ff9b3a5e304941020747e5c39b46c88"
+
+	tests := []struct {
+		desc     string
+		flags    []string
+		repoTest func(t *testing.T, repoPath string)
+	}{
+		{
+			desc: "full_clone",
+			repoTest: func(t *testing.T, repoPath string) {
+				testhelper.GitObjectMustExist(t, repoPath, blobGreaterThanLimit)
+			},
+		},
+		{
+			desc:  "partial_clone",
+			flags: []string{featureflag.UploadPackFilter},
+			repoTest: func(t *testing.T, repoPath string) {
+				testhelper.GitObjectMustNotExist(t, repoPath, blobGreaterThanLimit)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.desc, func(t *testing.T) {
+			// Run the clone with filtering enabled in both runs. The only
+			// difference is that in the first run, we have the
+			// UploadPackFilter flag disabled.
+			localPath := path.Join(testRepoRoot, fmt.Sprintf("gitlab-test-upload-pack-local-%s", tc.desc))
+			cmd := cloneCommand{
+				repository:   testRepo,
+				command:      exec.Command("git", "clone", "--filter=blob:limit=2048", "git@localhost:test/test.git", localPath),
+				server:       serverSocketPath,
+				featureFlags: tc.flags,
+			}
+			err := cmd.execute(t)
+			defer os.RemoveAll(localPath)
+			require.NoError(t, err, "clone failed")
+
+			testhelper.GitObjectMustExist(t, localPath, blobLessThanLimit)
+			tc.repoTest(t, localPath)
+		})
+	}
+}
+
 func TestUploadPackCloneSuccessWithGitProtocol(t *testing.T) {
 	restore := testhelper.EnableGitProtocolV2Support()
 	defer restore()
@@ -179,8 +294,14 @@ func TestUploadPackCloneSuccessWithGitProtocol(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.desc, func(t *testing.T) {
-			lHead, rHead, _, _, err := testClone(t, serverSocketPath, testRepo.GetStorageName(), testRepo.GetRelativePath(), localRepoPath, "", tc.cmd, git.ProtocolV2)
-			require.NoError(t, err, "clone failed")
+			cmd := cloneCommand{
+				repository:  testRepo,
+				command:     tc.cmd,
+				server:      serverSocketPath,
+				gitProtocol: git.ProtocolV2,
+			}
+
+			lHead, rHead, _, _ := cmd.test(t, localRepoPath)
 			require.Equal(t, lHead, rHead, "local and remote head not equal")
 
 			envData, err := testhelper.GetGitEnvData()
@@ -197,13 +318,14 @@ func TestUploadPackCloneHideTags(t *testing.T) {
 
 	localRepoPath := path.Join(testRepoRoot, "gitlab-test-upload-pack-local-hide-tags")
 
-	cmd := exec.Command("git", "clone", "--mirror", "git@localhost:test/test.git", localRepoPath)
-
-	_, _, lTags, rTags, err := testClone(t, serverSocketPath, testRepo.GetStorageName(), testRepo.GetRelativePath(), localRepoPath, "transfer.hideRefs=refs/tags", cmd, "")
-
-	if err != nil {
-		t.Fatalf("clone failed: %v", err)
+	cmd := cloneCommand{
+		repository: testRepo,
+		command:    exec.Command("git", "clone", "--mirror", "git@localhost:test/test.git", localRepoPath),
+		server:     serverSocketPath,
+		gitConfig:  "transfer.hideRefs=refs/tags",
 	}
+	_, _, lTags, rTags := cmd.test(t, localRepoPath)
+
 	if lTags == rTags {
 		t.Fatalf("local and remote tags are equal. clone failed: %q != %q", lTags, rTags)
 	}
@@ -218,55 +340,16 @@ func TestUploadPackCloneFailure(t *testing.T) {
 
 	localRepoPath := path.Join(testRepoRoot, "gitlab-test-upload-pack-local-failure")
 
-	cmd := exec.Command("git", "clone", "git@localhost:test/test.git", localRepoPath)
-
-	_, _, _, _, err := testClone(t, serverSocketPath, "foobar", testRepo.GetRelativePath(), localRepoPath, "", cmd, "")
-	if err == nil {
-		t.Fatalf("clone didn't fail")
-	}
-}
-
-func testClone(t *testing.T, serverSocketPath, storageName, relativePath, localRepoPath string, gitConfig string, cmd *exec.Cmd, gitProtocol string) (string, string, string, string, error) {
-	defer os.RemoveAll(localRepoPath)
-
-	pbTempRepo := &gitalypb.Repository{StorageName: storageName, RelativePath: relativePath}
-	req := &gitalypb.SSHUploadPackRequest{
-		Repository:  pbTempRepo,
-		GitProtocol: gitProtocol,
-	}
-	if gitConfig != "" {
-		req.GitConfigOptions = strings.Split(gitConfig, " ")
-	}
-	pbMarshaler := &jsonpb.Marshaler{}
-	payload, err := pbMarshaler.MarshalToString(req)
-
-	require.NoError(t, err)
-
-	cmd.Env = []string{
-		fmt.Sprintf("GITALY_ADDRESS=%s", serverSocketPath),
-		fmt.Sprintf("GITALY_PAYLOAD=%s", payload),
-		fmt.Sprintf("PATH=%s", ".:"+os.Getenv("PATH")),
-		fmt.Sprintf(`GIT_SSH_COMMAND=%s upload-pack`, gitalySSHPath),
-	}
-
-	out, err := cmd.CombinedOutput()
-	if err != nil {
-		return "", "", "", "", fmt.Errorf("%v: %q", err, out)
-	}
-	if !cmd.ProcessState.Success() {
-		return "", "", "", "", fmt.Errorf("Failed to run `git clone`: %q", out)
+	cmd := cloneCommand{
+		repository: &gitalypb.Repository{
+			StorageName:  "foobar",
+			RelativePath: testRepo.GetRelativePath(),
+		},
+		command: exec.Command("git", "clone", "git@localhost:test/test.git", localRepoPath),
+		server:  serverSocketPath,
 	}
-
-	storagePath := testhelper.GitlabTestStoragePath()
-	testRepoPath := path.Join(storagePath, testRepo.GetRelativePath())
-
-	remoteHead := bytes.TrimSpace(testhelper.MustRunCommand(t, nil, "git", "-C", testRepoPath, "rev-parse", "master"))
-	localHead := bytes.TrimSpace(testhelper.MustRunCommand(t, nil, "git", "-C", localRepoPath, "rev-parse", "master"))
-
-	remoteTags := bytes.TrimSpace(testhelper.MustRunCommand(t, nil, "git", "-C", testRepoPath, "tag"))
-	localTags := bytes.TrimSpace(testhelper.MustRunCommand(t, nil, "git", "-C", localRepoPath, "tag"))
-
-	return string(localHead), string(remoteHead), string(localTags), string(remoteTags), nil
+	err := cmd.execute(t)
+	require.Error(t, err, "clone didn't fail")
 }
 
 func testPostUploadPackFailedResponse(t *testing.T, stream gitalypb.SSHService_SSHUploadPackClient) error {