udp: use nni_aio_start

This also moves the close of the UDP socket later, to avoid a
potential use after free while the aio's are still in-flight.
Unfortunately we cannot unbind cleanly without a hard close.

3 files changed, 14 insertions, 32 deletions

diff --git a/src/platform/posix/posix_udp.c b/src/platform/posix/posix_udp.c
index ade9e02e..4bb50656 100644
--- a/src/platform/posix/posix_udp.c
+++ b/src/platform/posix/posix_udp.c
@@ -298,13 +298,10 @@ void
 nni_plat_udp_recv(nni_plat_udp *udp, nni_aio *aio)
 {
 	int rv;
-	if (nni_aio_begin(aio) != 0) {
-		return;
-	}
+	nni_aio_reset(aio);
 	nni_mtx_lock(&udp->udp_mtx);
-	if ((rv = nni_aio_schedule(aio, nni_plat_udp_cancel, udp)) != 0) {
+	if (!nni_aio_start(aio, nni_plat_udp_cancel, udp)) {
 		nni_mtx_unlock(&udp->udp_mtx);
-		nni_aio_finish_error(aio, rv);
 		return;
 	}
 	nni_list_append(&udp->udp_recvq, aio);
@@ -322,13 +319,10 @@ void
 nni_plat_udp_send(nni_plat_udp *udp, nni_aio *aio)
 {
 	int rv;
-	if (nni_aio_begin(aio) != 0) {
-		return;
-	}
+	nni_aio_reset(aio);
 	nni_mtx_lock(&udp->udp_mtx);
-	if ((rv = nni_aio_schedule(aio, nni_plat_udp_cancel, udp)) != 0) {
+	if (!nni_aio_start(aio, nni_plat_udp_cancel, udp)) {
 		nni_mtx_unlock(&udp->udp_mtx);
-		nni_aio_finish_error(aio, rv);
 		return;
 	}
 	nni_list_append(&udp->udp_sendq, aio);
diff --git a/src/platform/windows/win_udp.c b/src/platform/windows/win_udp.c
index c1c7ef21..c74b3f77 100644
--- a/src/platform/windows/win_udp.c
+++ b/src/platform/windows/win_udp.c
@@ -122,9 +122,7 @@ nni_plat_udp_send(nni_plat_udp *u, nni_aio *aio)
 	int              rv;
 	DWORD            nsent;
 
-	if (nni_aio_begin(aio) != 0) {
-		return;
-	}
+	nni_aio_reset(aio);
 	sa = nni_aio_get_input(aio, 0);
 	if ((tolen = nni_win_nn2sockaddr(&to, sa)) < 0) {
 		nni_aio_finish_error(aio, NNG_EADDRINVAL);
@@ -280,19 +278,15 @@ again:
 void
 nni_plat_udp_recv(nni_plat_udp *u, nni_aio *aio)
 {
-	int rv;
-	if (nni_aio_begin(aio) != 0) {
-		return;
-	}
+	nni_aio_reset(aio);
 	nni_mtx_lock(&u->lk);
 	if (u->closed) {
 		nni_mtx_unlock(&u->lk);
 		nni_aio_finish_error(aio, NNG_ECLOSED);
 		return;
 	}
-	if ((rv = nni_aio_schedule(aio, udp_recv_cancel, u)) != 0) {
+	if (!nni_aio_start(aio, udp_recv_cancel, u)) {
 		nni_mtx_unlock(&u->lk);
-		nni_aio_finish_error(aio, rv);
 		return;
 	}
 	nni_list_append(&u->rxq, aio);
diff --git a/src/sp/transport/udp/udp.c b/src/sp/transport/udp/udp.c
index 97ffdc04..f2fc6889 100644
--- a/src/sp/transport/udp/udp.c
+++ b/src/sp/transport/udp/udp.c
@@ -1108,6 +1108,10 @@ udp_ep_fini(void *arg)
 	nni_aio_fini(&ep->tx_aio);
 	nni_aio_fini(&ep->rx_aio);
 
+	if (ep->udp != NULL) {
+		nni_udp_close(ep->udp);
+	}
+
 	for (int i = 0; i < ep->tx_ring.size; i++) {
 		nni_msg_free(ep->tx_ring.descs[i].payload);
 		ep->tx_ring.descs[i].payload = NULL;
@@ -1167,10 +1171,6 @@ udp_ep_stop(void *arg)
 
 	// finally close the tx channel
 	nni_aio_stop(&ep->tx_aio);
-
-	if (ep->udp != NULL) {
-		nni_udp_close(ep->udp);
-	}
 }
 
 // timer handler - sends out additional creqs as needed,
@@ -1517,12 +1517,12 @@ static void
 udp_ep_connect(void *arg, nni_aio *aio)
 {
 	udp_ep *ep = arg;
-	int     rv;
 
-	if (nni_aio_begin(aio) != 0) {
+	nni_mtx_lock(&ep->mtx);
+	if (!nni_aio_start(aio, udp_ep_cancel, ep)) {
+		nni_mtx_unlock(&ep->mtx);
 		return;
 	}
-	nni_mtx_lock(&ep->mtx);
 	if (ep->closed) {
 		nni_mtx_unlock(&ep->mtx);
 		nni_aio_finish_error(aio, NNG_ECLOSED);
@@ -1536,12 +1536,6 @@ udp_ep_connect(void *arg, nni_aio *aio)
 	NNI_ASSERT(nni_list_empty(&ep->connaios));
 	ep->dialer = true;
 
-	if ((rv = nni_aio_schedule(aio, udp_ep_cancel, ep)) != 0) {
-		nni_mtx_unlock(&ep->mtx);
-		nni_aio_finish_error(aio, rv);
-		return;
-	}
-
 	nni_list_append(&ep->connaios, aio);
 
 	// lookup the IP address