Merge 9ceecf9 into cee3667

app.go

@@ -40,6 +40,7 @@ import (
 	"github.com/topfreegames/pitaya/constants"
 	pcontext "github.com/topfreegames/pitaya/context"
 	"github.com/topfreegames/pitaya/defaultpipelines"
+	"github.com/topfreegames/pitaya/docgenerator"
 	"github.com/topfreegames/pitaya/errors"
 	"github.com/topfreegames/pitaya/internal/codec"
 	"github.com/topfreegames/pitaya/internal/message"
@@ -585,3 +586,8 @@ func AddGRPCInfoToMetadata(
 	metadata[constants.RegionKey] = region
 	return metadata
 }
+
+// Descriptor returns the protobuf message descriptor for a given message name
+func Descriptor(protoName string) ([]byte, error) {
+	return docgenerator.ProtoDescriptors(protoName)
+}

app_test.go

@@ -494,6 +494,16 @@ func TestExtractSpan(t *testing.T) {
 	assert.Equal(t, span.Context(), spanCtx)
 }
 
+func TestDescriptor(t *testing.T) {
+	bts, err := Descriptor("kick.proto")
+	assert.NoError(t, err)
+	assert.NotNil(t, bts)
+
+	bts, err = Descriptor("not_exists.proto")
+	assert.Nil(t, bts)
+	assert.EqualError(t, constants.ErrProtodescriptor, err.Error())
+}
+
 func TestDocumentation(t *testing.T) {
 	doc, err := Documentation(false)
 	assert.NoError(t, err)

constants/errors.go

@@ -54,6 +54,7 @@ var (
 	ErrNonsenseRPC                    = errors.New("you are making a rpc that may be processed locally, either specify a different server type or specify a server id")
 	ErrNotifyOnRequest                = errors.New("tried to notify a request route")
 	ErrOnCloseBackend                 = errors.New("onclose callbacks are not allowed on backend servers")
+	ErrProtodescriptor                = errors.New("failed to get protobuf message descriptor")
 	ErrRPCClientNotInitialized        = errors.New("RPC client is not running")
 	ErrRPCLocal                       = errors.New("RPC must be to a different server type")
 	ErrRPCServerNotInitialized        = errors.New("RPC server is not running")

docgenerator/descriptors.go

@@ -0,0 +1,39 @@
+package docgenerator
+
+import (
+	"reflect"
+	"strings"
+
+	"github.com/gogo/protobuf/proto"
+	"github.com/topfreegames/pitaya/constants"
+)
+
+// ProtoDescriptors returns the descriptor for a given message name or .proto file
+func ProtoDescriptors(protoName string) ([]byte, error) {
+	if strings.HasSuffix(protoName, ".proto") {
+		descriptor := proto.FileDescriptor(protoName)
+		if descriptor == nil {
+			return nil, constants.ErrProtodescriptor
+		}
+		return descriptor, nil
+	}
+
+	protoReflectTypePointer := proto.MessageType(protoName)
+
+	if protoReflectTypePointer == nil {
+		return nil, constants.ErrProtodescriptor
+	}
+
+	protoReflectType := protoReflectTypePointer.Elem()
+	protoValue := reflect.New(protoReflectType)
+	descriptorMethod, ok := protoReflectTypePointer.MethodByName("Descriptor")
+
+	if !ok {
+		return nil, constants.ErrProtodescriptor
+	}
+
+	descriptorValue := descriptorMethod.Func.Call([]reflect.Value{protoValue})
+	protoDescriptor := descriptorValue[0].Bytes()
+
+	return protoDescriptor, nil
+}

docgenerator/descriptors_test.go

@@ -0,0 +1,35 @@
+package docgenerator
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/topfreegames/pitaya/constants"
+	_ "github.com/topfreegames/pitaya/protos"
+)
+
+func TestProtoDescriptors(t *testing.T) {
+	t.Parallel()
+	tables := []struct {
+		name        string
+		messageName string
+		err         error
+	}{
+		{"fail filename", "not_exists.proto", constants.ErrProtodescriptor},
+		{"success filename", "kick.proto", nil},
+		{"success message", "protos.Push", nil},
+		{"fail message", "protos.DoNotExist", constants.ErrProtodescriptor},
+	}
+	for _, table := range tables {
+		t.Run(table.name, func(t *testing.T) {
+			bts, err := ProtoDescriptors(table.messageName)
+			if table.err != nil {
+				assert.EqualError(t, table.err, err.Error())
+				assert.Nil(t, bts)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, bts)
+			}
+		})
+	}
+}