github.com/imannamdari/v2ray-core/v5@v5.0.5/app/router/weight.go (about)

     1  package router
     2  
     3  import (
     4  	"regexp"
     5  	"strconv"
     6  	"strings"
     7  )
     8  
     9  type weightScaler func(value, weight float64) float64
    10  
    11  var numberFinder = regexp.MustCompile(`\d+(\.\d+)?`)
    12  
    13  // NewWeightManager creates a new WeightManager with settings
    14  func NewWeightManager(s []*StrategyWeight, defaultWeight float64, scaler weightScaler) *WeightManager {
    15  	return &WeightManager{
    16  		settings:      s,
    17  		cache:         make(map[string]float64),
    18  		scaler:        scaler,
    19  		defaultWeight: defaultWeight,
    20  	}
    21  }
    22  
    23  // WeightManager manages weights for specific settings
    24  type WeightManager struct {
    25  	settings      []*StrategyWeight
    26  	cache         map[string]float64
    27  	scaler        weightScaler
    28  	defaultWeight float64
    29  }
    30  
    31  // Get gets the weight of specified tag
    32  func (s *WeightManager) Get(tag string) float64 {
    33  	weight, ok := s.cache[tag]
    34  	if ok {
    35  		return weight
    36  	}
    37  	weight = s.findValue(tag)
    38  	s.cache[tag] = weight
    39  	return weight
    40  }
    41  
    42  // Apply applies weight to the value
    43  func (s *WeightManager) Apply(tag string, value float64) float64 {
    44  	return s.scaler(value, s.Get(tag))
    45  }
    46  
    47  func (s *WeightManager) findValue(tag string) float64 {
    48  	for _, w := range s.settings {
    49  		matched := s.getMatch(tag, w.Match, w.Regexp)
    50  		if matched == "" {
    51  			continue
    52  		}
    53  		if w.Value > 0 {
    54  			return float64(w.Value)
    55  		}
    56  		// auto weight from matched
    57  		numStr := numberFinder.FindString(matched)
    58  		if numStr == "" {
    59  			return s.defaultWeight
    60  		}
    61  		weight, err := strconv.ParseFloat(numStr, 64)
    62  		if err != nil {
    63  			newError("unexpected error from ParseFloat: ", err).AtError().WriteToLog()
    64  			return s.defaultWeight
    65  		}
    66  		return weight
    67  	}
    68  	return s.defaultWeight
    69  }
    70  
    71  func (s *WeightManager) getMatch(tag, find string, isRegexp bool) string {
    72  	if !isRegexp {
    73  		idx := strings.Index(tag, find)
    74  		if idx < 0 {
    75  			return ""
    76  		}
    77  		return find
    78  	}
    79  	r, err := regexp.Compile(find)
    80  	if err != nil {
    81  		newError("invalid regexp: ", find, "err: ", err).AtError().WriteToLog()
    82  		return ""
    83  	}
    84  	return r.FindString(tag)
    85  }