Add tests for wepoll to debug

Author: nekohasekai

common/bufio/fd_demux_windows.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"sync"
 	"sync/atomic"
+	"unsafe"
 
 	"github.com/sagernet/sing/common/wepoll"
 
@@ -36,7 +37,7 @@ type FDDemultiplexer struct {
 	mutex               sync.Mutex
 	entries             map[int]*fdDemuxEntry
 	registrationCounter uint64
-	registrationToFD    map[uint64]int
+	iosbToFD            map[uintptr]int
 	running             bool
 	closed              atomic.Bool
 	wg                  sync.WaitGroup
@@ -56,12 +57,12 @@ func NewFDDemultiplexer(ctx context.Context) (*FDDemultiplexer, error) {
 
 	ctx, cancel := context.WithCancel(ctx)
 	demux := &FDDemultiplexer{
-		ctx:              ctx,
-		cancel:           cancel,
-		iocp:             iocp,
-		afd:              afd,
-		entries:          make(map[int]*fdDemuxEntry),
-		registrationToFD: make(map[uint64]int),
+		ctx:      ctx,
+		cancel:   cancel,
+		iocp:     iocp,
+		afd:      afd,
+		entries:  make(map[int]*fdDemuxEntry),
+		iosbToFD: make(map[uintptr]int),
 	}
 	return demux, nil
 }
@@ -94,14 +95,14 @@ func (p *FDDemultiplexer) Add(stream *reactorStream, fd int) error {
 	entry.pinner.Pin(&entry.state)
 
 	events := uint32(wepoll.AFD_POLL_RECEIVE | wepoll.AFD_POLL_DISCONNECT | wepoll.AFD_POLL_ABORT | wepoll.AFD_POLL_LOCAL_CLOSE)
-	err = p.afd.Poll(baseHandle, events, &entry.state.iosb, &entry.state.pollInfo, uintptr(regID))
+	err = p.afd.Poll(baseHandle, events, &entry.state.iosb, &entry.state.pollInfo)
 	if err != nil {
 		entry.pinner.Unpin()
 		return err
 	}
 
 	p.entries[fd] = entry
-	p.registrationToFD[regID] = fd
+	p.iosbToFD[uintptr(unsafe.Pointer(&entry.state.iosb))] = fd
 
 	if !p.running {
 		p.running = true
@@ -178,27 +179,28 @@ func (p *FDDemultiplexer) run() {
 
 		for i := uint32(0); i < numRemoved; i++ {
 			ev := entries[i]
-			regID := uint64(ev.CompletionKey)
 
-			if regID == 0 {
+			if ev.Overlapped == nil {
 				continue
 			}
 
+			iosbPtr := uintptr(unsafe.Pointer(ev.Overlapped))
+
 			p.mutex.Lock()
-			fd, ok := p.registrationToFD[regID]
+			fd, ok := p.iosbToFD[iosbPtr]
 			if !ok {
 				p.mutex.Unlock()
 				continue
 			}
 
 			entry := p.entries[fd]
-			if entry == nil || entry.registrationID != regID {
+			if entry == nil {
 				p.mutex.Unlock()
 				continue
 			}
 
 			entry.pinner.Unpin()
-			delete(p.registrationToFD, regID)
+			delete(p.iosbToFD, iosbPtr)
 			delete(p.entries, fd)
 
 			if entry.cancelled {

common/bufio/fd_demux_windows_test.go

@@ -0,0 +1,305 @@
+//go:build windows
+
+package bufio
+
+import (
+	"context"
+	"net"
+	"sync"
+	"sync/atomic"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func getSocketFD(t *testing.T, conn net.PacketConn) int {
+	syscallConn, ok := conn.(syscall.Conn)
+	require.True(t, ok)
+	rawConn, err := syscallConn.SyscallConn()
+	require.NoError(t, err)
+	var fd int
+	err = rawConn.Control(func(f uintptr) { fd = int(f) })
+	require.NoError(t, err)
+	return fd
+}
+
+func TestFDDemultiplexer_Create(t *testing.T) {
+	t.Parallel()
+
+	demux, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+
+	err = demux.Close()
+	require.NoError(t, err)
+}
+
+func TestFDDemultiplexer_CreateMultiple(t *testing.T) {
+	t.Parallel()
+
+	demux1, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+	defer demux1.Close()
+
+	demux2, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+	defer demux2.Close()
+}
+
+func TestFDDemultiplexer_AddRemove(t *testing.T) {
+	t.Parallel()
+
+	demux, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+	defer demux.Close()
+
+	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer conn.Close()
+
+	fd := getSocketFD(t, conn)
+
+	stream := &reactorStream{}
+
+	err = demux.Add(stream, fd)
+	require.NoError(t, err)
+
+	demux.Remove(fd)
+}
+
+func TestFDDemultiplexer_RapidAddRemove(t *testing.T) {
+	t.Parallel()
+
+	demux, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+	defer demux.Close()
+
+	const iterations = 50
+
+	for i := 0; i < iterations; i++ {
+		conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+		require.NoError(t, err)
+
+		fd := getSocketFD(t, conn)
+		stream := &reactorStream{}
+
+		err = demux.Add(stream, fd)
+		require.NoError(t, err)
+
+		demux.Remove(fd)
+		conn.Close()
+	}
+}
+
+func TestFDDemultiplexer_ConcurrentAccess(t *testing.T) {
+	t.Parallel()
+
+	demux, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+	defer demux.Close()
+
+	const numGoroutines = 10
+	const iterations = 20
+
+	var wg sync.WaitGroup
+	wg.Add(numGoroutines)
+
+	for g := 0; g < numGoroutines; g++ {
+		go func() {
+			defer wg.Done()
+
+			for i := 0; i < iterations; i++ {
+				conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+				if err != nil {
+					continue
+				}
+
+				fd := getSocketFD(t, conn)
+				stream := &reactorStream{}
+
+				err = demux.Add(stream, fd)
+				if err == nil {
+					demux.Remove(fd)
+				}
+				conn.Close()
+			}
+		}()
+	}
+
+	wg.Wait()
+}
+
+func TestFDDemultiplexer_ReceiveEvent(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	demux, err := NewFDDemultiplexer(ctx)
+	require.NoError(t, err)
+	defer demux.Close()
+
+	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer conn.Close()
+
+	fd := getSocketFD(t, conn)
+
+	triggered := make(chan struct{}, 1)
+	stream := &reactorStream{
+		state: atomic.Int32{},
+	}
+	stream.connection = &reactorConnection{
+		upload:   stream,
+		download: stream,
+		done:     make(chan struct{}),
+	}
+
+	originalRunActiveLoop := stream.runActiveLoop
+	_ = originalRunActiveLoop
+
+	err = demux.Add(stream, fd)
+	require.NoError(t, err)
+
+	sender, err := net.ListenPacket("udp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer sender.Close()
+
+	_, err = sender.WriteTo([]byte("test data"), conn.LocalAddr())
+	require.NoError(t, err)
+
+	time.Sleep(200 * time.Millisecond)
+
+	select {
+	case <-triggered:
+	default:
+	}
+
+	demux.Remove(fd)
+}
+
+func TestFDDemultiplexer_CloseWhilePolling(t *testing.T) {
+	t.Parallel()
+
+	demux, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+
+	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer conn.Close()
+
+	fd := getSocketFD(t, conn)
+	stream := &reactorStream{}
+
+	err = demux.Add(stream, fd)
+	require.NoError(t, err)
+
+	time.Sleep(50 * time.Millisecond)
+
+	done := make(chan struct{})
+	go func() {
+		demux.Close()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Close blocked - possible deadlock")
+	}
+}
+
+func TestFDDemultiplexer_RemoveNonExistent(t *testing.T) {
+	t.Parallel()
+
+	demux, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+	defer demux.Close()
+
+	demux.Remove(99999)
+}
+
+func TestFDDemultiplexer_AddAfterClose(t *testing.T) {
+	t.Parallel()
+
+	demux, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+
+	err = demux.Close()
+	require.NoError(t, err)
+
+	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer conn.Close()
+
+	fd := getSocketFD(t, conn)
+	stream := &reactorStream{}
+
+	err = demux.Add(stream, fd)
+	require.Error(t, err)
+}
+
+func TestFDDemultiplexer_MultipleSocketsSimultaneous(t *testing.T) {
+	t.Parallel()
+
+	demux, err := NewFDDemultiplexer(context.Background())
+	require.NoError(t, err)
+	defer demux.Close()
+
+	const numSockets = 5
+	conns := make([]net.PacketConn, numSockets)
+	fds := make([]int, numSockets)
+
+	for i := 0; i < numSockets; i++ {
+		conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+		require.NoError(t, err)
+		defer conn.Close()
+		conns[i] = conn
+
+		fd := getSocketFD(t, conn)
+		fds[i] = fd
+
+		stream := &reactorStream{}
+		err = demux.Add(stream, fd)
+		require.NoError(t, err)
+	}
+
+	for i := 0; i < numSockets; i++ {
+		demux.Remove(fds[i])
+	}
+}
+
+func TestFDDemultiplexer_ContextCancellation(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	demux, err := NewFDDemultiplexer(ctx)
+	require.NoError(t, err)
+
+	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer conn.Close()
+
+	fd := getSocketFD(t, conn)
+	stream := &reactorStream{}
+
+	err = demux.Add(stream, fd)
+	require.NoError(t, err)
+
+	cancel()
+
+	time.Sleep(100 * time.Millisecond)
+
+	done := make(chan struct{})
+	go func() {
+		demux.Close()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Close blocked after context cancellation")
+	}
+}