fixes #1837 IPC - Use After Free

This fixes a problem only found on Windows, that affected both IPC and TCP.

3 files changed, 59 insertions, 11 deletions

diff --git a/src/platform/windows/win_ipcconn.c b/src/platform/windows/win_ipcconn.c
index b2136a89..ef6ceb66 100644
--- a/src/platform/windows/win_ipcconn.c
+++ b/src/platform/windows/win_ipcconn.c
@@ -30,6 +30,8 @@ typedef struct ipc_conn {
 	int           send_rv;
 	int           conn_rv;
 	bool          closed;
+	bool          sending;
+	bool          recving;
 	nni_mtx       mtx;
 	nni_cv        cv;
 	nni_reap_node reap;
@@ -76,9 +78,11 @@ ipc_recv_start(ipc_conn *c)
 			len = 0x1000000;
 		}
 
+		c->recving = true;
 		if ((!ReadFile(c->f, buf, len, NULL, &c->recv_io.olpd)) &&
 		    ((rv = GetLastError()) != ERROR_IO_PENDING)) {
 			// Synchronous failure.
+			c->recving = false;
 			nni_aio_list_remove(aio);
 			nni_aio_finish_error(aio, nni_win_error(rv));
 		} else {
@@ -104,6 +108,7 @@ ipc_recv_cb(nni_win_io *io, int rv, size_t num)
 		// A zero byte receive is a remote close from the peer.
 		rv = NNG_ECONNSHUT;
 	}
+	c->recving = false;
 	nni_aio_list_remove(aio);
 	ipc_recv_start(c);
 	nni_mtx_unlock(&c->mtx);
@@ -196,9 +201,11 @@ ipc_send_start(ipc_conn *c)
 			len = 0x1000000;
 		}
 
+		c->sending = true;
 		if ((!WriteFile(c->f, buf, len, NULL, &c->send_io.olpd)) &&
 		    ((rv = GetLastError()) != ERROR_IO_PENDING)) {
 			// Synchronous failure.
+			c->sending = false;
 			nni_aio_list_remove(aio);
 			nni_aio_finish_error(aio, nni_win_error(rv));
 		} else {
@@ -217,6 +224,7 @@ ipc_send_cb(nni_win_io *io, int rv, size_t num)
 	aio = nni_list_first(&c->send_aios);
 	NNI_ASSERT(aio != NULL);
 	nni_aio_list_remove(aio);
+	c->sending = false;
 	if (c->send_rv != 0) {
 		rv         = c->send_rv;
 		c->send_rv = 0;
@@ -275,17 +283,32 @@ static void
 ipc_close(void *arg)
 {
 	ipc_conn *c = arg;
+	nni_time  now;
 	nni_mtx_lock(&c->mtx);
 	if (!c->closed) {
 		HANDLE f  = c->f;
 		c->closed = true;
-		c->f      = INVALID_HANDLE_VALUE;
+
+		c->f = INVALID_HANDLE_VALUE;
 
 		if (f != INVALID_HANDLE_VALUE) {
+			CancelIoEx(f, &c->send_io.olpd);
+			CancelIoEx(f, &c->recv_io.olpd);
 			DisconnectNamedPipe(f);
 			CloseHandle(f);
 		}
 	}
+	now = nni_clock();
+	// wait up to a maximum of 10 seconds before assuming something is
+	// badly amiss. from what we can tell, this doesn't happen, and we do
+	// see the timer expire properly, but this safeguard can prevent a
+	// hang.
+	while ((c->recving || c->sending) &&
+	    ((nni_clock() - now) < (NNI_SECOND * 10))) {
+		nni_mtx_unlock(&c->mtx);
+		nni_msleep(1);
+		nni_mtx_lock(&c->mtx);
+	}
 	nni_mtx_unlock(&c->mtx);
 }
 
diff --git a/src/platform/windows/win_tcp.h b/src/platform/windows/win_tcp.h
index e54f853c..d151b9a3 100644
--- a/src/platform/windows/win_tcp.h
+++ b/src/platform/windows/win_tcp.h
@@ -1,5 +1,5 @@
 //
-// Copyright 2019 Staysail Systems, Inc. <info@staysail.tech>
+// Copyright 2024 Staysail Systems, Inc. <info@staysail.tech>
 // Copyright 2018 Capitar IT Group BV <info@capitar.com>
 // Copyright 2019 Devolutions <info@devolutions.net>
 //
@@ -24,16 +24,18 @@ struct nni_tcp_conn {
 	nni_win_io        conn_io;
 	nni_list          recv_aios;
 	nni_list          send_aios;
-	nni_aio *         conn_aio;
+	nni_aio          *conn_aio;
 	SOCKADDR_STORAGE  sockname;
 	SOCKADDR_STORAGE  peername;
-	nni_tcp_dialer *  dialer;
+	nni_tcp_dialer   *dialer;
 	nni_tcp_listener *listener;
 	int               recv_rv;
 	int               send_rv;
 	int               conn_rv;
 	bool              closed;
 	char              buf[512]; // to hold acceptex results
+	bool              sending;
+	bool              recving;
 	nni_mtx           mtx;
 	nni_cv            cv;
 };
diff --git a/src/platform/windows/win_tcpconn.c b/src/platform/windows/win_tcpconn.c
index 24605505..0e74ccfd 100644
--- a/src/platform/windows/win_tcpconn.c
+++ b/src/platform/windows/win_tcpconn.c
@@ -48,13 +48,15 @@ tcp_recv_start(nni_tcp_conn *c)
 			}
 		}
 
