github.com/xmplusdev/xmcore@v1.8.11-0.20240412132628-5518b55526af/app/router/command/command.go (about)

     1  package command
     2  
     3  //go:generate go run github.com/xmplusdev/xmcore/common/errors/errorgen
     4  
     5  import (
     6  	"context"
     7  	"time"
     8  
     9  	"github.com/xmplusdev/xmcore/common"
    10  	"github.com/xmplusdev/xmcore/core"
    11  	"github.com/xmplusdev/xmcore/features/routing"
    12  	"github.com/xmplusdev/xmcore/features/stats"
    13  	"google.golang.org/grpc"
    14  )
    15  
    16  // routingServer is an implementation of RoutingService.
    17  type routingServer struct {
    18  	router       routing.Router
    19  	routingStats stats.Channel
    20  }
    21  
    22  func (s *routingServer) GetBalancerInfo(ctx context.Context, request *GetBalancerInfoRequest) (*GetBalancerInfoResponse, error) {
    23  	var ret GetBalancerInfoResponse
    24  	ret.Balancer = &BalancerMsg{}
    25  	if bo, ok := s.router.(routing.BalancerOverrider); ok {
    26  		{
    27  			res, err := bo.GetOverrideTarget(request.GetTag())
    28  			if err != nil {
    29  				return nil, err
    30  			}
    31  			ret.Balancer.Override = &OverrideInfo{
    32  				Target: res,
    33  			}
    34  		}
    35  	}
    36  
    37  	if pt, ok := s.router.(routing.BalancerPrincipleTarget); ok {
    38  		{
    39  			res, err := pt.GetPrincipleTarget(request.GetTag())
    40  			if err != nil {
    41  				newError("unable to obtain principle target").Base(err).AtInfo().WriteToLog()
    42  			} else {
    43  				ret.Balancer.PrincipleTarget = &PrincipleTargetInfo{Tag: res}
    44  			}
    45  		}
    46  	}
    47  	return &ret, nil
    48  }
    49  
    50  func (s *routingServer) OverrideBalancerTarget(ctx context.Context, request *OverrideBalancerTargetRequest) (*OverrideBalancerTargetResponse, error) {
    51  	if bo, ok := s.router.(routing.BalancerOverrider); ok {
    52  		return &OverrideBalancerTargetResponse{}, bo.SetOverrideTarget(request.BalancerTag, request.Target)
    53  	}
    54  	return nil, newError("unsupported router implementation")
    55  }
    56  
    57  func (s *routingServer) AddRule(ctx context.Context, request *AddRuleRequest) (*AddRuleResponse, error) {
    58  	if bo, ok := s.router.(routing.Router); ok {
    59  		return &AddRuleResponse{}, bo.AddRule(request.Config, request.ShouldAppend)
    60  	}
    61  	return nil, newError("unsupported router implementation")
    62  
    63  }
    64  func (s *routingServer) RemoveRule(ctx context.Context, request *RemoveRuleRequest) (*RemoveRuleResponse, error) {
    65  	if bo, ok := s.router.(routing.Router); ok {
    66  		return &RemoveRuleResponse{}, bo.RemoveRule(request.RuleTag)
    67  	}
    68  	return nil, newError("unsupported router implementation")
    69  }
    70  
    71  // NewRoutingServer creates a statistics service with statistics manager.
    72  func NewRoutingServer(router routing.Router, routingStats stats.Channel) RoutingServiceServer {
    73  	return &routingServer{
    74  		router:       router,
    75  		routingStats: routingStats,
    76  	}
    77  }
    78  
    79  func (s *routingServer) TestRoute(ctx context.Context, request *TestRouteRequest) (*RoutingContext, error) {
    80  	if request.RoutingContext == nil {
    81  		return nil, newError("Invalid routing request.")
    82  	}
    83  	route, err := s.router.PickRoute(AsRoutingContext(request.RoutingContext))
    84  	if err != nil {
    85  		return nil, err
    86  	}
    87  	if request.PublishResult && s.routingStats != nil {
    88  		ctx, _ := context.WithTimeout(context.Background(), 4*time.Second)
    89  		s.routingStats.Publish(ctx, route)
    90  	}
    91  	return AsProtobufMessage(request.FieldSelectors)(route), nil
    92  }
    93  
    94  func (s *routingServer) SubscribeRoutingStats(request *SubscribeRoutingStatsRequest, stream RoutingService_SubscribeRoutingStatsServer) error {
    95  	if s.routingStats == nil {
    96  		return newError("Routing statistics not enabled.")
    97  	}
    98  	genMessage := AsProtobufMessage(request.FieldSelectors)
    99  	subscriber, err := stats.SubscribeRunnableChannel(s.routingStats)
   100  	if err != nil {
   101  		return err
   102  	}
   103  	defer stats.UnsubscribeClosableChannel(s.routingStats, subscriber)
   104  	for {
   105  		select {
   106  		case value, ok := <-subscriber:
   107  			if !ok {
   108  				return newError("Upstream closed the subscriber channel.")
   109  			}
   110  			route, ok := value.(routing.Route)
   111  			if !ok {
   112  				return newError("Upstream sent malformed statistics.")
   113  			}
   114  			err := stream.Send(genMessage(route))
   115  			if err != nil {
   116  				return err
   117  			}
   118  		case <-stream.Context().Done():
   119  			return stream.Context().Err()
   120  		}
   121  	}
   122  }
   123  
   124  func (s *routingServer) mustEmbedUnimplementedRoutingServiceServer() {}
   125  
   126  type service struct {
   127  	v *core.Instance
   128  }
   129  
   130  func (s *service) Register(server *grpc.Server) {
   131  	common.Must(s.v.RequireFeatures(func(router routing.Router, stats stats.Manager) {
   132  		rs := NewRoutingServer(router, nil)
   133  		RegisterRoutingServiceServer(server, rs)
   134  
   135  		// For compatibility purposes
   136  		vCoreDesc := RoutingService_ServiceDesc
   137  		vCoreDesc.ServiceName = "v2ray.core.app.router.command.RoutingService"
   138  		server.RegisterService(&vCoreDesc, rs)
   139  	}))
   140  }
   141  
   142  func init() {
   143  	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) {
   144  		s := core.MustFromContext(ctx)
   145  		return &service{v: s}, nil
   146  	}))
   147  }