Merge branch 'smh-extract-snapshot-path' into 'master'

Restore snapshot's relative path in quarantined requests

See merge request https://gitlab.com/gitlab-org/gitaly/-/merge_requests/6550

Merged-by: Sami Hiltunen <shiltunen@gitlab.com>
Approved-by: James Fargher <jfargher@gitlab.com>
Approved-by: Justin Tobler <jtobler@gitlab.com>
Reviewed-by: Sami Hiltunen <shiltunen@gitlab.com>
Reviewed-by: Justin Tobler <jtobler@gitlab.com>

9 files changed, 224 insertions, 11 deletions

diff --git a/internal/gitaly/service/blob/blobs_test.go b/internal/gitaly/service/blob/blobs_test.go
index e0345c9a3..150f064cb 100644
--- a/internal/gitaly/service/blob/blobs_test.go
+++ b/internal/gitaly/service/blob/blobs_test.go
@@ -10,10 +10,12 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/quarantine"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/config"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagemgr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/structerr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitaly/v16/streamio"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -436,6 +438,19 @@ func TestListAllBlobs(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := ctx
+			if tc.request.Repository.GitObjectDirectory != "" {
+				// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+				// this is the snapshot's relative path. Include the metadata in the test as well as we're testing requests with quarantine
+				// as if they were coming from access checks.
+				ctx = metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath,
+					// Gitaly sends the snapshot's relative path to Rails from `pre-receive` and Rails
+					// sends it back to Gitaly when it performs requests in the access checks. The repository
+					// would have already been rewritten by Praefect, so we have to adjust for that as well.
+					gittest.RewrittenRepository(t, ctx, cfg, tc.request.Repository).RelativePath,
+				)
+			}
+
 			stream, err := client.ListAllBlobs(ctx, tc.request)
 			require.NoError(t, err)
 
