From 01d7a5bf47e5bc228ce637ee4dd7f16a8c55d0c4 Mon Sep 17 00:00:00 2001
From: Dipin Hora <dipinhora+github@gmail.com>
Date: Thu, 21 Nov 2024 22:19:39 -0500
Subject: [PATCH] Fix mpmcq use-after-free bug

according to address sanitizer, prior to this commit, the mpmcq node
free could cause a use-after-free error from another thread that
might still be reading from `tail->next`. This was technically safe
in this context since the logic will end up looping due to the CAS
failing in the while condition due to the aba counter being out of
sync.

This commit defers freeing of the nodes and instead recycles them
for future pushes onto the mpmcq from the same thread to ensure
that there are no use-after-free concerns any longer.
---
 src/libponyrt/sched/mpmcq.c     | 35 +++++++++++++++++++++++++++++++--
 src/libponyrt/sched/mpmcq.h     |  2 ++
 src/libponyrt/sched/scheduler.c |  2 ++
 3 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/src/libponyrt/sched/mpmcq.c b/src/libponyrt/sched/mpmcq.c
index 884d988fc1..7f747a389b 100644
--- a/src/libponyrt/sched/mpmcq.c
+++ b/src/libponyrt/sched/mpmcq.c
@@ -1,6 +1,7 @@
 #define PONY_WANT_ATOMIC_DEFS
 
 #include "mpmcq.h"
+#include "ponyassert.h"
 #include "../mem/pool.h"
 #include "../sched/cpu.h"
 
@@ -16,9 +17,19 @@ struct mpmcq_node_t
   PONY_ATOMIC(void*) data;
 };
 
+static __pony_thread_local mpmcq_node_t* nodes_to_recycle;
+
 static mpmcq_node_t* node_alloc(void* data)
 {
-  mpmcq_node_t* node = POOL_ALLOC(mpmcq_node_t);
+  mpmcq_node_t* node = nodes_to_recycle;
+
+  if (NULL != node)
+  {
+    nodes_to_recycle = atomic_load_explicit(&node->next, memory_order_relaxed);
+  } else {
+    node = POOL_ALLOC(mpmcq_node_t);
+  }
+
   atomic_store_explicit(&node->next, NULL, memory_order_relaxed);
   atomic_store_explicit(&node->data, data, memory_order_relaxed);
 
@@ -55,6 +66,7 @@ static size_t mpmcq_size_debug(mpmcq_t* q)
 
 void ponyint_mpmcq_init(mpmcq_t* q)
 {
+  nodes_to_recycle = NULL;
   mpmcq_node_t* node = node_alloc(NULL);
 
   atomic_store_explicit(&q->head, node, memory_order_relaxed);
@@ -66,8 +78,20 @@ void ponyint_mpmcq_init(mpmcq_t* q)
 #endif
 }
 
+void ponyint_mpmcq_cleanup()
+{
+  while(nodes_to_recycle != NULL)
+  {
+    mpmcq_node_t* node = nodes_to_recycle;
+    nodes_to_recycle = atomic_load_explicit(&node->next, memory_order_relaxed);
+    node_free(node);
+  }
+}
+
 void ponyint_mpmcq_destroy(mpmcq_t* q)
 {
+  pony_assert(atomic_load_explicit(&q->head, memory_order_relaxed) == q->tail.object);
+
   atomic_store_explicit(&q->head, NULL, memory_order_relaxed);
   node_free(q->tail.object);
   q->tail.object = NULL;
@@ -177,7 +201,14 @@ void* ponyint_mpmcq_pop(mpmcq_t* q)
   ANNOTATE_HAPPENS_AFTER(&tail->data);
 #endif
 
-  node_free(tail);
+  // if tail has already been consumed and freed by this thread, the read of
+  // `tail->next` above in another thread would be a use-after-free error but is
+  // technically safe in this context since we will end up looping due to the CAS
+  // failing in the while condition due to the aba counter being out of sync
+  // we cannot free the tail node if we want to ensure there are no use-after-free
+  // concerns so we instead add it to the list of nodes to recycle
+  atomic_store_explicit(&tail->next, nodes_to_recycle, memory_order_relaxed);
+  nodes_to_recycle = tail;
 
   return data;
 }
diff --git a/src/libponyrt/sched/mpmcq.h b/src/libponyrt/sched/mpmcq.h
index b08fcb5ab4..71e30ad184 100644
--- a/src/libponyrt/sched/mpmcq.h
+++ b/src/libponyrt/sched/mpmcq.h
@@ -24,6 +24,8 @@ void ponyint_mpmcq_init(mpmcq_t* q);
 
 void ponyint_mpmcq_destroy(mpmcq_t* q);
 
+void ponyint_mpmcq_cleanup();
+
 void ponyint_mpmcq_push(mpmcq_t* q, void* data);
 
 void ponyint_mpmcq_push_single(mpmcq_t* q, void* data);
diff --git a/src/libponyrt/sched/scheduler.c b/src/libponyrt/sched/scheduler.c
index 363a62b782..ef2dc85517 100644
--- a/src/libponyrt/sched/scheduler.c
+++ b/src/libponyrt/sched/scheduler.c
@@ -1125,6 +1125,8 @@ static DECLARE_THREAD_FN(run_thread)
 #endif
 
   run(sched);
+
+  ponyint_mpmcq_cleanup();
   ponyint_pool_thread_cleanup();
 
   return 0;