fix traffic limiter

go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/gorilla/websocket v1.5.0
 	github.com/lucas-clemente/quic-go v0.30.0
 	github.com/miekg/dns v1.1.50
+	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/pion/dtls/v2 v2.1.5
 	github.com/pires/go-proxyproto v0.6.2
 	github.com/prometheus/client_golang v1.12.1

go.sum

@@ -280,6 +280,8 @@ github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
 github.com/onsi/ginkgo/v2 v2.2.0 h1:3ZNA3L1c5FYDFTTxbFeVGGD8jYvjYauHD30YgLxVsNI=
 github.com/onsi/ginkgo/v2 v2.2.0/go.mod h1:MEH45j8TBi6u9BMogfbp0stKC5cdGjumZj5Y7AG4VIk=
 github.com/onsi/gomega v1.20.1 h1:PA/3qinGoukvymdIDV8pii6tiZgC8kbmJO6Z5+b002Q=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pelletier/go-toml v1.9.4 h1:tjENF6MfZAg8e4ZmZTeWaWiT2vXtsoO6+iuOjFhECwM=
 github.com/pelletier/go-toml v1.9.4/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
 github.com/pion/dtls/v2 v2.1.5 h1:jlh2vtIyUBShchoTDqpCCqiYCyRFJ/lvf/gQ8TALs+c=

limiter/traffic/generator.go

@@ -4,58 +4,28 @@ import (
 	limiter "github.com/go-gost/core/limiter/traffic"
 )
 
