fix #946 Use after free in TLS

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

2 files changed, 53 insertions, 6 deletions

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;