package router import ( "context" sync "sync" "github.com/xtls/xray-core/features/extension" "github.com/xtls/xray-core/features/outbound" ) type BalancingStrategy interface { PickOutbound([]string) string } type BalancingPrincipleTarget interface { GetPrincipleTarget([]string) []string } type RoundRobinStrategy struct { mu sync.Mutex index int } func (s *RoundRobinStrategy) PickOutbound(tags []string) string { n := len(tags) if n == 0 { panic("0 tags") } s.mu.Lock() defer s.mu.Unlock() tag := tags[s.index%n] s.index = (s.index + 1) % n return tag } type Balancer struct { selectors []string strategy BalancingStrategy ohm outbound.Manager fallbackTag string override override } // PickOutbound picks the tag of a outbound func (b *Balancer) PickOutbound() (string, error) { candidates, err := b.SelectOutbounds() if err != nil { if b.fallbackTag != "" { newError("fallback to [", b.fallbackTag, "], due to error: ", err).AtInfo().WriteToLog() return b.fallbackTag, nil } return "", err } var tag string if o := b.override.Get(); o != "" { tag = o } else { tag = b.strategy.PickOutbound(candidates) } if tag == "" { if b.fallbackTag != "" { newError("fallback to [", b.fallbackTag, "], due to empty tag returned").AtInfo().WriteToLog() return b.fallbackTag, nil } // will use default handler return "", newError("balancing strategy returns empty tag") } return tag, nil } func (b *Balancer) InjectContext(ctx context.Context) { if contextReceiver, ok := b.strategy.(extension.ContextReceiver); ok { contextReceiver.InjectContext(ctx) } } // SelectOutbounds select outbounds with selectors of the Balancer func (b *Balancer) SelectOutbounds() ([]string, error) { hs, ok := b.ohm.(outbound.HandlerSelector) if !ok { return nil, newError("outbound.Manager is not a HandlerSelector") } tags := hs.Select(b.selectors) return tags, nil } // GetPrincipleTarget implements routing.BalancerPrincipleTarget func (r *Router) GetPrincipleTarget(tag string) ([]string, error) { if b, ok := r.balancers[tag]; ok { if s, ok := b.strategy.(BalancingPrincipleTarget); ok { candidates, err := b.SelectOutbounds() if err != nil { return nil, newError("unable to select outbounds").Base(err) } return s.GetPrincipleTarget(candidates), nil } return nil, newError("unsupported GetPrincipleTarget") } return nil, newError("cannot find tag") } // SetOverrideTarget implements routing.BalancerOverrider func (r *Router) SetOverrideTarget(tag, target string) error { if b, ok := r.balancers[tag]; ok { b.override.Put(target) return nil } return newError("cannot find tag") } // GetOverrideTarget implements routing.BalancerOverrider func (r *Router) GetOverrideTarget(tag string) (string, error) { if b, ok := r.balancers[tag]; ok { return b.override.Get(), nil } return "", newError("cannot find tag") }