Browse Source

added custom loader for components

Shelikhoo 4 years ago
parent
commit
7e9f4acd9b
2 changed files with 20 additions and 12 deletions
  1. 11 5
      common/registry/implementation_set.go
  2. 9 7
      common/registry/registry.go

+ 11 - 5
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 {

+ 9 - 7
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)
 }