diff --git a/common/registry/implementation_set.go b/common/registry/implementation_set.go index 7b509620d..e2958e2a4 100644 --- a/common/registry/implementation_set.go +++ b/common/registry/implementation_set.go @@ -1,6 +1,9 @@ package registry -import "github.com/v2fly/v2ray-core/v4/common/protoext" +import ( + "github.com/golang/protobuf/proto" + "github.com/v2fly/v2ray-core/v4/common/protoext" +) //go:generate go run github.com/v2fly/v2ray-core/v4/common/errors/errorgen @@ -8,12 +11,15 @@ type implementationSet struct { AliasLookup map[string]*implementation } +type CustomLoader func(data []byte, dataType string) (proto.Message, error) + type implementation struct { FullName string Alias []string + Loader CustomLoader } -func (i *implementationSet) RegisterImplementation(name string, opt *protoext.MessageOpt) { +func (i *implementationSet) RegisterImplementation(name string, opt *protoext.MessageOpt, loader CustomLoader) { alias := opt.GetShortName() impl := &implementation{ @@ -26,12 +32,12 @@ func (i *implementationSet) RegisterImplementation(name string, opt *protoext.Me } } -func (i *implementationSet) FindImplementationByAlias(alias string) (string, error) { +func (i *implementationSet) FindImplementationByAlias(alias string) (string, CustomLoader, error) { impl, found := i.AliasLookup[alias] if found { - return impl.FullName, nil + return impl.FullName, impl.Loader, nil } - return "", newError("cannot find implementation by alias") + return "", nil, newError("cannot find implementation by alias") } func newImplementationSet() *implementationSet { diff --git a/common/registry/registry.go b/common/registry/registry.go index 28ffa9038..6d721654a 100644 --- a/common/registry/registry.go +++ b/common/registry/registry.go @@ -9,20 +9,20 @@ type implementationRegistry struct { implSet map[string]*implementationSet } -func (i *implementationRegistry) RegisterImplementation(name string, opt *protoext.MessageOpt) { +func (i *implementationRegistry) RegisterImplementation(name string, opt *protoext.MessageOpt, loader CustomLoader) { interfaceType := opt.GetType()[0] implSet, found := i.implSet[interfaceType] if !found { implSet = newImplementationSet() i.implSet[interfaceType] = implSet } - implSet.RegisterImplementation(name, opt) + implSet.RegisterImplementation(name, opt, loader) } -func (i *implementationRegistry) FindImplementationByAlias(interfaceType, alias string) (string, error) { +func (i *implementationRegistry) FindImplementationByAlias(interfaceType, alias string) (string, CustomLoader, error) { implSet, found := i.implSet[interfaceType] if !found { - return "", newError("cannot find implemention unknown interface type") + return "", nil, newError("cannot find implemention unknown interface type") } return implSet.FindImplementationByAlias(alias) } @@ -33,17 +33,19 @@ func newImplementationRegistry() *implementationRegistry { var globalImplementationRegistry = newImplementationRegistry() -func RegisterImplementation(proto proto.Message) error { +// RegisterImplementation register an implementation of a type of interface +// loader(CustomLoader) is a private API, its interface is subject to breaking changes +func RegisterImplementation(proto proto.Message, loader CustomLoader) error { msgDesc := proto.ProtoReflect().Type().Descriptor() fullName := string(msgDesc.FullName()) msgOpts, err := protoext.GetMessageOptions(msgDesc) if err != nil { return newError("unable to find message options").Base(err) } - globalImplementationRegistry.RegisterImplementation(fullName, msgOpts) + globalImplementationRegistry.RegisterImplementation(fullName, msgOpts, loader) return nil } -func FindImplementationByAlias(interfaceType, alias string) (string, error) { +func FindImplementationByAlias(interfaceType, alias string) (string, CustomLoader, error) { return globalImplementationRegistry.FindImplementationByAlias(interfaceType, alias) }