add conn limiter

limiter/conn/conn.go

@@ -0,0 +1,360 @@
+package conn
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"net"
+	"sort"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	limiter "github.com/go-gost/core/limiter/conn"
+	"github.com/go-gost/core/logger"
+	"github.com/go-gost/x/internal/loader"
+	"github.com/yl2chen/cidranger"
+)
+
+const (
+	GlobalLimitKey = "$"
+	IPLimitKey     = "$$"
+)
+
+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) Allow(n int) (b bool) {
+	var i int
+
+	for i = range l.limiters {
+		if b = l.limiters[i].Allow(n); !b {
+			break
+		}
+	}
+	if !b && i > 0 && n > 0 {
+		for i := range l.limiters[:i] {
+			l.limiters[i].Allow(-n)
+		}
+	}
+
+	return
+}
+
+func (l *limiterGroup) Limit() int {
+	if len(l.limiters) == 0 {
+		return 0
+	}
+
+	return l.limiters[0].Limit()
+}
+
+type options struct {
+	limits      []string
+	fileLoader  loader.Loader
+	redisLoader loader.Loader
+	period      time.Duration
+	logger      logger.Logger
+}
+
+type Option func(opts *options)
+
+func LimitsOption(limits ...string) Option {
+	return func(opts *options) {
+		opts.limits = limits
+	}
+}
+
+func ReloadPeriodOption(period time.Duration) Option {
+	return func(opts *options) {
+		opts.period = period
+	}
+}
+
+func FileLoaderOption(fileLoader loader.Loader) Option {
+	return func(opts *options) {
+		opts.fileLoader = fileLoader
+	}
+}
+
+func RedisLoaderOption(redisLoader loader.Loader) Option {
+	return func(opts *options) {
+		opts.redisLoader = redisLoader
+	}
+}
+
+func LoggerOption(logger logger.Logger) Option {
+	return func(opts *options) {
+		opts.logger = logger
+	}
+}
+
+type connLimiter struct {
+	ipLimits   map[string]ConnLimitGenerator
+	cidrLimits cidranger.Ranger
+	limits     map[string]limiter.Limiter
+	mu         sync.RWMutex
+	cancelFunc context.CancelFunc
+	options    options
+}
+
+func NewConnLimiter(opts ...Option) limiter.ConnLimiter {
+	var options options
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	ctx, cancel := context.WithCancel(context.TODO())
+	lim := &connLimiter{
+		ipLimits:   make(map[string]ConnLimitGenerator),
+		cidrLimits: cidranger.NewPCTrieRanger(),
+		limits:     make(map[string]limiter.Limiter),
+		options:    options,
+		cancelFunc: cancel,
+	}
+
+	if err := lim.reload(ctx); err != nil {
+		options.logger.Warnf("reload: %v", err)
+	}
+	if lim.options.period > 0 {
+		go lim.periodReload(ctx)
+	}
+	return lim
+}
+
+func (l *connLimiter) Limiter(key string) limiter.Limiter {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	if lim, ok := l.limits[key]; ok {
+		return lim
+	}
+
+	var lims []limiter.Limiter
+
+	if ip := net.ParseIP(key); ip != nil {
+		found := false
+		if p := l.ipLimits[key]; p != nil {
+			if lim := p.Limiter(); 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.Limiter(); lim != nil {
+						lims = append(lims, lim)
+					}
+				}
+			}
+		}
+	}
+
+	if len(lims) == 0 {
+		if p := l.ipLimits[IPLimitKey]; p != nil {
+			if lim := p.Limiter(); lim != nil {
+				lims = append(lims, lim)
+			}
+		}
+	}
+
+	if p := l.ipLimits[GlobalLimitKey]; p != nil {
+		if lim := p.Limiter(); lim != nil {
+			lims = append(lims, lim)
+		}
+	}
+
+	var lim limiter.Limiter
+	if len(lims) > 0 {
+		lim = newLimiterGroup(lims...)
+	}
+	l.limits[key] = lim
+
+	if lim != nil && l.options.logger != nil {
+		l.options.logger.Debugf("conn limit for %s: %d", key, lim.Limit())
+	}
+
+	return lim
+}
+
+func (l *connLimiter) periodReload(ctx context.Context) error {
+	period := l.options.period
+	if period < time.Second {
+		period = time.Second
+	}
+	ticker := time.NewTicker(period)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			if err := l.reload(ctx); err != nil {
+				l.options.logger.Warnf("reload: %v", err)
+				// return err
+			}
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+}
+
+func (l *connLimiter) reload(ctx context.Context) error {
+	v, err := l.load(ctx)
+	if err != nil {
+		return err
+	}
+
+	lines := append(l.options.limits, v...)
+
+	ipLimits := make(map[string]ConnLimitGenerator)
+	cidrLimits := cidranger.NewPCTrieRanger()
+
+	for _, s := range lines {
+		key, limit := l.parseLimit(s)
+		if key == "" || limit <= 0 {
+			continue
+		}
+		switch key {
+		case GlobalLimitKey:
+			ipLimits[key] = NewConnLimitSingleGenerator(limit)
+		case IPLimitKey:
+			ipLimits[key] = NewConnLimitGenerator(limit)
+		default:
+			if ip := net.ParseIP(key); ip != nil {
+				ipLimits[key] = NewConnLimitGenerator(limit)
+				break
+			}
+			if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
+				cidrLimits.Insert(&cidrLimitEntry{
+					ipNet: *ipNet,
+					limit: NewConnLimitGenerator(limit),
+				})
+			}
+		}
+	}
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	l.ipLimits = ipLimits
+	l.cidrLimits = cidrLimits
+	l.limits = make(map[string]limiter.Limiter)
+
+	return nil
+}
+
+func (l *connLimiter) load(ctx context.Context) (patterns []string, err error) {
+	if l.options.fileLoader != nil {
+		if lister, ok := l.options.fileLoader.(loader.Lister); ok {
+			list, er := lister.List(ctx)
+			if er != nil {
+				l.options.logger.Warnf("file loader: %v", er)
+			}
+			for _, s := range list {
+				if line := l.parseLine(s); line != "" {
+					patterns = append(patterns, line)
+				}
+			}
+		} 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 = append(patterns, v...)
+			}
+		}
+	}
+	if l.options.redisLoader != nil {
+		if lister, ok := l.options.redisLoader.(loader.Lister); ok {
+			list, er := lister.List(ctx)
+			if er != nil {
+				l.options.logger.Warnf("redis loader: %v", er)
+			}
+			patterns = append(patterns, list...)
+		} 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 = append(patterns, v...)
+			}
+		}
+	}
+
+	l.options.logger.Debugf("load items %d", len(patterns))
+	return
+}
+
+func (l *connLimiter) parsePatterns(r io.Reader) (patterns []string, err error) {
+	if r == nil {
+		return
+	}
+
+	scanner := bufio.NewScanner(r)
+	for scanner.Scan() {
+		if line := l.parseLine(scanner.Text()); line != "" {
+			patterns = append(patterns, line)
+		}
+	}
+
+	err = scanner.Err()
+	return
+}
+
+func (l *connLimiter) parseLine(s string) string {
+	if n := strings.IndexByte(s, '#'); n >= 0 {
+		s = s[:n]
+	}
+	return strings.TrimSpace(s)
+}
+
+func (l *connLimiter) parseLimit(s string) (key string, limit int) {
+	s = strings.Replace(s, "\t", " ", -1)
+	s = strings.TrimSpace(s)
+	var ss []string
+	for _, v := range strings.Split(s, " ") {
+		if v != "" {
+			ss = append(ss, v)
+		}
+	}
+	if len(ss) < 2 {
+		return
+	}
+
+	key = ss[0]
+	limit, _ = strconv.Atoi(ss[1])
+
+	return
+}
+
+func (l *connLimiter) Close() error {
+	l.cancelFunc()
+	if l.options.fileLoader != nil {
+		l.options.fileLoader.Close()
+	}
+	if l.options.redisLoader != nil {
+		l.options.redisLoader.Close()
+	}
+	return nil
+}
+
+type cidrLimitEntry struct {
+	ipNet net.IPNet
+	limit ConnLimitGenerator
+}
+
+func (p *cidrLimitEntry) Network() net.IPNet {
+	return p.ipNet
+}

