repository: Implement UpdateHead option for FetchBundle

Allows updating HEAD based off of a HEAD reference in a bundle file.
This is intended to allow updating the default branch from incremental
backups.

Changelog: added

2 files changed, 61 insertions, 1 deletions

diff --git a/internal/gitaly/service/repository/fetch_bundle.go b/internal/gitaly/service/repository/fetch_bundle.go
index 05204d064..af6389d5b 100644
--- a/internal/gitaly/service/repository/fetch_bundle.go
+++ b/internal/gitaly/service/repository/fetch_bundle.go
@@ -1,6 +1,8 @@
 package repository
 
 import (
+	"context"
+	"fmt"
 	"io"
 	"os"
 	"path/filepath"
@@ -8,6 +10,7 @@ import (
 	gitalyerrors "gitlab.com/gitlab-org/gitaly/v14/internal/errors"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git/localrepo"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/service/ref"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/helper"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/tempdir"
 	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
@@ -41,6 +44,7 @@ func (s *server) FetchBundle(stream gitalypb.RepositoryService_FetchBundleServer
 
 	ctx := stream.Context()
 	repo := s.localrepo(firstRequest.GetRepository())
+	updateHead := firstRequest.GetUpdateHead()
 
 	tmpDir, err := tempdir.New(ctx, repo.GetStorageName(), s.locator)
 	if err != nil {
@@ -71,5 +75,60 @@ func (s *server) FetchBundle(stream gitalypb.RepositoryService_FetchBundleServer
 		return helper.ErrInternal(err)
 	}
 
+	if updateHead {
+		if err := s.updateHeadFromBundle(ctx, repo, bundlePath); err != nil {
+			return helper.ErrInternal(err)
+		}
+	}
+
 	return stream.SendAndClose(&gitalypb.FetchBundleResponse{})
 }
+
+// updateHeadFromBundle updates HEAD from a bundle file
+func (s *server) updateHeadFromBundle(ctx context.Context, repo *localrepo.Repo, bundlePath string) error {
+	head, err := s.findBundleHead(ctx, repo, bundlePath)
+	if err != nil {
+		return fmt.Errorf("update head from bundle: %w", err)
+	}
+	if head == nil {
+		return nil
+	}
+
+	branch, err := repo.GuessHead(ctx, *head)
+	if err != nil {
+		return fmt.Errorf("update head from bundle: %w", err)
+	}
+
+	if err := ref.SetDefaultBranchRef(ctx, repo, branch.String(), s.cfg); err != nil {
+		return fmt.Errorf("update head from bundle: %w", err)
+	}
+	return nil
+}
+
+// findBundleHead tries to extract HEAD and its target from a bundle. Returns
+// nil when HEAD is not found.
+func (s *server) findBundleHead(ctx context.Context, repo git.RepositoryExecutor, bundlePath string) (*git.Reference, error) {
+	cmd, err := repo.Exec(ctx, git.SubSubCmd{
+		Name:   "bundle",
+		Action: "list-heads",
+		Args:   []string{bundlePath, "HEAD"},
+	})
+	if err != nil {
+		return nil, err
+	}
+	decoder := git.NewShowRefDecoder(cmd)
+	for {
+		var ref git.Reference
+		err := decoder.Decode(&ref)
+		if err == io.EOF {
+			break
+		} else if err != nil {
+			return nil, err
+		}
+		if ref.Name != "HEAD" {
+			continue
+		}
+		return &ref, nil
+	}
+	return nil, nil
+}
diff --git a/internal/gitaly/service/repository/fetch_bundle_test.go b/internal/gitaly/service/repository/fetch_bundle_test.go
index 25b2d9961..c60eb8149 100644
--- a/internal/gitaly/service/repository/fetch_bundle_test.go
+++ b/internal/gitaly/service/repository/fetch_bundle_test.go
@@ -31,6 +31,7 @@ func TestServer_FetchBundle_success(t *testing.T) {
 	tmp := testhelper.TempDir(t)
 	bundlePath := filepath.Join(tmp, "test.bundle")
 
+	gittest.Exec(t, cfg, "-C", repoPath, "symbolic-ref", "HEAD", "refs/heads/feature")
 	gittest.Exec(t, cfg, "-C", repoPath, "bundle", "create", bundlePath, "--all")
 	expectedRefs := gittest.Exec(t, cfg, "-C", repoPath, "show-ref", "--head")
 
@@ -42,7 +43,7 @@ func TestServer_FetchBundle_success(t *testing.T) {
 	stream, err := client.FetchBundle(ctx)
 	require.NoError(t, err)
 
-	request := &gitalypb.FetchBundleRequest{Repository: targetRepo}
+	request := &gitalypb.FetchBundleRequest{Repository: targetRepo, UpdateHead: true}
 	writer := streamio.NewWriter(func(p []byte) error {
 		request.Data = p