github.com/eagleql/xray-core@v1.4.4/app/router/command/command.go (about)

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