diff --git a/internal/gitaly/service/blob/lfs_pointers_test.go b/internal/gitaly/service/blob/lfs_pointers_test.go
index 74ea55948..8149c52a9 100644
--- a/internal/gitaly/service/blob/lfs_pointers_test.go
+++ b/internal/gitaly/service/blob/lfs_pointers_test.go
@@ -17,6 +17,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/quarantine"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/config"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagemgr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/helper/text"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/structerr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
@@ -283,6 +284,19 @@ size 12345`
 
 			setup := tc.setup(t)
 
+			ctx := ctx
+			if setup.repo.GetGitObjectDirectory() != "" {
+				// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+				// this is the snapshot's relative path. Include the metadata in the test as well as we're testing requests with quarantine
+				// as if they were coming from access checks.
+				ctx = metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath,
+					// Gitaly sends the snapshot's relative path to Rails from `pre-receive` and Rails
+					// sends it back to Gitaly when it performs requests in the access checks. The repository
+					// would have already been rewritten by Praefect, so we have to adjust for that as well.
+					gittest.RewrittenRepository(t, ctx, cfg, setup.repo).RelativePath,
+				)
+			}
+
 			stream, err := client.ListAllLFSPointers(ctx, &gitalypb.ListAllLFSPointersRequest{
 				Repository: setup.repo,
 			})
diff --git a/internal/gitaly/service/commit/find_commits_test.go b/internal/gitaly/service/commit/find_commits_test.go
index 8f7b5492c..21a92e71a 100644
--- a/internal/gitaly/service/commit/find_commits_test.go
+++ b/internal/gitaly/service/commit/find_commits_test.go
@@ -14,11 +14,13 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/localrepo"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagemgr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/structerr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
 	"golang.org/x/text/encoding/charmap"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -664,6 +666,16 @@ func TestFindCommits_quarantine(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			repo.GitAlternateObjectDirectories = tc.altDirs
 
+			// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+			// this is the snapshot's relative path. Include the metadata in the test as well as we're testing requests with quarantine
+			// as if they were coming from access checks.
+			ctx := metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath,
+				// Gitaly sends the snapshot's relative path to Rails from `pre-receive` and Rails
+				// sends it back to Gitaly when it performs requests in the access checks. The repository
+				// would have already been rewritten by Praefect, so we have to adjust for that as well.
+				gittest.RewrittenRepository(t, ctx, cfg, repo).RelativePath,
+			)
+
 			commits, err := getCommits(t, ctx, client, &gitalypb.FindCommitsRequest{
 				Repository: repo,
 				Revision:   []byte(commitID.String()),
diff --git a/internal/gitaly/service/commit/isancestor_test.go b/internal/gitaly/service/commit/isancestor_test.go
index 93b3229f4..3bc6ac676 100644
--- a/internal/gitaly/service/commit/isancestor_test.go
+++ b/internal/gitaly/service/commit/isancestor_test.go
@@ -6,10 +6,12 @@ import (
 
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagemgr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/structerr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
@@ -308,6 +310,19 @@ func TestCommitIsAncestor(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			setup := tc.setup(t)
 
+			ctx := ctx
+			if setup.request.GetRepository().GetGitObjectDirectory() != "" || len(setup.request.GetRepository().GetGitAlternateObjectDirectories()) > 0 {
+				// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+				// this is the snapshot's relative path. Include the metadata in the test as well as we're testing requests with quarantine
+				// as if they were coming from access checks.
+				ctx = metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath,
+					// Gitaly sends the snapshot's relative path to Rails from `pre-receive` and Rails
+					// sends it back to Gitaly when it performs requests in the access checks. The repository
+					// would have already been rewritten by Praefect, so we have to adjust for that as well.
+					gittest.RewrittenRepository(t, ctx, cfg, setup.request.GetRepository()).RelativePath,
+				)
+			}
+
 			resp, err := client.CommitIsAncestor(ctx, setup.request)
 			testhelper.ProtoEqual(t, setup.expectedResponse, resp)
 			testhelper.RequireGrpcError(t, setup.expectedErr, err)
diff --git a/internal/gitaly/service/commit/list_all_commits_test.go b/internal/gitaly/service/commit/list_all_commits_test.go
index d95b74568..862a5081d 100644
--- a/internal/gitaly/service/commit/list_all_commits_test.go
+++ b/internal/gitaly/service/commit/list_all_commits_test.go
@@ -13,10 +13,12 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/localrepo"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagemgr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/helper/perm"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/structerr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
+	"google.golang.org/grpc/metadata"
 )
 
 func TestListAllCommits(t *testing.T) {
@@ -161,6 +163,19 @@ func TestListAllCommits(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			setup := tc.setup(t)
 
+			ctx := ctx
+			if setup.request.GetRepository().GetGitObjectDirectory() != "" {
+				// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+				// this is the snapshot's relative path. Include the metadata in the test as well as we're testing requests with quarantine
+				// as if they were coming from access checks.
+				ctx = metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath,
+					// Gitaly sends the snapshot's relative path to Rails from `pre-receive` and Rails
+					// sends it back to Gitaly when it performs requests in the access checks. The repository
+					// would have already been rewritten by Praefect, so we have to adjust for that as well.
+					gittest.RewrittenRepository(t, ctx, cfg, setup.request.Repository).RelativePath,
+				)
+			}
+
 			stream, err := client.ListAllCommits(ctx, setup.request)
 			require.NoError(t, err)
 
diff --git a/internal/gitaly/service/hook/pre_receive_test.go b/internal/gitaly/service/hook/pre_receive_test.go
index 82017dee2..ca5bdd618 100644
--- a/internal/gitaly/service/hook/pre_receive_test.go
+++ b/internal/gitaly/service/hook/pre_receive_test.go
@@ -18,6 +18,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/config"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/config/prometheus"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagemgr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/transaction"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitlab"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/helper/text"
@@ -28,6 +29,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/transaction/txinfo"
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitaly/v16/streamio"
+	"google.golang.org/grpc/metadata"
 )
 
 func TestPreReceiveInvalidArgument(t *testing.T) {
@@ -170,6 +172,19 @@ func TestPreReceiveHook_GitlabAPIAccess(t *testing.T) {
 		},
 	}
 
+	// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+	// this is the snapshot's relative path.
+	//
+	// Transaction middleware fails if this metadata is not present but this is not correct. We only start transactions when
+	// they come through the external API, not when it comes through the internal socket used by hooks to call into Gitaly.
+	// This test setup however is calling the HookService through the external API.
+	//
+	// For now, include the header so the test runs. In the longer term, we should stop serving HookService on the external
+	// socket given it is a service intended only to be used internally by Gitaly for hook callbacks.
+	//
+	// Related issue: https://gitlab.com/gitlab-org/gitaly/-/issues/3746
+	ctx = metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath, repo.RelativePath)
+
 	stream, err := client.PreReceiveHook(ctx)
 	require.NoError(t, err)
 	require.NoError(t, stream.Send(&req))
diff --git a/internal/gitaly/service/repository/size_test.go b/internal/gitaly/service/repository/size_test.go
index 0e0ef55d4..d955d33ad 100644
--- a/internal/gitaly/service/repository/size_test.go
+++ b/internal/gitaly/service/repository/size_test.go
@@ -13,10 +13,12 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/quarantine"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/config"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagemgr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/helper/perm"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/structerr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -151,6 +153,18 @@ func TestGetObjectDirectorySize_successful(t *testing.T) {
 	repo, repoPath := gittest.CreateRepository(t, ctx, cfg)
 	repo.GitObjectDirectory = "objects/"
 
+	// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+	// this is the snapshot's relative path. Include the metadata in the test as well as we're testing requests with quarantine
+	// as if they were coming from access checks. The RPC is also a special case as it only works with a quarantine set.
+	//
+	// Related issue: https://gitlab.com/gitlab-org/gitaly/-/issues/5710
+	ctx = metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath,
+		// Gitaly sends the snapshot's relative path to Rails from `pre-receive` and Rails
+		// sends it back to Gitaly when it performs requests in the access checks. The repository
+		// would have already been rewritten by Praefect, so we have to adjust for that as well.
+		gittest.RewrittenRepository(t, ctx, cfg, repo).RelativePath,
+	)
+
 	// Initially, the object directory should be empty and thus have a size of zero.
 	requireObjectDirectorySize(t, ctx, client, repo, 0)
 
@@ -171,6 +185,17 @@ func TestGetObjectDirectorySize_quarantine(t *testing.T) {
 		repo, repoPath := gittest.CreateRepository(t, ctx, cfg)
 		repo.GitObjectDirectory = "objects/"
 		gittest.WriteBlob(t, cfg, repoPath, uncompressibleData(16*1024))
+
+		// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+		// this is the snapshot's relative path. Include the metadata in the test as well as we're testing requests with quarantine
+		// as if they were coming from access checks. The RPC is also a special case as it only works with a quarantine set.
+		ctx := metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath,
+			// Gitaly sends the snapshot's relative path to Rails from `pre-receive` and Rails
+			// sends it back to Gitaly when it performs requests in the access checks. The repository
+			// would have already been rewritten by Praefect, so we have to adjust for that as well.
+			gittest.RewrittenRepository(t, ctx, cfg, repo).RelativePath,
+		)
+
 		requireObjectDirectorySize(t, ctx, client, repo, 16)
 
 		quarantine, err := quarantine.New(ctx, gittest.RewrittenRepository(t, ctx, cfg, repo), logger, locator)
@@ -211,6 +236,16 @@ func TestGetObjectDirectorySize_quarantine(t *testing.T) {
 		// it back through the API.
 		repo.RelativePath = repo1.RelativePath
 
+		// Rails sends the repository's relative path from the access checks as provided by Gitaly. If transactions are enabled,
+		// this is the snapshot's relative path. Include the metadata in the test as well as we're testing requests with quarantine
+		// as if they were coming from access checks. The RPC is also a special case as it only works with a quarantine set.
+		ctx := metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath,
+			// Gitaly sends the snapshot's relative path to Rails from `pre-receive` and Rails
+			// sends it back to Gitaly when it performs requests in the access checks. The repository
+			// would have already been rewritten by Praefect, so we have to adjust for that as well.
+			gittest.RewrittenRepository(t, ctx, cfg, repo).RelativePath,
+		)
+
 		response, err := client.GetObjectDirectorySize(ctx, &gitalypb.GetObjectDirectorySizeRequest{
 			Repository: repo,
 		})
diff --git a/internal/gitaly/storage/storagemgr/middleware.go b/internal/gitaly/storage/storagemgr/middleware.go
index 2dc9eed3b..21fbd6d81 100644
--- a/internal/gitaly/storage/storagemgr/middleware.go
+++ b/internal/gitaly/storage/storagemgr/middleware.go
@@ -7,15 +7,26 @@ import (
 
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagectx"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/grpc/metadata"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/grpc/protoregistry"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/log"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/structerr"
+	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
 )
 
-// ErrQuarantineConfiguredOnMutator is returned when a mutator request is received with a quarantine configured.
-var ErrQuarantineConfiguredOnMutator = errors.New("quarantine configured on a mutator request")
+// MetadataKeySnapshotRelativePath is the header key that contains the snapshot's relative path. Rails relays
+// the relative path in the RPC calls performed as part of access checks.
+const MetadataKeySnapshotRelativePath = "relative-path-bin"
+
+var (
+	// ErrQuarantineConfiguredOnMutator is returned when a mutator request is received with a quarantine configured.
+	ErrQuarantineConfiguredOnMutator = errors.New("quarantine configured on a mutator request")
+	// ErrQuarantineWithoutSnapshotRelativePath is returned when a request is configured with a quarantine but the snapshot's
+	// relative path was not sent in a header.
+	ErrQuarantineWithoutSnapshotRelativePath = errors.New("quarantined request did not contain snapshot relative path")
+)
 
 // NonTransactionalRPCs are the RPCs that do not support transactions.
 var NonTransactionalRPCs = map[string]struct{}{
@@ -223,10 +234,6 @@ func transactionalizeRequest(ctx context.Context, logger log.Logger, txRegistry
 		// object directory. This allows for circumventing the transaction management by configuring the either
 		// of the object directories. We'll leave this unaddressed for now and later address this by removing
 		// the options to configure object directories and alternates in a request.
-		//
-		// The relative path in quarantined requests is currently still pointing to the original repository.
-		// https://gitlab.com/gitlab-org/gitaly/-/issues/5483 tracks having Rails send the snapshot's relative
-		// path instead.
 
 		if methodInfo.Operation == protoregistry.OpMutator {
 			// Accessor requests may come with quarantine configured from Rails' access checks. Since the
@@ -236,7 +243,12 @@ func transactionalizeRequest(ctx context.Context, logger log.Logger, txRegistry
 			return transactionalizedRequest{}, ErrQuarantineConfiguredOnMutator
 		}
 
-		return nonTransactionalRequest(ctx, req), nil
+		rewrittenReq, err := restoreSnapshotRelativePath(ctx, methodInfo, req)
+		if err != nil {
+			return transactionalizedRequest{}, fmt.Errorf("restore snapshot relative path: %w", err)
+		}
+
+		return nonTransactionalRequest(ctx, rewrittenReq), nil
 	}
 
 	// While the PartitionManager already verifies the repository's storage and relative path, it does not
@@ -307,12 +319,41 @@ func transactionalizeRequest(ctx context.Context, logger log.Logger, txRegistry
 	}, nil
 }
 
-func rewriteRequest(tx *finalizableTransaction, methodInfo protoregistry.MethodInfo, req proto.Message) (proto.Message, error) {
+func rewritableRequest(methodInfo protoregistry.MethodInfo, req proto.Message) (proto.Message, *gitalypb.Repository, error) {
 	// Clone the request in order to not rewrite the request in the earlier interceptors.
 	rewrittenReq := proto.Clone(req)
 	targetRepo, err := methodInfo.TargetRepo(rewrittenReq)
 	if err != nil {
-		return nil, fmt.Errorf("extract target repository: %w", err)
+		return nil, nil, fmt.Errorf("extract target repository: %w", err)
+	}
+
+	return rewrittenReq, targetRepo, nil
+}
+
+func restoreSnapshotRelativePath(ctx context.Context, methodInfo protoregistry.MethodInfo, req proto.Message) (proto.Message, error) {
+	// Rails sends RPCs from its access checks with quarantine applied. The quarantine paths are relative to the
+	// snapshot repository of the original transaction. While the relative path of the request is the original relative
+	// path of the repository before snapshotting, Rails sends the snapshot repository's path in a header. For the
+	// quarantine paths to apply correctly, we must thus rewrite the request to point to the snapshot repository here.
+	snapshotRelativePath := metadata.GetValue(ctx, MetadataKeySnapshotRelativePath)
+	if snapshotRelativePath == "" {
+		return nil, ErrQuarantineWithoutSnapshotRelativePath
+	}
+
+	rewrittenReq, targetRepo, err := rewritableRequest(methodInfo, req)
+	if err != nil {
+		return nil, fmt.Errorf("rewritable request: %w", err)
+	}
+
+	targetRepo.RelativePath = snapshotRelativePath
+
+	return rewrittenReq, nil
+}
+
+func rewriteRequest(tx *finalizableTransaction, methodInfo protoregistry.MethodInfo, req proto.Message) (proto.Message, error) {
+	rewrittenReq, targetRepo, err := rewritableRequest(methodInfo, req)
+	if err != nil {
+		return nil, fmt.Errorf("rewritable request: %w", err)
 	}
 
 	*targetRepo = *tx.RewriteRepository(targetRepo)
diff --git a/internal/gitaly/storage/storagemgr/middleware_test.go b/internal/gitaly/storage/storagemgr/middleware_test.go
index c7203d437..d731e389b 100644
--- a/internal/gitaly/storage/storagemgr/middleware_test.go
+++ b/internal/gitaly/storage/storagemgr/middleware_test.go
@@ -23,7 +23,10 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper/testserver"
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/health/grpc_health_v1"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -562,8 +565,9 @@ messages and behavior by erroring out the requests before they even hit this int
 	}
 
 	for _, tc := range []struct {
-		desc           string
-		performRequest func(*testing.T, context.Context, *grpc.ClientConn)
+		desc               string
+		performRequest     func(*testing.T, context.Context, *grpc.ClientConn)
+		expectedRepository *gitalypb.Repository
 	}{
 		{
 			desc: "health service",
@@ -574,8 +578,42 @@ messages and behavior by erroring out the requests before they even hit this int
 			},
 		},
 		{
+			desc: "repository with object directory missing snapshot relative path header",
+			performRequest: func(t *testing.T, ctx context.Context, cc *grpc.ClientConn) {
+				resp, err := gitalypb.NewRepositoryServiceClient(cc).ObjectFormat(ctx, &gitalypb.ObjectFormatRequest{
+					Repository: &gitalypb.Repository{
+						StorageName:        "default",
+						GitObjectDirectory: "non-default",
+					},
+				})
+				testhelper.RequireGrpcError(t,
+					status.Error(codes.Internal, "restore snapshot relative path: "+storagemgr.ErrQuarantineWithoutSnapshotRelativePath.Error()),
+					err,
+				)
+				require.Nil(t, resp)
+			},
+		},
+		{
+			desc: "repository with alternate object directory missing snapshot relative path header",
+			performRequest: func(t *testing.T, ctx context.Context, cc *grpc.ClientConn) {
+				resp, err := gitalypb.NewRepositoryServiceClient(cc).ObjectFormat(ctx, &gitalypb.ObjectFormatRequest{
+					Repository: &gitalypb.Repository{
+						StorageName:                   "default",
+						GitAlternateObjectDirectories: []string{"non-default"},
+					},
+				})
+				testhelper.RequireGrpcError(t,
+					status.Error(codes.Internal, "restore snapshot relative path: "+storagemgr.ErrQuarantineWithoutSnapshotRelativePath.Error()),
+					err,
+				)
+				require.Nil(t, resp)
+			},
+		},
+		{
 			desc: "repository with object directory does not start a transaction",
 			performRequest: func(t *testing.T, ctx context.Context, cc *grpc.ClientConn) {
+				ctx = metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath, "snapshot-relative-path")
+
 				resp, err := gitalypb.NewRepositoryServiceClient(cc).ObjectFormat(ctx, &gitalypb.ObjectFormatRequest{
 					Repository: &gitalypb.Repository{
 						StorageName:        "default",
@@ -585,10 +623,17 @@ messages and behavior by erroring out the requests before they even hit this int
 				require.NoError(t, err)
 				testhelper.ProtoEqual(t, &gitalypb.ObjectFormatResponse{}, resp)
 			},
+			expectedRepository: &gitalypb.Repository{
+				StorageName:        "default",
+				RelativePath:       "snapshot-relative-path",
+				GitObjectDirectory: "non-default",
+			},
 		},
 		{
 			desc: "repository with alternate object directory does not start a transaction",
 			performRequest: func(t *testing.T, ctx context.Context, cc *grpc.ClientConn) {
+				ctx = metadata.AppendToOutgoingContext(ctx, storagemgr.MetadataKeySnapshotRelativePath, "snapshot-relative-path")
+
 				resp, err := gitalypb.NewRepositoryServiceClient(cc).ObjectFormat(ctx, &gitalypb.ObjectFormatRequest{
 					Repository: &gitalypb.Repository{
 						StorageName:                   "default",
@@ -598,6 +643,11 @@ messages and behavior by erroring out the requests before they even hit this int
 				require.NoError(t, err)
 				testhelper.ProtoEqual(t, &gitalypb.ObjectFormatResponse{}, resp)
 			},
+			expectedRepository: &gitalypb.Repository{
+				StorageName:                   "default",
+				RelativePath:                  "snapshot-relative-path",
+				GitAlternateObjectDirectories: []string{"non-default"},
+			},
 		},
 		{
 			desc: "streaming rpc without first message",
@@ -642,6 +692,7 @@ messages and behavior by erroring out the requests before they even hit this int
 				gitalypb.RegisterRepositoryServiceServer(server, mockRepositoryService{
 					objectFormatFunc: func(ctx context.Context, req *gitalypb.ObjectFormatRequest) (*gitalypb.ObjectFormatResponse, error) {
 						assertHandler(ctx)
+						testhelper.ProtoEqual(t, tc.expectedRepository, req.Repository)
 						return &gitalypb.ObjectFormatResponse{}, nil
 					},
 					setCustomHooksFunc: func(stream gitalypb.RepositoryService_SetCustomHooksServer) error {