docs(limiter/traffic): add doc comments, unit tests

limiter/traffic/cache/limiter.go

@@ -1,3 +1,6 @@
+// Package limiter provides a cached traffic limiter that wraps a
+// TrafficLimiter with in-memory caching to reduce calls to the underlying
+// limiter (e.g., plugin-based limiters).
 package limiter
 
 import (
@@ -20,20 +23,27 @@ type options struct {
 	scope           string
 }
 
+// Option configures a cachedTrafficLimiter.
 type Option func(*options)
 
+// RefreshIntervalOption sets the cache refresh interval. Cached limiters
+// are refreshed after this duration. The minimum value is 1 second.
 func RefreshIntervalOption(interval time.Duration) Option {
 	return func(o *options) {
 		o.refreshInterval = interval
 	}
 }
 
+// CleanupIntervalOption sets the interval for cleaning up expired cache
+// entries. The minimum value is 1 second.
 func CleanupIntervalOption(interval time.Duration) Option {
 	return func(o *options) {
 		o.cleanupInterval = interval
 	}
 }
 
+// ScopeOption filters cached limiters by scope. When set, only requests
+// with a matching scope return a limiter; non-matching scopes return nil.
 func ScopeOption(scope string) Option {
 	return func(o *options) {
 		o.scope = scope
@@ -47,6 +57,9 @@ type cachedTrafficLimiter struct {
 	options   options
 }
 
+// NewCachedTrafficLimiter wraps a TrafficLimiter with an in-memory cache.
+// Cached limiters are refreshed after the configured refresh interval.
+// Returns nil if the inner limiter is nil.
 func NewCachedTrafficLimiter(limiter traffic.TrafficLimiter, opts ...Option) traffic.TrafficLimiter {
 	if limiter == nil {
 		return nil

limiter/traffic/cache/limiter_test.go

@@ -0,0 +1,175 @@
+package limiter
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+)
+
+// --- mock types ---
+
+type mockLimiter struct {
+	limit int
+}
+
+func (m *mockLimiter) Wait(ctx context.Context, n int) int { return n }
+func (m *mockLimiter) Limit() int                          { return m.limit }
+func (m *mockLimiter) Set(n int)                           { m.limit = n }
+
+type mockTrafficLimiter struct {
+	inCalls  int
+	outCalls int
+	inLim    traffic.Limiter
+	outLim   traffic.Limiter
+}
+
+func (m *mockTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	m.inCalls++
+	return m.inLim
+}
+
+func (m *mockTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	m.outCalls++
+	return m.outLim
+}
+
+// --- tests ---
+
+func TestCachedTrafficLimiter_In(t *testing.T) {
+	inner := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}}
+	cached := NewCachedTrafficLimiter(inner, RefreshIntervalOption(50*time.Millisecond))
+	if cached == nil {
+		t.Fatal("NewCachedTrafficLimiter should not return nil")
+	}
+
+	ctx := context.Background()
+	v1 := cached.In(ctx, "key1")
+	if v1 == nil {
+		t.Fatal("first In() should return a limiter")
+	}
+	if v1.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", v1.Limit())
+	}
+	if inner.inCalls != 1 {
+		t.Fatalf("expected 1 In call, got %d", inner.inCalls)
+	}
+
+	// Second call: should return cached result without calling inner again.
+	v2 := cached.In(ctx, "key1")
+	if v2 == nil {
+		t.Fatal("cached In() should return a limiter")
+	}
+	if inner.inCalls != 1 {
+		t.Fatalf("cached call should not invoke inner, got %d calls", inner.inCalls)
+	}
+}
+
+func TestCachedTrafficLimiter_Out(t *testing.T) {
+	inner := &mockTrafficLimiter{outLim: &mockLimiter{limit: 200}}
+	cached := NewCachedTrafficLimiter(inner, RefreshIntervalOption(50*time.Millisecond))
+
+	ctx := context.Background()
+	v1 := cached.Out(ctx, "key1")
+	if v1 == nil {
+		t.Fatal("first Out() should return a limiter")
+	}
+	if v1.Limit() != 200 {
+		t.Fatalf("expected limit 200, got %d", v1.Limit())
+	}
+	if inner.outCalls != 1 {
+		t.Fatalf("expected 1 Out call, got %d", inner.outCalls)
+	}
+
+	cached.Out(ctx, "key1")
+	if inner.outCalls != 1 {
+		t.Fatalf("cached call should not invoke inner, got %d calls", inner.outCalls)
+	}
+}
+
+func TestCachedTrafficLimiter_Expiration(t *testing.T) {
+	inner := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}}
+	// Bypass constructor to avoid 1s minimum refresh interval.
+	cached := NewCachedTrafficLimiter(inner, RefreshIntervalOption(10*time.Millisecond))
+	if cached == nil {
+		t.Fatal("NewCachedTrafficLimiter should not return nil")
+	}
+	ct := cached.(*cachedTrafficLimiter)
+	ct.options.refreshInterval = 10 * time.Millisecond
+
+	ctx := context.Background()
+	ct.In(ctx, "key1")
+	if inner.inCalls != 1 {
+		t.Fatalf("expected 1 In call, got %d", inner.inCalls)
+	}
+
+	// Wait for cache to expire.
+	time.Sleep(30 * time.Millisecond)
+
+	ct.In(ctx, "key1")
+	if inner.inCalls != 2 {
+		t.Fatalf("after expiration, inner In should be called again, got %d calls", inner.inCalls)
+	}
+}
+
+func TestCachedTrafficLimiter_ScopeFilter(t *testing.T) {
+	inner := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}}
+	cached := NewCachedTrafficLimiter(inner,
+		RefreshIntervalOption(50*time.Millisecond),
+		ScopeOption("service"),
+	)
+
+	ctx := context.Background()
+
+	// Non-matching scope: returns nil.
+	v := cached.In(ctx, "key1", limiter.ScopeOption(limiter.ScopeConn))
+	if v != nil {
+		t.Fatal("non-matching scope should return nil")
+	}
+
+	// Matching scope: returns limiter.
+	v = cached.In(ctx, "key1", limiter.ScopeOption(limiter.ScopeService))
+	if v == nil {
+		t.Fatal("matching scope should return limiter")
+	}
+}
+
+func TestCachedTrafficLimiter_NilInner(t *testing.T) {
+	cached := NewCachedTrafficLimiter(nil)
+	if cached != nil {
+		t.Fatal("NewCachedTrafficLimiter(nil) should return nil")
+	}
+}
+
+func TestCachedTrafficLimiter_NilNew(t *testing.T) {
+	inner := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}}
+	cached := NewCachedTrafficLimiter(inner, RefreshIntervalOption(10*time.Millisecond))
+	if cached == nil {
+		t.Fatal("NewCachedTrafficLimiter should not return nil")
+	}
+	ct := cached.(*cachedTrafficLimiter)
+	ct.options.refreshInterval = 10 * time.Millisecond
+
+	ctx := context.Background()
+	v1 := ct.In(ctx, "key1")
+	if v1 == nil || v1.Limit() != 100 {
+		t.Fatal("first In() should return limiter with limit 100")
+	}
+
+	// Now make inner return nil.
+	inner.inLim = nil
+
+	// Wait for cache to expire.
+	time.Sleep(30 * time.Millisecond)
+
+	// After expiration: inner returns nil, falls back to cached limiter.
+	v2 := ct.In(ctx, "key1")
+	if v2 == nil {
+		t.Fatal("should fall back to cached limiter when inner returns nil")
+	}
+	if v2.Limit() != 100 {
+		t.Fatalf("expected fallback limit 100, got %d", v2.Limit())
+	}
+}

