github.com/Uhtred009/v2ray-core-1@v4.31.2+incompatible/app/router/command/command.go (about)

     1  // +build !confonly
     2  
     3  package command
     4  
     5  //go:generate go run v2ray.com/core/common/errors/errorgen
     6  
     7  import (
     8  	"context"
     9  	"time"
    10  
    11  	"google.golang.org/grpc"
    12  
    13  	"v2ray.com/core"
    14  	"v2ray.com/core/common"
    15  	"v2ray.com/core/features/routing"
    16  	"v2ray.com/core/features/stats"
    17  )
    18  
    19  // routingServer is an implementation of RoutingService.
    20  type routingServer struct {
    21  	router       routing.Router
    22  	routingStats stats.Channel
    23  }
    24  
    25  // NewRoutingServer creates a statistics service with statistics manager.
    26  func NewRoutingServer(router routing.Router, routingStats stats.Channel) RoutingServiceServer {
    27  	return &routingServer{
    28  		router:       router,
    29  		routingStats: routingStats,
    30  	}
    31  }
    32  
    33  func (s *routingServer) TestRoute(ctx context.Context, request *TestRouteRequest) (*RoutingContext, error) {
    34  	if request.RoutingContext == nil {
    35  		return nil, newError("Invalid routing request.")
    36  	}
    37  	route, err := s.router.PickRoute(AsRoutingContext(request.RoutingContext))
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  	if request.PublishResult && s.routingStats != nil {
    42  		ctx, _ := context.WithTimeout(context.Background(), 4*time.Second) // nolint: govet
    43  		s.routingStats.Publish(ctx, route)
    44  	}
    45  	return AsProtobufMessage(request.FieldSelectors)(route), nil
    46  }
    47  
    48  func (s *routingServer) SubscribeRoutingStats(request *SubscribeRoutingStatsRequest, stream RoutingService_SubscribeRoutingStatsServer) error {
    49  	if s.routingStats == nil {
    50  		return newError("Routing statistics not enabled.")
    51  	}
    52  	genMessage := AsProtobufMessage(request.FieldSelectors)
    53  	subscriber, err := stats.SubscribeRunnableChannel(s.routingStats)
    54  	if err != nil {
    55  		return err
    56  	}
    57  	defer stats.UnsubscribeClosableChannel(s.routingStats, subscriber) // nolint: errcheck
    58  	for {
    59  		select {
    60  		case value, ok := <-subscriber:
    61  			if !ok {
    62  				return newError("Upstream closed the subscriber channel.")
    63  			}
    64  			route, ok := value.(routing.Route)
    65  			if !ok {
    66  				return newError("Upstream sent malformed statistics.")
    67  			}
    68  			err := stream.Send(genMessage(route))
    69  			if err != nil {
    70  				return err
    71  			}
    72  		case <-stream.Context().Done():
    73  			return stream.Context().Err()
    74  		}
    75  	}
    76  }
    77  
    78  func (s *routingServer) mustEmbedUnimplementedRoutingServiceServer() {}
    79  
    80  type service struct {
    81  	v *core.Instance
    82  }
    83  
    84  func (s *service) Register(server *grpc.Server) {
    85  	common.Must(s.v.RequireFeatures(func(router routing.Router, stats stats.Manager) {
    86  		RegisterRoutingServiceServer(server, NewRoutingServer(router, nil))
    87  	}))
    88  }
    89  
    90  func init() {
    91  	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) {
    92  		s := core.MustFromContext(ctx)
    93  		return &service{v: s}, nil
    94  	}))
    95  }