github.com/imannamdari/v2ray-core/v5@v5.0.5/common/registry/registry.go (about)

     1  package registry
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"reflect"
     7  	"strings"
     8  	"sync"
     9  
    10  	"github.com/golang/protobuf/jsonpb"
    11  	"github.com/golang/protobuf/proto"
    12  	protov2 "google.golang.org/protobuf/proto"
    13  
    14  	"github.com/imannamdari/v2ray-core/v5/common/protoext"
    15  	"github.com/imannamdari/v2ray-core/v5/common/protofilter"
    16  	"github.com/imannamdari/v2ray-core/v5/common/serial"
    17  )
    18  
    19  type implementationRegistry struct {
    20  	implSet map[string]*implementationSet
    21  }
    22  
    23  func (i *implementationRegistry) RegisterImplementation(name string, opt *protoext.MessageOpt, loader CustomLoader) {
    24  	interfaceType := opt.GetType()
    25  	for _, v := range interfaceType {
    26  		i.registerSingleImplementation(v, name, opt, loader)
    27  	}
    28  }
    29  
    30  func (i *implementationRegistry) registerSingleImplementation(interfaceType, name string, opt *protoext.MessageOpt, loader CustomLoader) {
    31  	implSet, found := i.implSet[interfaceType]
    32  	if !found {
    33  		implSet = newImplementationSet()
    34  		i.implSet[interfaceType] = implSet
    35  	}
    36  	implSet.RegisterImplementation(name, opt, loader)
    37  }
    38  
    39  func (i *implementationRegistry) findImplementationByAlias(interfaceType, alias string) (string, CustomLoader, error) {
    40  	implSet, found := i.implSet[interfaceType]
    41  	if !found {
    42  		return "", nil, newError("cannot find implemention unknown interface type")
    43  	}
    44  	return implSet.findImplementationByAlias(alias)
    45  }
    46  
    47  func (i *implementationRegistry) LoadImplementationByAlias(ctx context.Context, interfaceType, alias string, data []byte) (proto.Message, error) {
    48  	var implementationFullName string
    49  
    50  	if strings.HasPrefix(alias, "#") {
    51  		// skip resolution for full name
    52  		implementationFullName = alias
    53  	} else {
    54  		registryResult, customLoader, err := i.findImplementationByAlias(interfaceType, alias)
    55  		if err != nil {
    56  			return nil, newError("unable to find implementation").Base(err)
    57  		}
    58  		if customLoader != nil {
    59  			return customLoader(data, i)
    60  		}
    61  		implementationFullName = registryResult
    62  	}
    63  	implementationConfigInstance, err := serial.GetInstance(implementationFullName)
    64  	if err != nil {
    65  		return nil, newError("unable to create implementation config instance").Base(err)
    66  	}
    67  
    68  	unmarshaler := jsonpb.Unmarshaler{AllowUnknownFields: false}
    69  	err = unmarshaler.Unmarshal(bytes.NewReader(data), implementationConfigInstance.(proto.Message))
    70  	if err != nil {
    71  		return nil, newError("unable to parse json content").Base(err)
    72  	}
    73  
    74  	implementationConfigInstancev2 := proto.MessageV2(implementationConfigInstance)
    75  	if err := protofilter.FilterProtoConfig(ctx, implementationConfigInstancev2); err != nil {
    76  		return nil, err
    77  	}
    78  
    79  	return implementationConfigInstance.(proto.Message), nil
    80  }
    81  
    82  func newImplementationRegistry() *implementationRegistry {
    83  	return &implementationRegistry{implSet: map[string]*implementationSet{}}
    84  }
    85  
    86  var globalImplementationRegistry = newImplementationRegistry()
    87  
    88  var initialized = &sync.Once{}
    89  
    90  type registerRequest struct {
    91  	proto  interface{}
    92  	loader CustomLoader
    93  }
    94  
    95  var registerRequests []registerRequest
    96  
    97  // RegisterImplementation register an implementation of a type of interface
    98  // loader(CustomLoader) is a private API, its interface is subject to breaking changes
    99  func RegisterImplementation(proto interface{}, loader CustomLoader) error {
   100  	registerRequests = append(registerRequests, registerRequest{
   101  		proto:  proto,
   102  		loader: loader,
   103  	})
   104  	return nil
   105  }
   106  
   107  func registerImplementation(proto interface{}, loader CustomLoader) error {
   108  	protoReflect := reflect.New(reflect.TypeOf(proto).Elem())
   109  	proto2 := protoReflect.Interface().(protov2.Message)
   110  	msgDesc := proto2.ProtoReflect().Descriptor()
   111  	fullName := string(msgDesc.FullName())
   112  	msgOpts, err := protoext.GetMessageOptions(msgDesc)
   113  	if err != nil {
   114  		return newError("unable to find message options").Base(err)
   115  	}
   116  	globalImplementationRegistry.RegisterImplementation(fullName, msgOpts, loader)
   117  	return nil
   118  }
   119  
   120  type LoadByAlias interface {
   121  	LoadImplementationByAlias(ctx context.Context, interfaceType, alias string, data []byte) (proto.Message, error)
   122  }
   123  
   124  func LoadImplementationByAlias(ctx context.Context, interfaceType, alias string, data []byte) (proto.Message, error) {
   125  	initialized.Do(func() {
   126  		for _, v := range registerRequests {
   127  			registerImplementation(v.proto, v.loader)
   128  		}
   129  	})
   130  	return globalImplementationRegistry.LoadImplementationByAlias(ctx, interfaceType, alias, data)
   131  }