limiter/traffic/generator_test.go

@@ -0,0 +1,63 @@
+package traffic
+
+import (
+	"testing"
+)
+
+func TestLimitGenerator_In(t *testing.T) {
+	g := newLimitGenerator(100, 0)
+	lim := g.In()
+	if lim == nil {
+		t.Fatal("In() should not be nil for positive in")
+	}
+	if lim.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", lim.Limit())
+	}
+
+	// Each In() call creates a new instance.
+	lim2 := g.In()
+	if lim2 == lim {
+		t.Fatal("each In() call should create a new instance")
+	}
+}
+
+func TestLimitGenerator_In_Zero(t *testing.T) {
+	g := newLimitGenerator(0, 200)
+	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)
+	lim := g.Out()
+	if lim == nil {
+		t.Fatal("Out() should not be nil for positive out")
+	}
+	if lim.Limit() != 200 {
+		t.Fatalf("expected limit 200, got %d", lim.Limit())
+	}
+
+	// Each Out() call creates a new instance.
+	lim2 := g.Out()
+	if lim2 == lim {
+		t.Fatal("each Out() call should create a new instance")
+	}
+}
+
+func TestLimitGenerator_Out_Zero(t *testing.T) {
+	g := newLimitGenerator(100, 0)
+	if lim := g.Out(); lim != nil {
+		t.Fatal("Out() should be nil for zero out")
+	}
+}
+
+func TestLimitGenerator_NilReceiver(t *testing.T) {
+	var g *limitGenerator
+	if lim := g.In(); lim != nil {
+		t.Fatal("nil receiver In() should return nil")
+	}
+	if lim := g.Out(); lim != nil {
+		t.Fatal("nil receiver Out() should return nil")
+	}
+}

limiter/traffic/limiter.go

@@ -14,6 +14,8 @@ type llimiter struct {
 	limiter *rate.Limiter
 }
 
