go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix traffic limiter for CIDR

Browse Source
This commit is contained in:
ginuerzh
2022-11-05 16:10:26 +08:00
parent 2b9ea187b8
commit 84079e32d2
3 changed files with 59 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
limiter/conn/conn.go
View File

@ -238,7 +238,7 @@ func (l *connLimiter) reload(ctx context.Context) error {
ipLimits[key] = NewConnLimitGenerator(limit)
default:
if ip := net.ParseIP(key); ip != nil {
ipLimits[key] = NewConnLimitGenerator(limit)
ipLimits[key] = NewConnLimitSingleGenerator(limit)
break
}
if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {

2
limiter/rate/rate.go
View File

@ -231,7 +231,7 @@ func (l *rateLimiter) reload(ctx context.Context) error {
ipLimits[key] = NewRateLimitGenerator(limit)
default:
if ip := net.ParseIP(key); ip != nil {
ipLimits[key] = NewRateLimitGenerator(limit)
ipLimits[key] = NewRateLimitSingleGenerator(limit)
break
}
if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {

96
limiter/traffic/traffic.go
View File

@ -100,6 +100,8 @@ func LoggerOption(logger logger.Logger) Option {
type trafficLimiter struct {
limits map[string]TrafficLimitGenerator
cidrLimits cidranger.Ranger
inLimits map[string]limiter.Limiter
outLimits map[string]limiter.Limiter
mu sync.Mutex
cancelFunc context.CancelFunc
options options
@ -115,6 +117,8 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
lim := &trafficLimiter{
limits: make(map[string]TrafficLimitGenerator),
cidrLimits: cidranger.NewPCTrieRanger(),
inLimits: make(map[string]limiter.Limiter),
outLimits: make(map[string]limiter.Limiter),
options: options,
cancelFunc: cancel,
}
@ -134,25 +138,6 @@ func (l *trafficLimiter) In(key string) limiter.Limiter {
var lims []limiter.Limiter
if ip := net.ParseIP(key); ip != nil {
found := false
if p := l.limits[key]; p != nil {
if lim := p.In(); lim != nil {
lims = append(lims, lim)
found = true
}
}
if !found {
if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 {
if v, _ := p[0].(*cidrLimitEntry); v != nil {
if lim := v.limit.In(); lim != nil {
lims = append(lims, lim)
}
}
}
}
}
if p := l.limits[ConnLimitKey]; p != nil {
if lim := p.In(); lim != nil {
lims = append(lims, lim)
@ -164,6 +149,31 @@ func (l *trafficLimiter) In(key string) limiter.Limiter {
}
}
// IP limiter
if lim, ok := l.inLimits[key]; ok {
if lim != nil {
lims = append(lims, lim)
}
} else {
if ip := net.ParseIP(key); ip != nil {
if p := l.limits[key]; p != nil {
if lim = p.In(); lim != nil {
lims = append(lims, lim)
}
}
if lim == nil {
if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 {
if v, _ := p[0].(*cidrLimitEntry); v != nil {
if lim = v.limit.In(); lim != nil {
lims = append(lims, lim)
}
}
}
}
l.inLimits[key] = lim
}
}
var lim limiter.Limiter
if len(lims) > 0 {
lim = newLimiterGroup(lims...)
@ -182,25 +192,6 @@ func (l *trafficLimiter) Out(key string) limiter.Limiter {
var lims []limiter.Limiter
if ip := net.ParseIP(key); ip != nil {
found := false
if p := l.limits[key]; p != nil {
if lim := p.Out(); lim != nil {
lims = append(lims, lim)
found = true
}
}
if !found {
if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 {
if v, _ := p[0].(*cidrLimitEntry); v != nil {
if lim := v.limit.Out(); lim != nil {
lims = append(lims, lim)
}
}
}
}
}
if p := l.limits[ConnLimitKey]; p != nil {
if lim := p.Out(); lim != nil {
lims = append(lims, lim)
@ -212,6 +203,31 @@ func (l *trafficLimiter) Out(key string) limiter.Limiter {
}
}
// IP limiter
if lim, ok := l.outLimits[key]; ok {
if lim != nil {
lims = append(lims, lim)
}
} else {
if ip := net.ParseIP(key); ip != nil {
if p := l.limits[key]; p != nil {
if lim = p.Out(); lim != nil {
lims = append(lims, lim)
}
}
if lim == nil {
if p, _ := l.cidrLimits.ContainingNetworks(ip); len(p) > 0 {
if v, _ := p[0].(*cidrLimitEntry); v != nil {
if lim = v.limit.Out(); lim != nil {
lims = append(lims, lim)
}
}
}
}
l.outLimits[key] = lim
}
}
var lim limiter.Limiter
if len(lims) > 0 {
lim = newLimiterGroup(lims...)
@ -268,7 +284,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
limits[key] = NewTrafficLimitGenerator(in, out)
default:
if ip := net.ParseIP(key); ip != nil {
limits[key] = NewTrafficLimitGenerator(in, out)
limits[key] = NewTrafficLimitSingleGenerator(in, out)
break
}
if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
@ -285,6 +301,8 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
l.limits = limits
l.cidrLimits = cidrLimits
l.inLimits = make(map[string]limiter.Limiter)
l.outLimits = make(map[string]limiter.Limiter)
return nil
}