fix #946 Use after free in TLS

This also introduces a more efficient reference counting usage based
on atomics, rather than locks.

6 files changed, 78 insertions, 32 deletions

diff --git a/src/core/platform.h b/src/core/platform.h
index 1418ba4f..65d9af82 100644
--- a/src/core/platform.h
+++ b/src/core/platform.h
@@ -172,11 +172,13 @@ extern void nni_atomic_flag_reset(nni_atomic_flag *);
 typedef struct nni_atomic_u64 nni_atomic_u64;
 
 extern void     nni_atomic_init64(nni_atomic_u64 *);
-extern void     nni_atomic_inc64(nni_atomic_u64 *, uint64_t);
-extern void     nni_atomic_dec64(nni_atomic_u64 *, uint64_t);
+extern void     nni_atomic_add64(nni_atomic_u64 *, uint64_t);
+extern void     nni_atomic_sub64(nni_atomic_u64 *, uint64_t);
 extern uint64_t nni_atomic_get64(nni_atomic_u64 *);
 extern void     nni_atomic_set64(nni_atomic_u64 *, uint64_t);
 extern uint64_t nni_atomic_swap64(nni_atomic_u64 *, uint64_t);
+extern uint64_t nni_atomic_dec64_nv(nni_atomic_u64 *);
+extern void     nni_atomic_inc64(nni_atomic_u64 *);
 
 //
 // Clock Support
diff --git a/src/core/stats.c b/src/core/stats.c
index f3f969c3..247e7ab4 100644
--- a/src/core/stats.c
+++ b/src/core/stats.c
@@ -152,13 +152,13 @@ nni_stat_init_atomic(nni_stat_item *stat, const char *name, const char *desc)
 void
 nni_stat_inc_atomic(nni_stat_item *stat, uint64_t inc)
 {
-	nni_atomic_inc64(&stat->si_atomic, inc);
+	nni_atomic_add64(&stat->si_atomic, inc);
 }
 
 void
 nni_stat_dec_atomic(nni_stat_item *stat, uint64_t inc)
 {
-	nni_atomic_dec64(&stat->si_atomic, inc);
+	nni_atomic_sub64(&stat->si_atomic, inc);
 }
 #endif
 
diff --git a/src/platform/posix/posix_atomic.c b/src/platform/posix/posix_atomic.c
index 7f735ec0..8f2284b5 100644
--- a/src/platform/posix/posix_atomic.c
+++ b/src/platform/posix/posix_atomic.c
@@ -30,13 +30,13 @@ nni_atomic_flag_reset(nni_atomic_flag *f)
 }
 
 void
-nni_atomic_inc64(nni_atomic_u64 *v, uint64_t bump)
+nni_atomic_add64(nni_atomic_u64 *v, uint64_t bump)
 {
 	(void) atomic_fetch_add_explicit(&v->v, bump, memory_order_relaxed);
 }
 
 void
-nni_atomic_dec64(nni_atomic_u64 *v, uint64_t bump)
+nni_atomic_sub64(nni_atomic_u64 *v, uint64_t bump)
 {
 	(void) atomic_fetch_sub_explicit(&v->v, bump, memory_order_relaxed);
 }
@@ -65,6 +65,22 @@ nni_atomic_init64(nni_atomic_u64 *v)
 	atomic_init(&v->v, 0);
 }
 
+void
+nni_atomic_inc64(nni_atomic_u64 *v)
+{
+	atomic_fetch_add(&v->v, 1);
+}
+
+uint64_t
+nni_atomic_dec64_nv(nni_atomic_u64 *v)
+{
+	uint64_t ov;
+
+	// C11 atomics give the old rather than new value.
+	ov = atomic_fetch_sub(&v->v, 1);
+	return (ov - 1);
+}
+
 #else
 
 #include <pthread.h>
@@ -91,7 +107,7 @@ nni_atomic_flag_reset(nni_atomic_flag *f)
 }
 
 void
