WireGuard inbound: Reduce memory held by idle UDP flows (#6856)

https://github.com/XTLS/Xray-core/pull/6856#issuecomment-5870639238

proxy/wireguard/tun_test.go

@@ -0,0 +1,91 @@
+package wireguard
+
+import (
+	"testing"
+	"time"
+
+	"github.com/xtls/xray-core/common/net"
+)
+
+// BenchmarkUDPManagerNewSession measures what one new UDP flow costs the
+// inbound while it stays open: QUIC and DNS open many short flows, and each
+// one lives until the connection idle timeout.
+func BenchmarkUDPManagerNewSession(b *testing.B) {
+	m := &udpManager{
+		handler: func(conn net.Conn, dest net.Destination) {},
+		m:       make(map[string]*udpConn),
+	}
+	dst := net.UDPDestination(net.ParseAddress("1.1.1.1"), 443)
+	payload := make([]byte, 1200)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		src := net.UDPDestination(net.IPAddress([]byte{10, byte(i >> 16), byte(i >> 8), byte(i)}), net.Port(1024+i%60000))
+		m.feed(src, dst, payload)
+	}
+}
+
+func TestPacketQueueOrderAndClose(t *testing.T) {
+	q := newPacketQueue(udpQueueLimit)
+	for i := 0; i < 3; i++ {
+		if !q.push(&packet{p: []byte{byte(i)}}) {
+			t.Fatalf("push %d rejected", i)
+		}
+	}
+	for i := 0; i < 3; i++ {
+		p, ok := q.pop()
+		if !ok || p.p[0] != byte(i) {
+			t.Fatalf("pop %d: got %v, %v", i, p, ok)
+		}
+	}
+	q.close()
+	if _, ok := q.pop(); ok {
+		t.Fatal("pop after close returned a packet")
+	}
+	if q.push(&packet{}) {
+		t.Fatal("push after close accepted")
+	}
+}
+
+func TestPacketQueueLimit(t *testing.T) {
+	q := newPacketQueue(udpQueueLimit)
+	for i := 0; i < udpQueueLimit; i++ {
+		if !q.push(&packet{}) {
+			t.Fatalf("push %d rejected below the limit", i)
+		}
+	}
+	if q.push(&packet{}) {
+		t.Fatal("push above the limit accepted")
+	}
+}
+
+func TestPacketQueueCloseUnblocksReader(t *testing.T) {
+	q := newPacketQueue(udpQueueLimit)
+	done := make(chan bool)
+	go func() {
+		_, ok := q.pop()
+		done <- ok
+	}()
+	q.close()
+	select {
+	case ok := <-done:
+		if ok {
+			t.Fatal("blocked pop returned a packet after close")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("close did not wake the reader")
+	}
+}
+
+func TestPacketQueueDropsDrainedStorage(t *testing.T) {
+	q := newPacketQueue(udpQueueLimit)
+	for i := 0; i < 100; i++ {
+		q.push(&packet{})
+	}
+	for i := 0; i < 100; i++ {
+		q.pop()
+	}
+	if q.items != nil {
+		t.Fatalf("drained queue still holds %d slots", cap(q.items))
+	}
+}