test(limiter/conn): add unit tests for conn limiter and wrappers

limiter/conn/conn_test.go

@@ -0,0 +1,458 @@
+package conn
+
+import (
+	"context"
+	"io"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	limiter "github.com/go-gost/core/limiter/conn"
+	"github.com/go-gost/core/logger"
+	xlogger "github.com/go-gost/x/logger"
+	"github.com/yl2chen/cidranger"
+)
+
+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
+}
+
+func nopLogger() logger.Logger {
+	return xlogger.Nop()
+}
+
+func TestParseLine(t *testing.T) {
+	cl := &connLimiter{}
+
+	tests := []struct {
+		input    string
+		expected string
+	}{
+		{"192.168.1.1 100", "192.168.1.1 100"},
+		{"  192.168.1.1 100  ", "192.168.1.1 100"},
+		{"192.168.1.1 100 # comment", "192.168.1.1 100"},
+		{"# only comment", ""},
+		{"   # indented comment", ""},
+		{"", ""},
+	}
+	for _, tt := range tests {
+		got := cl.parseLine(tt.input)
+		if got != tt.expected {
+			t.Errorf("parseLine(%q) = %q, want %q", tt.input, got, tt.expected)
+		}
+	}
+}
+
+func TestParseLimit(t *testing.T) {
+	cl := &connLimiter{}
+
+	tests := []struct {
+		input string
+		key   string
+		limit int
+	}{
+		{"192.168.1.1 100", "192.168.1.1", 100},
+		{"192.168.1.1\t100", "192.168.1.1", 100},
+		{"192.168.1.1     100", "192.168.1.1", 100},
+		{"$ 50", "$", 50},
+		{"$$ 20", "$$", 20},
+		{"10.0.0.0/8 500", "10.0.0.0/8", 500},
+		{"key only", "key", 0},
+		{"key invalid", "key", 0},
+		{"", "", 0},
+	}
+	for _, tt := range tests {
+		key, limit := cl.parseLimit(tt.input)
+		if key != tt.key || limit != tt.limit {
+			t.Errorf("parseLimit(%q) = (%q, %d), want (%q, %d)", tt.input, key, limit, tt.key, tt.limit)
+		}
+	}
+}
+
+func TestParsePatterns(t *testing.T) {
+	cl := &connLimiter{}
+
+	patterns, err := cl.parsePatterns(strings.NewReader("192.168.1.1 10\n10.0.0.1 20\n# comment\n\n  172.16.0.1 30  \n"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(patterns) != 3 {
+		t.Fatalf("expected 3 patterns, got %d: %v", len(patterns), patterns)
+	}
+}
+
+func TestParsePatterns_NilReader(t *testing.T) {
+	cl := &connLimiter{}
+	patterns, err := cl.parsePatterns(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(patterns) != 0 {
+		t.Fatal("expected no patterns from nil reader")
+	}
+}
+
+func TestLoad_NoLoaders(t *testing.T) {
+	cl := &connLimiter{options: options{}, logger: nopLogger()}
+	patterns, err := cl.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(patterns) != 0 {
+		t.Fatalf("expected 0 patterns, got %d", len(patterns))
+	}
+}
+
+func TestLoad_FileLister(t *testing.T) {
+	fl := &fakeLoader{data: "10.0.0.1 10\n# ignored\n10.0.0.2 20"}
+	cl := &connLimiter{
+		options: options{fileLoader: fl},
+		logger:  nopLogger(),
+	}
+	patterns, err := cl.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(patterns) != 2 {
+		t.Fatalf("expected 2 patterns, got %d: %v", len(patterns), patterns)
+	}
+}
+
+func TestLoad_RedisLoader(t *testing.T) {
+	rl := &fakeLoader{data: "10.0.0.1 5"}
+	cl := &connLimiter{
+		options: options{redisLoader: rl},
+		logger:  nopLogger(),
+	}
+	patterns, err := cl.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	// redis list items are NOT parseLine-filtered (raw).
+	if len(patterns) != 1 {
+		t.Fatalf("expected 1 pattern, got %d", len(patterns))
+	}
+}
+
+func TestLoad_HTTPLoader(t *testing.T) {
+	hl := &fakeLoader{data: "10.0.0.1 5\n10.0.0.2 3"}
+	cl := &connLimiter{
+		options: options{httpLoader: hl},
+		logger:  nopLogger(),
+	}
+	patterns, err := cl.load(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(patterns) != 2 {
+		t.Fatalf("expected 2 patterns, got %d", len(patterns))
+	}
+}
+
+func TestLoad_FileListerError(t *testing.T) {
+	fl := &fakeLoader{data: "", listErr: io.ErrUnexpectedEOF}
+	cl := &connLimiter{
+		options: options{fileLoader: fl},
+		logger:  nopLogger(),
+	}
+	patterns, err := cl.load(context.Background())
+	if err != nil {
+		t.Fatal("load should not return error on loader failure")
+	}
+	if len(patterns) != 0 {
+		t.Fatalf("expected 0 patterns, got %d", len(patterns))
+	}
+}
+
+func TestReload_StaticLimits(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cl := &connLimiter{
+		options: options{
+			limits: []string{"192.168.1.1 3", "10.0.0.0/8 10", "$ 100"},
+		},
+		ipLimits:   make(map[string]ConnLimitGenerator),
+		cidrLimits: cidranger.NewPCTrieRanger(),
+		limits:     make(map[string]limiter.Limiter),
+		logger:     nopLogger(),
+		cancelFunc: cancel,
+	}
+
+	if err := cl.reload(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	// Global limit applies to any key.
+	lim := cl.Limiter("any-key")
+	if lim == nil {
+		t.Fatal("global limiter should exist")
+	}
+	if lim.Limit() != 100 {
+		t.Fatalf("expected global limit 100, got %d", lim.Limit())
+	}
+}
+
+// newTestConnLimiter creates a connLimiter with sync reload, no goroutine.
+func newTestConnLimiter(limits ...string) *connLimiter {
+	ctx, cancel := context.WithCancel(context.Background())
+	cl := &connLimiter{
+		options:    options{limits: limits},
+		ipLimits:   make(map[string]ConnLimitGenerator),
+		cidrLimits: cidranger.NewPCTrieRanger(),
+		limits:     make(map[string]limiter.Limiter),
+		logger:     nopLogger(),
+		cancelFunc: cancel,
+	}
+	_ = cl.reload(ctx)
+	return cl
+}
+
+func TestReload_ClearsCache(t *testing.T) {
+	cl := newTestConnLimiter("192.168.1.1 5")
+	defer cl.Close()
+
+	lim := cl.Limiter("192.168.1.1")
+	if lim == nil || lim.Limit() != 5 {
+		t.Fatal("limiter should exist with limit 5")
+	}
+
+	// Reload clears cache and re-parses static limits.
+	if err := cl.reload(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	lim2 := cl.Limiter("192.168.1.1")
+	if lim2 == nil {
+		t.Fatal("limiter should still exist after reload")
+	}
+	if lim2.Limit() != 5 {
+		t.Fatalf("expected limit 5 after reload, got %d", lim2.Limit())
+	}
+}
+
+func TestLimiter_NilCacheDoesNotStick(t *testing.T) {
+	cl := newTestConnLimiter()
+	defer cl.Close()
+
+	// Lookup for a key with no matching limit should NOT cache nil.
+	lim := cl.Limiter("10.0.0.1")
+	if lim != nil {
+		t.Fatal("should be no limiter for 10.0.0.1")
+	}
+
+	// Manually add a limit via reload.
+	cl.options.limits = []string{"10.0.0.1 3"}
+	if err := cl.reload(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	// Should now find the limit because nil was not cached.
+	lim = cl.Limiter("10.0.0.1")
+	if lim == nil {
+		t.Fatal("limiter should exist after reload when nil was not cached")
+	}
+	if lim.Limit() != 3 {
+		t.Fatalf("expected limit 3, got %d", lim.Limit())
+	}
+}
+
+func TestLimiter_IPLimitKey(t *testing.T) {
+	cl := newTestConnLimiter("$$ 5")
+	defer cl.Close()
+
+	lim := cl.Limiter("10.0.0.1")
+	if lim == nil {
+		t.Fatal("IP-level limiter should exist")
+	}
+	if lim.Limit() != 5 {
+		t.Fatalf("expected limit 5, got %d", lim.Limit())
+	}
+}
+
+func TestLimiter_CIDRMatch(t *testing.T) {
+	cl := newTestConnLimiter("10.0.0.0/8 7")
+	defer cl.Close()
+
+	lim := cl.Limiter("10.1.2.3")
+	if lim == nil {
+		t.Fatal("CIDR-based limiter should exist for 10.1.2.3")
+	}
+	if lim.Limit() != 7 {
+		t.Fatalf("expected limit 7, got %d", lim.Limit())
+	}
+}
+
+func TestLimiter_CIDRNoMatch(t *testing.T) {
+	cl := newTestConnLimiter("10.0.0.0/8 7")
+	defer cl.Close()
+
+	lim := cl.Limiter("192.168.1.1")
+	if lim != nil {
+		t.Fatal("CIDR-based limiter should not match 192.168.1.1")
+	}
+}
+
+func TestLimiter_SpecificIPTakesPrecedence(t *testing.T) {
+	cl := newTestConnLimiter("10.0.0.0/8 5", "10.1.2.3 2")
+	defer cl.Close()
+
+	lim := cl.Limiter("10.1.2.3")
+	if lim == nil {
+		t.Fatal("specific IP limiter should exist")
+	}
+	if lim.Limit() != 2 {
+		t.Fatalf("expected specific IP limit 2, got %d", lim.Limit())
+	}
+}
+
+func TestLimiter_GlobalAndIP(t *testing.T) {
+	cl := newTestConnLimiter("$$ 5", "$ 10")
+	defer cl.Close()
+
+	lim := cl.Limiter("10.0.0.1")
+	if lim == nil {
+		t.Fatal("should have IP+global limiter")
+	}
+	if lim.Limit() != 5 {
+		t.Fatalf("expected limit 5 (min of 5,10), got %d", lim.Limit())
+	}
+}
+
+func TestLimiter_NonIPKeyHasGlobalOnly(t *testing.T) {
+	cl := newTestConnLimiter("$$ 5", "$ 10")
+	defer cl.Close()
+
+	lim := cl.Limiter("not-an-ip")
+	if lim == nil {
+		t.Fatal("should have limiters")
+	}
+	// IPLimitKey ($$=5) acts as fallback for any key when no specific IP match.
+	// Global ($=10) also applies. Group picks the minimum: 5.
+	if lim.Limit() != 5 {
+		t.Fatalf("expected limit 5 (min of IP fallback 5 and global 10), got %d", lim.Limit())
+	}
+}
+
+func TestClose(t *testing.T) {
+	fl := &fakeLoader{}
+	rl := &fakeLoader{}
+	hl := &fakeLoader{}
+
+	cl := NewConnLimiter(
+		FileLoaderOption(fl),
+		RedisLoaderOption(rl),
+		HTTPLoaderOption(hl),
+	).(*connLimiter)
+
+	cl.Close()
+
+	if !fl.closed {
+		t.Error("fileLoader was not closed")
+	}
+	if !rl.closed {
+		t.Error("redisLoader was not closed")
+	}
+	if !hl.closed {
+		t.Error("httpLoader was not closed")
+	}
+}
+
+func TestClose_Idempotent(t *testing.T) {
+	fl := &fakeLoader{}
+	cl := NewConnLimiter(FileLoaderOption(fl)).(*connLimiter)
+	cl.Close()
+	cl.Close()
+}
+
+func TestLimiter_CachedResult(t *testing.T) {
+	cl := newTestConnLimiter("192.168.1.1 5")
+	defer cl.Close()
+
+	lim1 := cl.Limiter("192.168.1.1")
+	if lim1 == nil {
+		t.Fatal("first lookup should return a limiter")
+	}
+
+	lim2 := cl.Limiter("192.168.1.1")
+	if lim2 != lim1 {
+		t.Fatal("second lookup should return cached limiter")
+	}
+}
+
+func TestLimiter_ConcurrentAccess(t *testing.T) {
+	cl := newTestConnLimiter("192.168.1.1 1000")
+	defer cl.Close()
+
+	var wg sync.WaitGroup
+	n := 50
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_ = cl.Limiter("192.168.1.1")
+		}()
+	}
+	wg.Wait()
+}
+
+func TestConnLimiter_Interface(t *testing.T) {
+	var cl limiter.ConnLimiter = NewConnLimiter()
+	if cl == nil {
+		t.Fatal("ConnLimiter should not be nil")
+	}
+}
+
+func TestPeriodReload_ZeroPeriod(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cl := &connLimiter{
+		options: options{period: 0},
+		logger:  nopLogger(),
+	}
+
+	err := cl.periodReload(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestPeriodReload_ContextCancel(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cl := &connLimiter{
+		options: options{period: 10 * time.Second},
+		logger:  nopLogger(),
+	}
+	cancel()
+	err := cl.periodReload(ctx)
+	if err != context.Canceled {
+		t.Fatalf("expected context.Canceled, got %v", err)
+	}
+}

limiter/conn/generator_test.go

@@ -0,0 +1,84 @@
+package conn
+
+import (
+	"testing"
+)
+
+func TestConnLimitGenerator(t *testing.T) {
+	g := NewConnLimitGenerator(5)
+	lim := g.Limiter()
+	if lim == nil {
+		t.Fatal("Limiter() should not be nil for non-zero limit")
+	}
+	if lim.Limit() != 5 {
+		t.Fatalf("expected limit 5, got %d", lim.Limit())
+	}
+
+	// Each Limiter() call creates a new instance.
+	lim2 := g.Limiter()
+	if lim2 == lim {
+		t.Fatal("each Limiter() call should create a new instance")
+	}
+}
+
+func TestConnLimitGenerator_Zero(t *testing.T) {
+	g := NewConnLimitGenerator(0)
+	if g == nil {
+		t.Fatal("generator should not be nil")
+	}
+	if lim := g.Limiter(); lim != nil {
+		t.Fatal("Limiter() should be nil for zero limit")
+	}
+}
+
+func TestConnLimitGenerator_Negative(t *testing.T) {
+	g := NewConnLimitGenerator(-5)
+	if lim := g.Limiter(); lim != nil {
+		t.Fatal("Limiter() should be nil for negative limit")
+	}
+}
+
+func TestConnLimitGenerator_NilReceiver(t *testing.T) {
+	var g *connLimitGenerator
+	if lim := g.Limiter(); lim != nil {
+		t.Fatal("nil receiver Limiter() should return nil")
+	}
+}
+
+func TestConnLimitSingleGenerator(t *testing.T) {
+	g := NewConnLimitSingleGenerator(5)
+	lim := g.Limiter()
+	if lim == nil {
+		t.Fatal("Limiter() should not be nil")
+	}
+	if lim.Limit() != 5 {
+		t.Fatalf("expected limit 5, got %d", lim.Limit())
+	}
+
+	// Same instance every call.
+	lim2 := g.Limiter()
+	if lim2 != lim {
+		t.Fatal("single generator should return same instance")
+	}
+}
+
+func TestConnLimitSingleGenerator_Zero(t *testing.T) {
+	g := NewConnLimitSingleGenerator(0)
+	if lim := g.Limiter(); lim != nil {
+		t.Fatal("Limiter() should be nil for zero limit")
+	}
+}
+
+func TestConnLimitSingleGenerator_Negative(t *testing.T) {
+	g := NewConnLimitSingleGenerator(-1)
+	if lim := g.Limiter(); lim != nil {
+		t.Fatal("Limiter() should be nil for negative limit")
+	}
+}
+
+func TestConnLimitSingleGenerator_NilReceiver(t *testing.T) {
+	var g *connLimitSingleGenerator
+	if lim := g.Limiter(); lim != nil {
+		t.Fatal("nil receiver Limiter() should return nil")
+	}
+}

limiter/conn/limiter_test.go

@@ -0,0 +1,155 @@
+package conn
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	limiter "github.com/go-gost/core/limiter/conn"
+)
+
+func TestLimiter_Allow(t *testing.T) {
+	l := NewLimiter(3).(*llimiter)
+
+	if l.Limit() != 3 {
+		t.Fatalf("expected limit 3, got %d", l.Limit())
+	}
+
+	// Basic acquire/release within limit.
+	for i := 0; i < 3; i++ {
+		if !l.Allow(1) {
+			t.Fatalf("Allow(1) #%d should succeed", i)
+		}
+	}
+	if l.Allow(1) {
+		t.Fatal("Allow(1) should fail after reaching limit")
+	}
+
+	// Release one and re-acquire.
+	if !l.Allow(-1) {
+		t.Fatal("Allow(-1) should always succeed")
+	}
+	if !l.Allow(1) {
+		t.Fatal("Allow(1) should succeed after release")
+	}
+}
+
+func TestLimiter_AllowZero(t *testing.T) {
+	l := NewLimiter(1)
+	if !l.Allow(0) {
+		t.Fatal("Allow(0) should not fail with no connections")
+	}
+}
+
+func TestLimiter_AllowNegative(t *testing.T) {
+	l := NewLimiter(1)
+	// Allow(-1) should always return true even without prior acquire.
+	if !l.Allow(-1) {
+		t.Fatal("Allow(-1) should succeed")
+	}
+	// Allow(-1) again — still succeeds.
+	if !l.Allow(-1) {
+		t.Fatal("Allow(-1) should succeed again")
+	}
+}
+
+func TestLimiter_Concurrent(t *testing.T) {
+	l := NewLimiter(100)
+	var wg sync.WaitGroup
+	n := 200
+	var success int64
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if l.Allow(1) {
+				atomic.AddInt64(&success, 1)
+			}
+		}()
+	}
+	wg.Wait()
+	if v := atomic.LoadInt64(&success); v != 100 {
+		t.Errorf("expected 100 successes, got %d", v)
+	}
+}
+
+func TestLimiterGroup_Allow(t *testing.T) {
+	lg := newLimiterGroup(NewLimiter(2), NewLimiter(5), NewLimiter(10))
+
+	if lg.Limit() != 2 {
+		t.Fatalf("expected group limit 2 (smallest), got %d", lg.Limit())
+	}
+
+	// Each Allow must pass all limiters.
+	for i := 0; i < 2; i++ {
+		if !lg.Allow(1) {
+			t.Fatalf("lg.Allow(1) #%d should succeed", i)
+		}
+	}
+	if lg.Allow(1) {
+		t.Fatal("lg.Allow(1) should fail after 2 connections (tightest limit=2)")
+	}
+}
+
+func TestLimiterGroup_AllowRelease(t *testing.T) {
+	lg := newLimiterGroup(NewLimiter(1), NewLimiter(10))
+
+	if !lg.Allow(1) {
+		t.Fatal("Allow(1) should succeed")
+	}
+	if lg.Allow(1) {
+		t.Fatal("Allow(1) should fail")
+	}
+
+	// Release on group must propagate to all limiters.
+	if !lg.Allow(-1) {
+		t.Fatal("Allow(-1) should succeed")
+	}
+
+	// Should be able to acquire again.
+	if !lg.Allow(1) {
+		t.Fatal("Allow(1) should succeed after release")
+	}
+}
+
+func TestLimiterGroup_AllowRollback(t *testing.T) {
+	// limiters sorted: [1, 5] (ascending by limit).
+	lg := newLimiterGroup(NewLimiter(5), NewLimiter(1))
+
+	// First Allow passes both limiters.
+	if !lg.Allow(1) {
+		t.Fatal("first Allow(1) should succeed")
+	}
+
+	// Second Allow fails on the tightest limiter (limit=1) and rolls back.
+	if lg.Allow(1) {
+		t.Fatal("second Allow(1) should fail after limit reached")
+	}
+
+	// The tightest limiter should still be at 1 (rollback worked).
+	ll := lg.limiters[0].(*llimiter)
+	if ll.current != 1 {
+		t.Errorf("first limiter should be at 1 after rollback, got %d", ll.current)
+	}
+}
+
+func TestLimiterGroup_BadInput(t *testing.T) {
+	lg := newLimiterGroup(NewLimiter(-5), NewLimiter(10))
+	if lg.Limit() != -5 {
+		t.Fatalf("expected limit -5 for negative-limit limiter")
+	}
+}
+
+func TestLimiterGroup_Empty(t *testing.T) {
+	lg := newLimiterGroup()
+	if lg.Limit() != 0 {
+		t.Fatal("empty group limit should be 0")
+	}
+	// Empty group returns false (no limiters to consult).
+	if lg.Allow(1) {
+		t.Fatal("Allow(1) on empty group should return false")
+	}
+}
+
+var _ limiter.Limiter = (*llimiter)(nil)
+var _ limiter.Limiter = (*limiterGroup)(nil)

limiter/conn/wrapper/conn_test.go

@@ -0,0 +1,193 @@
+package wrapper
+
+import (
+	"context"
+	"io"
+	"net"
+	"syscall"
+	"testing"
+	"time"
+
+	xio "github.com/go-gost/x/internal/io"
+)
+
+type connWithCloseRead struct {
+	net.Conn
+	closeReadCalled bool
+}
+
+func (c *connWithCloseRead) CloseRead() error {
+	c.closeReadCalled = true
+	return nil
+}
+
+type connWithCloseWrite struct {
+	net.Conn
+	closeWriteCalled bool
+}
+
+func (c *connWithCloseWrite) CloseWrite() error {
+	c.closeWriteCalled = true
+	return nil
+}
+
+type connWithSyscall struct {
+	net.Conn
+}
+
+func (c *connWithSyscall) SyscallConn() (syscall.RawConn, error) {
+	return nil, nil
+}
+
+type connWithContext struct {
+	net.Conn
+}
+
+func (c *connWithContext) Context() context.Context {
+	return context.Background()
+}
+
+func TestWrapConn_NilLimiter(t *testing.T) {
+	c := &mockConn{}
+	wc := WrapConn(nil, c)
+	if wc != c {
+		t.Fatal("WrapConn should return original conn when limiter is nil")
+	}
+}
+
+func TestWrapConn_WithLimiter(t *testing.T) {
+	c := &mockConn{}
+	lim := &allowLimiter{limit: 5}
+	wc := WrapConn(lim, c)
+
+	sc, ok := wc.(*serverConn)
+	if !ok {
+		t.Fatal("WrapConn should return *serverConn when limiter is not nil")
+	}
+	if sc.Conn != c {
+		t.Fatal("inner conn should be the original")
+	}
+	if sc.limiter != lim {
+		t.Fatal("limiter should be set")
+	}
+}
+
+func TestServerConn_Close_ReleasesLimiter(t *testing.T) {
+	c := &mockConn{}
+	lim := &allowLimiter{limit: 1}
+	lim.Allow(1) // acquire one
+
+	sc := &serverConn{Conn: c, limiter: lim}
+	if err := sc.Close(); err != nil {
+		t.Fatal(err)
+	}
+	if !c.closed {
+		t.Fatal("inner conn should be closed")
+	}
+	// After Close, the limiter should have one slot released.
+	if !lim.Allow(1) {
+		t.Fatal("limiter should allow after close released a slot")
+	}
+}
+
+func TestServerConn_CloseRead_Supported(t *testing.T) {
+	inner := &connWithCloseRead{Conn: &mockConn{}}
+	sc := &serverConn{Conn: inner}
+	if err := sc.CloseRead(); err != nil {
+		t.Fatal(err)
+	}
+	if !inner.closeReadCalled {
+		t.Fatal("CloseRead should delegate to inner conn")
+	}
+}
+
+func TestServerConn_CloseRead_Unsupported(t *testing.T) {
+	inner := &mockConn{}
+	sc := &serverConn{Conn: inner}
+	err := sc.CloseRead()
+	if err != xio.ErrUnsupported {
+		t.Fatalf("expected ErrUnsupported, got %v", err)
+	}
+}
+
+func TestServerConn_CloseWrite_Supported(t *testing.T) {
+	inner := &connWithCloseWrite{Conn: &mockConn{}}
+	sc := &serverConn{Conn: inner}
+	if err := sc.CloseWrite(); err != nil {
+		t.Fatal(err)
+	}
+	if !inner.closeWriteCalled {
+		t.Fatal("CloseWrite should delegate to inner conn")
+	}
+}
+
+func TestServerConn_CloseWrite_Unsupported(t *testing.T) {
+	inner := &mockConn{}
+	sc := &serverConn{Conn: inner}
+	err := sc.CloseWrite()
+	if err != xio.ErrUnsupported {
+		t.Fatalf("expected ErrUnsupported, got %v", err)
+	}
+}
+
+func TestServerConn_SyscallConn_Supported(t *testing.T) {
+	inner := &connWithSyscall{Conn: &mockConn{}}
+	sc := &serverConn{Conn: inner}
+	_, err := sc.SyscallConn()
+	if err != nil {
+		t.Fatalf("expected success, got %v", err)
+	}
+}
+
+func TestServerConn_SyscallConn_Unsupported(t *testing.T) {
+	inner := &mockConn{}
+	sc := &serverConn{Conn: inner}
+	_, err := sc.SyscallConn()
+	if err != errUnsupport {
+		t.Fatalf("expected errUnsupport, got %v", err)
+	}
+}
+
+func TestServerConn_Context_Supported(t *testing.T) {
+	inner := &connWithContext{Conn: &mockConn{}}
+	sc := &serverConn{Conn: inner}
+	ctx := sc.Context()
+	if ctx == nil {
+		t.Fatal("Context should not be nil when inner conn supports it")
+	}
+}
+
+func TestServerConn_Context_Unsupported(t *testing.T) {
+	inner := &mockConn{}
+	sc := &serverConn{Conn: inner}
+	ctx := sc.Context()
+	if ctx == nil {
+		t.Fatal("Context should return context.Background(), not nil")
+	}
+}
+
+func TestServerConn_ReadWrite(t *testing.T) {
+	inner := &mockConn{}
+	sc := &serverConn{Conn: inner}
+
+	n, err := sc.Read(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 0 {
+		t.Fatalf("expected 0, got %d", n)
+	}
+
+	n, err = sc.Write([]byte("hello"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 5 {
+		t.Fatalf("expected 5, got %d", n)
+	}
+}
+
+// Ensure unused imports are fine:
+var _ io.ReadWriter = (*serverConn)(nil)
+var _ net.Conn = (*serverConn)(nil)
+var _ = time.Now

limiter/conn/wrapper/listener_test.go

@@ -0,0 +1,147 @@
+package wrapper
+
+import (
+	"errors"
+	"net"
+	"strings"
+	"testing"
+	"time"
+
+	limiter "github.com/go-gost/core/limiter/conn"
+)
+
+type mockConn struct {
+	closed bool
+}
+
+func (c *mockConn) Read(b []byte) (int, error)         { return 0, nil }
+func (c *mockConn) Write(b []byte) (int, error)        { return len(b), nil }
+func (c *mockConn) Close() error                       { c.closed = true; return nil }
+func (c *mockConn) LocalAddr() net.Addr                { return nil }
+func (c *mockConn) RemoteAddr() net.Addr               { return &mockAddr{addr: "192.168.1.1:1234"} }
+func (c *mockConn) SetDeadline(t time.Time) error      { return nil }
+func (c *mockConn) SetReadDeadline(t time.Time) error  { return nil }
+func (c *mockConn) SetWriteDeadline(t time.Time) error { return nil }
+
+type mockAddr struct {
+	addr string
+}
+
+func (a *mockAddr) Network() string { return "tcp" }
+func (a *mockAddr) String() string  { return a.addr }
+
+type mockListener struct {
+	conns   []net.Conn
+	current int
+}
+
+func (l *mockListener) Accept() (net.Conn, error) {
+	if l.current >= len(l.conns) {
+		return nil, errors.New("no more conns")
+	}
+	c := l.conns[l.current]
+	l.current++
+	return c, nil
+}
+
+func (l *mockListener) Close() error   { return nil }
+func (l *mockListener) Addr() net.Addr { return nil }
+
+type allowLimiter struct {
+	limit int
+	count int
+}
+
+func (l *allowLimiter) Allow(n int) bool {
+	if l.count+n > l.limit {
+		return false
+	}
+	l.count += n
+	return true
+}
+func (l *allowLimiter) Limit() int { return l.limit }
+
+type connLimiter struct {
+	lims map[string]limiter.Limiter
+}
+
+func (cl *connLimiter) Limiter(key string) limiter.Limiter {
+	return cl.lims[key]
+}
+
+func TestWrapListener_NilLimiter(t *testing.T) {
+	inner := &mockListener{}
+	ln := WrapListener(nil, inner)
+	if ln != inner {
+		t.Fatal("WrapListener should return the inner listener when limiter is nil")
+	}
+}
+
+func TestAccept_NoLimiterForKey(t *testing.T) {
+	inner := &mockListener{conns: []net.Conn{&mockConn{}}}
+	cl := &connLimiter{lims: map[string]limiter.Limiter{}}
+	ln := WrapListener(cl, inner)
+
+	c, err := ln.Accept()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if c == nil {
+		t.Fatal("conn should not be nil")
+	}
+}
+
+func TestAccept_AllowSuccess(t *testing.T) {
+	inner := &mockListener{conns: []net.Conn{&mockConn{}}}
+	cl := &connLimiter{lims: map[string]limiter.Limiter{
+		"192.168.1.1": &allowLimiter{limit: 1},
+	}}
+	ln := WrapListener(cl, inner)
+
+	c, err := ln.Accept()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if c == nil {
+		t.Fatal("conn should not be nil")
+	}
+
+	// Should be wrapped in serverConn.
+	if _, ok := c.(*serverConn); !ok {
+		t.Fatal("conn should be wrapped as *serverConn")
+	}
+}
+
+func TestAccept_LimitExceeded(t *testing.T) {
+	mc := &mockConn{}
+	inner := &mockListener{conns: []net.Conn{mc}}
+	cl := &connLimiter{lims: map[string]limiter.Limiter{
+		"192.168.1.1": &allowLimiter{limit: 0}, // always deny
+	}}
+	ln := WrapListener(cl, inner)
+
+	c, err := ln.Accept()
+	if err == nil {
+		t.Fatal("expected error when limit exceeded")
+	}
+	if !strings.Contains(err.Error(), "limit exceeded") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if c != nil {
+		t.Fatal("conn should be nil when limit exceeded")
+	}
+	if !mc.closed {
+		t.Fatal("conn should be closed when rejected")
+	}
+}
+
+func TestAccept_ListenerError(t *testing.T) {
+	inner := &mockListener{} // no conns, Accept returns error
+	cl := &connLimiter{lims: map[string]limiter.Limiter{}}
+	ln := WrapListener(cl, inner)
+
+	_, err := ln.Accept()
+	if err == nil {
+		t.Fatal("expected error from listener")
+	}
+}