github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/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/v2fly/v2ray-core/v5/common/protoext"
    15  	"github.com/v2fly/v2ray-core/v5/common/protofilter"
    16  	"github.com/v2fly/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, _ = strings.CutPrefix(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  
    76  	if isRestrictedModeContext(ctx) {
    77  		if err := enforceRestriction(implementationConfigInstancev2); err != nil {
    78  			return nil, err
    79  		}
    80  	}
    81  
    82  	if err := protofilter.FilterProtoConfig(ctx, implementationConfigInstancev2); err != nil {
    83  		return nil, err
    84  	}
    85  
    86  	return implementationConfigInstance.(proto.Message), nil
    87  }
    88  
    89  func newImplementationRegistry() *implementationRegistry {
    90  	return &implementationRegistry{implSet: map[string]*implementationSet{}}
    91  }
    92  
    93  var globalImplementationRegistry = newImplementationRegistry()
    94  
    95  var initialized = &sync.Once{}
    96  
    97  type registerRequest struct {
    98  	proto  interface{}
    99  	loader CustomLoader
   100  }
   101  
   102  var registerRequests []registerRequest
   103  
   104  // RegisterImplementation register an implementation of a type of interface
   105  // loader(CustomLoader) is a private API, its interface is subject to breaking changes
   106  func RegisterImplementation(proto interface{}, loader CustomLoader) error {
   107  	registerRequests = append(registerRequests, registerRequest{
   108  		proto:  proto,
   109  		loader: loader,
   110  	})
   111  	return nil
   112  }
   113  
   114  func registerImplementation(proto interface{}, loader CustomLoader) error {
   115  	protoReflect := reflect.New(reflect.TypeOf(proto).Elem())
   116  	proto2 := protoReflect.Interface().(protov2.Message)
   117  	msgDesc := proto2.ProtoReflect().Descriptor()
   118  	fullName := string(msgDesc.FullName())
   119  	msgOpts, err := protoext.GetMessageOptions(msgDesc)
   120  	if err != nil {
   121  		return newError("unable to find message options").Base(err)
   122  	}
   123  	globalImplementationRegistry.RegisterImplementation(fullName, msgOpts, loader)
   124  	return nil
   125  }
   126  
   127  type LoadByAlias interface {
   128  	LoadImplementationByAlias(ctx context.Context, interfaceType, alias string, data []byte) (proto.Message, error)
   129  }
   130  
   131  func LoadImplementationByAlias(ctx context.Context, interfaceType, alias string, data []byte) (proto.Message, error) {
   132  	initialized.Do(func() {
   133  		for _, v := range registerRequests {
   134  			registerImplementation(v.proto, v.loader)
   135  		}
   136  	})
   137  	return globalImplementationRegistry.LoadImplementationByAlias(ctx, interfaceType, alias, data)
   138  }