feat(limiter/traffic): add burst support, dropped-packet counters, and cache fixes

- Add NewLimiterWithBurst(rate, burst) for independent burst control
- Extend parseLimit to accept optional 4th field as burst size
- Add DroppedPacketCounter interface with DroppedPackets() for packet/udp conns
- Export ErrRateLimited sentinel for rate-limit error checks
- Fix IP-level cache expiration (NoExpiration → defaultExpiration) to prevent
  unbounded growth; reset TTL on access in In/Out
- Short-circuit single-limiter path to avoid limiterGroup allocation
- Add nil guard in cachedTrafficLimiter when both old and new values are nil
- Remove unused ctx field from wrapper structs

limiter/traffic/cache/limiter.go

@@ -115,6 +115,9 @@ func (p *cachedTrafficLimiter) In(ctx context.Context, key string, opts ...limit
 
 	limNew := p.limiter.In(ctx, key, opts...)
 	if limNew == nil {
+		if lim == nil {
+			return nil
+		}
 		limNew = lim
 	}
 	if item == nil || !p.equal(lim, limNew) {
@@ -149,6 +152,9 @@ func (p *cachedTrafficLimiter) Out(ctx context.Context, key string, opts ...limi
 
 	limNew := p.limiter.Out(ctx, key, opts...)
 	if limNew == nil {
+		if lim == nil {
+			return nil
+		}
 		limNew = lim
 	}
 	if item == nil || !p.equal(lim, limNew) {

limiter/traffic/generator.go

@@ -5,14 +5,16 @@ import (
 )
 
 type limitGenerator struct {
-	in  int
+	in    int
-	out int
+	out   int
+	burst int
 }
 
-func newLimitGenerator(in, out int) *limitGenerator {
+func newLimitGenerator(in, out, burst int) *limitGenerator {
 	return &limitGenerator{
-		in:  in,
+		in:    in,
-		out: out,
+		out:   out,
+		burst: burst,
 	}
 }
 
@@ -20,6 +22,9 @@ func (p *limitGenerator) In() limiter.Limiter {
 	if p == nil || p.in <= 0 {
 		return nil
 	}
+	if p.burst > 0 {
+		return NewLimiterWithBurst(p.in, p.burst)
+	}
 	return NewLimiter(p.in)
 }
 
@@ -27,5 +32,8 @@ func (p *limitGenerator) Out() limiter.Limiter {
 	if p == nil || p.out <= 0 {
 		return nil
 	}
+	if p.burst > 0 {
+		return NewLimiterWithBurst(p.out, p.burst)
+	}
 	return NewLimiter(p.out)
 }

limiter/traffic/generator_test.go

@@ -1,11 +1,12 @@
 package traffic
 
 import (
+	"context"
 	"testing"
 )
 
 func TestLimitGenerator_In(t *testing.T) {
-	g := newLimitGenerator(100, 0)
+	g := newLimitGenerator(100, 0, 0)
 	lim := g.In()
 	if lim == nil {
 		t.Fatal("In() should not be nil for positive in")
@@ -22,14 +23,14 @@ func TestLimitGenerator_In(t *testing.T) {
 }
 
 func TestLimitGenerator_In_Zero(t *testing.T) {
-	g := newLimitGenerator(0, 200)
+	g := newLimitGenerator(0, 200, 0)
 	if lim := g.In(); lim != nil {
 		t.Fatal("In() should be nil for zero in")
 	}
 }
 
 func TestLimitGenerator_Out(t *testing.T) {
-	g := newLimitGenerator(0, 200)
+	g := newLimitGenerator(0, 200, 0)
 	lim := g.Out()
 	if lim == nil {
 		t.Fatal("Out() should not be nil for positive out")
@@ -46,7 +47,7 @@ func TestLimitGenerator_Out(t *testing.T) {
 }
 
 func TestLimitGenerator_Out_Zero(t *testing.T) {
-	g := newLimitGenerator(100, 0)
+	g := newLimitGenerator(100, 0, 0)
 	if lim := g.Out(); lim != nil {
 		t.Fatal("Out() should be nil for zero out")
 	}
@@ -61,3 +62,27 @@ func TestLimitGenerator_NilReceiver(t *testing.T) {
 		t.Fatal("nil receiver Out() should return nil")
 	}
 }
+
+func TestLimitGenerator_Burst(t *testing.T) {
+	// Burst specified: should use NewLimiterWithBurst.
+	g := newLimitGenerator(100, 200, 500)
+	limIn := g.In()
+	if limIn == nil {
+		t.Fatal("In() should not be nil")
+	}
+	if limIn.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", limIn.Limit())
+	}
+	// With burst=500, Wait for 500 should succeed immediately.
+	if n := limIn.Wait(context.Background(), 500); n != 500 {
+		t.Fatalf("expected 500 with burst, got %d", n)
+	}
+
+	limOut := g.Out()
+	if limOut == nil {
+		t.Fatal("Out() should not be nil")
+	}
+	if limOut.Limit() != 200 {
+		t.Fatalf("expected limit 200, got %d", limOut.Limit())
+	}
+}

limiter/traffic/limiter.go

@@ -15,13 +15,24 @@ type llimiter struct {
 }
 
 // NewLimiter creates a rate-based traffic Limiter with the given rate
-// (bytes per second) and burst size.
+// (bytes per second). The burst size is set equal to the rate.
 func NewLimiter(r int) limiter.Limiter {
 	return &llimiter{
 		limiter: rate.NewLimiter(rate.Limit(r), r),
 	}
 }
 
+// NewLimiterWithBurst creates a rate-based traffic Limiter with the given rate
+// (bytes per second) and burst size. If burst <= 0, it defaults to r.
+func NewLimiterWithBurst(r, burst int) limiter.Limiter {
+	if burst <= 0 {
+		burst = r
+	}
+	return &llimiter{
+		limiter: rate.NewLimiter(rate.Limit(r), burst),
+	}
+}
+
 func (l *llimiter) Wait(ctx context.Context, n int) int {
 	if l.limiter.Burst() < n {
 		n = l.limiter.Burst()

limiter/traffic/limiter_test.go

@@ -17,6 +17,40 @@ func TestNewLimiter(t *testing.T) {
 	}
 }
 
+func TestNewLimiterWithBurst(t *testing.T) {
+	l := NewLimiterWithBurst(100, 500)
+	if l == nil {
+		t.Fatal("NewLimiterWithBurst should not be nil")
+	}
+	if l.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", l.Limit())
+	}
+	// With burst=500, Wait for 500 should succeed immediately.
+	if n := l.Wait(context.Background(), 500); n != 500 {
+		t.Fatalf("expected 500 with burst, got %d", n)
+	}
+	// With burst=500, Wait for 501 clamps to burst.
+	if n := l.Wait(context.Background(), 501); n != 500 {
+		t.Fatalf("expected burst 500, got %d", n)
+	}
+}
+
+func TestNewLimiterWithBurst_Defaults(t *testing.T) {
+	// burst <= 0 defaults to rate.
+	l := NewLimiterWithBurst(100, 0)
+	if l.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", l.Limit())
+	}
+	if n := l.Wait(context.Background(), 100); n != 100 {
+		t.Fatalf("expected burst=rate=100, got %d", n)
+	}
+
+	l2 := NewLimiterWithBurst(50, -1)
+	if n := l2.Wait(context.Background(), 50); n != 50 {
+		t.Fatalf("expected burst=rate=50, got %d", n)
+	}
+}
+
 func TestLimiter_Wait(t *testing.T) {
 	// High burst so Wait does not block.
 	l := NewLimiter(1000000)

limiter/traffic/traffic.go

@@ -92,8 +92,9 @@ func LoggerOption(logger logger.Logger) Option {
 }
 
 type limitValue struct {
-	in  int
+	in    int
-	out int
+	out   int
+	burst int
 }
 
 type trafficLimiter struct {
@@ -193,6 +194,8 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt
 		// cached IP limiter
 		if lim != nil {
 			lims = append(lims, lim.(traffic.Limiter))
+			// reset expiration to prevent unbounded cache growth
+			l.inLimits.Set(host, lim, defaultExpiration)
 		}
 	} else {
 		l.mu.RLock()
@@ -204,14 +207,16 @@ func (l *trafficLimiter) In(ctx context.Context, key string, opts ...limiter.Opt
 			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)
+					l.inLimits.Set(host, lim, defaultExpiration)
 				}
 			}
 		}
 	}
 
 	var lim traffic.Limiter
-	if len(lims) > 0 {
+	if len(lims) == 1 {
+		lim = lims[0]
+	} else if len(lims) > 1 {
 		lim = newLimiterGroup(lims...)
 	}
 
@@ -275,6 +280,8 @@ func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Op
 		if lim != nil {
 			// cached IP level limiter
 			lims = append(lims, lim.(traffic.Limiter))
+			// reset expiration to prevent unbounded cache growth
+			l.outLimits.Set(host, lim, defaultExpiration)
 		}
 	} else {
 		l.mu.RLock()
@@ -286,14 +293,16 @@ func (l *trafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Op
 			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)
+					l.outLimits.Set(host, lim, defaultExpiration)
 				}
 			}
 		}
 	}
 
 	var lim traffic.Limiter
-	if len(lims) > 0 {
+	if len(lims) == 1 {
+		lim = lims[0]
+	} else if len(lims) > 1 {
 		lim = newLimiterGroup(lims...)
 	}
 
@@ -351,7 +360,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 			}
 		} else {
 			if value.in > 0 {
-				l.inLimits.Set(ServiceLimitKey, NewLimiter(value.in), cache.NoExpiration)
+				l.inLimits.Set(ServiceLimitKey, NewLimiterWithBurst(value.in, value.burst), cache.NoExpiration)
 			}
 		}
 
@@ -364,7 +373,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 			}
 		} else {
 			if value.out > 0 {
-				l.outLimits.Set(ServiceLimitKey, NewLimiter(value.out), cache.NoExpiration)
+				l.outLimits.Set(ServiceLimitKey, NewLimiterWithBurst(value.out, value.burst), cache.NoExpiration)
 			}
 		}
 		delete(values, ServiceLimitKey)
