diff options
author | Paul Okstad <pokstad@gitlab.com> | 2019-07-20 18:32:38 +0300 |
---|---|---|
committer | John Cai <jcai@gitlab.com> | 2019-07-20 18:32:38 +0300 |
commit | a6b17e3736c149f716b82c65058929b999269707 (patch) | |
tree | 7aac5b20cffe1f027a53eb1a9f1e6f283ef2a6ee /internal/praefect/protoregistry | |
parent | b566483bda61ae48c58efc40606036276edc3df6 (diff) |
Gitaly proto method request factories
Diffstat (limited to 'internal/praefect/protoregistry')
-rw-r--r-- | internal/praefect/protoregistry/protoregistry.go | 55 | ||||
-rw-r--r-- | internal/praefect/protoregistry/protoregistry_test.go | 14 |
2 files changed, 62 insertions, 7 deletions
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) +} |