limiter/conn/generator.go

@@ -0,0 +1,42 @@
+package conn
+
+import (
+	limiter "github.com/go-gost/core/limiter/conn"
+)
+
+type ConnLimitGenerator interface {
+	Limiter() limiter.Limiter
+}
+
+type connLimitGenerator struct {
+	n int
+}
+
+func NewConnLimitGenerator(n int) ConnLimitGenerator {
+	return &connLimitGenerator{
+		n: n,
+	}
+}
+
+func (p *connLimitGenerator) Limiter() limiter.Limiter {
+	if p == nil || p.n <= 0 {
+		return nil
+	}
+	return NewLimiter(p.n)
+}
+
+type connLimitSingleGenerator struct {
+	limiter limiter.Limiter
+}
+
+func NewConnLimitSingleGenerator(n int) ConnLimitGenerator {
+	p := &connLimitSingleGenerator{}
+	if n > 0 {
+		p.limiter = NewLimiter(n)
+	}
+	return p
+}
+
+func (p *connLimitSingleGenerator) Limiter() limiter.Limiter {
+	return p.limiter
+}

limiter/conn/limiter.go

@@ -0,0 +1,30 @@
+package conn
+
+import (
+	"sync/atomic"
+
+	limiter "github.com/go-gost/core/limiter/conn"
+)
+
+type llimiter struct {
+	limit   int
+	current int64
+}
+
+func NewLimiter(n int) limiter.Limiter {
+	return &llimiter{limit: n}
+}
+
+func (l *llimiter) Limit() int {
+	return l.limit
+}
+
+func (l *llimiter) Allow(n int) bool {
+	if atomic.AddInt64(&l.current, int64(n)) >= int64(l.limit) {
+		if n > 0 {
+			atomic.AddInt64(&l.current, -int64(n))
+		}
+		return false
+	}
+	return true
+}