@@ -378,7 +387,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 		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))
+		l.generators.Store(ConnLimitKey, newLimitGenerator(value.in, value.out, value.burst))
 
 		if value.in <= 0 {
 			l.connInLimits.Flush()
@@ -420,7 +429,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 			if _, ipNet, _ := net.ParseCIDR(key); ipNet != nil {
 				cidrGenerators.Insert(&cidrLimitEntry{
 					ipNet:     *ipNet,
-					generator: newLimitGenerator(value.in, value.out),
+					generator: newLimitGenerator(value.in, value.out, value.burst),
 				})
 				continue
 			}
@@ -435,7 +444,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 				delete(inLimits, key)
 			} else {
 				if value.in > 0 {
-					l.inLimits.Set(key, NewLimiter(value.in), cache.NoExpiration)
+					l.inLimits.Set(key, NewLimiterWithBurst(value.in, value.burst), defaultExpiration)
 				}
 			}
 
@@ -449,7 +458,7 @@ func (l *trafficLimiter) reload(ctx context.Context) error {
 				delete(outLimits, key)
 			} else {
 				if value.out > 0 {
-					l.outLimits.Set(key, NewLimiter(value.out), cache.NoExpiration)
+					l.outLimits.Set(key, NewLimiterWithBurst(value.out, value.burst), defaultExpiration)
 				}
 			}
 		}
