github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/app/router/command/command.go (about)

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