From c5c1d4557bf63d9a9343ea76a1ae3f1be5d36639 Mon Sep 17 00:00:00 2001
From: rdb <git@rdb.name>
Date: Tue, 10 Mar 2020 12:01:51 +0100
Subject: [PATCH] pgraph: fix double free if weak ptr to state is locked while
 being gc'ed

Fixes #499
---
 panda/src/express/referenceCount.I  | 15 +++++++++++++++
 panda/src/express/referenceCount.h  |  1 +
 panda/src/pgraph/renderAttrib.cxx   |  7 ++++---
 panda/src/pgraph/renderState.cxx    |  8 +++++---
 panda/src/pgraph/transformState.cxx |  7 ++++---
 5 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/panda/src/express/referenceCount.I b/panda/src/express/referenceCount.I
index bb4ac1fc18..6b5847407c 100644
--- a/panda/src/express/referenceCount.I
+++ b/panda/src/express/referenceCount.I
@@ -317,6 +317,21 @@ ref_if_nonzero() const {
   return true;
 }
 
+/**
+ * Atomically decreases the reference count of this object if it is one.
+ * Do not use this.  This exists only to implement a special case with the
+ * state cache.
+ * @return true if the reference count was decremented to zero.
+ */
+INLINE bool ReferenceCount::
+unref_if_one() const {
+#ifdef _DEBUG
+  nassertr(test_ref_count_integrity(), 0);
+  nassertr(_ref_count > 0, 0);
+#endif
+  return (AtomicAdjust::compare_and_exchange(_ref_count, 1, 0) == 1);
+}
+
 /**
  * This global helper function will unref the given ReferenceCount object, and
  * if the reference count reaches zero, automatically delete it.  It can't be
diff --git a/panda/src/express/referenceCount.h b/panda/src/express/referenceCount.h
index 46f0231117..09ef7b4501 100644
--- a/panda/src/express/referenceCount.h
+++ b/panda/src/express/referenceCount.h
@@ -64,6 +64,7 @@ public:
   INLINE void weak_unref();
 
   INLINE bool ref_if_nonzero() const;
+  INLINE bool unref_if_one() const;
 
 protected:
   bool do_test_ref_count_integrity() const;
diff --git a/panda/src/pgraph/renderAttrib.cxx b/panda/src/pgraph/renderAttrib.cxx
index 9bac4e666a..9b5bbd6303 100644
--- a/panda/src/pgraph/renderAttrib.cxx
+++ b/panda/src/pgraph/renderAttrib.cxx
@@ -215,14 +215,15 @@ garbage_collect() {
 
   do {
     RenderAttrib *attrib = (RenderAttrib *)_attribs->get_key(si);
-    if (attrib->get_ref_count() == 1) {
+    if (attrib->unref_if_one()) {
       // This attrib has recently been unreffed to 1 (the one we added when
       // we stored it in the cache).  Now it's time to delete it.  This is
       // safe, because we're holding the _attribs_lock, so it's not possible
       // for some other thread to find the attrib in the cache and ref it
-      // while we're doing this.
+      // while we're doing this.  Also, we've just made sure to unref it to 0,
+      // to ensure that another thread can't get it via a weak pointer.
       attrib->release_new();
-      unref_delete(attrib);
+      delete attrib;
 
       // When we removed it from the hash map, it swapped the last element
       // with the one we just removed.  So the current index contains one we
diff --git a/panda/src/pgraph/renderState.cxx b/panda/src/pgraph/renderState.cxx
index 5ce1527247..5098f9dc34 100644
--- a/panda/src/pgraph/renderState.cxx
+++ b/panda/src/pgraph/renderState.cxx
@@ -934,15 +934,17 @@ garbage_collect() {
       }
     }
 
-    if (state->get_ref_count() == 1) {
+    if (state->unref_if_one()) {
       // This state has recently been unreffed to 1 (the one we added when
       // we stored it in the cache).  Now it's time to delete it.  This is
       // safe, because we're holding the _states_lock, so it's not possible
       // for some other thread to find the state in the cache and ref it
-      // while we're doing this.
+      // while we're doing this.  Also, we've just made sure to unref it to 0,
+      // to ensure that another thread can't get it via a weak pointer.
+
       state->release_new();
       state->remove_cache_pointers();
-      state->cache_unref();
+      state->cache_unref_only();
       delete state;
 
       // When we removed it from the hash map, it swapped the last element
diff --git a/panda/src/pgraph/transformState.cxx b/panda/src/pgraph/transformState.cxx
index fc8b3de606..220c9c5d8b 100644
--- a/panda/src/pgraph/transformState.cxx
+++ b/panda/src/pgraph/transformState.cxx
@@ -1204,15 +1204,16 @@ garbage_collect() {
       }
     }
 
-    if (state->get_ref_count() == 1) {
+    if (state->unref_if_one()) {
       // This state has recently been unreffed to 1 (the one we added when
       // we stored it in the cache).  Now it's time to delete it.  This is
       // safe, because we're holding the _states_lock, so it's not possible
       // for some other thread to find the state in the cache and ref it
-      // while we're doing this.
+      // while we're doing this.  Also, we've just made sure to unref it to 0,
+      // to ensure that another thread can't get it via a weak pointer.
       state->release_new();
       state->remove_cache_pointers();
-      state->cache_unref();
+      state->cache_unref_only();
       delete state;
 
       // When we removed it from the hash map, it swapped the last element