@@ -504,11 +513,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue
 	values = make(map[string]limitValue)
 
 	for _, v := range l.options.limits {
-		key, in, out := l.parseLimit(v)
+		key, in, out, burst := l.parseLimit(v)
 		if key == "" {
 			continue
 		}
-		values[key] = limitValue{in: in, out: out}
+		values[key] = limitValue{in: in, out: out, burst: burst}
 	}
 
 	if l.options.fileLoader != nil {
@@ -518,11 +527,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue
 				l.logger.Warnf("file loader: %v", er)
 			}
 			for _, s := range list {
-				key, in, out := l.parseLimit(l.parseLine(s))
+				key, in, out, burst := l.parseLimit(l.parseLine(s))
 				if key == "" {
 					continue
 				}
-				values[key] = limitValue{in: in, out: out}
+				values[key] = limitValue{in: in, out: out, burst: burst}
 			}
 		} else {
 			r, er := l.options.fileLoader.Load(ctx)
@@ -531,11 +540,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue
 			}
 			patterns, _ := l.parsePatterns(r)
 			for _, s := range patterns {
-				key, in, out := l.parseLimit(l.parseLine(s))
+				key, in, out, burst := l.parseLimit(l.parseLine(s))
 				if key == "" {
 					continue
 				}
-				values[key] = limitValue{in: in, out: out}
+				values[key] = limitValue{in: in, out: out, burst: burst}
 			}
 		}
 	}
@@ -546,11 +555,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue
 				l.logger.Warnf("redis loader: %v", er)
 			}
 			for _, s := range list {
-				key, in, out := l.parseLimit(l.parseLine(s))
+				key, in, out, burst := l.parseLimit(l.parseLine(s))
 				if key == "" {
 					continue
 				}
-				values[key] = limitValue{in: in, out: out}
+				values[key] = limitValue{in: in, out: out, burst: burst}
 			}
 		} else {
 			r, er := l.options.redisLoader.Load(ctx)
@@ -559,11 +568,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue
 			}
 			patterns, _ := l.parsePatterns(r)
 			for _, s := range patterns {
-				key, in, out := l.parseLimit(l.parseLine(s))
+				key, in, out, burst := l.parseLimit(l.parseLine(s))
 				if key == "" {
 					continue
 				}
-				values[key] = limitValue{in: in, out: out}
+				values[key] = limitValue{in: in, out: out, burst: burst}
 			}
 		}
 	}
