add extension type register function

.gitignore

@@ -1,3 +1,4 @@
 /bin
 /coverage.txt
 /dist
+/.vscode

go.mod

@@ -3,7 +3,7 @@ module github.com/pseudomuto/protokit
 go 1.17
 
 require (
-	github.com/golang/protobuf v1.0.0
+	github.com/golang/protobuf v1.5.2
 	github.com/stretchr/testify v1.2.1
 	google.golang.org/genproto v0.0.0-20180427144745-86e600f69ee4
 )
@@ -12,4 +12,5 @@ require (
 	github.com/davecgh/go-spew v1.1.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c // indirect
+	google.golang.org/protobuf v1.28.0
 )

go.sum

@@ -2,11 +2,20 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/golang/protobuf v1.0.0 h1:lsek0oXi8iFE9L+EXARyHIjU5rlWIhhTkjDz3vHhWWQ=
 github.com/golang/protobuf v1.0.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
+github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
+github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
+github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
+github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.2.1 h1:52QO5WkIUcHGIR7EnGagH88x1bUzqGXTC5/1bDTUQ7U=
 github.com/stretchr/testify v1.2.1/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/genproto v0.0.0-20180427144745-86e600f69ee4 h1:0rk3/gV3HbvCeUzVMhdxV3TEVKMVPDnayjN7sYRmcxY=
 google.golang.org/genproto v0.0.0-20180427144745-86e600f69ee4/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
+google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
+google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
+google.golang.org/protobuf v1.28.0 h1:w43yiav+6bVFTBQFZX0r7ipe9JQ1QsbMgHwbBziscLw=
+google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=

utils/protobuf.go

@@ -1,9 +1,16 @@
 package utils
 
 import (
-	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/protoc-gen-go/descriptor"
-	"github.com/golang/protobuf/protoc-gen-go/plugin"
+	plugin_go "github.com/golang/protobuf/protoc-gen-go/plugin"
+
+	"google.golang.org/protobuf/reflect/protodesc"
+	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/protobuf/reflect/protoregistry"
+
+	"google.golang.org/protobuf/types/descriptorpb"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/dynamicpb"
 
 	"errors"
 	"io/ioutil"
@@ -12,8 +19,9 @@ import (
 
 // CreateGenRequest creates a codegen request from a `FileDescriptorSet`
 func CreateGenRequest(fds *descriptor.FileDescriptorSet, filesToGen ...string) *plugin_go.CodeGeneratorRequest {
+	files := RegisterExtensions(fds)
 	req := new(plugin_go.CodeGeneratorRequest)
-	req.ProtoFile = fds.GetFile()
+	req.ProtoFile = files
 
 	for _, f := range req.GetProtoFile() {
 		if InStringSlice(filesToGen, f.GetName()) {
@@ -87,3 +95,46 @@ func LoadDescriptor(name string, pathSegments ...string) (*descriptor.FileDescri
 
 	return nil, errors.New("FileDescriptor not found")
 }
+
+func RegisterExtensions(fds *descriptor.FileDescriptorSet) []*descriptorpb.FileDescriptorProto{
+	extTypes := new(protoregistry.Types)
+
+	files,err := protodesc.NewFiles(fds)
+	if err != nil {
+		panic(err)
+	}
+
+	files.RangeFiles(func(fileDescriptor protoreflect.FileDescriptor) bool {
+		registerAllExtensions(extTypes, fileDescriptor)
+		return true
+	} )
+
+	new_files := make([]*descriptorpb.FileDescriptorProto,0)
+	for _,fd := range fds.GetFile() {
+		b, _ := proto.Marshal(fd)
+		err := proto.UnmarshalOptions{Resolver: extTypes}.Unmarshal(b, fd)
+		if err != nil {
+			panic(err)
+		}
+		new_files = append(new_files,fd)
+	} 
+
+	return new_files
+}
+
+func registerAllExtensions(extTypes *protoregistry.Types, descs interface {
+	Messages() protoreflect.MessageDescriptors
+	Extensions() protoreflect.ExtensionDescriptors
+}) error {
+	mds := descs.Messages()
+	for i := 0; i < mds.Len(); i++ {
+		registerAllExtensions(extTypes, mds.Get(i))
+	}
+	xds := descs.Extensions()
+	for i := 0; i < xds.Len(); i++ {
+		if err := extTypes.RegisterExtension(dynamicpb.NewExtensionType(xds.Get(i))); err != nil {
+			return err
+		}
+	}
+	return nil
+}