github.com/moqsien/xraycore@v1.8.5/app/router/command/command.go (about)

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