limiter/conn/wrapper/conn.go

@@ -0,0 +1,43 @@
+package wrapper
+
+import (
+	"errors"
+	"net"
+	"syscall"
+
+	limiter "github.com/go-gost/core/limiter/conn"
+)
+
+var (
+	errUnsupport = errors.New("unsupported operation")
+)
+
+// serverConn is a server side Conn with metrics supported.
+type serverConn struct {
+	net.Conn
+	limiter limiter.Limiter
+}
+
+func WrapConn(limiter limiter.Limiter, c net.Conn) net.Conn {
+	if limiter == nil {
+		return c
+	}
+	return &serverConn{
+		Conn:    c,
+		limiter: limiter,
+	}
+}
+
+func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
+	if sc, ok := c.Conn.(syscall.Conn); ok {
+		rc, err = sc.SyscallConn()
+		return
+	}
+	err = errUnsupport
+	return
+}
+
+func (c *serverConn) Close() error {
+	c.limiter.Allow(-1)
+	return c.Conn.Close()
+}

limiter/conn/wrapper/listener.go

@@ -0,0 +1,40 @@
+package wrapper
+
+import (
+	"net"
+
+	limiter "github.com/go-gost/core/limiter/conn"
+)
+
+type listener struct {
+	net.Listener
+	limiter limiter.ConnLimiter
+}
+
+func WrapListener(limiter limiter.ConnLimiter, ln net.Listener) net.Listener {
+	if limiter == nil {
+		return ln
+	}
+
+	return &listener{
+		limiter:  limiter,
+		Listener: ln,
+	}
+}
+
+func (ln *listener) Accept() (net.Conn, error) {
+	c, err := ln.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+
+	host, _, _ := net.SplitHostPort(c.RemoteAddr().String())
+	if lim := ln.limiter.Limiter(host); lim != nil {
+		if lim.Allow(1) {
+			return WrapConn(lim, c), nil
+		}
+		c.Close()
+	}
+
+	return c, nil
+}