@@ -574,11 +583,11 @@ func (l *trafficLimiter) load(ctx context.Context) (values map[string]limitValue
 		}
 		patterns, _ := l.parsePatterns(r)
 		for _, s := range patterns {
-			key, in, out := l.parseLimit(l.parseLine(s))
+			key, in, out, burst := l.parseLimit(l.parseLine(s))
 			if key == "" {
 				continue
 			}
-			values[key] = limitValue{in: in, out: out}
+			values[key] = limitValue{in: in, out: out, burst: burst}
 		}
 	}
 
@@ -609,7 +618,7 @@ func (l *trafficLimiter) parseLine(s string) string {
 	return strings.TrimSpace(s)
 }
 
-func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) {
+func (l *trafficLimiter) parseLimit(s string) (key string, in, out, burst int) {
 	s = strings.Replace(s, "\t", " ", -1)
 	s = strings.TrimSpace(s)
 	if s == "" {
@@ -635,6 +644,11 @@ func (l *trafficLimiter) parseLimit(s string) (key string, in, out int) {
 			out = int(v)
 		}
 	}
+	if len(ss) > 3 {
+		if v, _ := units.ParseBase2Bytes(ss[3]); v > 0 {
+			burst = int(v)
+		}
+	}
 
 	return
 }

limiter/traffic/traffic_test.go

@@ -110,7 +110,10 @@ func TestParseLimit(t *testing.T) {
 		{"key invalid_bytes", "key", 0, 0},
 	}
 	for _, tt := range tests {
-		key, in, out := l.parseLimit(tt.input)
+		key, in, out, burst := l.parseLimit(tt.input)
+		if burst != 0 {
+			t.Errorf("parseLimit(%q) unexpected burst: %d", tt.input, burst)
+		}
 		if key != tt.key || in != tt.in || out != tt.out {
 			t.Errorf("parseLimit(%q) = (%q, %d, %d), want (%q, %d, %d)",
 				tt.input, key, in, out, tt.key, tt.in, tt.out)

limiter/traffic/wrapper/conn.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"io"
 	"net"
+	"sync/atomic"
 	"syscall"
 
 	"github.com/go-gost/core/limiter"
@@ -23,6 +24,16 @@ var (
 	errRateLimited = errors.New("rate limited")
 )
 
+// ErrRateLimited is returned when a write exceeds the rate limit.
+var ErrRateLimited = errRateLimited
+
+// DroppedPacketCounter is an optional interface implemented by rate-limited
+// packet connections that report the number of packets discarded due to rate
+// limiting.
+type DroppedPacketCounter interface {
+	DroppedPackets() int64
+}
+
 // limitConn wraps a net.Conn with traffic rate limiting applied to reads and writes.
 type limitConn struct {
 	net.Conn
@@ -135,6 +146,12 @@ type packetConn struct {
 	limiter traffic.TrafficLimiter
 	opts    []limiter.Option
 	key     string
+	dropped atomic.Int64
+}
+
+// DroppedPackets returns the number of packets discarded due to rate limiting.
+func (c *packetConn) DroppedPackets() int64 {
+	return c.dropped.Load()
 }
 
 // WrapPacketConn wraps a net.PacketConn with traffic rate limiting. Packets
@@ -165,6 +182,7 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 
 		// discard when exceed the limit size.
 		if limiter.Wait(context.Background(), n) < n {
+			c.dropped.Add(1)
 			continue
 		}
 
@@ -176,7 +194,7 @@ func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 	limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 	if limiter != nil && limiter.Limit() > 0 &&
 		limiter.Wait(context.Background(), len(p)) < len(p) {
-		return 0, errRateLimited
+		return 0, ErrRateLimited
 	}
 
 	return c.PacketConn.WriteTo(p, addr)
@@ -194,6 +212,12 @@ type udpConn struct {
 	limiter traffic.TrafficLimiter
 	opts    []limiter.Option
 	key     string
+	dropped atomic.Int64
+}
+
+// DroppedPackets returns the number of packets discarded due to rate limiting.
+func (c *udpConn) DroppedPackets() int64 {
+	return c.dropped.Load()
 }
 
 // WrapUDPConn wraps a net.PacketConn as a udp.Conn with traffic rate limiting.
@@ -251,6 +275,7 @@ func (c *udpConn) Read(b []byte) (n int, err error) {
 
 		// discard when exceed the limit size.
 		if limiter.Wait(context.Background(), n) < n {
+			c.dropped.Add(1)
 			continue
 		}
 
@@ -276,6 +301,7 @@ func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 
 		// discard when exceed the limit size.
 		if limiter.Wait(context.Background(), n) < n {
+			c.dropped.Add(1)
 			continue
 		}
 
@@ -307,6 +333,7 @@ func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
 
 		// discard when exceed the limit size.
 		if limiter.Wait(context.Background(), n) < n {
+			c.dropped.Add(1)
 			continue
 		}
 
@@ -338,6 +365,7 @@ func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAd
 
 		// discard when exceed the limit size.
 		if limiter.Wait(context.Background(), n) < n {
+			c.dropped.Add(1)
 			continue
 		}
 		return
@@ -355,7 +383,7 @@ func (c *udpConn) Write(p []byte) (n int, err error) {
 		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 		if limiter != nil && limiter.Limit() > 0 &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
-			return 0, errRateLimited
+			return 0, ErrRateLimited
 		}
 	}
 
@@ -368,7 +396,7 @@ func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 		if limiter != nil && limiter.Limit() > 0 &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
-			return 0, errRateLimited
+			return 0, ErrRateLimited
 		}
 	}
 
@@ -387,7 +415,7 @@ func (c *udpConn) WriteToUDP(p []byte, addr *net.UDPAddr) (n int, err error) {
 		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 		if limiter != nil && limiter.Limit() > 0 &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
-			return 0, errRateLimited
+			return 0, ErrRateLimited
 		}
 	}
 
@@ -406,7 +434,7 @@ func (c *udpConn) WriteMsgUDP(p, oob []byte, addr *net.UDPAddr) (n, oobn int, er
 		limiter := c.limiter.Out(context.Background(), c.key, c.opts...)
 		if limiter != nil && limiter.Limit() > 0 &&
 			limiter.Wait(context.Background(), len(p)) < len(p) {
-			return 0, 0, errRateLimited
+			return 0, 0, ErrRateLimited
 		}
 	}
 

limiter/traffic/wrapper/conn_test.go

@@ -362,6 +362,35 @@ func TestWrapUDPConn(t *testing.T) {
 	}
 }
 
