From 6ee74f20322250f10698405e15a85ad2e44fc5cb Mon Sep 17 00:00:00 2001
From: Sanaei <ho3ein.sanaei@gmail.com>
Date: Wed, 9 Sep 2026 09:37:12 +0200
Subject: [PATCH] fix(amneziawgnet): wait for the client netstack goroutines
 before closing its device

TestPortForwardRoundTripTCPAndUDP flakes in the race job: closing the test's
client WireGuard device races the goroutines still writing into its netstack.

amneziawg-go's device.Close() calls tun.Close() before it stops the routine
draining the tun, and netTun.Close() closes the unbuffered incomingPacket
channel that WriteNotify sends on. A goroutine still inside a netstack write
when the deferred clientDev.Close() runs therefore closes and sends on the same
channel -- reported as a data race, and on a bad interleaving a "send on closed
channel" panic.

The TCP echo listener, its per-connection copies and the UDP echo all write into
clientNet, and teardown only closed the two listeners before the device: nothing
waited for the goroutines themselves. A WaitGroup deferred right after
clientDev.Close() supplies the missing edge, since LIFO then puts the wait
between the listener closes and the device close.

Confirmed by flooding the existing UDP echo goroutine under GOMAXPROCS=1 and 2,
which failed 3/6 and 2/6 runs with the stack CI reported and 0/12 with the fix.
---
 internal/amneziawgnet/portfwd_test.go | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/internal/amneziawgnet/portfwd_test.go b/internal/amneziawgnet/portfwd_test.go
index e83924a69..1c8f81811 100644
--- a/internal/amneziawgnet/portfwd_test.go
+++ b/internal/amneziawgnet/portfwd_test.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"net"
 	"net/netip"
+	"sync"
 	"testing"
 	"time"
 
@@ -285,6 +286,10 @@ func TestPortForwardRoundTripTCPAndUDP(t *testing.T) {
 	}
 	clientDev := device.NewDevice(clientTun, awgconn.NewDefaultBind(), device.NewLogger(device.LogLevelSilent, ""))
 	defer clientDev.Close()
+	// clientDev.Close() closes the tun's packet channel without waiting for
+	// writers, so every goroutine writing into clientNet must be gone first.
+	var clientSvc sync.WaitGroup
+	defer clientSvc.Wait()
 
 	clientPrivHex, err := wireguard.KeyToHex(clientPriv)
 	if err != nil {
@@ -338,13 +343,16 @@ primed:
 		t.Fatalf("client ListenTCP: %v", err)
 	}
 	defer tcpSvc.Close()
+	clientSvc.Add(1)
 	go func() {
+		defer clientSvc.Done()
 		for {
 			c, err := tcpSvc.Accept()
 			if err != nil {
 				return
 			}
-			go func() { io.Copy(c, c); c.Close() }()
+			clientSvc.Add(1)
+			go func() { defer clientSvc.Done(); io.Copy(c, c); c.Close() }()
 		}
 	}()
 
@@ -353,7 +361,9 @@ primed:
 		t.Fatalf("client ListenUDP: %v", err)
 	}
 	defer udpSvc.Close()
+	clientSvc.Add(1)
 	go func() {
+		defer clientSvc.Done()
 		buf := make([]byte, 1500)
 		for {
 			n, addr, err := udpSvc.ReadFrom(buf)