-nni_atomic_inc64(nni_atomic_u64 *v, uint64_t bump)
+nni_atomic_add64(nni_atomic_u64 *v, uint64_t bump)
 {
 	pthread_mutex_lock(&plat_atomic_lock);
 	v += bump;
@@ -99,7 +115,7 @@ nni_atomic_inc64(nni_atomic_u64 *v, uint64_t bump)
 }
 
 void
-nni_atomic_dec64(nni_atomic_u64 *v, uint64_t bump)
+nni_atomic_sub64(nni_atomic_u64 *v, uint64_t bump)
 {
 	pthread_mutex_lock(&plat_atomic_lock);
 	v -= bump;
@@ -141,6 +157,25 @@ nni_atomic_init64(nni_atomic_u64 *v)
 	v->v = 0;
 }
 
+void
+nni_atomic_inc64(nni_atomic_u64 *v)
+{
+	pthread_mutex_lock(&plat_atomic_lock);
+	v->v++;
+	pthread_mutex_unlock(&plat_atomic_lock);
+}
+
+void
+nni_atomic_dec64_nv(nni_atomic_u64 *v)
+{
+	uint64_t nv;
+	pthread_mutex_lock(&plat_atomic_lock);
+	v->v--;
+	nv = v->v;
+	pthread_mutex_unlock(&plat_atomic_lock);
+	return (nv);
+}
+
 #endif
 
 #endif // NNG_PLATFORM_POSIX
diff --git a/src/platform/windows/win_thread.c b/src/platform/windows/win_thread.c
index 158b8133..146520e9 100644
--- a/src/platform/windows/win_thread.c
+++ b/src/platform/windows/win_thread.c
@@ -128,13 +128,13 @@ nni_atomic_flag_reset(nni_atomic_flag *f)
 }
 
 void
-nni_atomic_inc64(nni_atomic_u64 *v, uint64_t bump)
+nni_atomic_add64(nni_atomic_u64 *v, uint64_t bump)
 {
 	InterlockedAddNoFence64(&v->v, (LONGLONG) bump);
 }
 
 void
-nni_atomic_dec64(nni_atomic_u64 *v, uint64_t bump)
+nni_atomic_sub64(nni_atomic_u64 *v, uint64_t bump)
 {
 	// Windows lacks a sub, so we add the negative.
 	InterlockedAddNoFence64(&v->v, (0ll - (LONGLONG) bump));
@@ -165,6 +165,18 @@ nni_atomic_init64(nni_atomic_u64 *v)
 	InterlockedExchange64(&v->v, 0);
 }
 
