Gitaly proto method request factories

3 files changed, 67 insertions, 7 deletions

diff --git a/changelogs/unreleased/po-proto-factories.yml b/changelogs/unreleased/po-proto-factories.yml
new file mode 100644
index 000000000..b22c84a73
--- /dev/null
+++ b/changelogs/unreleased/po-proto-factories.yml
@@ -0,0 +1,5 @@
+---
+title: Gitaly proto method request factories
+merge_request: 1375
+author:
+type: added
diff --git a/internal/praefect/protoregistry/protoregistry.go b/internal/praefect/protoregistry/protoregistry.go
index c0418893d..909b2e2b2 100644
--- a/internal/praefect/protoregistry/protoregistry.go
+++ b/internal/praefect/protoregistry/protoregistry.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"reflect"
 	"strconv"
 	"strings"
 	"sync"
@@ -46,10 +47,10 @@ const (
 // for message type "OperationMsg" shared.proto in gitlab-org/gitaly-proto for
 // more documentation.
 type MethodInfo struct {
-	Operation   OpType
-	targetRepo  []int
-	requestName string // protobuf message name for input type
-
+	Operation      OpType
+	targetRepo     []int
+	requestName    string // protobuf message name for input type
+	requestFactory protoFactory
 }
 
 // TargetRepo returns the target repository for a protobuf message if it exists
@@ -64,6 +65,12 @@ func (mi MethodInfo) TargetRepo(msg proto.Message) (*gitalypb.Repository, error)
 	return reflectFindRepoTarget(msg, mi.targetRepo)
 }
 
+// UnmarshalRequestProto will unmarshal the bytes into the method's request
+// message type
+func (mi MethodInfo) UnmarshalRequestProto(b []byte) (proto.Message, error) {
+	return mi.requestFactory(b)
+}
+
 // Registry contains info about RPC methods
 type Registry struct {
 	sync.RWMutex
@@ -122,6 +129,34 @@ func getOpExtension(m *descriptor.MethodDescriptorProto) (*gitalypb.OperationMsg
 	return opMsg, nil
 }
 
+type protoFactory func([]byte) (proto.Message, error)
+
+func methodReqFactory(method *descriptor.MethodDescriptorProto) (protoFactory, error) {
+	// for some reason, the descriptor prepends a dot not expected in Go
+	inputTypeName := strings.TrimPrefix(method.GetInputType(), ".")
+
+	inputType := proto.MessageType(inputTypeName)
+	if inputType == nil {
+		return nil, fmt.Errorf("no message type found for %s", inputType)
+	}
+
+	f := func(buf []byte) (proto.Message, error) {
+		v := reflect.New(inputType.Elem())
+		pb, ok := v.Interface().(proto.Message)
+		if !ok {
+			return nil, fmt.Errorf("factory function expected protobuf message but got %T", v.Interface())
+		}
+
+		if err := proto.Unmarshal(buf, pb); err != nil {
+			return nil, err
+		}
+
+		return pb, nil
+	}
+
+	return f, nil
+}
+
 func parseMethodInfo(methodDesc *descriptor.MethodDescriptorProto) (MethodInfo, error) {
 	opMsg, err := getOpExtension(methodDesc)
 	if err != nil {
@@ -149,10 +184,16 @@ func parseMethodInfo(methodDesc *descriptor.MethodDescriptorProto) (MethodInfo,
 	// the two copies consistent for comparisons.
 	requestName := strings.TrimLeft(methodDesc.GetInputType(), ".")
 
+	reqFactory, err := methodReqFactory(methodDesc)
+	if err != nil {
+		return MethodInfo{}, err
+	}
+
 	return MethodInfo{
-		Operation:   opCode,
-		targetRepo:  targetRepo,
-		requestName: requestName,
+		Operation:      opCode,
+		targetRepo:     targetRepo,
+		requestName:    requestName,
+		requestFactory: reqFactory,
 	}, nil
 }
 
diff --git a/internal/praefect/protoregistry/protoregistry_test.go b/internal/praefect/protoregistry/protoregistry_test.go
index 3d896d6d7..42671b5a8 100644
--- a/internal/praefect/protoregistry/protoregistry_test.go
+++ b/internal/praefect/protoregistry/protoregistry_test.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/protoregistry"
 )
 
@@ -194,3 +195,16 @@ func TestPopulatesProtoRegistry(t *testing.T) {
 		}
 	}
 }
+
+func TestRequestFactory(t *testing.T) {
+	r := protoregistry.New()
+	require.NoError(t, r.RegisterFiles(protoregistry.GitalyProtoFileDescriptors...))
+
+	mInfo, err := r.LookupMethod("/gitaly.RepositoryService/RepositoryExists")
+	require.NoError(t, err)
+
+	pb, err := mInfo.UnmarshalRequestProto([]byte{})
+	require.NoError(t, err)
+
+	require.Exactly(t, &gitalypb.RepositoryExistsRequest{}, pb)
+}