fix(limiter/traffic): infinite Write loop on zero burst, missing httpLoader.Close, stale reload values, ignored WaitN error, silent UDP drops, SplitHostPort fallback

limiter/traffic/limiter.go

@@ -24,7 +24,9 @@ func (l *llimiter) Wait(ctx context.Context, n int) int {
 	if l.limiter.Burst() < n {
 		n = l.limiter.Burst()
 	}
-	l.limiter.WaitN(ctx, n)
+	if err := l.limiter.WaitN(ctx, n); err != nil {
+		return 0
+	}
 	return n
 }
 

limiter/traffic/traffic.go

@@ -166,7 +166,10 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt
 		}
 	}
 
-	host, _, _ := net.SplitHostPort(key)
+	host, _, err := net.SplitHostPort(key)
+	if err != nil {
+		host = key
+	}
 	// IP level limiter
 	if lim, ok := l.inLimits.Get(host); ok {
 		// cached IP limiter
@@ -245,7 +248,10 @@ func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Op
 		}
 	}
 
-	host, _, _ := net.SplitHostPort(key)
+	host, _, err := net.SplitHostPort(key)
+	if err != nil {
+		host = key
+	}
 	// IP level limiter
 	if lim, ok := l.outLimits.Get(host); ok {
 		if lim != nil {
@@ -362,7 +368,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 			if in != value.in {
 				for _, item := range l.connInLimits.Items() {
 					if v := item.Object; v != nil {
-						v.(traffic.Limiter).Set(in)
+						v.(traffic.Limiter).Set(value.in)
 					}
 				}
 			}
@@ -374,7 +380,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 			if out != value.out {
 				for _, item := range l.connOutLimits.Items() {
 					if v := item.Object; v != nil {
-						v.(traffic.Limiter).Set(out)
+						v.(traffic.Limiter).Set(value.out)
 					}
 				}
 			}
@@ -623,6 +629,9 @@ func (l *trafficLimiter) Close() error {
 	if l.options.redisLoader != nil {
 		l.options.redisLoader.Close()
 	}
+	if l.options.httpLoader != nil {
+		l.options.httpLoader.Close()
+	}
 	return nil
 }
 

limiter/traffic/wrapper/conn.go

@@ -17,7 +17,8 @@ import (
 )
 
 var (
-	errUnsupport = errors.New("unsupported operation")
+	errUnsupport   = errors.New("unsupported operation")
+	errRateLimited = errors.New("rate limited")
 )
 
 // limitConn is a Conn with traffic limiter supported.
@@ -80,7 +81,11 @@ func (c *limitConn) Write(b []byte) (n int, err error) {
 
 	nn := 0
 	for len(b) > 0 {
-		nn, err = c.Conn.Write(b[:limiter.Wait(context.Background(), len(b))])
+		burst := limiter.Wait(context.Background(), len(b))
+		if burst == 0 {
+			return
+		}
+		nn, err = c.Conn.Write(b[:burst])
 		n += nn
 		if err != nil {
 			return
@@ -162,12 +167,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) {
-	// discard when exceed the limit size.
 	limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 	if limiter != nil && limiter.Limit() > 0 &&
 		limiter.Wait(context.Background(), len(p)) < len(p) {
-		n = len(p)
-		return
+		return 0, errRateLimited
 	}
 
 	return c.PacketConn.WriteTo(p, addr)
@@ -342,12 +345,10 @@ func (c *udpConn) Write(p []byte) (n int, err error) {
 	}
 
 	if c.limiter != nil {
-		// discard when exceed the limit size.
 		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 		if limiter != nil && limiter.Limit() > 0 &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
-			n = len(p)
-			return
+			return 0, errRateLimited
 		}
 	}
 
@@ -357,12 +358,10 @@ func (c *udpConn) Write(p []byte) (n int, err error) {
 
 func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 	if c.limiter != nil {
-		// discard when exceed the limit size.
 		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 		if limiter != nil && limiter.Limit() > 0 &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
-			n = len(p)
-			return
+			return 0, errRateLimited
 		}
 	}
 
@@ -378,12 +377,10 @@ func (c *udpConn) WriteToUDP(p []byte, addr *net.UDPAddr) (n int, err error) {
 	}
 
 	if c.limiter != nil {
-		// discard when exceed the limit size.
 		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 		if limiter != nil && limiter.Limit() > 0 &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
-			n = len(p)
-			return
+			return 0, errRateLimited
 		}
 	}
 
@@ -399,12 +396,10 @@ func (c *udpConn) WriteMsgUDP(p, oob []byte, addr *net.UDPAddr) (n, oobn int, er
 	}
 
 	if c.limiter != nil {
-		// discard when exceed the limit size.
 		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 		if limiter != nil && limiter.Limit() > 0 &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
-			n = len(p)
-			return
+			return 0, 0, errRateLimited
 		}
 	}
 

limiter/traffic/wrapper/io.go

@@ -69,7 +69,11 @@ func (p *readWriter) Write(b []byte) (n int, err error) {
 
 	nn := 0
 	for len(b) > 0 {
-		nn, err = p.ReadWriter.Write(b[:limiter.Wait(context.Background(), len(b))])
+		burst := limiter.Wait(context.Background(), len(b))
+		if burst == 0 {
+			return
+		}
+		nn, err = p.ReadWriter.Write(b[:burst])
 		n += nn
 		if err != nil {
 			return