+void
+nni_atomic_inc64(nni_atomic_u64 *v)
+{
+	(void) InterlockedIncrementAcquire64(&v->v);
+}
+
+uint64_t
+nni_atomic_dec64_nv(nni_atomic_u64 *v)
+{
+	return ((uint64_t)(InterlockedDecrementRelease64(&v->v)));
+}
+
 static unsigned int __stdcall nni_plat_thr_main(void *arg)
 {
 	nni_plat_thr *thr = arg;
diff --git a/src/supplemental/tls/mbedtls/tls.c b/src/supplemental/tls/mbedtls/tls.c
index d3816747..b7ed0575 100644
--- a/src/supplemental/tls/mbedtls/tls.c
+++ b/src/supplemental/tls/mbedtls/tls.c
@@ -95,7 +95,7 @@ struct nng_tls_config {
 	mbedtls_x509_crt ca_certs;
 	mbedtls_x509_crl crl;
 
-	int refcnt; // servers increment the reference
+	nni_atomic_u64 refcnt;
 
 	nni_list certkeys;
 };
@@ -159,13 +159,9 @@ nni_tls_config_fini(nng_tls_config *cfg)
 	nni_tls_certkey *ck;
 
 	if (cfg != NULL) {
-		nni_mtx_lock(&cfg->lk);
-		cfg->refcnt--;
-		if (cfg->refcnt != 0) {
-			nni_mtx_unlock(&cfg->lk);
+		if (nni_atomic_dec64_nv(&cfg->refcnt) != 0) {
 			return;
 		}
-		nni_mtx_unlock(&cfg->lk);
 
 		mbedtls_ssl_config_free(&cfg->cfg_ctx);
 #ifdef NNG_TLS_USE_CTR_DRBG
@@ -199,7 +195,8 @@ nni_tls_config_init(nng_tls_config **cpp, enum nng_tls_mode mode)
 	if ((cfg = NNI_ALLOC_STRUCT(cfg)) == NULL) {
 		return (NNG_ENOMEM);
 	}
-	cfg->refcnt = 1;
+	nni_atomic_init64(&cfg->refcnt);
+	nni_atomic_inc64(&cfg->refcnt);
 	nni_mtx_init(&cfg->lk);
 	if (mode == NNG_TLS_MODE_SERVER) {
 		sslmode  = MBEDTLS_SSL_IS_SERVER;
@@ -247,9 +244,7 @@ nni_tls_config_init(nng_tls_config **cpp, enum nng_tls_mode mode)
 void
 nni_tls_config_hold(nng_tls_config *cfg)
 {
-	nni_mtx_lock(&cfg->lk);
-	cfg->refcnt++;
-	nni_mtx_unlock(&cfg->lk);
+	nni_atomic_inc64(&cfg->refcnt);
 }
 
 // tls_mkerr converts an mbed error to an NNG error.  In all cases
@@ -295,15 +290,15 @@ tls_free(void *arg)
 		}
 		nni_aio_stop(tls->tcp_send);
 		nni_aio_stop(tls->tcp_recv);
-
 		nni_aio_fini(tls->com.aio);
-		nng_tls_config_free(tls->com.cfg);
 
 		// And finalize / free everything.
 		nng_stream_free(tls->tcp);
 		nni_aio_fini(tls->tcp_send);
 		nni_aio_fini(tls->tcp_recv);
 		mbedtls_ssl_free(&tls->ctx);
+		nng_tls_config_free(tls->com.cfg);
+
 		if (tls->recvbuf != NULL) {
 			nni_free(tls->recvbuf, NNG_TLS_MAX_RECV_SIZE);
 		}
@@ -311,6 +306,7 @@ tls_free(void *arg)
 			nni_free(tls->sendbuf, NNG_TLS_MAX_RECV_SIZE);
 		}
 		nni_mtx_fini(&tls->lk);
+		memset(tls, 0xff, sizeof(*tls));
 		NNI_FREE_STRUCT(tls);
 	}
 }
diff --git a/src/supplemental/tls/tls_common.c b/src/supplemental/tls/tls_common.c
index 990d3add..97bcce26 100644
--- a/src/supplemental/tls/tls_common.c
+++ b/src/supplemental/tls/tls_common.c
@@ -160,14 +160,14 @@ tls_dialer_set_config(void *arg, const void *buf, size_t sz, nni_type t)
 	if (cfg == NULL) {
 		return (NNG_EINVAL);
 	}
-	nni_mtx_lock(&d->lk);
-	old = d->cfg;
 	nng_tls_config_hold(cfg);
+
+	nni_mtx_lock(&d->lk);
+	old    = d->cfg;
 	d->cfg = cfg;
 	nni_mtx_unlock(&d->lk);
-	if (old != NULL) {
-		nng_tls_config_free(old);
-	}
+
+	nng_tls_config_free(old);
 	return (0);
 }
 
@@ -432,14 +432,15 @@ tls_listener_set_config(void *arg, const void *buf, size_t sz, nni_type t)
 		return (NNG_EINVAL);
 	}
 
-	nni_mtx_lock(&l->lk);
-	old = l->cfg;
 	nng_tls_config_hold(cfg);
+
+	nni_mtx_lock(&l->lk);
+	old    = l->cfg;
 	l->cfg = cfg;
 	nni_mtx_unlock(&l->lk);
-	if (old != NULL) {
-		nng_tls_config_free(old);
-	}
+
+	nng_tls_config_free(old);
+
 	return (0);
 }