mirror of https://github.com/v2fly/v2ray-core.git
104 lines
2.6 KiB
Go
104 lines
2.6 KiB
Go
//go:build !confonly
|
|
// +build !confonly
|
|
|
|
package router
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/v2fly/v2ray-core/v5/features/extension"
|
|
"github.com/v2fly/v2ray-core/v5/features/outbound"
|
|
)
|
|
|
|
type BalancingStrategy interface {
|
|
PickOutbound([]string) string
|
|
}
|
|
|
|
type BalancingPrincipleTarget interface {
|
|
GetPrincipleTarget([]string) []string
|
|
}
|
|
|
|
type Balancer struct {
|
|
selectors []string
|
|
strategy BalancingStrategy
|
|
ohm outbound.Manager
|
|
fallbackTag string
|
|
|
|
override override
|
|
}
|
|
|
|
// PickOutbound picks the tag of an 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")
|
|
}
|