Praefect: storage scoped RPCs not handled properly

Scope of the FindRemoteRepository RPC call changed to STORAGE.
Storage value of the incoming message can be changed by new
'SetStorage' method.
Human-readable string for 'Scope' type with default string formatting.

Closes: https://gitlab.com/gitlab-org/gitaly/-/issues/2442

3 files changed, 80 insertions, 0 deletions

diff --git a/internal/praefect/protoregistry/find_oid.go b/internal/praefect/protoregistry/find_oid.go
index 5a110c93c..abc87ccb0 100644
--- a/internal/praefect/protoregistry/find_oid.go
+++ b/internal/praefect/protoregistry/find_oid.go
@@ -50,6 +50,16 @@ func reflectFindStorage(pbMsg proto.Message, targetOID []int) (string, error) {
 	return targetRepo, nil
 }
 
+func reflectSetStorage(pbMsg proto.Message, targetOID []int, storage string) error {
+	msgV, err := reflectFindOID(pbMsg, targetOID)
+	if err != nil {
+		return err
+	}
+
+	msgV.Set(reflect.ValueOf(storage))
+	return nil
+}
+
 // ErrProtoFieldEmpty indicates the protobuf field is empty
 var ErrProtoFieldEmpty = errors.New("proto field is empty")
 
diff --git a/internal/praefect/protoregistry/find_oid_test.go b/internal/praefect/protoregistry/find_oid_test.go
index 4a87db9c7..a040bb17b 100644
--- a/internal/praefect/protoregistry/find_oid_test.go
+++ b/internal/praefect/protoregistry/find_oid_test.go
@@ -159,3 +159,48 @@ func TestProtoRegistryStorage(t *testing.T) {
 		})
 	}
 }
+
+func TestMethodInfo_SetStorage(t *testing.T) {
+	testCases := []struct {
+		desc      string
+		service   string
+		method    string
+		pbMsg     proto.Message
+		storage   string
+		expectErr error
+	}{
+		{
+			desc:    "valid request type",
+			service: "NamespaceService",
+			method:  "AddNamespace",
+			pbMsg: &gitalypb.AddNamespaceRequest{
+				StorageName: "old_storage",
+			},
+			storage: "new_storage",
+		},
+		{
+			desc:      "incorrect request type",
+			service:   "RepositoryService",
+			method:    "RepackIncremental",
+			pbMsg:     &gitalypb.RepackIncrementalResponse{},
+			expectErr: errors.New("proto message gitaly.RepackIncrementalResponse does not match expected RPC request message gitaly.RepackIncrementalRequest"),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			info, err := protoregistry.GitalyProtoPreregistered.LookupMethod("/gitaly." + tc.service + "/" + tc.method)
+			require.NoError(t, err)
+
+			err = info.SetStorage(tc.pbMsg, tc.storage)
+			if tc.expectErr == nil {
+				require.NoError(t, err)
+				changed, err := info.Storage(tc.pbMsg)
+				require.NoError(t, err)
+				require.Equal(t, tc.storage, changed)
+			} else {
+				require.Equal(t, tc.expectErr, err)
+			}
+		})
+	}
+}
diff --git a/internal/praefect/protoregistry/protoregistry.go b/internal/praefect/protoregistry/protoregistry.go
index cabc38b19..7acc94d03 100644
--- a/internal/praefect/protoregistry/protoregistry.go
+++ b/internal/praefect/protoregistry/protoregistry.go
@@ -67,6 +67,19 @@ const (
 	ScopeServer
 )
 
+func (s Scope) String() string {
+	switch s {
+	case ScopeServer:
+		return "server"
+	case ScopeStorage:
+		return "storage"
+	case ScopeRepository:
+		return "repository"
+	default:
+		return fmt.Sprintf("N/A: %d", s)
+	}
+}
+
 var protoScope = map[gitalypb.OperationMsg_Scope]Scope{
 	gitalypb.OperationMsg_SERVER:     ScopeServer,
 	gitalypb.OperationMsg_REPOSITORY: ScopeRepository,
@@ -137,6 +150,18 @@ func (mi MethodInfo) Storage(msg proto.Message) (string, error) {
 	return reflectFindStorage(msg, mi.storage)
 }
 
+// SetStorage sets the storage name for a protobuf message
+func (mi MethodInfo) SetStorage(msg proto.Message, storage string) error {
+	if mi.requestName != proto.MessageName(msg) {
+		return fmt.Errorf(
+			"proto message %s does not match expected RPC request message %s",
+			proto.MessageName(msg), mi.requestName,
+		)
+	}
+
+	return reflectSetStorage(msg, mi.storage, storage)
+}
+
 // UnmarshalRequestProto will unmarshal the bytes into the method's request
 // message type
 func (mi MethodInfo) UnmarshalRequestProto(b []byte) (proto.Message, error) {