fixes #1657 Use after free in listener (data race)

3 files changed, 27 insertions, 4 deletions

diff --git a/src/core/listener.c b/src/core/listener.c
index 09fddf03..0ec8e507 100644
--- a/src/core/listener.c
+++ b/src/core/listener.c
@@ -306,6 +306,7 @@ nni_listener_rele(nni_listener *l)
 	bool reap;
 
 	nni_mtx_lock(&listeners_lk);
+	NNI_ASSERT(l->l_ref > 0);
 	l->l_ref--;
 	reap = ((l->l_ref == 0) && (l->l_closed));
 	nni_mtx_unlock(&listeners_lk);
@@ -329,6 +330,7 @@ nni_listener_close(nni_listener *l)
 
 	nni_listener_shutdown(l);
 
+	nni_sock_remove_listener(l);
 	nni_listener_rele(l); // This will reap if reference count is zero.
 }
 
diff --git a/src/core/socket.c b/src/core/socket.c
index 1e7d978e..316f3603 100644
--- a/src/core/socket.c
+++ b/src/core/socket.c
@@ -9,6 +9,7 @@
 //
 
 #include "core/nng_impl.h"
+#include "list.h"
 #include "sockimpl.h"
 
 #include <stdio.h>
@@ -691,7 +692,6 @@ nni_sock_shutdown(nni_sock *sock)
 
 	while ((l = nni_list_first(&sock->s_listeners)) != NULL) {
 		nni_listener_hold(l);
-		nni_list_node_remove(&l->l_node);
 		nni_mtx_unlock(&sock->s_mx);
 		nni_listener_close(l);
 		nni_mtx_lock(&sock->s_mx);
@@ -890,10 +890,17 @@ int
 nni_sock_add_listener(nni_sock *s, nni_listener *l)
 {
 	nni_sockopt *sopt;
+	int          rv;
+
+	// grab a hold on the listener for the socket
+	if ((rv = nni_listener_hold(l)) != 0) {
+		return (rv);
+	}
 
 	nni_mtx_lock(&s->s_mx);
 	if (s->s_closing) {
 		nni_mtx_unlock(&s->s_mx);
+		nni_listener_rele(l);
 		return (NNG_ECLOSED);
 	}
 
@@ -917,6 +924,19 @@ nni_sock_add_listener(nni_sock *s, nni_listener *l)
 	return (0);
 }
 
+void
+nni_sock_remove_listener(nni_listener *l)
+{
+	nni_sock *s = l->l_sock;
+	nni_mtx_lock(&s->s_mx);
+	NNI_ASSERT(nni_list_node_active(&l->l_node));
+	nni_list_node_remove(&l->l_node);
+	nni_mtx_unlock(&s->s_mx);
+
+	// also drop the hold from the socket
+	nni_listener_rele(l);
+}
+
 int
 nni_sock_add_dialer(nni_sock *s, nni_dialer *d)
 {
diff --git a/src/core/sockimpl.h b/src/core/sockimpl.h
index 83cbb877..9734d94a 100644
--- a/src/core/sockimpl.h
+++ b/src/core/sockimpl.h
@@ -18,7 +18,7 @@
 struct nni_dialer {
 	nni_sp_dialer_ops d_ops;  // transport ops
 	nni_sp_tran      *d_tran; // transport pointer
-	void	     *d_data; // transport private
+	void             *d_data; // transport private
 	uint32_t          d_id;   // endpoint id
 	nni_list_node     d_node; // per socket list
 	nni_sock         *d_sock;
@@ -123,8 +123,9 @@ struct nni_pipe {
 #endif
 };
 
-extern int nni_sock_add_dialer(nni_sock *, nni_dialer *);
-extern int nni_sock_add_listener(nni_sock *, nni_listener *);
+extern int  nni_sock_add_dialer(nni_sock *, nni_dialer *);
+extern int  nni_sock_add_listener(nni_sock *, nni_listener *);
+extern void nni_sock_remove_listener(nni_listener *);
 
 extern void nni_dialer_add_pipe(nni_dialer *, void *);
 extern void nni_dialer_shutdown(nni_dialer *);