protoregistry: Allow `additional_repository`s to be unset

When fields are tagged with the `additional_repository` extension then
we will always try to extract those fields and rewrite them in Praefect.
If this fails, we return an error to the caller and reject the request.

This is needlessly rigid though: the coordinator really only needs info
about the target repository so that it can properly route the request to
the correct Gitaly node. The case where additional repositories are not
set shouldn't necessarily cause us to fail the request. In fact, there
may be (and indeed are) perfectly valid usecases where you want to only
set a repository field conditionally depending on certain parameters.

Refactor the code to allow the case where additional repositories are
not set. We will simply forward these requests to the corresponding
Gitaly node, which can then verify by itself whether it is okay for the
repository to not be set.

Changelog: changed

3 files changed, 50 insertions, 48 deletions

diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index 5b88eaf87..48fbdd460 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -373,9 +373,15 @@ func (c *Coordinator) mutatorStreamParameters(ctx context.Context, call grpcCall
 	}
 
 	var additionalRepoRelativePath string
-	if additionalRepo, ok, err := call.methodInfo.AdditionalRepo(call.msg); err != nil {
+	if additionalRepo, err := call.methodInfo.AdditionalRepo(call.msg); errors.Is(err, protoregistry.ErrTargetRepoMissing) {
+		// We can land here in two cases: either the message doesn't have an additional
+		// repository, or the repository wasn't set. The former case is obviously fine, but
+		// the latter case is fine, too, given that the additional repository may be an
+		// optional field for some RPC calls. The Gitaly-side RPC handlers should know to
+		// handle this case anyway, so we just leave the field unset in that case.
+	} else if err != nil {
 		return nil, structerr.NewInvalidArgument("%w", err)
-	} else if ok {
+	} else {
 		additionalRepoRelativePath = additionalRepo.GetRelativePath()
 	}
 
@@ -794,12 +800,12 @@ func rewrittenRepositoryMessage(mi protoregistry.MethodInfo, m proto.Message, st
 	targetRepo.StorageName = storage
 	targetRepo.RelativePath = relativePath
 
-	additionalRepo, ok, err := mi.AdditionalRepo(m)
-	if err != nil {
+	if additionalRepo, err := mi.AdditionalRepo(m); errors.Is(err, protoregistry.ErrTargetRepoMissing) {
+		// Nothing to rewrite in case the additional repository either doesn't exist in the
+		// message or wasn't set by the caller.
+	} else if err != nil {
 		return nil, structerr.NewInvalidArgument("%w", err)
-	}
-
-	if ok {
+	} else {
 		additionalRepo.StorageName = storage
 		additionalRepo.RelativePath = additionalRelativePath
 	}
diff --git a/internal/praefect/protoregistry/method_info.go b/internal/praefect/protoregistry/method_info.go
index 3887e51d3..0ae8bda37 100644
--- a/internal/praefect/protoregistry/method_info.go
+++ b/internal/praefect/protoregistry/method_info.go
@@ -79,16 +79,10 @@ func (mi MethodInfo) TargetRepo(msg proto.Message) (*gitalypb.Repository, error)
 	return mi.getRepo(msg, gitalypb.E_TargetRepository)
 }
 
-// AdditionalRepo returns the additional repository for a protobuf message that needs a storage rewritten
-// if it exists
-func (mi MethodInfo) AdditionalRepo(msg proto.Message) (*gitalypb.Repository, bool, error) {
-	if mi.additionalRepo == nil {
-		return nil, false, nil
-	}
-
-	repo, err := mi.getRepo(msg, gitalypb.E_AdditionalRepository)
-
-	return repo, true, err
+// AdditionalRepo returns the additional repository for a Protobuf message that needs a storage
+// rewritten if it exists.
+func (mi MethodInfo) AdditionalRepo(msg proto.Message) (*gitalypb.Repository, error) {
+	return mi.getRepo(msg, gitalypb.E_AdditionalRepository)
 }
 
 //nolint:revive // This is unintentionally missing documentation.
diff --git a/internal/praefect/protoregistry/method_info_test.go b/internal/praefect/protoregistry/method_info_test.go
index d97d171aa..1532a4d99 100644
--- a/internal/praefect/protoregistry/method_info_test.go
+++ b/internal/praefect/protoregistry/method_info_test.go
@@ -40,8 +40,9 @@ func TestMethodInfo_getRepo(t *testing.T) {
 		method               string
 		pbMsg                proto.Message
 		expectRepo           *gitalypb.Repository
-		expectAdditionalRepo *gitalypb.Repository
 		expectErr            error
+		expectAdditionalRepo *gitalypb.Repository
+		expectAdditionalErr  error
 	}{
 		{
 			desc:   "valid request type single depth",
@@ -50,14 +51,16 @@ func TestMethodInfo_getRepo(t *testing.T) {
 			pbMsg: &gitalypb.OptimizeRepositoryRequest{
 				Repository: testRepos[0],
 			},
-			expectRepo: testRepos[0],
+			expectRepo:          testRepos[0],
+			expectAdditionalErr: ErrTargetRepoMissing,
 		},
 		{
-			desc:      "unset oneof",
-			svc:       "OperationService",
-			method:    "UserCommitFiles",
-			pbMsg:     &gitalypb.UserCommitFilesRequest{},
-			expectErr: ErrTargetRepoMissing,
+			desc:                "unset oneof",
+			svc:                 "OperationService",
+			method:              "UserCommitFiles",
+			pbMsg:               &gitalypb.UserCommitFilesRequest{},
+			expectErr:           ErrTargetRepoMissing,
+			expectAdditionalErr: ErrTargetRepoMissing,
 		},
 		{
 			desc:   "unset value in oneof",
@@ -66,7 +69,8 @@ func TestMethodInfo_getRepo(t *testing.T) {
 			pbMsg: &gitalypb.UserCommitFilesRequest{
 				UserCommitFilesRequestPayload: &gitalypb.UserCommitFilesRequest_Header{},
 			},
-			expectErr: ErrTargetRepoMissing,
+			expectErr:           ErrTargetRepoMissing,
+			expectAdditionalErr: ErrTargetRepoMissing,
 		},
 		{
 			desc:   "unset repository in oneof",
@@ -77,7 +81,8 @@ func TestMethodInfo_getRepo(t *testing.T) {
 					Header: &gitalypb.UserCommitFilesRequestHeader{},
 				},
 			},
-			expectErr: ErrTargetRepoMissing,
+			expectErr:           ErrTargetRepoMissing,
+			expectAdditionalErr: ErrTargetRepoMissing,
 		},
 		{
 			desc:   "target nested in oneOf",
@@ -90,7 +95,8 @@ func TestMethodInfo_getRepo(t *testing.T) {
 					},
 				},
 			},
-			expectRepo: testRepos[1],
+			expectRepo:          testRepos[1],
+			expectAdditionalErr: ErrTargetRepoMissing,
 		},
 		{
 			desc:   "target nested, includes additional repository",
@@ -104,35 +110,31 @@ func TestMethodInfo_getRepo(t *testing.T) {
 			expectAdditionalRepo: testRepos[0],
 		},
 		{
-			desc:      "target repo is nil",
-			svc:       "RepositoryService",
-			method:    "OptimizeRepository",
-			pbMsg:     &gitalypb.OptimizeRepositoryRequest{Repository: nil},
-			expectErr: ErrTargetRepoMissing,
+			desc:                "target repo is nil",
+			svc:                 "RepositoryService",
+			method:              "OptimizeRepository",
+			pbMsg:               &gitalypb.OptimizeRepositoryRequest{Repository: nil},
+			expectErr:           ErrTargetRepoMissing,
+			expectAdditionalErr: ErrTargetRepoMissing,
 		},
 	}
 
 	for _, tc := range testcases {
-		desc := fmt.Sprintf("%s:%s %s", tc.svc, tc.method, tc.desc)
-		t.Run(desc, func(t *testing.T) {
+		t.Run(tc.desc, func(t *testing.T) {
 			info, err := GitalyProtoPreregistered.LookupMethod(fmt.Sprintf("/gitaly.%s/%s", tc.svc, tc.method))
 			require.NoError(t, err)
 
-			actualTarget, actualErr := info.TargetRepo(tc.pbMsg)
-			require.Equal(t, tc.expectErr, actualErr)
-
-			// not only do we want the value to be the same, but we actually want the
-			// exact same instance to be returned
-			if tc.expectRepo != actualTarget {
-				t.Fatal("pointers do not match")
-			}
+			t.Run("TargetRepo", func(t *testing.T) {
+				repo, err := info.TargetRepo(tc.pbMsg)
+				require.Equal(t, tc.expectErr, err)
+				require.Same(t, tc.expectRepo, repo)
+			})
 
-			if tc.expectAdditionalRepo != nil {
-				additionalRepo, ok, err := info.AdditionalRepo(tc.pbMsg)
-				require.True(t, ok)
-				require.NoError(t, err)
-				require.Equal(t, tc.expectAdditionalRepo, additionalRepo)
-			}
+			t.Run("AdditionalRepo", func(t *testing.T) {
+				additionalRepo, err := info.AdditionalRepo(tc.pbMsg)
+				require.Equal(t, tc.expectAdditionalErr, err)
+				require.Same(t, tc.expectAdditionalRepo, additionalRepo)
+			})
 		})
 	}
 }