limiter/generator.go

@@ -1,61 +0,0 @@
-package limiter
-
-import (
-	"github.com/go-gost/core/limiter"
-)
-
-type RateLimitGenerator interface {
-	In() limiter.Limiter
-	Out() limiter.Limiter
-}
-
-type rateLimitGenerator struct {
-	in  int
-	out int
-}
-
-func NewRateLimitGenerator(in, out int) RateLimitGenerator {
-	return &rateLimitGenerator{
-		in:  in,
-		out: out,
-	}
-}
-
-func (p *rateLimitGenerator) In() limiter.Limiter {
-	if p == nil || p.in <= 0 {
-		return nil
-	}
-	return NewLimiter(p.in)
-}
-
-func (p *rateLimitGenerator) Out() limiter.Limiter {
-	if p == nil || p.out <= 0 {
-		return nil
-	}
-	return NewLimiter(p.out)
-}
-
-type rateLimitSingleGenerator struct {
-	in  limiter.Limiter
-	out limiter.Limiter
-}
-
-func NewRateLimitSingleGenerator(in, out int) RateLimitGenerator {
-	p := &rateLimitSingleGenerator{}
-	if in > 0 {
-		p.in = NewLimiter(in)
-	}
-	if out > 0 {
-		p.out = NewLimiter(out)
-	}
-
-	return p
-}
-
-func (p *rateLimitSingleGenerator) In() limiter.Limiter {
-	return p.in
-}
-
-func (p *rateLimitSingleGenerator) Out() limiter.Limiter {
-	return p.out
-}

limiter/traffic/generator.go

@@ -0,0 +1,61 @@
+package traffic
+
+import (
+	limiter "github.com/go-gost/core/limiter/traffic"
+)
+
+type TrafficLimitGenerator interface {
+	In() limiter.Limiter
+	Out() limiter.Limiter
+}
+
+type trafficLimitGenerator struct {
+	in  int
+	out int
+}
+
+func NewTrafficLimitGenerator(in, out int) TrafficLimitGenerator {
+	return &trafficLimitGenerator{
+		in:  in,
+		out: out,
+	}
+}
+
+func (p *trafficLimitGenerator) In() limiter.Limiter {
+	if p == nil || p.in <= 0 {
+		return nil
+	}
+	return NewLimiter(p.in)
+}
+
+func (p *trafficLimitGenerator) 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

@@ -1,9 +1,9 @@
-package limiter
+package traffic
 
 import (
 	"context"
 
-	"github.com/go-gost/core/limiter"
+	limiter "github.com/go-gost/core/limiter/traffic"
 	"golang.org/x/time/rate"
 )
 

limiter/traffic/traffic.go