-type TrafficLimitGenerator interface {
+type limitGenerator struct {
-	In() limiter.Limiter
-	Out() limiter.Limiter
-}
-
-type trafficLimitGenerator struct {
 	in  int
 	out int
 }
 
-func NewTrafficLimitGenerator(in, out int) TrafficLimitGenerator {
+func newLimitGenerator(in, out int) *limitGenerator {
-	return &trafficLimitGenerator{
+	return &limitGenerator{
 		in:  in,
 		out: out,
 	}
 }
 
-func (p *trafficLimitGenerator) In() limiter.Limiter {
+func (p *limitGenerator) In() limiter.Limiter {
 	if p == nil || p.in <= 0 {
 		return nil
 	}
 	return NewLimiter(p.in)
 }
 
-func (p *trafficLimitGenerator) Out() limiter.Limiter {
+func (p *limitGenerator) Out() limiter.Limiter {
 	if p == nil || p.out <= 0 {
 		return nil
 	}
 	return NewLimiter(p.out)
 }
-
-type trafficLimitSingleGenerator struct {
-	in  limiter.Limiter
-	out limiter.Limiter
-}
-
-func NewTrafficLimitSingleGenerator(in, out int) TrafficLimitGenerator {
-	p := &trafficLimitSingleGenerator{}
-	if in > 0 {
-		p.in = NewLimiter(in)
-	}
-	if out > 0 {
-		p.out = NewLimiter(out)
-	}
-
-	return p
-}
-
-func (p *trafficLimitSingleGenerator) In() limiter.Limiter {
-	return p.in
-}
-
-func (p *trafficLimitSingleGenerator) Out() limiter.Limiter {
-	return p.out
-}

limiter/traffic/limiter.go

@@ -2,6 +2,9 @@ package traffic
 
 import (
 	"context"
+	"fmt"
+	"sort"
+	"strconv"
 
 	limiter "github.com/go-gost/core/limiter/traffic"
 	"golang.org/x/time/rate"
@@ -26,5 +29,48 @@ func (l *llimiter) Wait(ctx context.Context, n int) int {
 }
 
 func (l *llimiter) Limit() int {
-	return l.limiter.Burst()
+	return int(l.limiter.Limit())
+}
+
+func (l *llimiter) Set(n int) {
+	l.limiter.SetLimit(rate.Limit(n))
+	l.limiter.SetBurst(n)
+}
+
+func (l *llimiter) String() string {
+	return strconv.Itoa(int(l.limiter.Limit()))
+}
+
+type limiterGroup struct {
+	limiters []limiter.Limiter
+}
+
+func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup {
+	sort.Slice(limiters, func(i, j int) bool {
+		return limiters[i].Limit() < limiters[j].Limit()
+	})
+	return &limiterGroup{limiters: limiters}
+}
+
+func (l *limiterGroup) Wait(ctx context.Context, n int) int {
+	for i := range l.limiters {
+		if v := l.limiters[i].Wait(ctx, n); v < n {
+			n = v
+		}
+	}
+	return n
+}
+
+func (l *limiterGroup) Limit() int {
+	if len(l.limiters) == 0 {
+		return 0
+	}
+
+	return l.limiters[0].Limit()
+}
+
+func (l *limiterGroup) Set(n int) {}
+
+func (l *limiterGroup) String() string {
+	return fmt.Sprintf("%v", l.limiters)
 }

limiter/traffic/traffic.go

@@ -5,7 +5,6 @@ import (
 	"context"
 	"io"
 	"net"
-	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -14,6 +13,7 @@ import (
 	limiter "github.com/go-gost/core/limiter/traffic"
 	"github.com/go-gost/core/logger"
 	"github.com/go-gost/x/internal/loader"
+	"github.com/patrickmn/go-cache"
 	"github.com/yl2chen/cidranger"
 )
 
@@ -22,33 +22,10 @@ const (
 	ConnLimitKey   = "$$"
 )
 
-type limiterGroup struct {
+const (
-	limiters []limiter.Limiter
+	defaultExpiration = 15 * time.Second
-}
+	cleanupInterval   = 30 * time.Second
-
+)
-func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup {
-	sort.Slice(limiters, func(i, j int) bool {
-		return limiters[i].Limit() < limiters[j].Limit()
-	})
-	return &limiterGroup{limiters: limiters}
-}
-
-func (l *limiterGroup) Wait(ctx context.Context, n int) int {
-	for i := range l.limiters {
-		if v := l.limiters[i].Wait(ctx, n); v < n {
-			n = v
-		}
-	}
-	return n
-}
-
-func (l *limiterGroup) Limit() int {
-	if len(l.limiters) == 0 {
-		return 0
-	}
-
-	return l.limiters[0].Limit()
-}
 
 type options struct {
 	limits      []string
@@ -97,14 +74,21 @@ func LoggerOption(logger logger.Logger) Option {
 	}
 }
 
+type limitValue struct {
+	in  int
+	out int
+}
+
 type trafficLimiter struct {
-	limits     map[string]TrafficLimitGenerator
+	generators     sync.Map
-	cidrLimits cidranger.Ranger
+	cidrGenerators cidranger.Ranger
-	inLimits   map[string]limiter.Limiter
+	connInLimits   *cache.Cache
-	outLimits  map[string]limiter.Limiter
+	connOutLimits  *cache.Cache
-	mu         sync.Mutex
+	inLimits       *cache.Cache
-	cancelFunc context.CancelFunc
+	outLimits      *cache.Cache
-	options    options
+	mu             sync.RWMutex
+	cancelFunc     context.CancelFunc
+	options        options
 }
 
 func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
@@ -115,12 +99,13 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
 
 	ctx, cancel := context.WithCancel(context.TODO())
 	lim := &trafficLimiter{
-		limits:     make(map[string]TrafficLimitGenerator),
+		cidrGenerators: cidranger.NewPCTrieRanger(),
-		cidrLimits: cidranger.NewPCTrieRanger(),
+		connInLimits:   cache.New(defaultExpiration, cleanupInterval),
-		inLimits:   make(map[string]limiter.Limiter),
+		connOutLimits:  cache.New(defaultExpiration, cleanupInterval),
-		outLimits:  make(map[string]limiter.Limiter),
+		inLimits:       cache.New(defaultExpiration, cleanupInterval),
-		options:    options,
+		outLimits:      cache.New(defaultExpiration, cleanupInterval),
-		cancelFunc: cancel,
+		options:        options,
+		cancelFunc:     cancel,
 	}
 
 	if err := lim.reload(ctx); err != nil {
@@ -132,45 +117,55 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
 	return lim
 }
 
+// In obtains a traffic input limiter based on key.
+// The key should be client connection address.
 func (l *trafficLimiter) In(key string) limiter.Limiter {
-	l.mu.Lock()
-	defer l.mu.Unlock()
-
 	var lims []limiter.Limiter
 
-	if p := l.limits[ConnLimitKey]; p != nil {
+	// service level limiter
-		if lim := p.In(); lim != nil {
+	if lim, ok := l.inLimits.Get(GlobalLimitKey); ok && lim != nil {
-			lims = append(lims, lim)
+		lims = append(lims, lim.(limiter.Limiter))
-		}
 	}
-	if p := l.limits[GlobalLimitKey]; p != nil {
+
-		if lim := p.In(); lim != nil {
+	// connection level limiter
-			lims = append(lims, lim)
+	if lim, ok := l.connInLimits.Get(key); ok {
+		if lim != nil {
+			// cached connection level limiter
+			lims = append(lims, lim.(limiter.Limiter))
+			// reset expiration
+			l.connInLimits.Set(key, lim, defaultExpiration)
+		}
+	} else {
+		// generate a new connection level limiter and cache it
+		if v, ok := l.generators.Load(ConnLimitKey); ok && v != nil {
+			lim := v.(*limitGenerator).In()
+			if lim != nil {
+				lims = append(lims, lim)
+				l.connInLimits.Set(key, lim, defaultExpiration)
+			}
 		}
 	}
 
-	// IP limiter
+	host, _, _ := net.SplitHostPort(key)
-	if lim, ok := l.inLimits[key]; ok {
+	// IP level limiter
+	if lim, ok := l.inLimits.Get(host); ok {
+		// cached IP limiter
 		if lim != nil {
-			lims = append(lims, lim)
+			lims = append(lims, lim.(limiter.Limiter))
 		}
 	} else {
-		if ip := net.ParseIP(key); ip != nil {
+		l.mu.RLock()
-			if p := l.limits[key]; p != nil {
+		ranger := l.cidrGenerators
-				if lim = p.In(); lim != nil {
+		l.mu.RUnlock()
+
+		// CIDR level limiter
+		if p, _ := ranger.ContainingNetworks(net.ParseIP(host)); len(p) > 0 {
+			if v, _ := p[0].(*cidrLimitEntry); v != nil {
+				if lim := v.generator.In(); lim != nil {
 					lims = append(lims, lim)
+					l.inLimits.Set(host, lim, cache.NoExpiration)
 				}
 			}
-			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
 		}
 	}
 
@@ -180,51 +175,61 @@ func (l *trafficLimiter) In(key string) limiter.Limiter {
 	}
 
 	if lim != nil && l.options.logger != nil {
-		l.options.logger.Debugf("input limit for %s: %d", key, lim.Limit())
+		l.options.logger.Debugf("input limit for %s: %s", key, lim)
 	}
 
 	return lim
 }
 
+// Out obtains a traffic output limiter based on key.
+// The key should be client connection address.
 func (l *trafficLimiter) Out(key string) limiter.Limiter {
-	l.mu.Lock()
-	defer l.mu.Unlock()
-
 	var lims []limiter.Limiter
 
-	if p := l.limits[ConnLimitKey]; p != nil {
+	// service level limiter
-		if lim := p.Out(); lim != nil {
+	if lim, ok := l.outLimits.Get(GlobalLimitKey); ok && lim != nil {
-			lims = append(lims, lim)
+		lims = append(lims, lim.(limiter.Limiter))
-		}
 	}
-	if p := l.limits[GlobalLimitKey]; p != nil {
+
-		if lim := p.Out(); lim != nil {
+	// connection level limiter
-			lims = append(lims, lim)
+	if lim, ok := l.connOutLimits.Get(key); ok {
+		if lim != nil {
+			// cached connection level limiter
+			lims = append(lims, lim.(limiter.Limiter))
+			// reset expiration
+			l.connOutLimits.Set(key, lim, defaultExpiration)
+		}
+	} else {
+		// generate a new connection level limiter
+		if v, ok := l.generators.Load(ConnLimitKey); ok && v != nil {
+			lim := v.(*limitGenerator).Out()
+			if lim != nil {
+				lims = append(lims, lim)
+				l.connOutLimits.Set(key, lim, defaultExpiration)
+			}
 		}
 	}
 
-	// IP limiter
+	host, _, _ := net.SplitHostPort(key)
-	if lim, ok := l.outLimits[key]; ok {
+	// IP level limiter
+	if lim, ok := l.outLimits.Get(host); ok {
 		if lim != nil {
-			lims = append(lims, lim)
+			// cached IP level limiter
+			lims = append(lims, lim.(limiter.Limiter))
 		}
 	} else {
-		if ip := net.ParseIP(key); ip != nil {
+		l.mu.RLock()
-			if p := l.limits[key]; p != nil {
+		ranger := l.cidrGenerators
-				if lim = p.Out(); lim != nil {
+		l.mu.RUnlock()
+
+		// CIDR level limiter
+		if p, _ := ranger.ContainingNetworks(net.ParseIP(host)); len(p) > 0 {
+			if v, _ := p[0].(*cidrLimitEntry); v != nil {
+				if lim := v.generator.Out(); lim != nil {
 					lims = append(lims, lim)
+					l.outLimits.Set(host, lim, cache.NoExpiration)
 				}
 			}
-			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
 		}
 	}
 
@@ -234,7 +239,7 @@ func (l *trafficLimiter) Out(key string) limiter.Limiter {
 	}
 
 	if lim != nil && l.options.logger != nil {
-		l.options.logger.Debugf("output limit for %s: %d", key, lim.Limit())
+		l.options.logger.Debugf("output limit for %s: %s", key, lim)
 	}
 
 	return lim
@@ -262,36 +267,160 @@ func (l *trafficLimiter) periodReload(ctx context.Context) error {
 }
 
 func (l *trafficLimiter) reload(ctx context.Context) error {
-	v, err := l.load(ctx)
+	values, err := l.load(ctx)
 	if err != nil {
 		return err
 	}
 
-	lines := append(l.options.limits, v...)
+	// service level limiter, never expired
-
+	{
-	limits := make(map[string]TrafficLimitGenerator)
+		value := values[GlobalLimitKey]
-	cidrLimits := cidranger.NewPCTrieRanger()
+		if v, _ := l.inLimits.Get(GlobalLimitKey); v != nil {
-
+			lim := v.(limiter.Limiter)
-	for _, s := range lines {
+			if value.in <= 0 {
-		key, in, out := l.parseLimit(s)
+				l.inLimits.Delete(GlobalLimitKey)
-		if key == "" {
+			} else {
-			continue
+				lim.Set(value.in)
-		}
-		switch key {
-		case GlobalLimitKey:
-			limits[key] = NewTrafficLimitSingleGenerator(in, out)
-		case ConnLimitKey:
-			limits[key] = NewTrafficLimitGenerator(in, out)
-		default:
-			if ip := net.ParseIP(key); ip != nil {
-				limits[key] = NewTrafficLimitSingleGenerator(in, out)
-				break
 			}
+		} else {
+			if value.in > 0 {
+				l.inLimits.Set(GlobalLimitKey, NewLimiter(value.in), cache.NoExpiration)
+			}
+		}
+
+		if v, _ := l.outLimits.Get(GlobalLimitKey); v != nil {
+			lim := v.(limiter.Limiter)
+			if value.out <= 0 {
+				l.outLimits.Delete(GlobalLimitKey)
+			} else {
+				lim.Set(value.out)
+			}
+		} else {
+			if value.out > 0 {
+				l.outLimits.Set(GlobalLimitKey, NewLimiter(value.out), cache.NoExpiration)
+			}
+		}
+		delete(values, GlobalLimitKey)
+	}
+
+	// connection level limiters
+	{
+		value := values[ConnLimitKey]
+
+		var in, out int
+		if v, _ := l.generators.Load(ConnLimitKey); v != nil {
+			in, out = v.(*limitGenerator).in, v.(*limitGenerator).out
+		}
+		l.generators.Store(ConnLimitKey, newLimitGenerator(value.in, value.out))
+
+		if value.in <= 0 {
+			l.connInLimits.Flush()
+		} else {
+			if in != value.in {
+				for _, item := range l.connInLimits.Items() {
+					if v := item.Object; v != nil {
+						v.(limiter.Limiter).Set(in)
+					}
+				}
+			}
+		}
+
+		if value.out <= 0 {
+			l.connOutLimits.Flush()
+		} else {
+			if out != value.out {
+				for _, item := range l.connOutLimits.Items() {
+					if v := item.Object; v != nil {
+						v.(limiter.Limiter).Set(out)
+					}
+				}
+			}
+		}
+		delete(values, ConnLimitKey)
+	}
+
+	cidrGenerators := cidranger.NewPCTrieRanger()
+	// IP/CIDR level limiters
+	{
+		// snapshot of the current limiters
+		inLimits := l.inLimits.Items()
+		outLimits := l.outLimits.Items()
+
+		delete(inLimits, GlobalLimitKey)
+		delete(outLimits, GlobalLimitKey)
+
+		for key, value := range values {
 			if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
-				cidrLimits.Insert(&cidrLimitEntry{
+				cidrGenerators.Insert(&cidrLimitEntry{
-					ipNet: *ipNet,
+					ipNet:     *ipNet,
-					limit: NewTrafficLimitGenerator(in, out),
+					generator: newLimitGenerator(value.in, value.out),
 				})
+				continue
+			}
+
+			if v, _ := l.inLimits.Get(key); v != nil {
+				lim := v.(limiter.Limiter)
+				if value.in <= 0 {
+					l.inLimits.Delete(key)
+				} else {
+					lim.Set(value.in)
+				}
+				delete(inLimits, key)
+			} else {
+				if value.in > 0 {
+					l.inLimits.Set(key, NewLimiter(value.in), cache.NoExpiration)
+				}
+			}
+
+			if v, _ := l.outLimits.Get(key); v != nil {
+				lim := v.(limiter.Limiter)
+				if value.out <= 0 {
+					l.outLimits.Delete(key)
+				} else {
+					lim.Set(value.out)
+				}
+				delete(outLimits, key)
+			} else {
+				if value.out > 0 {
+					l.outLimits.Set(key, NewLimiter(value.out), cache.NoExpiration)
+				}
+			}
+		}
+
+		// check the CIDR for remain limiters, clean the unmatched ones.
+		for k, v := range inLimits {
+			if p, _ := cidrGenerators.ContainingNetworks(net.ParseIP(k)); len(p) > 0 {
+				if le, _ := p[0].(*cidrLimitEntry); le != nil {
+					in := le.generator.in
+					if in <= 0 {
+						l.inLimits.Delete(k)
+						continue
+					}
+					lim := v.Object.(limiter.Limiter)
+					if lim.Limit() != in {
+						lim.Set(in)
+					}
+				}
+			} else {
+				l.inLimits.Delete(k)
+			}
+		}
+		for k, v := range outLimits {
+			if p, _ := cidrGenerators.ContainingNetworks(net.ParseIP(k)); len(p) > 0 {
+				if le, _ := p[0].(*cidrLimitEntry); le != nil {
+					out := le.generator.out
+					if out <= 0 {
+						l.outLimits.Delete(k)
+						continue
+					}
+					lim := v.Object.(limiter.Limiter)
+					if lim.Limit() != out {
+						lim.Set(out)
+					}
+					delete(outLimits, k)
+				}
+			} else {
+				l.outLimits.Delete(k)
 			}
 		}
 	}
@@ -299,15 +428,22 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
-	l.limits = limits
+	l.cidrGenerators = cidrGenerators
-	l.cidrLimits = cidrLimits
-	l.inLimits = make(map[string]limiter.Limiter)
-	l.outLimits = make(map[string]limiter.Limiter)
 
 	return nil
 }
 
-func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error) {
+func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue, err error) {
+	values = make(map[string]limitValue)
+
+	for _, v := range l.options.limits {
+		key, in, out := l.parseLimit(v)
+		if key == "" {
+			continue
+		}
+		values[key] = limitValue{in: in, out: out}
+	}
+
 	if l.options.fileLoader != nil {
 		if lister, ok := l.options.fileLoader.(loader.Lister); ok {
 			list, er := lister.List(ctx)
@@ -315,17 +451,24 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error
 				l.options.logger.Warnf("file loader: %v", er)
 			}
 			for _, s := range list {
-				if line := l.parseLine(s); line != "" {
+				key, in, out := l.parseLimit(l.parseLine(s))
-					patterns = append(patterns, line)
+				if key == "" {
+					continue
 				}
+				values[key] = limitValue{in: in, out: out}
 			}
 		} else {
 			r, er := l.options.fileLoader.Load(ctx)
 			if er != nil {
 				l.options.logger.Warnf("file loader: %v", er)
 			}
-			if v, _ := l.parsePatterns(r); v != nil {
+			patterns, _ := l.parsePatterns(r)
-				patterns = append(patterns, v...)
+			for _, s := range patterns {
+				key, in, out := l.parseLimit(l.parseLine(s))
+				if key == "" {
+					continue
+				}
+				values[key] = limitValue{in: in, out: out}
 			}
 		}
 	}
@@ -335,14 +478,25 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error
 			if er != nil {
 				l.options.logger.Warnf("redis loader: %v", er)
 			}
-			patterns = append(patterns, list...)
+			for _, s := range list {
+				key, in, out := l.parseLimit(l.parseLine(s))
+				if key == "" {
+					continue
+				}
+				values[key] = limitValue{in: in, out: out}
+			}
 		} else {
 			r, er := l.options.redisLoader.Load(ctx)
 			if er != nil {
 				l.options.logger.Warnf("redis loader: %v", er)
 			}
-			if v, _ := l.parsePatterns(r); v != nil {
+			patterns, _ := l.parsePatterns(r)
-				patterns = append(patterns, v...)
+			for _, s := range patterns {
+				key, in, out := l.parseLimit(l.parseLine(s))
+				if key == "" {
+					continue
+				}
+				values[key] = limitValue{in: in, out: out}
 			}
 		}
 	}
@@ -351,12 +505,17 @@ func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error
 		if er != nil {
 			l.options.logger.Warnf("http loader: %v", er)
 		}
-		if v, _ := l.parsePatterns(r); v != nil {
+		patterns, _ := l.parsePatterns(r)
-			patterns = append(patterns, v...)
+		for _, s := range patterns {
+			key, in, out := l.parseLimit(l.parseLine(s))
+			if key == "" {
+				continue
+			}
+			values[key] = limitValue{in: in, out: out}
 		}
 	}
 
-	l.options.logger.Debugf("load items %d", len(patterns))
+	l.options.logger.Debugf("load items %d", len(values))
 	return
 }
 
@@ -386,6 +545,10 @@ func (l *trafficLimiter) parseLine(s string) string {
 func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) {
 	s = strings.Replace(s, "\t", " ", -1)
 	s = strings.TrimSpace(s)
+	if s == "" {
+		return
+	}
+
 	var ss []string
 	for _, v := range strings.Split(s, " ") {
 		if v != "" {
@@ -421,8 +584,8 @@ func (l *trafficLimiter) Close() error {
 }
 
 type cidrLimitEntry struct {
-	ipNet net.IPNet
+	ipNet     net.IPNet
-	limit TrafficLimitGenerator
+	generator *limitGenerator
 }
 
 func (p *cidrLimitEntry) Network() net.IPNet {

limiter/traffic/wrapper/conn.go

@@ -6,40 +6,63 @@ import (
 	"errors"
 	"io"
 	"net"
-	"sync"
 	"syscall"
+	"time"
 
 	limiter "github.com/go-gost/core/limiter/traffic"
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/net/udp"
+	"github.com/patrickmn/go-cache"
 )
 
 var (
 	errUnsupport = errors.New("unsupported operation")
 )
 
-// serverConn is a server side Conn with metrics supported.
+// serverConn is a server side Conn with traffic limiter supported.
 type serverConn struct {
 	net.Conn
 	rbuf       bytes.Buffer
+	limiter    limiter.TrafficLimiter
 	limiterIn  limiter.Limiter
+	expIn      int64
 	limiterOut limiter.Limiter
+	expOut     int64
 }
 
-func WrapConn(rlimiter limiter.TrafficLimiter, c net.Conn) net.Conn {
+func WrapConn(limiter limiter.TrafficLimiter, c net.Conn) net.Conn {
-	if rlimiter == nil {
+	if limiter == nil {
 		return c
 	}
-	host, _, _ := net.SplitHostPort(c.RemoteAddr().String())
 	return &serverConn{
-		Conn:       c,
+		Conn:    c,
-		limiterIn:  rlimiter.In(host),
+		limiter: limiter,
-		limiterOut: rlimiter.Out(host),
 	}
 }
 
+func (c *serverConn) getInLimiter(addr net.Addr) limiter.Limiter {
+	now := time.Now().UnixNano()
+	// cache the limiter for 1s
+	if c.limiter != nil && time.Duration(now-c.expIn) > time.Second {
+		c.limiterIn = c.limiter.In(addr.String())
+		c.expIn = now
+	}
+	return c.limiterIn
+}
+
+func (c *serverConn) getOutLimiter(addr net.Addr) limiter.Limiter {
+	now := time.Now().UnixNano()
+	// cache the limiter for 1s
+	if c.limiter != nil && time.Duration(now-c.expOut) > time.Second {
+		c.limiterOut = c.limiter.Out(addr.String())
+		c.expOut = now
+	}
+	return c.limiterOut
+}
+
 func (c *serverConn) Read(b []byte) (n int, err error) {
-	if c.limiterIn == nil {
+	limiter := c.getInLimiter(c.RemoteAddr())
+	if limiter == nil {
 		return c.Conn.Read(b)
 	}
 
@@ -48,7 +71,7 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
 		if c.rbuf.Len() < burst {
 			burst = c.rbuf.Len()
 		}
-		lim := c.limiterIn.Wait(context.Background(), burst)
+		lim := limiter.Wait(context.Background(), burst)
 		return c.rbuf.Read(b[:lim])
 	}
 
@@ -57,7 +80,7 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
 		return nn, err
 	}
 
-	n = c.limiterIn.Wait(context.Background(), nn)
+	n = limiter.Wait(context.Background(), nn)
 	if n < nn {
 		if _, err = c.rbuf.Write(b[n:nn]); err != nil {
 			return 0, err
@@ -68,13 +91,14 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
 }
 
 func (c *serverConn) Write(b []byte) (n int, err error) {
-	if c.limiterOut == nil {
+	limiter := c.getOutLimiter(c.RemoteAddr())
+	if limiter == nil {
 		return c.Conn.Write(b)
 	}
 
 	nn := 0
 	for len(b) > 0 {
-		nn, err = c.Conn.Write(b[:c.limiterOut.Wait(context.Background(), len(b))])
+		nn, err = c.Conn.Write(b[:limiter.Wait(context.Background(), len(b))])
 		n += nn
 		if err != nil {
 			return
@@ -97,10 +121,8 @@ func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
 type packetConn struct {
 	net.PacketConn
 	limiter   limiter.TrafficLimiter
-	inLimits  map[string]limiter.Limiter
+	inLimits  *cache.Cache
-	inMux     sync.RWMutex
+	outLimits *cache.Cache
-	outLimits map[string]limiter.Limiter
-	outMux    sync.RWMutex
 }
 
 func WrapPacketConn(lim limiter.TrafficLimiter, pc net.PacketConn) net.PacketConn {
@@ -110,8 +132,8 @@ func WrapPacketConn(lim limiter.TrafficLimiter, pc net.PacketConn) net.PacketCon
 	return &packetConn{
 		PacketConn: pc,
 		limiter:    lim,
-		inLimits:   make(map[string]limiter.Limiter),
+		inLimits:   cache.New(time.Second, 10*time.Second),
-		outLimits:  make(map[string]limiter.Limiter),
+		outLimits:  cache.New(time.Second, 10*time.Second),
 	}
 }
 
@@ -120,24 +142,21 @@ func (c *packetConn) getInLimiter(addr net.Addr) limiter.Limiter {
 		return nil
 	}
 
-	lim, ok := func() (limiter.Limiter, bool) {
+	lim, ok := func() (lim limiter.Limiter, ok bool) {
-		c.inMux.RLock()
+		v, ok := c.inLimits.Get(addr.String())
-		defer c.inMux.RUnlock()
+		if ok {
-
+			if v != nil {
-		lim, ok := c.inLimits[addr.String()]
+				lim = v.(limiter.Limiter)
-		return lim, ok
+			}
+		}
+		return
 	}()
 	if ok {
 		return lim
 	}
 
-	host, _, _ := net.SplitHostPort(addr.String())
+	lim = c.limiter.In(addr.String())
-	lim = c.limiter.In(host)
+	c.inLimits.Set(addr.String(), lim, 0)
-
-	c.inMux.Lock()
-	defer c.inMux.Unlock()
-
-	c.inLimits[addr.String()] = lim
 
 	return lim
 }
@@ -147,24 +166,21 @@ func (c *packetConn) getOutLimiter(addr net.Addr) limiter.Limiter {
 		return nil
 	}
 
-	lim, ok := func() (limiter.Limiter, bool) {
+	lim, ok := func() (lim limiter.Limiter, ok bool) {
-		c.outMux.RLock()
+		v, ok := c.outLimits.Get(addr.String())
-		defer c.outMux.RUnlock()
+		if ok {
-
+			if v != nil {
-		lim, ok := c.outLimits[addr.String()]
+				lim = v.(limiter.Limiter)
-		return lim, ok
+			}
+		}
+		return
 	}()
 	if ok {
 		return lim
 	}
 
-	host, _, _ := net.SplitHostPort(addr.String())
+	lim = c.limiter.Out(addr.String())
-	lim = c.limiter.Out(host)
+	c.outLimits.Set(addr.String(), lim, 0)
-
-	c.outMux.Lock()
-	defer c.outMux.Unlock()
-
-	c.outLimits[addr.String()] = lim
 
 	return lim
 }
@@ -204,16 +220,16 @@ func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 type udpConn struct {
 	net.PacketConn
 	limiter   limiter.TrafficLimiter
-	inLimits  map[string]limiter.Limiter
+	inLimits  *cache.Cache
-	inMux     sync.RWMutex
+	outLimits *cache.Cache
-	outLimits map[string]limiter.Limiter
-	outMux    sync.RWMutex
 }
 
 func WrapUDPConn(limiter limiter.TrafficLimiter, pc net.PacketConn) udp.Conn {
 	return &udpConn{
 		PacketConn: pc,
 		limiter:    limiter,
+		inLimits:   cache.New(time.Second, 10*time.Second),
+		outLimits:  cache.New(time.Second, 10*time.Second),
 	}
 }
 
@@ -222,24 +238,21 @@ func (c *udpConn) getInLimiter(addr net.Addr) limiter.Limiter {
 		return nil
 	}
 
-	lim, ok := func() (limiter.Limiter, bool) {
+	lim, ok := func() (lim limiter.Limiter, ok bool) {
-		c.inMux.RLock()
+		v, ok := c.inLimits.Get(addr.String())
-		defer c.inMux.RUnlock()
+		if ok {
-
+			if v != nil {
-		lim, ok := c.inLimits[addr.String()]
+				lim = v.(limiter.Limiter)
-		return lim, ok
+			}
+		}
+		return
 	}()
 	if ok {
 		return lim
 	}
 
-	host, _, _ := net.SplitHostPort(addr.String())
+	lim = c.limiter.In(addr.String())
-	lim = c.limiter.In(host)
+	c.inLimits.Set(addr.String(), lim, 0)
-
-	c.inMux.Lock()
-	defer c.inMux.Unlock()
-
-	c.inLimits[addr.String()] = lim
 
 	return lim
 }
@@ -249,24 +262,21 @@ func (c *udpConn) getOutLimiter(addr net.Addr) limiter.Limiter {
 		return nil
 	}
 
-	lim, ok := func() (limiter.Limiter, bool) {
+	lim, ok := func() (lim limiter.Limiter, ok bool) {
-		c.outMux.RLock()
+		v, ok := c.outLimits.Get(addr.String())
-		defer c.outMux.RUnlock()
+		if ok {
-
+			if v != nil {
-		lim, ok := c.outLimits[addr.String()]
+				lim = v.(limiter.Limiter)
-		return lim, ok
+			}
+		}
+		return
 	}()
 	if ok {
 		return lim
 	}
 
-	host, _, _ := net.SplitHostPort(addr.String())
+	lim = c.limiter.Out(addr.String())
-	lim = c.limiter.Out(host)
+	c.outLimits.Set(addr.String(), lim, 0)
-
-	c.outMux.Lock()
-	defer c.outMux.Unlock()
-
-	c.outLimits[addr.String()] = lim
 
 	return lim
 }