From a2cbb62b6e008fefacbd7e9d4aeb5d29ba455ae0 Mon Sep 17 00:00:00 2001
From: Ulrich Hornung <hornunguli@gmx.de>
Date: Mon, 16 Sep 2024 23:05:21 +0200
Subject: [PATCH] fix race condition where "CloseAndZero" is executed while
 still used

---
 udp/udp_rio_windows.go | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/udp/udp_rio_windows.go b/udp/udp_rio_windows.go
index ee7e1e002..d68957ab1 100644
--- a/udp/udp_rio_windows.go
+++ b/udp/udp_rio_windows.go
@@ -14,6 +14,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"syscall"
+	"time"
 	"unsafe"
 
 	"github.com/sirupsen/logrus"
@@ -178,10 +179,11 @@ func (u *RIOConn) receive(buf []byte) (int, windows.RawSockaddrInet6, error) {
 retry:
 	count = 0
 	for tries := 0; count == 0 && tries < receiveSpins; tries++ {
+		if !u.isOpen.Load() { // might have changes since first check before the mutex lock
+			return 0, windows.RawSockaddrInet6{}, net.ErrClosed
+		}
+
 		if tries > 0 {
-			if !u.isOpen.Load() {
-				return 0, windows.RawSockaddrInet6{}, net.ErrClosed
-			}
 			procyield(1)
 		}
 
@@ -247,6 +249,10 @@ func (u *RIOConn) WriteTo(buf []byte, ip netip.AddrPort) error {
 	u.tx.mu.Lock()
 	defer u.tx.mu.Unlock()
 
+	if !u.isOpen.Load() { // might have changes since first check before the mutex lock
+		return net.ErrClosed
+	}
+
 	count := winrio.DequeueCompletion(u.tx.cq, u.results[:])
 	if count == 0 && u.tx.isFull {
 		err := winrio.Notify(u.tx.cq)
@@ -323,6 +329,14 @@ func (u *RIOConn) Close() error {
 	windows.PostQueuedCompletionStatus(u.rx.iocp, 0, 0, nil)
 	windows.PostQueuedCompletionStatus(u.tx.iocp, 0, 0, nil)
 
+	u.rx.mu.Lock()                   // for waiting till active reader is done
+	time.Sleep(time.Millisecond * 0) // avoid warning about empty critical section
+	u.rx.mu.Unlock()
+
+	u.tx.mu.Lock()                   // for waiting till active writer is done
+	time.Sleep(time.Millisecond * 0) // avoid warning about empty critical section
+	u.tx.mu.Unlock()
+
 	u.rx.CloseAndZero()
 	u.tx.CloseAndZero()
 	if u.sock != 0 {