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

fix traffic limiter

Browse Source
This commit is contained in:
ginuerzh 2022-11-05 15:23:12 +08:00
parent 05ddda70e3
commit 2b9ea187b8
2 changed files with 176 additions and 106 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
limiter/traffic/traffic.go
View File

@ -98,10 +98,8 @@ func LoggerOption(logger logger.Logger) Option {
} }
type trafficLimiter struct { type trafficLimiter struct {
ipLimits map[string]TrafficLimitGenerator limits map[string]TrafficLimitGenerator
cidrLimits cidranger.Ranger cidrLimits cidranger.Ranger
inLimits map[string]limiter.Limiter
outLimits map[string]limiter.Limiter
mu sync.Mutex mu sync.Mutex
cancelFunc context.CancelFunc cancelFunc context.CancelFunc
options options options options
@ -115,10 +113,8 @@ func NewTrafficLimiter(opts ...Option) limiter.TrafficLimiter {
ctx, cancel := context.WithCancel(context.TODO()) ctx, cancel := context.WithCancel(context.TODO())
lim := &trafficLimiter{ lim := &trafficLimiter{
ipLimits: make(map[string]TrafficLimitGenerator), limits: make(map[string]TrafficLimitGenerator),
cidrLimits: cidranger.NewPCTrieRanger(), cidrLimits: cidranger.NewPCTrieRanger(),
inLimits: make(map[string]limiter.Limiter),
outLimits: make(map[string]limiter.Limiter),
options: options, options: options,
cancelFunc: cancel, cancelFunc: cancel,
} }
@ -136,15 +132,11 @@ func (l *trafficLimiter) In(key string) limiter.Limiter {
l.mu.Lock() l.mu.Lock()
defer l.mu.Unlock() defer l.mu.Unlock()
if lim, ok := l.inLimits[key]; ok {
return lim
}
var lims []limiter.Limiter var lims []limiter.Limiter
if ip := net.ParseIP(key); ip != nil { if ip := net.ParseIP(key); ip != nil {
found := false found := false
if p := l.ipLimits[key]; p != nil { if p := l.limits[key]; p != nil {
if lim := p.In(); lim != nil { if lim := p.In(); lim != nil {
lims = append(lims, lim) lims = append(lims, lim)
found = true found = true
@ -161,12 +153,12 @@ func (l *trafficLimiter) In(key string) limiter.Limiter {
} }
} }
if p := l.ipLimits[ConnLimitKey]; p != nil { if p := l.limits[ConnLimitKey]; p != nil {
if lim := p.In(); lim != nil { if lim := p.In(); lim != nil {
lims = append(lims, lim) lims = append(lims, lim)
} }
} }
if p := l.ipLimits[GlobalLimitKey]; p != nil { if p := l.limits[GlobalLimitKey]; p != nil {
if lim := p.In(); lim != nil { if lim := p.In(); lim != nil {
lims = append(lims, lim) lims = append(lims, lim)
} }
@ -176,7 +168,6 @@ func (l *trafficLimiter) In(key string) limiter.Limiter {
if len(lims) > 0 { if len(lims) > 0 {
lim = newLimiterGroup(lims...) lim = newLimiterGroup(lims...)
} }
l.inLimits[key] = lim
if lim != nil && l.options.logger != nil { 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: %d", key, lim.Limit())
@ -189,15 +180,11 @@ func (l *trafficLimiter) Out(key string) limiter.Limiter {
l.mu.Lock() l.mu.Lock()
defer l.mu.Unlock() defer l.mu.Unlock()
if lim, ok := l.outLimits[key]; ok {
return lim
}
var lims []limiter.Limiter var lims []limiter.Limiter
if ip := net.ParseIP(key); ip != nil { if ip := net.ParseIP(key); ip != nil {
found := false found := false
if p := l.ipLimits[key]; p != nil { if p := l.limits[key]; p != nil {
if lim := p.Out(); lim != nil { if lim := p.Out(); lim != nil {
lims = append(lims, lim) lims = append(lims, lim)
found = true found = true
@ -214,12 +201,12 @@ func (l *trafficLimiter) Out(key string) limiter.Limiter {
} }
} }
if p := l.ipLimits[ConnLimitKey]; p != nil { if p := l.limits[ConnLimitKey]; p != nil {
if lim := p.Out(); lim != nil { if lim := p.Out(); lim != nil {
lims = append(lims, lim) lims = append(lims, lim)
} }
} }
if p := l.ipLimits[GlobalLimitKey]; p != nil { if p := l.limits[GlobalLimitKey]; p != nil {
if lim := p.Out(); lim != nil { if lim := p.Out(); lim != nil {
lims = append(lims, lim) lims = append(lims, lim)
} }
@ -229,7 +216,6 @@ func (l *trafficLimiter) Out(key string) limiter.Limiter {
if len(lims) > 0 { if len(lims) > 0 {
lim = newLimiterGroup(lims...) lim = newLimiterGroup(lims...)
} }
l.outLimits[key] = lim
if lim != nil && l.options.logger != nil { 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: %d", key, lim.Limit())
@ -267,7 +253,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
lines := append(l.options.limits, v...) lines := append(l.options.limits, v...)
ipLimits := make(map[string]TrafficLimitGenerator) limits := make(map[string]TrafficLimitGenerator)
cidrLimits := cidranger.NewPCTrieRanger() cidrLimits := cidranger.NewPCTrieRanger()
for _, s := range lines { for _, s := range lines {
@ -277,12 +263,12 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
} }
switch key { switch key {
case GlobalLimitKey: case GlobalLimitKey:
ipLimits[key] = NewTrafficLimitSingleGenerator(in, out) limits[key] = NewTrafficLimitSingleGenerator(in, out)
case ConnLimitKey: case ConnLimitKey:
ipLimits[key] = NewTrafficLimitGenerator(in, out) limits[key] = NewTrafficLimitGenerator(in, out)
default: default:
if ip := net.ParseIP(key); ip != nil { if ip := net.ParseIP(key); ip != nil {
ipLimits[key] = NewTrafficLimitGenerator(in, out) limits[key] = NewTrafficLimitGenerator(in, out)
break break
} }
if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil { if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
@ -297,10 +283,8 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
l.mu.Lock() l.mu.Lock()
defer l.mu.Unlock() defer l.mu.Unlock()
l.ipLimits = ipLimits l.limits = limits
l.cidrLimits = cidrLimits l.cidrLimits = cidrLimits
l.inLimits = make(map[string]limiter.Limiter)
l.outLimits = make(map[string]limiter.Limiter)
return nil return nil
} }

200
limiter/traffic/wrapper/conn.go
View File

@ -6,6 +6,7 @@ import (
"errors" "errors"
"io" "io"
"net" "net"
"sync"
"syscall" "syscall"
limiter "github.com/go-gost/core/limiter/traffic" limiter "github.com/go-gost/core/limiter/traffic"
@ -21,8 +22,8 @@ var (
type serverConn struct { type serverConn struct {
net.Conn net.Conn
rbuf bytes.Buffer rbuf bytes.Buffer
raddr string limiterIn limiter.Limiter
limiter limiter.TrafficLimiter limiterOut limiter.Limiter
} }
func WrapConn(rlimiter limiter.TrafficLimiter, c net.Conn) net.Conn { func WrapConn(rlimiter limiter.TrafficLimiter, c net.Conn) net.Conn {
@ -32,25 +33,22 @@ func WrapConn(rlimiter limiter.TrafficLimiter, c net.Conn) net.Conn {
host, _, _ := net.SplitHostPort(c.RemoteAddr().String()) host, _, _ := net.SplitHostPort(c.RemoteAddr().String())
return &serverConn{ return &serverConn{
Conn: c, Conn: c,
limiter: rlimiter, limiterIn: rlimiter.In(host),
raddr: host, limiterOut: rlimiter.Out(host),
} }
} }
func (c *serverConn) Read(b []byte) (n int, err error) { func (c *serverConn) Read(b []byte) (n int, err error) {
if c.limiter == nil || if c.limiterIn == nil {
c.limiter.In(c.raddr) == nil {
return c.Conn.Read(b) return c.Conn.Read(b)
} }
limiter := c.limiter.In(c.raddr)
if c.rbuf.Len() > 0 { if c.rbuf.Len() > 0 {
burst := len(b) burst := len(b)
if c.rbuf.Len() < burst { if c.rbuf.Len() < burst {
burst = c.rbuf.Len() burst = c.rbuf.Len()
} }
lim := limiter.Wait(context.Background(), burst) lim := c.limiterIn.Wait(context.Background(), burst)
return c.rbuf.Read(b[:lim]) return c.rbuf.Read(b[:lim])
} }
@ -59,7 +57,7 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
return nn, err return nn, err
} }
n = limiter.Wait(context.Background(), nn) n = c.limiterIn.Wait(context.Background(), nn)
if n < nn { if n < nn {
if _, err = c.rbuf.Write(b[n:nn]); err != nil { if _, err = c.rbuf.Write(b[n:nn]); err != nil {
return 0, err return 0, err
@ -70,15 +68,13 @@ func (c *serverConn) Read(b []byte) (n int, err error) {
} }
func (c *serverConn) Write(b []byte) (n int, err error) { func (c *serverConn) Write(b []byte) (n int, err error) {
if c.limiter == nil || if c.limiterOut == nil {
c.limiter.Out(c.raddr) == nil {
return c.Conn.Write(b) return c.Conn.Write(b)
} }
limiter := c.limiter.Out(c.raddr)
nn := 0 nn := 0
for len(b) > 0 { for len(b) > 0 {
nn, err = c.Conn.Write(b[:limiter.Wait(context.Background(), len(b))]) nn, err = c.Conn.Write(b[:c.limiterOut.Wait(context.Background(), len(b))])
n += nn n += nn
if err != nil { if err != nil {
return return
@ -101,18 +97,78 @@ func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
type packetConn struct { type packetConn struct {
net.PacketConn net.PacketConn
limiter limiter.TrafficLimiter limiter limiter.TrafficLimiter
inLimits map[string]limiter.Limiter
inMux sync.RWMutex
outLimits map[string]limiter.Limiter
outMux sync.RWMutex
} }
func WrapPacketConn(limiter limiter.TrafficLimiter, pc net.PacketConn) net.PacketConn { func WrapPacketConn(lim limiter.TrafficLimiter, pc net.PacketConn) net.PacketConn {
if limiter == nil { if lim == nil {
return pc return pc
} }
return &packetConn{ return &packetConn{
PacketConn: pc, PacketConn: pc,
limiter: limiter, limiter: lim,
inLimits: make(map[string]limiter.Limiter),
outLimits: make(map[string]limiter.Limiter),
} }
} }
func (c *packetConn) getInLimiter(addr net.Addr) limiter.Limiter {
if c.limiter == nil {
return nil
}
lim, ok := func() (limiter.Limiter, bool) {
c.inMux.RLock()
defer c.inMux.RUnlock()
lim, ok := c.inLimits[addr.String()]
return lim, ok
}()
if ok {
return lim
}
host, _, _ := net.SplitHostPort(addr.String())
lim = c.limiter.In(host)
c.inMux.Lock()
defer c.inMux.Unlock()
c.inLimits[addr.String()] = lim
return lim
}
func (c *packetConn) getOutLimiter(addr net.Addr) limiter.Limiter {
if c.limiter == nil {
return nil
}
lim, ok := func() (limiter.Limiter, bool) {
c.outMux.RLock()
defer c.outMux.RUnlock()
lim, ok := c.outLimits[addr.String()]
return lim, ok
}()
if ok {
return lim
}
host, _, _ := net.SplitHostPort(addr.String())
lim = c.limiter.Out(host)
c.outMux.Lock()
defer c.outMux.Unlock()
c.outLimits[addr.String()] = lim
return lim
}
func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
for { for {
n, addr, err = c.PacketConn.ReadFrom(p) n, addr, err = c.PacketConn.ReadFrom(p)
@ -120,13 +176,11 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
return return
} }
host, _, _ := net.SplitHostPort(addr.String()) limiter := c.getInLimiter(addr)
if limiter == nil {
if c.limiter == nil || c.limiter.In(host) == nil {
return return
} }
limiter := c.limiter.In(host)
// discard when exceed the limit size. // discard when exceed the limit size.
if limiter.Wait(context.Background(), n) < n { if limiter.Wait(context.Background(), n) < n {
continue continue
@ -137,15 +191,12 @@ 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) { func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
if c.limiter != nil {
host, _, _ := net.SplitHostPort(addr.String())
// discard when exceed the limit size. // discard when exceed the limit size.
if limiter := c.limiter.Out(host); limiter != nil && if limiter := c.getOutLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), len(p)) < len(p) { limiter.Wait(context.Background(), len(p)) < len(p) {
n = len(p) n = len(p)
return return
} }
}
return c.PacketConn.WriteTo(p, addr) return c.PacketConn.WriteTo(p, addr)
} }
@ -153,6 +204,10 @@ func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
type udpConn struct { type udpConn struct {
net.PacketConn net.PacketConn
limiter limiter.TrafficLimiter limiter limiter.TrafficLimiter
inLimits map[string]limiter.Limiter
inMux sync.RWMutex
outLimits map[string]limiter.Limiter
outMux sync.RWMutex
} }
func WrapUDPConn(limiter limiter.TrafficLimiter, pc net.PacketConn) udp.Conn { func WrapUDPConn(limiter limiter.TrafficLimiter, pc net.PacketConn) udp.Conn {
@ -162,6 +217,60 @@ func WrapUDPConn(limiter limiter.TrafficLimiter, pc net.PacketConn) udp.Conn {
} }
} }
func (c *udpConn) getInLimiter(addr net.Addr) limiter.Limiter {
if c.limiter == nil {
return nil
}
lim, ok := func() (limiter.Limiter, bool) {
c.inMux.RLock()
defer c.inMux.RUnlock()
lim, ok := c.inLimits[addr.String()]
return lim, ok
}()
if ok {
return lim
}
host, _, _ := net.SplitHostPort(addr.String())
lim = c.limiter.In(host)
c.inMux.Lock()
defer c.inMux.Unlock()
c.inLimits[addr.String()] = lim
return lim
}
func (c *udpConn) getOutLimiter(addr net.Addr) limiter.Limiter {
if c.limiter == nil {
return nil
}
lim, ok := func() (limiter.Limiter, bool) {
c.outMux.RLock()
defer c.outMux.RUnlock()
lim, ok := c.outLimits[addr.String()]
return lim, ok
}()
if ok {
return lim
}
host, _, _ := net.SplitHostPort(addr.String())
lim = c.limiter.Out(host)
c.outMux.Lock()
defer c.outMux.Unlock()
c.outLimits[addr.String()] = lim
return lim
}
func (c *udpConn) RemoteAddr() net.Addr { func (c *udpConn) RemoteAddr() net.Addr {
if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok { if nc, ok := c.PacketConn.(xnet.RemoteAddr); ok {
return nc.RemoteAddr() return nc.RemoteAddr()
@ -198,14 +307,10 @@ func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
if err != nil { if err != nil {
return return
} }
host, _, _ := net.SplitHostPort(addr.String())
if c.limiter == nil || c.limiter.In(host) == nil {
return
}
limiter := c.limiter.In(host)
// discard when exceed the limit size. // discard when exceed the limit size.
if limiter.Wait(context.Background(), n) < n { if limiter := c.getInLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), n) < n {
continue continue
} }
return return
@ -220,14 +325,9 @@ func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
return return
} }
host, _, _ := net.SplitHostPort(addr.String())
if c.limiter == nil || c.limiter.In(host) == nil {
return
}
limiter := c.limiter.In(host)
// discard when exceed the limit size. // discard when exceed the limit size.
if limiter.Wait(context.Background(), n) < n { if limiter := c.getInLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), n) < n {
continue continue
} }
return return
@ -245,14 +345,9 @@ func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAd
return return
} }
host, _, _ := net.SplitHostPort(addr.String())
if c.limiter == nil || c.limiter.In(host) == nil {
return
}
limiter := c.limiter.In(host)
// discard when exceed the limit size. // discard when exceed the limit size.
if limiter.Wait(context.Background(), n) < n { if limiter := c.getInLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), n) < n {
continue continue
} }
return return
@ -272,30 +367,24 @@ func (c *udpConn) Write(b []byte) (n int, err error) {
} }
func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
if c.limiter != nil {
host, _, _ := net.SplitHostPort(addr.String())
// discard when exceed the limit size. // discard when exceed the limit size.
if limiter := c.limiter.Out(host); limiter != nil && if limiter := c.getOutLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), len(p)) < len(p) { limiter.Wait(context.Background(), len(p)) < len(p) {
n = len(p) n = len(p)
return return
} }
}
n, err = c.PacketConn.WriteTo(p, addr) n, err = c.PacketConn.WriteTo(p, addr)
return return
} }
func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) { func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
if c.limiter != nil {
host, _, _ := net.SplitHostPort(addr.String())
// discard when exceed the limit size. // discard when exceed the limit size.
if limiter := c.limiter.Out(host); limiter != nil && if limiter := c.getOutLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), len(b)) < len(b) { limiter.Wait(context.Background(), len(b)) < len(b) {
n = len(b) n = len(b)
return return
} }
}
if nc, ok := c.PacketConn.(udp.WriteUDP); ok { if nc, ok := c.PacketConn.(udp.WriteUDP); ok {
n, err = nc.WriteToUDP(b, addr) n, err = nc.WriteToUDP(b, addr)
@ -306,15 +395,12 @@ 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) { func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
if c.limiter != nil {
host, _, _ := net.SplitHostPort(addr.String())
// discard when exceed the limit size. // discard when exceed the limit size.
if limiter := c.limiter.Out(host); limiter != nil && if limiter := c.getOutLimiter(addr); limiter != nil &&
limiter.Wait(context.Background(), len(b)) < len(b) { limiter.Wait(context.Background(), len(b)) < len(b) {
n = len(b) n = len(b)
return return
} }
}
if nc, ok := c.PacketConn.(udp.WriteUDP); ok { if nc, ok := c.PacketConn.(udp.WriteUDP); ok {
n, oobn, err = nc.WriteMsgUDP(b, oob, addr) n, oobn, err = nc.WriteMsgUDP(b, oob, addr)