+func TestPacketConn_DroppedPackets(t *testing.T) {
+	pc := &mockPacketConn{writeBuf: make([]byte, 100)}
+	calls := 0
+	pc.readFromFunc = func(p []byte) (int, net.Addr, error) {
+		calls++
+		if calls == 1 {
+			return copy(p, "data-4"), nil, nil // 6 bytes, limiter allows 5 = drop
+		}
+		return copy(p, "ok"), nil, nil // 2 bytes, limiter allows 5 = pass
+	}
+	ml := &mockLimiter{
+		limit: 100,
+		waitFunc: func(ctx context.Context, n int) int {
+			return 5 // allow 5 bytes; 6-byte packet is dropped, 2-byte passes
+		},
+	}
+	tl := &mockTrafficLimiter{inLim: ml}
+	wrapped := WrapPacketConn(pc, tl, "test-key")
+
+	_, _, err := wrapped.ReadFrom(make([]byte, 10))
+	if err != nil {
+		t.Fatal(err)
+	}
+	counter := wrapped.(DroppedPacketCounter)
+	if n := counter.DroppedPackets(); n != 1 {
+		t.Fatalf("expected 1 dropped packet, got %d", n)
+	}
+}
+
 func TestUDPConn_Write_RateLimited(t *testing.T) {
 	pc := &mockPacketConn{writeBuf: make([]byte, 100)}
 	ml := &mockLimiter{
@@ -388,6 +417,9 @@ var (
 	_ xio.CloseWrite = (*limitConn)(nil)
 	_ syscall.Conn   = (*limitConn)(nil)
 	_ ctxutil.Context = (*limitConn)(nil)
+
+	_ DroppedPacketCounter = (*packetConn)(nil)
+	_ DroppedPacketCounter = (*udpConn)(nil)
 )
 
 // --- helpers ---
@@ -412,7 +444,8 @@ func (rw *bytesReadWriter) Write(b []byte) (int, error) {
 
 type mockPacketConn struct {
 	net.PacketConn
-	writeBuf []byte
+	writeBuf     []byte
+	readFromFunc func(p []byte) (int, net.Addr, error)
 }
 
 func (pc *mockPacketConn) Write(p []byte) (int, error) {
@@ -426,6 +459,9 @@ func (pc *mockPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
 }
 
 func (pc *mockPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
+	if pc.readFromFunc != nil {
+		return pc.readFromFunc(p)
+	}
 	return copy(p, "test"), nil, nil
 }