-		flags = 0;
-		rv    = WSARecv(
+		c->recving = true;
+		flags      = 0;
+		rv         = WSARecv(
                     c->s, iov, niov, &nrecv, &flags, &c->recv_io.olpd, NULL);
 
 		if ((rv == SOCKET_ERROR) &&
 		    ((rv = GetLastError()) != ERROR_IO_PENDING)) {
 			// Synchronous error.
+			c->recving = false;
 			nni_aio_list_remove(aio);
 			nni_aio_finish_error(aio, nni_win_error(rv));
 		} else {
@@ -76,7 +78,6 @@ tcp_recv_cb(nni_win_io *io, int rv, size_t num)
 	nni_mtx_lock(&c->mtx);
 	aio = nni_list_first(&c->recv_aios);
 	NNI_ASSERT(aio != NULL);
-	nni_aio_list_remove(aio);
 
 	if (c->recv_rv != 0) {
 		rv         = c->recv_rv;
@@ -86,6 +87,8 @@ tcp_recv_cb(nni_win_io *io, int rv, size_t num)
 		// A zero byte receive is a remote close from the peer.
 		rv = NNG_ECONNSHUT;
 	}
+	c->recving = false;
+	nni_aio_list_remove(aio);
 	tcp_recv_start(c);
 	nni_mtx_unlock(&c->mtx);
 
@@ -164,11 +167,13 @@ tcp_send_start(nni_tcp_conn *c)
 			}
 		}
 
+		c->sending = true;
 		rv = WSASend(c->s, iov, niov, NULL, 0, &c->send_io.olpd, NULL);
 
 		if ((rv == SOCKET_ERROR) &&
 		    ((rv = GetLastError()) != ERROR_IO_PENDING)) {
 			// Synchronous failure.
+			c->sending = false;
 			nni_aio_list_remove(aio);
 			nni_aio_finish_error(aio, nni_win_error(rv));
 		} else {
@@ -208,6 +213,7 @@ tcp_send_cb(nni_win_io *io, int rv, size_t num)
 	aio = nni_list_first(&c->send_aios);
 	NNI_ASSERT(aio != NULL);
 	nni_aio_list_remove(aio); // should always be at head
+	c->sending = false;
 
 	if (c->send_rv != 0) {
 		rv         = c->send_rv;
@@ -246,14 +252,31 @@ tcp_close(void *arg)
 {
 	nni_tcp_conn *c = arg;
 	nni_mtx_lock(&c->mtx);
+	nni_time now;
 	if (!c->closed) {
+		SOCKET s = c->s;
+
 		c->closed = true;
-		if (c->s != INVALID_SOCKET) {
-			shutdown(c->s, SD_BOTH);
-			closesocket(c->s);
-			c->s = INVALID_SOCKET;
+		c->s      = INVALID_SOCKET;
+
+		if (s != INVALID_SOCKET) {
+			CancelIoEx(s, &c->send_io.olpd);
+			CancelIoEx(s, &c->recv_io.olpd);
+			shutdown(s, SD_BOTH);
+			closesocket(s);
 		}
 	}
+	now = nni_clock();
+	// wait up to a maximum of 10 seconds before assuming something is
+	// badly amiss. from what we can tell, this doesn't happen, and we do
+	// see the timer expire properly, but this safeguard can prevent a
+	// hang.
+	while ((c->recving || c->sending) &&
+	    ((nni_clock() - now) < (NNI_SECOND * 10))) {
+		nni_mtx_unlock(&c->mtx);
+		nni_msleep(1);
+		nni_mtx_lock(&c->mtx);
+	}
 	nni_mtx_unlock(&c->mtx);
 }