+// NewLimiter creates a rate-based traffic Limiter with the given rate
+// (bytes per second) and burst size.
 func NewLimiter(r int) limiter.Limiter {
 	return &llimiter{
 		limiter: rate.NewLimiter(rate.Limit(r), r),

limiter/traffic/limiter_test.go

@@ -0,0 +1,135 @@
+package traffic
+
+import (
+	"context"
+	"testing"
+
+	limiter "github.com/go-gost/core/limiter/traffic"
+)
+
+func TestNewLimiter(t *testing.T) {
+	l := NewLimiter(100)
+	if l == nil {
+		t.Fatal("NewLimiter should not be nil")
+	}
+	if l.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", l.Limit())
+	}
+}
+
+func TestLimiter_Wait(t *testing.T) {
+	// High burst so Wait does not block.
+	l := NewLimiter(1000000)
+	n := l.Wait(context.Background(), 3)
+	if n != 3 {
+		t.Fatalf("expected 3, got %d", n)
+	}
+}
+
+func TestLimiter_Wait_BurstClamped(t *testing.T) {
+	// Burst is 5, request 100 -> clamped to burst.
+	l := NewLimiter(5)
+	n := l.Wait(context.Background(), 100)
+	if n != 5 {
+		t.Fatalf("expected burst 5, got %d", n)
+	}
+}
+
+func TestLimiter_Wait_ContextCancelled(t *testing.T) {
+	l := NewLimiter(1000)
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	n := l.Wait(ctx, 10)
+	if n != 0 {
+		t.Fatalf("cancelled context should return 0, got %d", n)
+	}
+}
+
+func TestLimiter_Set(t *testing.T) {
+	l := NewLimiter(5)
+	l.Set(10)
+	if l.Limit() != 10 {
+		t.Fatalf("expected limit 10 after Set, got %d", l.Limit())
+	}
+	// Verify burst also updated: Wait should allow up to 10.
+	n := l.Wait(context.Background(), 10)
+	if n != 10 {
+		t.Fatalf("expected burst 10 after Set, got %d", n)
+	}
+}
+
+func TestLimiter_String(t *testing.T) {
+	l := NewLimiter(42)
+	// llimiter implements fmt.Stringer (not part of core interface).
+	if s := l.(*llimiter).String(); s != "42" {
+		t.Fatalf("expected '42', got '%s'", s)
+	}
+}
+
+func TestLimiterGroup_SortedByLimit(t *testing.T) {
+	g := newLimiterGroup(
+		NewLimiter(200),
+		NewLimiter(50),
+		NewLimiter(100),
+	)
+	if g.Limit() != 50 {
+		t.Fatalf("expected smallest limit 50, got %d", g.Limit())
+	}
+}
+
+func TestLimiterGroup_Wait(t *testing.T) {
+	// Group: [3, 100000] — smallest limit is 3.
+	small := NewLimiter(3)
+	large := NewLimiter(100000)
+	g := newLimiterGroup(large, small)
+
+	// Wait for 3: small returns 3, large returns 3, min = 3.
+	n := g.Wait(context.Background(), 3)
+	if n != 3 {
+		t.Fatalf("expected 3, got %d", n)
+	}
+
+	// Wait for 5: small clamps to burst=3, returns 3, large returns 5, min = 3.
+	n = g.Wait(context.Background(), 5)
+	if n != 3 {
+		t.Fatalf("expected min=3, got %d", n)
+	}
+}
+
+func TestLimiterGroup_Limit(t *testing.T) {
+	g := newLimiterGroup(NewLimiter(10), NewLimiter(5), NewLimiter(20))
+	if g.Limit() != 5 {
+		t.Fatalf("expected 5, got %d", g.Limit())
+	}
+}
+
+func TestLimiterGroup_Set(t *testing.T) {
+	g := newLimiterGroup(NewLimiter(5))
+	g.Set(100)
+	if g.Limit() != 5 {
+		t.Fatal("Set should be a no-op on limiterGroup")
+	}
+}
+
+func TestLimiterGroup_Empty(t *testing.T) {
+	g := newLimiterGroup()
+	if g.Limit() != 0 {
+		t.Fatalf("empty group limit should be 0, got %d", g.Limit())
+	}
+	n := g.Wait(context.Background(), 100)
+	if n != 100 {
+		t.Fatalf("empty group Wait should return n unchanged, got %d", n)
+	}
+}
+
+func TestLimiterGroup_String(t *testing.T) {
+	g := newLimiterGroup(NewLimiter(5))
+	if g.String() != "[5]" {
+		t.Fatalf("expected '[5]', got '%s'", g.String())
+	}
+}
+
+var (
+	_ limiter.Limiter = (*llimiter)(nil)
+	_ limiter.Limiter = (*limiterGroup)(nil)
+)

limiter/traffic/traffic.go

@@ -1,3 +1,6 @@
+// Package traffic provides a traffic/bandwidth limiter implementation that
+// supports per-service, per-connection, per-IP, and per-CIDR rate limiting
+// for both ingress (In) and egress (Out) directions.
 package traffic
 
 import (
@@ -20,8 +23,10 @@ import (
 )
 
 const (
+	// ServiceLimitKey is the special key used for service-level rate limits.
 	ServiceLimitKey = "$"
-	ConnLimitKey    = "$$"
+	// ConnLimitKey is the special key used for connection-level rate limits.
+	ConnLimitKey = "$$"
 )
 
 const (
@@ -38,38 +43,48 @@ type options struct {
 	logger      logger.Logger
 }
 
+// Option configures a TrafficLimiter.
 type Option func(opts *options)
 
+// LimitsOption sets the static limit rules. Each string is a limit definition
+// in the format "<key> <in> [out]" where in/out are byte quantities with
+// optional SI/IEC suffixes (e.g., "100B", "1MB", "512KB").
 func LimitsOption(limits ...string) Option {
 	return func(opts *options) {
 		opts.limits = limits
 	}
 }
 
+// ReloadPeriodOption sets the period for reloading limit rules from external
+// loaders. A period of zero disables periodic reloading.
 func ReloadPeriodOption(period time.Duration) Option {
 	return func(opts *options) {
 		opts.period = period
 	}
 }
 
+// FileLoaderOption sets a file-based loader for limit rules.
 func FileLoaderOption(fileLoader loader.Loader) Option {
 	return func(opts *options) {
 		opts.fileLoader = fileLoader
 	}
 }
 
+// RedisLoaderOption sets a Redis-based loader for limit rules.
 func RedisLoaderOption(redisLoader loader.Loader) Option {
 	return func(opts *options) {
 		opts.redisLoader = redisLoader
 	}
 }
 
+// HTTPLoaderOption sets an HTTP-based loader for limit rules.
 func HTTPLoaderOption(httpLoader loader.Loader) Option {
 	return func(opts *options) {
 		opts.httpLoader = httpLoader
 	}
 }
 
+// LoggerOption sets the logger for the traffic limiter.
 func LoggerOption(logger logger.Logger) Option {
 	return func(opts *options) {
 		opts.logger = logger
@@ -96,6 +111,9 @@ type trafficLimiter struct {
 	cancelFunc context.CancelFunc
 }
 
+// NewTrafficLimiter creates a new TrafficLimiter with the given options.
+// It starts a background goroutine for periodic reload of limit rules from
+// external loaders.
 func NewTrafficLimiter(opts ...Option) traffic.TrafficLimiter {
 	var options options
 	for _, opt := range opts {

limiter/traffic/traffic_test.go

@@ -0,0 +1,578 @@
+package traffic
+
+import (
+	"context"
+	"io"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/go-gost/core/limiter"
+	traffic "github.com/go-gost/core/limiter/traffic"
+	"github.com/go-gost/x/internal/loader"
+	xlogger "github.com/go-gost/x/logger"
+	"github.com/patrickmn/go-cache"
+	"github.com/yl2chen/cidranger"
+)
+
+// --- test helpers ---
+
+func newTestTrafficLimiter(limits ...string) *trafficLimiter {
+	ctx, cancel := context.WithCancel(context.Background())
+	lim := &trafficLimiter{
+		cidrGenerators: cidranger.NewPCTrieRanger(),
+		connInLimits:   cache.New(defaultExpiration, cleanupInterval),
+		connOutLimits:  cache.New(defaultExpiration, cleanupInterval),
+		inLimits:       cache.New(defaultExpiration, cleanupInterval),
+		outLimits:      cache.New(defaultExpiration, cleanupInterval),
+		options:        options{limits: limits},
+		cancelFunc:     cancel,
+		logger:         xlogger.Nop(),
+	}
+	_ = lim.reload(ctx)
+	return lim
+}
+
+type fakeLoader struct {
+	data    string
+	listErr error
+	loadErr error
+	closed  bool
+}
+
+func (l *fakeLoader) List(ctx context.Context) ([]string, error) {
+	if l.listErr != nil {
+		return nil, l.listErr
+	}
+	if l.data == "" {
+		return nil, nil
+	}
+	return strings.Split(l.data, "\n"), nil
+}
+
+func (l *fakeLoader) Load(ctx context.Context) (io.Reader, error) {
+	if l.loadErr != nil {
+		return nil, l.loadErr
+	}
+	return strings.NewReader(l.data), nil
+}
+
+func (l *fakeLoader) Close() error {
+	l.closed = true
+	return nil
+}
+
+type errorLoader struct{}
+
+func (e *errorLoader) Load(ctx context.Context) (io.Reader, error) { return nil, io.ErrUnexpectedEOF }
+func (e *errorLoader) Close() error                                { return nil }
+
+// --- parsing tests ---
+
+func TestParseLine(t *testing.T) {
+	l := newTestTrafficLimiter()
+	tests := []struct {
+		input, expected string
+	}{
+		{"192.168.1.1 100B 200B", "192.168.1.1 100B 200B"},
+		{"  192.168.1.1 100B 200B  ", "192.168.1.1 100B 200B"},
+		{"192.168.1.1 100B 200B # comment", "192.168.1.1 100B 200B"},
+		{"# only comment", ""},
+		{"   # indented comment", ""},
+		{"", ""},
+	}
+	for _, tt := range tests {
+		got := l.parseLine(tt.input)
+		if got != tt.expected {
+			t.Errorf("parseLine(%q) = %q, want %q", tt.input, got, tt.expected)
+		}
+	}
+}
+
+func TestParseLimit(t *testing.T) {
+	l := newTestTrafficLimiter()
+	tests := []struct {
+		input   string
+		key     string
+		in, out int
+	}{
+		{"192.168.1.1 100B 200B", "192.168.1.1", 100, 200},
+		{"192.168.1.1\t100B\t200B", "192.168.1.1", 100, 200},
+		{"$ 1KB", "$", 1024, 0},
+		{"$$ 500B", "$$", 500, 0},
+		{"10.0.0.0/8 1KB 2KB", "10.0.0.0/8", 1024, 2048},
+		{"key   100B   200B", "key", 100, 200},
+		{"key 1MB", "key", 1048576, 0},
+		{"key 512KB 1MB", "key", 524288, 1048576},
+		{"key_only", "", 0, 0},
+		{"", "", 0, 0},
+		{"key invalid_bytes", "key", 0, 0},
+	}
+	for _, tt := range tests {
+		key, in, out := l.parseLimit(tt.input)
+		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)
+		}
+	}
+}
+
+func TestParsePatterns(t *testing.T) {
+	l := newTestTrafficLimiter()
+	r := strings.NewReader("192.168.1.1 100B\n# comment\n\n10.0.0.1 200B\n")
+	patterns, err := l.parsePatterns(r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(patterns) != 2 {
+		t.Fatalf("expected 2 patterns, got %d: %v", len(patterns), patterns)
+	}
+	if patterns[0] != "192.168.1.1 100B" {
+		t.Errorf("patterns[0] = %q", patterns[0])
+	}
+	if patterns[1] != "10.0.0.1 200B" {
+		t.Errorf("patterns[1] = %q", patterns[1])
+	}
+}
+
+func TestParsePatterns_NilReader(t *testing.T) {
+	l := newTestTrafficLimiter()
+	patterns, err := l.parsePatterns(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(patterns) != 0 {
+		t.Fatalf("expected 0 patterns, got %d", len(patterns))
+	}
+}
+
+// --- load tests ---
+
+func TestLoad_NoLoaders(t *testing.T) {
+	l := newTestTrafficLimiter("$ 100B")
+	values, err := l.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	v, ok := values["$"]
+	if !ok {
+		t.Fatal("expected $ key in loaded values")
+	}
+	if v.in != 100 || v.out != 0 {
+		t.Fatalf("expected in=100 out=0, got in=%d out=%d", v.in, v.out)
+	}
+}
+
+func TestLoad_FileLister(t *testing.T) {
+	l := newTestTrafficLimiter()
+	l.options.fileLoader = &fakeLoader{data: "$ 50B\n192.168.1.1 200B 300B"}
+	values, err := l.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if v := values["$"]; v.in != 50 {
+		t.Fatalf("expected $ in=50, got %d", v.in)
+	}
+	v := values["192.168.1.1"]
+	if v.in != 200 || v.out != 300 {
+		t.Fatalf("expected in=200 out=300, got in=%d out=%d", v.in, v.out)
+	}
+}
+
+func TestLoad_FileLoader(t *testing.T) {
+	l := newTestTrafficLimiter()
+	l.options.fileLoader = &fakeLoader{data: "$ 30B"}
+	values, err := l.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if v := values["$"]; v.in != 30 {
+		t.Fatalf("expected $ in=30, got %d", v.in)
+	}
+}
+
+func TestLoad_RedisLoader(t *testing.T) {
+	l := newTestTrafficLimiter()
+	l.options.redisLoader = &fakeLoader{data: "$ 40B"}
+	values, err := l.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if v := values["$"]; v.in != 40 {
+		t.Fatalf("expected $ in=40, got %d", v.in)
+	}
+}
+
+func TestLoad_HTTPLoader(t *testing.T) {
+	l := newTestTrafficLimiter()
+	l.options.httpLoader = &fakeLoader{data: "$ 35B"}
+	values, err := l.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if v := values["$"]; v.in != 35 {
+		t.Fatalf("expected $ in=35, got %d", v.in)
+	}
+}
+
+func TestLoad_ErrorLoader(t *testing.T) {
+	l := newTestTrafficLimiter("$ 100B")
+	l.options.fileLoader = &errorLoader{}
+	values, err := l.load(context.Background())
+	if err != nil {
+		t.Fatal("error loader should not propagate error")
+	}
+	if v := values["$"]; v.in != 100 {
+		t.Fatalf("expected $ in=100 from static limits, got %d", v.in)
+	}
+}
+
+func TestLoad_MergesLoaders(t *testing.T) {
+	l := newTestTrafficLimiter("$ 100B")
+	l.options.fileLoader = &fakeLoader{data: "192.168.1.1 50B"}
+	l.options.httpLoader = &fakeLoader{data: "$ 200B"}
+	values, err := l.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if v := values["$"]; v.in != 200 {
+		t.Fatalf("expected $ in=200 (HTTP overrides static), got %d", v.in)
+	}
+	if v := values["192.168.1.1"]; v.in != 50 {
+		t.Fatalf("expected 192.168.1.1 in=50, got %d", v.in)
+	}
+}
+
+// --- In/Out tests ---
+
+func TestIn_ServiceScope(t *testing.T) {
+	l := newTestTrafficLimiter("$ 100B")
+	v := l.In(context.Background(), "any-key", limiter.ScopeOption(limiter.ScopeService))
+	if v == nil {
+		t.Fatal("expected service-level In limiter")
+	}
+	if v.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", v.Limit())
+	}
+}
+
+func TestIn_ClientScope(t *testing.T) {
+	l := newTestTrafficLimiter("$ 100B")
+	v := l.In(context.Background(), "any-key", limiter.ScopeOption(limiter.ScopeClient))
+	if v != nil {
+		t.Fatal("ScopeClient should return nil")
+	}
+}
+
+func TestIn_ConnScope_GeneratesLimiter(t *testing.T) {
+	l := newTestTrafficLimiter("$$ 100B")
+	v := l.In(context.Background(), "10.0.0.1:12345")
+	if v == nil {
+		t.Fatal("expected conn-level In limiter")
+	}
+	if v.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", v.Limit())
+	}
+}
+
+func TestIn_ConnScope_CachedResult(t *testing.T) {
+	l := newTestTrafficLimiter("$$ 100B")
+	v1 := l.In(context.Background(), "10.0.0.1:12345")
+	v2 := l.In(context.Background(), "10.0.0.1:12345")
+	// Both should have the correct limit (cached limiter reused in group).
+	if v1 == nil || v2 == nil {
+		t.Fatal("limiters should not be nil")
+	}
+	if v1.Limit() != 100 || v2.Limit() != 100 {
+		t.Fatalf("expected limit 100, got v1=%d v2=%d", v1.Limit(), v2.Limit())
+	}
+	// The cache entry should exist for the key.
+	if _, ok := l.connInLimits.Get("10.0.0.1:12345"); !ok {
+		t.Fatal("expected cached conn-level In limiter")
+	}
+}
+
+func TestIn_IPMatch(t *testing.T) {
+	l := newTestTrafficLimiter("192.168.1.1 100B")
+	v := l.In(context.Background(), "192.168.1.1:12345")
+	if v == nil {
+		t.Fatal("expected IP-level In limiter")
+	}
+	if v.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", v.Limit())
+	}
+}
+
+func TestIn_CIDRMatch(t *testing.T) {
+	l := newTestTrafficLimiter("10.0.0.0/8 100B")
+	v := l.In(context.Background(), "10.0.0.1:12345")
+	if v == nil {
+		t.Fatal("expected CIDR-level In limiter")
+	}
+	if v.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", v.Limit())
+	}
+}
+
+func TestIn_CIDRNoMatch(t *testing.T) {
+	l := newTestTrafficLimiter("10.0.0.0/8 100B")
+	v := l.In(context.Background(), "192.168.1.1:12345")
+	if v != nil {
+		t.Fatalf("CIDR should not match different IP range, got limit %d", v.Limit())
+	}
+}
+
+func TestIn_IPTakesPrecedence(t *testing.T) {
+	l := newTestTrafficLimiter("10.0.0.0/8 5B", "10.0.0.1 2B")
+	v := l.In(context.Background(), "10.0.0.1:12345")
+	if v == nil {
+		t.Fatal("expected limiter for exact IP + CIDR")
+	}
+	if v.Limit() != 2 {
+		t.Fatalf("exact IP should take precedence, expected limit 2, got %d", v.Limit())
+	}
+}
+
+func TestOut_ServiceScope(t *testing.T) {
+	l := newTestTrafficLimiter("$ 100B 100B")
+	v := l.Out(context.Background(), "any-key", limiter.ScopeOption(limiter.ScopeService))
+	if v == nil {
+		t.Fatal("expected service-level Out limiter")
+	}
+	if v.Limit() != 100 {
+		t.Fatalf("expected limit 100, got %d", v.Limit())
+	}
+}
+
+func TestOut_ClientScope(t *testing.T) {
+	l := newTestTrafficLimiter("$ 100B")
+	v := l.Out(context.Background(), "any-key", limiter.ScopeOption(limiter.ScopeClient))
+	if v != nil {
+		t.Fatal("ScopeClient Out should return nil")
+	}
+}
+
+func TestOut_ConnScope(t *testing.T) {
+	l := newTestTrafficLimiter("$$ 200B 200B")
+	v := l.Out(context.Background(), "10.0.0.1:12345")
+	if v == nil {
+		t.Fatal("expected conn-level Out limiter")
+	}
+	if v.Limit() != 200 {
+		t.Fatalf("expected limit 200, got %d", v.Limit())
+	}
+}
+
+// --- reload tests ---
+
+func TestReload_ServiceLimit(t *testing.T) {
+	ctx := context.Background()
+	l := newTestTrafficLimiter("$ 100B")
+	v1 := l.In(ctx, "x", limiter.ScopeOption(limiter.ScopeService))
+	if v1.Limit() != 100 {
+		t.Fatalf("expected 100, got %d", v1.Limit())
+	}
+
+	l.options.limits = []string{"$ 200B"}
+	_ = l.reload(ctx)
+
+	v2 := l.In(ctx, "x", limiter.ScopeOption(limiter.ScopeService))
+	if v2.Limit() != 200 {
+		t.Fatalf("expected 200 after reload, got %d", v2.Limit())
+	}
+}
+
+func TestReload_ConnLimit(t *testing.T) {
+	ctx := context.Background()
+	l := newTestTrafficLimiter("$$ 100B")
+	v1 := l.In(ctx, "10.0.0.1:12345")
+	if v1.Limit() != 100 {
+		t.Fatalf("expected 100, got %d", v1.Limit())
+	}
+
+	l.options.limits = []string{"$$ 200B"}
+	_ = l.reload(ctx)
+
+	v2 := l.In(ctx, "10.0.0.2:12345")
+	if v2.Limit() != 200 {
+		t.Fatalf("expected 200 after reload, got %d", v2.Limit())
+	}
+}
+
+func TestReload_ConnLimit_UpdatesExisting(t *testing.T) {
+	ctx := context.Background()
+	l := newTestTrafficLimiter("$$ 100B")
+	v1 := l.In(ctx, "10.0.0.1:12345")
+	if v1.Limit() != 100 {
+		t.Fatalf("expected 100, got %d", v1.Limit())
+	}
+
+	l.options.limits = []string{"$$ 200B"}
+	_ = l.reload(ctx)
+
+	// Same cached limiter should have been updated in-place.
+	if v1.Limit() != 200 {
+		t.Fatalf("existing cached limiter should be updated to 200, got %d", v1.Limit())
+	}
+}
+
+func TestReload_IPLimit(t *testing.T) {
+	ctx := context.Background()
+	l := newTestTrafficLimiter("192.168.1.1 100B")
+	v := l.In(ctx, "192.168.1.1:12345")
+	if v.Limit() != 100 {
+		t.Fatalf("expected 100, got %d", v.Limit())
+	}
+
+	l.options.limits = []string{"192.168.1.1 200B"}
+	_ = l.reload(ctx)
+
+	if v.Limit() != 200 {
+		t.Fatalf("expected 200 after reload, got %d", v.Limit())
+	}
+}
+
+func TestReload_CIDRLimit(t *testing.T) {
+	ctx := context.Background()
+	l := newTestTrafficLimiter("10.0.0.0/8 100B")
+	v := l.In(ctx, "10.0.0.1:12345")
+	if v.Limit() != 100 {
+		t.Fatalf("expected 100, got %d", v.Limit())
+	}
+
+	l.options.limits = []string{"10.0.0.0/8 200B"}
+	_ = l.reload(ctx)
+
+	if v.Limit() != 200 {
+		t.Fatalf("expected 200 after CIDR reload, got %d", v.Limit())
+	}
+}
+
+func TestReload_RemovesObsoleteLimit(t *testing.T) {
+	ctx := context.Background()
+	l := newTestTrafficLimiter("192.168.1.1 100B")
+	_ = l.In(ctx, "192.168.1.1:12345")
+
+	l.options.limits = nil
+	_ = l.reload(ctx)
+
+	v := l.In(ctx, "192.168.1.1:12345")
+	if v != nil {
+		t.Fatalf("obsolete IP limit should be removed, got limit %d", v.Limit())
+	}
+}
+
+func TestReload_ClearsConnCacheWhenRemoved(t *testing.T) {
+	ctx := context.Background()
+	l := newTestTrafficLimiter("$$ 100B")
+	_ = l.In(ctx, "10.0.0.1:12345")
+
+	l.options.limits = []string{"$$ 0B"}
+	_ = l.reload(ctx)
+
+	v := l.In(ctx, "10.0.0.2:12345")
+	if v != nil {
+		t.Fatalf("conn cache should be flushed, got limiter with limit %d", v.Limit())
+	}
+}
+
+// --- close tests ---
+
+func TestClose(t *testing.T) {
+	fl := &fakeLoader{}
+	rl := &fakeLoader{}
+	hl := &fakeLoader{}
+
+	l := newTestTrafficLimiter()
+	l.options.fileLoader = fl
+	l.options.redisLoader = rl
+	l.options.httpLoader = hl
+
+	l.Close()
+
+	if !fl.closed {
+		t.Error("fileLoader should be closed")
+	}
+	if !rl.closed {
+		t.Error("redisLoader should be closed")
+	}
+	if !hl.closed {
+		t.Error("httpLoader should be closed")
+	}
+}
+
+func TestClose_Idempotent(t *testing.T) {
+	l := newTestTrafficLimiter()
+	l.Close()
+	l.Close()
+}
+
+func TestClose_NilLoaders(t *testing.T) {
+	l := newTestTrafficLimiter()
+	if err := l.Close(); err != nil {
+		t.Fatal("Close with nil loaders should return nil error")
+	}
+}
+
+// --- lifecycle tests ---
+
+func TestPeriodReload_ZeroPeriod(t *testing.T) {
+	l := newTestTrafficLimiter()
+	l.options.period = 0
+
+	ctx := context.Background()
+	err := l.periodReload(ctx)
+	if err != nil {
+		t.Fatalf("zero period should return nil, got %v", err)
+	}
+}
+
+func TestPeriodReload_ContextCancel(t *testing.T) {
+	l := newTestTrafficLimiter("$ 100B")
+	l.options.period = time.Second
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err := l.periodReload(ctx)
+	if err != context.Canceled {
+		t.Fatalf("expected context.Canceled, got %v", err)
+	}
+}
+
+func TestNewTrafficLimiter_Defaults(t *testing.T) {
+	lim := NewTrafficLimiter()
+	if lim == nil {
+		t.Fatal("NewTrafficLimiter should not return nil")
+	}
+	if closer, ok := lim.(io.Closer); ok {
+		closer.Close()
+	}
+}
+
+// --- concurrency ---
+
+func TestIn_ConcurrentAccess(t *testing.T) {
+	l := newTestTrafficLimiter("$$ 1KB", "192.168.1.1 500B", "10.0.0.0/8 200B", "$ 100B")
+	ctx := context.Background()
+	var wg sync.WaitGroup
+	for i := 0; i < 50; i++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+			_ = l.In(ctx, "10.0.0.1:12345")
+			_ = l.In(ctx, "192.168.1.1:8080")
+			_ = l.Out(ctx, "10.0.0.2:9090")
+		}(i)
+	}
+	wg.Wait()
+}
+
+// --- interface compliance ---
+
+var (
+	_ traffic.TrafficLimiter = (*trafficLimiter)(nil)
+	_ io.Closer             = (*trafficLimiter)(nil)
+	_ loader.Loader         = (*fakeLoader)(nil)
+	_ loader.Lister         = (*fakeLoader)(nil)
+)

limiter/traffic/wrapper/conn.go

@@ -1,3 +1,5 @@
+// Package wrapper provides net.Conn, net.PacketConn, and io.ReadWriter
+// wrappers that apply traffic rate limiting to reads and writes.
 package wrapper
 
 import (
@@ -21,7 +23,7 @@ var (
 	errRateLimited = errors.New("rate limited")
 )
 
-// limitConn is a Conn with traffic limiter supported.
+// limitConn wraps a net.Conn with traffic rate limiting applied to reads and writes.
 type limitConn struct {
 	net.Conn
 	rbuf    bytes.Buffer
@@ -30,6 +32,8 @@ type limitConn struct {
 	key     string
 }
 
+// WrapConn wraps a net.Conn with traffic rate limiting. If tlimiter is nil,
+// the original conn is returned unchanged.
 func WrapConn(c net.Conn, tlimiter traffic.TrafficLimiter, key string, opts ...limiter.Option) net.Conn {
 	if tlimiter == nil {
 		return c
@@ -133,6 +137,8 @@ type packetConn struct {
 	key     string
 }
 
+// WrapPacketConn wraps a net.PacketConn with traffic rate limiting. Packets
+// exceeding the rate limit are discarded on read or rejected with an error on write.
 func WrapPacketConn(pc net.PacketConn, lim traffic.TrafficLimiter, key string, opts ...limiter.Option) net.PacketConn {
 	if lim == nil {
 		return pc
@@ -190,6 +196,7 @@ type udpConn struct {
 	key     string
 }
 
+// WrapUDPConn wraps a net.PacketConn as a udp.Conn with traffic rate limiting.
 func WrapUDPConn(pc net.PacketConn, limiter traffic.TrafficLimiter, key string, opts ...limiter.Option) udp.Conn {
 	return &udpConn{
 		PacketConn: pc,

limiter/traffic/wrapper/conn_test.go

@@ -0,0 +1,448 @@
+package wrapper
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/go-gost/core/limiter"
+	"github.com/go-gost/core/limiter/traffic"
+	xio "github.com/go-gost/x/internal/io"
+	ctxutil "github.com/go-gost/x/ctx"
+)
+
+// --- mock types ---
+
+type mockLimiter struct {
+	waitFunc func(ctx context.Context, n int) int
+	limit    int
+}
+
+func (m *mockLimiter) Wait(ctx context.Context, n int) int {
+	if m.waitFunc != nil {
+		return m.waitFunc(ctx, n)
+	}
+	return n
+}
+func (m *mockLimiter) Limit() int { return m.limit }
+func (m *mockLimiter) Set(n int)  { m.limit = n }
+
+type mockTrafficLimiter struct {
+	inLim  traffic.Limiter
+	outLim traffic.Limiter
+}
+
+func (m *mockTrafficLimiter) In(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	return m.inLim
+}
+func (m *mockTrafficLimiter) Out(ctx context.Context, key string, opts ...limiter.Option) traffic.Limiter {
+	return m.outLim
+}
+
+// --- WrapConn tests ---
+
+func TestWrapConn_NilLimiter(t *testing.T) {
+	c := &struct{ net.Conn }{}
+	result := WrapConn(c, nil, "key")
+	if result != c {
+		t.Fatal("nil limiter should return original conn")
+	}
+}
+
+func TestWrapConn_WithLimiter(t *testing.T) {
+	c := &struct{ net.Conn }{}
+	tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}}
+	result := WrapConn(c, tl, "key")
+	lc, ok := result.(*limitConn)
+	if !ok {
+		t.Fatalf("expected *limitConn, got %T", result)
+	}
+	if lc.Conn != c {
+		t.Fatal("wrapped conn should retain inner conn")
+	}
+}
+
+func TestLimitConn_Read_NoLimiter(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	tl := &mockTrafficLimiter{inLim: nil}
+	wrapped := WrapConn(client, tl, "test-key")
+
+	go server.Write([]byte("hello"))
+
+	buf := make([]byte, 10)
+	n, err := wrapped.Read(buf)
+	if err != nil && err != io.EOF {
+		t.Fatal(err)
+	}
+	if n != 5 || string(buf[:5]) != "hello" {
+		t.Fatalf("expected 'hello', got %q", string(buf[:n]))
+	}
+}
+
+func TestLimitConn_Read_WithLimiter(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	ml := &mockLimiter{limit: 100}
+	tl := &mockTrafficLimiter{inLim: ml}
+	wrapped := WrapConn(client, tl, "test-key")
+
+	go func() {
+		server.Write([]byte("hello world"))
+		server.Close()
+	}()
+
+	buf := make([]byte, 100)
+	n, err := wrapped.Read(buf)
+	if err != nil && err != io.EOF {
+		t.Fatal(err)
+	}
+	if string(buf[:n]) != "hello world" {
+		t.Fatalf("expected 'hello world', got %q", string(buf[:n]))
+	}
+}
+
+func TestLimitConn_Read_BufferedData(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	callCount := 0
+	ml := &mockLimiter{
+		limit: 100,
+		waitFunc: func(ctx context.Context, n int) int {
+			callCount++
+			if callCount == 1 {
+				return 5 // only allow 5 of 11 bytes
+			}
+			return 10
+		},
+	}
+	tl := &mockTrafficLimiter{inLim: ml}
+	wrapped := WrapConn(client, tl, "test-key")
+
+	go func() {
+		server.Write([]byte("hello world"))
+		server.Close()
+	}()
+
+	buf := make([]byte, 100)
+	n, err := readFull(wrapped, buf, 5)
+	if err != nil {
+		t.Fatalf("first read: %v", err)
+	}
+	if n != 5 {
+		t.Fatalf("expected 5 bytes, got %d", n)
+	}
+
+	// Remaining bytes from rbuf.
+	n, err = readFull(wrapped, buf[5:], 6)
+	if err != nil {
+		t.Fatalf("second read (buffered): %v", err)
+	}
+	if n != 6 {
+		t.Fatalf("expected 6 buffered bytes, got %d", n)
+	}
+}
+
+func TestLimitConn_Write_NoLimiter(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	tl := &mockTrafficLimiter{outLim: nil}
+	wrapped := WrapConn(client, tl, "test-key")
+
+	go func() {
+		buf := make([]byte, 5)
+		server.Read(buf)
+	}()
+
+	n, err := wrapped.Write([]byte("hello"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 5 {
+		t.Fatalf("expected 5, got %d", n)
+	}
+}
+
+func TestLimitConn_Write_ZeroBurst(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	ml := &mockLimiter{
+		limit: 100,
+		waitFunc: func(ctx context.Context, n int) int {
+			return 0
+		},
+	}
+	tl := &mockTrafficLimiter{outLim: ml}
+	wrapped := WrapConn(client, tl, "test-key")
+
+	n, err := wrapped.Write([]byte("data"))
+	if err != nil {
+		t.Fatalf("should not error on zero burst: %v", err)
+	}
+	if n != 0 {
+		t.Fatalf("expected 0 bytes, got %d", n)
+	}
+}
+
+type connWithSyscall struct {
+	net.Conn
+}
+
+func (c *connWithSyscall) SyscallConn() (syscall.RawConn, error) {
+	return nil, nil
+}
+
+func TestLimitConn_SyscallConn(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	c := &connWithSyscall{Conn: client}
+	tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}}
+	wrapped := WrapConn(c, tl, "test-key")
+
+	_, err := wrapped.(*limitConn).SyscallConn()
+	if err != nil {
+		t.Fatal("SyscallConn should succeed", err)
+	}
+}
+
+type connWithCloseRead struct {
+	net.Conn
+	closed bool
+}
+
+func (c *connWithCloseRead) CloseRead() error {
+	c.closed = true
+	return nil
+}
+
+func TestLimitConn_CloseRead(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	c := &connWithCloseRead{Conn: client}
+	tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}}
+	wrapped := WrapConn(c, tl, "test-key")
+
+	lc := wrapped.(*limitConn)
+	if err := lc.CloseRead(); err != nil {
+		t.Fatal("CloseRead should succeed", err)
+	}
+	if !c.closed {
+		t.Fatal("CloseRead should delegate to inner")
+	}
+}
+
+type connWithCloseWrite struct {
+	net.Conn
+	closed bool
+}
+
+func (c *connWithCloseWrite) CloseWrite() error {
+	c.closed = true
+	return nil
+}
+
+func TestLimitConn_CloseWrite(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	c := &connWithCloseWrite{Conn: client}
+	tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 100}}
+	wrapped := WrapConn(c, tl, "test-key")
+
+	lc := wrapped.(*limitConn)
+	if err := lc.CloseWrite(); err != nil {
+		t.Fatal("CloseWrite should succeed", err)
+	}
+	if !c.closed {
+		t.Fatal("CloseWrite should delegate to inner")
+	}
+}
+
+// --- WrapReadWriter tests ---
+
+func TestWrapReadWriter_NilLimiter(t *testing.T) {
+	rw := &bytesReadWriter{}
+	result := WrapReadWriter(nil, rw, "key")
+	if result != rw {
+		t.Fatal("nil limiter should return original ReadWriter")
+	}
+}
+
+func TestReadWriter_Read_WithLimiter(t *testing.T) {
+	data := []byte("hello world")
+	rw := &bytesReadWriter{buf: append([]byte{}, data...)}
+
+	ml := &mockLimiter{limit: 100}
+	tl := &mockTrafficLimiter{inLim: ml}
+	wrapped := WrapReadWriter(tl, rw, "test-key")
+
+	buf := make([]byte, 100)
+	n, err := wrapped.Read(buf)
+	if err != nil && err != io.EOF {
+		t.Fatal(err)
+	}
+	if string(buf[:n]) != "hello world" {
+		t.Fatalf("expected 'hello world', got %q", string(buf[:n]))
+	}
+}
+
+func TestReadWriter_Write_ZeroBurst(t *testing.T) {
+	rw := &bytesReadWriter{}
+	ml := &mockLimiter{
+		limit: 100,
+		waitFunc: func(ctx context.Context, n int) int {
+			return 0
+		},
+	}
+	tl := &mockTrafficLimiter{outLim: ml}
+	wrapped := WrapReadWriter(tl, rw, "test-key")
+
+	n, err := wrapped.Write([]byte("data"))
+	if err != nil {
+		t.Fatalf("should not error on zero burst: %v", err)
+	}
+	if n != 0 {
+		t.Fatalf("expected 0 bytes, got %d", n)
+	}
+}
+
+// --- WrapPacketConn tests ---
+
+func TestWrapPacketConn_NilLimiter(t *testing.T) {
+	pc := &struct{ net.PacketConn }{}
+	result := WrapPacketConn(pc, nil, "key")
+	if result != pc {
+		t.Fatal("nil limiter should return original PacketConn")
+	}
+}
+
+func TestPacketConn_WriteTo_RateLimited(t *testing.T) {
+	pc := &mockPacketConn{writeBuf: make([]byte, 100)}
+	ml := &mockLimiter{
+		limit: 100,
+		waitFunc: func(ctx context.Context, n int) int {
+			return 3 // allow only 3, but packet is 4 bytes
+		},
+	}
+	tl := &mockTrafficLimiter{outLim: ml}
+	wrapped := WrapPacketConn(pc, tl, "test-key")
+
+	_, err := wrapped.WriteTo([]byte("data"), nil)
+	if !errors.Is(err, errRateLimited) {
+		t.Fatalf("expected errRateLimited, got %v", err)
+	}
+}
+
+// --- WrapUDPConn tests ---
+
+func TestWrapUDPConn(t *testing.T) {
+	pc := &mockPacketConn{writeBuf: make([]byte, 100)}
+	result := WrapUDPConn(pc, &mockTrafficLimiter{}, "key")
+	if _, ok := result.(*udpConn); !ok {
+		t.Fatalf("expected *udpConn, got %T", result)
+	}
+}
+
+func TestUDPConn_Write_RateLimited(t *testing.T) {
+	pc := &mockPacketConn{writeBuf: make([]byte, 100)}
+	ml := &mockLimiter{
+		limit: 100,
+		waitFunc: func(ctx context.Context, n int) int {
+			return 2 // allow fewer bytes than packet
+		},
+	}
+	tl := &mockTrafficLimiter{outLim: ml}
+	wrapped := WrapUDPConn(pc, tl, "test-key")
+
+	_, err := wrapped.Write([]byte("dat"))
+	if !errors.Is(err, errRateLimited) {
+		t.Fatalf("expected errRateLimited, got %v", err)
+	}
+}
+
+// --- interface compliance ---
+
+var (
+	_ net.Conn     = (*limitConn)(nil)
+	_ net.Conn     = (*struct{ net.Conn })(nil)
+	_ xio.CloseRead  = (*limitConn)(nil)
+	_ xio.CloseWrite = (*limitConn)(nil)
+	_ syscall.Conn   = (*limitConn)(nil)
+	_ ctxutil.Context = (*limitConn)(nil)
+)
+
+// --- helpers ---
+
+type bytesReadWriter struct {
+	buf []byte
+}
+
+func (rw *bytesReadWriter) Read(b []byte) (int, error) {
+	if len(rw.buf) == 0 {
+		return 0, io.EOF
+	}
+	n := copy(b, rw.buf)
+	rw.buf = rw.buf[n:]
+	return n, nil
+}
+
+func (rw *bytesReadWriter) Write(b []byte) (int, error) {
+	rw.buf = append(rw.buf, b...)
+	return len(b), nil
+}
+
+type mockPacketConn struct {
+	net.PacketConn
+	writeBuf []byte
+}
+
+func (pc *mockPacketConn) Write(p []byte) (int, error) {
+	pc.writeBuf = append(pc.writeBuf[:0], p...)
+	return len(p), nil
+}
+
+func (pc *mockPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
+	pc.writeBuf = append(pc.writeBuf[:0], p...)
+	return len(p), nil
+}
+
+func (pc *mockPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
+	return copy(p, "test"), nil, nil
+}
+
+func (pc *mockPacketConn) Close() error                       { return nil }
+func (pc *mockPacketConn) LocalAddr() net.Addr                { return nil }
+func (pc *mockPacketConn) SetDeadline(t time.Time) error      { return nil }
+func (pc *mockPacketConn) SetReadDeadline(t time.Time) error  { return nil }
+func (pc *mockPacketConn) SetWriteDeadline(t time.Time) error { return nil }
+
+func readFull(c net.Conn, buf []byte, want int) (int, error) {
+	total := 0
+	for total < want {
+		n, err := c.Read(buf[total:])
+		if err != nil {
+			return total, err
+		}
+		total += n
+	}
+	return total, nil
+}

limiter/traffic/wrapper/io.go

@@ -9,7 +9,7 @@ import (
 	"github.com/go-gost/core/limiter/traffic"
 )
 
-// readWriter is an io.ReadWriter with traffic limiter supported.
+// readWriter wraps an io.ReadWriter with traffic rate limiting applied to reads and writes.
 type readWriter struct {
 	io.ReadWriter
 	rbuf    bytes.Buffer
@@ -18,6 +18,8 @@ type readWriter struct {
 	key     string
 }
 
+// WrapReadWriter wraps an io.ReadWriter with traffic rate limiting. If limiter
+// is nil, the original ReadWriter is returned unchanged.
 func WrapReadWriter(limiter traffic.TrafficLimiter, rw io.ReadWriter, key string, opts ...limiter.Option) io.ReadWriter {
 	if limiter == nil {
 		return rw

limiter/traffic/wrapper/listener.go

@@ -14,6 +14,8 @@ type listener struct {
 	service string
 }
 
+// WrapListener wraps a net.Listener so that accepted connections are
+// automatically wrapped with service-level traffic rate limiting.
 func WrapListener(service string, ln net.Listener, limiter traffic.TrafficLimiter) net.Listener {
 	if limiter == nil {
 		return ln

limiter/traffic/wrapper/listener_test.go

@@ -0,0 +1,65 @@
+package wrapper
+
+import (
+	"errors"
+	"net"
+	"testing"
+
+	traffic "github.com/go-gost/core/limiter/traffic"
+)
+
+type testListener struct {
+	connCh chan net.Conn
+	addr   net.Addr
+}
+
+func (l *testListener) Accept() (net.Conn, error) {
+	c, ok := <-l.connCh
+	if !ok {
+		return nil, errors.New("listener closed")
+	}
+	return c, nil
+}
+
+func (l *testListener) Close() error {
+	close(l.connCh)
+	return nil
+}
+
+func (l *testListener) Addr() net.Addr { return l.addr }
+
+func TestWrapListener_NilLimiter(t *testing.T) {
+	ln := &testListener{}
+	result := WrapListener("svc", ln, nil)
+	if result != ln {
+		t.Fatal("nil limiter should return original listener")
+	}
+}
+
+func TestAccept_WrapsConn(t *testing.T) {
+	server, client := net.Pipe()
+	defer server.Close()
+	defer client.Close()
+
+	ln := &testListener{
+		connCh: make(chan net.Conn, 1),
+		addr:   &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080},
+	}
+	ln.connCh <- client
+
+	tl := &mockTrafficLimiter{inLim: &mockLimiter{limit: 1000}}
+	wrappedLn := WrapListener("test-svc", ln, tl)
+
+	conn, err := wrappedLn.Accept()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if _, ok := conn.(*limitConn); !ok {
+		t.Fatalf("accepted conn should be *limitConn, got %T", conn)
+	}
+}
+
+var (
+	_ traffic.TrafficLimiter = (*mockTrafficLimiter)(nil)
+)