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

2026-05-24 19:38:54 +08:00
parent fa708f4b5f
commit 05fe690f67
5 changed files with 1037 additions and 0 deletions
@@ -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
@@ -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")
+	}
+}