@@ -1,4 +1,4 @@
-package limiter
+package traffic
 
 import (
 	"bufio"
@@ -11,7 +11,7 @@ import (
 	"time"
 
 	"github.com/alecthomas/units"
-	"github.com/go-gost/core/limiter"
+	limiter "github.com/go-gost/core/limiter/traffic"
 	"github.com/go-gost/core/logger"
 	"github.com/go-gost/x/internal/loader"
 	"github.com/yl2chen/cidranger"
@@ -90,8 +90,8 @@ func LoggerOption(logger logger.Logger) Option {
 	}
 }
 
-type rateLimiter struct {
-	ipLimits   map[string]RateLimitGenerator
+type trafficLimiter struct {
+	ipLimits   map[string]TrafficLimitGenerator
 	cidrLimits cidranger.Ranger
 	inLimits   map[string]limiter.Limiter
 	outLimits  map[string]limiter.Limiter
@@ -100,15 +100,15 @@ type rateLimiter struct {
 	options    options
 }
 
-func NewRateLimiter(opts ...Option) limiter.RateLimiter {
+func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
 	var options options
 	for _, opt := range opts {
 		opt(&options)
 	}
 
 	ctx, cancel := context.WithCancel(context.TODO())
-	lim := &rateLimiter{
-		ipLimits:   make(map[string]RateLimitGenerator),
+	lim := &trafficLimiter{
+		ipLimits:   make(map[string]TrafficLimitGenerator),
 		cidrLimits: cidranger.NewPCTrieRanger(),
 		inLimits:   make(map[string]limiter.Limiter),
 		outLimits:  make(map[string]limiter.Limiter),
@@ -125,7 +125,7 @@ func NewRateLimiter(opts ...Option) limiter.RateLimiter {
 	return lim
 }
 
-func (l *rateLimiter) In(key string) limiter.Limiter {
+func (l *trafficLimiter) In(key string) limiter.Limiter {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
@@ -178,7 +178,7 @@ func (l *rateLimiter) In(key string) limiter.Limiter {
 	return lim
 }
 
-func (l *rateLimiter) Out(key string) limiter.Limiter {
+func (l *trafficLimiter) Out(key string) limiter.Limiter {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
@@ -231,7 +231,7 @@ func (l *rateLimiter) Out(key string) limiter.Limiter {
 	return lim
 }
 
-func (l *rateLimiter) periodReload(ctx context.Context) error {
+func (l *trafficLimiter) periodReload(ctx context.Context) error {
 	period := l.options.period
 	if period < time.Second {
 		period = time.Second
@@ -252,7 +252,7 @@ func (l *rateLimiter) periodReload(ctx context.Context) error {
 	}
 }
 
-func (l *rateLimiter) reload(ctx context.Context) error {
+func (l *trafficLimiter) reload(ctx context.Context) error {
 	v, err := l.load(ctx)
 	if err != nil {
 		return err
@@ -260,7 +260,7 @@ func (l *rateLimiter) reload(ctx context.Context) error {
 
 	lines := append(l.options.limits, v...)
 
-	ipLimits := make(map[string]RateLimitGenerator)
+	ipLimits := make(map[string]TrafficLimitGenerator)
 	cidrLimits := cidranger.NewPCTrieRanger()
 
 	for _, s := range lines {
@@ -270,18 +270,18 @@ func (l *rateLimiter) reload(ctx context.Context) error {
 		}
 		switch key {
 		case GlobalLimitKey:
-			ipLimits[key] = NewRateLimitSingleGenerator(in, out)
+			ipLimits[key] = NewTrafficLimitSingleGenerator(in, out)
 		case ConnLimitKey:
-			ipLimits[key] = NewRateLimitGenerator(in, out)
+			ipLimits[key] = NewTrafficLimitGenerator(in, out)
 		default:
 			if ip := net.ParseIP(key); ip != nil {
-				ipLimits[key] = NewRateLimitGenerator(in, out)
+				ipLimits[key] = NewTrafficLimitGenerator(in, out)
 				break
 			}
 			if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
 				cidrLimits.Insert(&cidrLimitEntry{
 					ipNet: *ipNet,
-					limit: NewRateLimitGenerator(in, out),
+					limit: NewTrafficLimitGenerator(in, out),
 				})
 			}
 		}
@@ -298,7 +298,7 @@ func (l *rateLimiter) reload(ctx context.Context) error {
 	return nil
 }
 
-func (l *rateLimiter) load(ctx context.Context) (patterns []string, err error) {
+func (l *trafficLimiter) load(ctx context.Context) (patterns []string, err error) {
 	if l.options.fileLoader != nil {
 		if lister, ok := l.options.fileLoader.(loader.Lister); ok {
 			list, er := lister.List(ctx)
@@ -342,7 +342,7 @@ func (l *rateLimiter) load(ctx context.Context) (patterns []string, err error) {
 	return
 }
 
-func (l *rateLimiter) parsePatterns(r io.Reader) (patterns []string, err error) {
+func (l *trafficLimiter) parsePatterns(r io.Reader) (patterns []string, err error) {
 	if r == nil {
 		return
 	}
@@ -358,14 +358,14 @@ func (l *rateLimiter) parsePatterns(r io.Reader) (patterns []string, err error)
 	return
 }
 
-func (l *rateLimiter) parseLine(s string) string {
+func (l *trafficLimiter) parseLine(s string) string {
 	if n := strings.IndexByte(s, '#'); n >= 0 {
 		s = s[:n]
 	}
 	return strings.TrimSpace(s)
 }
 
-func (l *rateLimiter) parseLimit(s string) (key string, in, out int) {
+func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) {
 	s = strings.Replace(s, "\t", " ", -1)
 	s = strings.TrimSpace(s)
 	var ss []string
@@ -391,7 +391,7 @@ func (l *rateLimiter) parseLimit(s string) (key string, in, out int) {
 	return
 }
 
-func (l *rateLimiter) Close() error {
+func (l *trafficLimiter) Close() error {
 	l.cancelFunc()
 	if l.options.fileLoader != nil {
 		l.options.fileLoader.Close()
@@ -404,7 +404,7 @@ func (l *rateLimiter) Close() error {
 
 type cidrLimitEntry struct {
 	ipNet net.IPNet
-	limit RateLimitGenerator
+	limit TrafficLimitGenerator
 }
 
 func (p *cidrLimitEntry) Network() net.IPNet {

limiter/traffic/wrapper/conn.go

@@ -8,7 +8,7 @@ import (
 	"net"
 	"syscall"
 
-	"github.com/go-gost/core/limiter"
+	limiter "github.com/go-gost/core/limiter/traffic"
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/go-gost/x/internal/net/udp"
 )
@@ -22,10 +22,10 @@ type serverConn struct {
 	net.Conn
 	rbuf     bytes.Buffer
 	raddr    string
-	rlimiter limiter.RateLimiter
+	rlimiter limiter.TrafficLimiter
 }
 
-func WrapConn(rlimiter limiter.RateLimiter, c net.Conn) net.Conn {
+func WrapConn(rlimiter limiter.TrafficLimiter, c net.Conn) net.Conn {
 	if rlimiter == nil {
 		return c
 	}
@@ -100,16 +100,16 @@ func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
 
 type packetConn struct {
 	net.PacketConn
-	rlimiter limiter.RateLimiter
+	limiter limiter.TrafficLimiter
 }
 
-func WrapPacketConn(rlimiter limiter.RateLimiter, pc net.PacketConn) net.PacketConn {
-	if rlimiter == nil {
+func WrapPacketConn(limiter limiter.TrafficLimiter, pc net.PacketConn) net.PacketConn {
+	if limiter == nil {
 		return pc
 	}
 	return &packetConn{
 		PacketConn: pc,
-		rlimiter:   rlimiter,
+		limiter:    limiter,
 	}
 }
 
@@ -122,11 +122,11 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 
 		host, _, _ := net.SplitHostPort(addr.String())
 
-		if c.rlimiter == nil || c.rlimiter.In(host) == nil {
+		if c.limiter == nil || c.limiter.In(host) == nil {
 			return
 		}
 
-		limiter := c.rlimiter.In(host)
+		limiter := c.limiter.In(host)
 		// discard when exceed the limit size.
 		if limiter.Wait(context.Background(), n) < n {
 			continue
@@ -137,10 +137,10 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 }
 
 func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
-	if c.rlimiter != nil {
+	if c.limiter != nil {
 		host, _, _ := net.SplitHostPort(addr.String())
 		// discard when exceed the limit size.
-		if limiter := c.rlimiter.Out(host); limiter != nil &&
+		if limiter := c.limiter.Out(host); limiter != nil &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
 			n = len(p)
 			return
@@ -152,13 +152,13 @@ func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 
 type udpConn struct {
 	net.PacketConn
-	rlimiter limiter.RateLimiter
+	limiter limiter.TrafficLimiter
 }
 
-func WrapUDPConn(rlimiter limiter.RateLimiter, pc net.PacketConn) udp.Conn {
+func WrapUDPConn(limiter limiter.TrafficLimiter, pc net.PacketConn) udp.Conn {
 	return &udpConn{
 		PacketConn: pc,
-		rlimiter:   rlimiter,
+		limiter:    limiter,
 	}
 }
 
@@ -200,10 +200,10 @@ func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 		}
 		host, _, _ := net.SplitHostPort(addr.String())
 
-		if c.rlimiter == nil || c.rlimiter.In(host) == nil {
+		if c.limiter == nil || c.limiter.In(host) == nil {
 			return
 		}
-		limiter := c.rlimiter.In(host)
+		limiter := c.limiter.In(host)
 		// discard when exceed the limit size.
 		if limiter.Wait(context.Background(), n) < n {
 			continue
@@ -222,10 +222,10 @@ func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
 
 			host, _, _ := net.SplitHostPort(addr.String())
 
-			if c.rlimiter == nil || c.rlimiter.In(host) == nil {
+			if c.limiter == nil || c.limiter.In(host) == nil {
 				return
 			}
-			limiter := c.rlimiter.In(host)
+			limiter := c.limiter.In(host)
 			// discard when exceed the limit size.
 			if limiter.Wait(context.Background(), n) < n {
 				continue
@@ -247,10 +247,10 @@ func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAd
 
 			host, _, _ := net.SplitHostPort(addr.String())
 
-			if c.rlimiter == nil || c.rlimiter.In(host) == nil {
+			if c.limiter == nil || c.limiter.In(host) == nil {
 				return
 			}
-			limiter := c.rlimiter.In(host)
+			limiter := c.limiter.In(host)
 			// discard when exceed the limit size.
 			if limiter.Wait(context.Background(), n) < n {
 				continue
@@ -272,10 +272,10 @@ func (c *udpConn) Write(b []byte) (n int, err error) {
 }
 
 func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
-	if c.rlimiter != nil {
+	if c.limiter != nil {
 		host, _, _ := net.SplitHostPort(addr.String())
 		// discard when exceed the limit size.
-		if limiter := c.rlimiter.Out(host); limiter != nil &&
+		if limiter := c.limiter.Out(host); limiter != nil &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
 			n = len(p)
 			return
@@ -287,10 +287,10 @@ func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 }
 
 func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
-	if c.rlimiter != nil {
+	if c.limiter != nil {
 		host, _, _ := net.SplitHostPort(addr.String())
 		// discard when exceed the limit size.
-		if limiter := c.rlimiter.Out(host); limiter != nil &&
+		if limiter := c.limiter.Out(host); limiter != nil &&
 			limiter.Wait(context.Background(), len(b)) < len(b) {
 			n = len(b)
 			return
@@ -306,10 +306,10 @@ func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
 }
 
 func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
-	if c.rlimiter != nil {
+	if c.limiter != nil {
 		host, _, _ := net.SplitHostPort(addr.String())
 		// discard when exceed the limit size.
-		if limiter := c.rlimiter.Out(host); limiter != nil &&
+		if limiter := c.limiter.Out(host); limiter != nil &&
 			limiter.Wait(context.Background(), len(b)) < len(b) {
 			n = len(b)
 			return

limiter/traffic/wrapper/listener.go

@@ -3,21 +3,21 @@ package wrapper
 import (
 	"net"
 
-	"github.com/go-gost/core/limiter"
+	limiter "github.com/go-gost/core/limiter/traffic"
 )
 
 type listener struct {
 	net.Listener
-	rlimiter limiter.RateLimiter
+	limiter limiter.TrafficLimiter
 }
 
-func WrapListener(rlimiter limiter.RateLimiter, ln net.Listener) net.Listener {
-	if rlimiter == nil {
+func WrapListener(limiter limiter.TrafficLimiter, ln net.Listener) net.Listener {
+	if limiter == nil {
 		return ln
 	}
 
 	return &listener{
-		rlimiter: rlimiter,
+		limiter:  limiter,
 		Listener: ln,
 	}
 }
@@ -28,5 +28,5 @@ func (ln *listener) Accept() (net.Conn, error) {
 		return nil, err
 	}
 
-	return WrapConn(ln.rlimiter, c), nil
+	return WrapConn(ln.limiter, c), nil
 }