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

2026-05-24 20:54:34 +08:00
parent e9372ea44a
commit dc99f4731f
12 changed files with 1511 additions and 3 deletions
@@ -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,
@@ -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
+}
@@ -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
@@ -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
@@ -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)
+)