From 95d48b8d131e77dd04056ad07525c2855095655c Mon Sep 17 00:00:00 2001
From: Autumn <autumnontape@tutanota.com>
Date: Wed, 30 Sep 2020 12:37:16 -0700
Subject: [PATCH 1/3] use drop guard to prevent leaving a dangling reference in
 run_as_context

---
 src/internals/serialize/id.rs | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/internals/serialize/id.rs b/src/internals/serialize/id.rs
index 993901bf..54ff9cff 100644
--- a/src/internals/serialize/id.rs
+++ b/src/internals/serialize/id.rs
@@ -27,34 +27,39 @@ thread_local! {
 
 /// Runs the provided function with this canon as context for Entity (de)serialization.
 pub fn run_as_context<F: FnOnce() -> R, R>(context: &dyn EntitySerializer, f: F) -> R {
+    struct SerializerGuard(Option<&'static dyn EntitySerializer>);
+
+    impl Drop for SerializerGuard {
+        fn drop(&mut self) {
+            // swap context back out of TLS, putting back what we took out.
+            SERIALIZER.with(|cell| {
+                let mut existing = cell.borrow_mut();
+                *existing = self.0.take();
+            });
+        }
+    }
+
     SERIALIZER.with(|cell| {
         // swap context into TLS
-        let prev = {
+        let _serializer_guard = {
             let mut existing = cell.borrow_mut();
             // Safety? Hmm
             // This &'static reference is only stored in the TLS for the duration of this function,
             // meaning it will be valid for all TLS uses as long as it's not copied anywhere else.
             // Can't express this lifetime in a TLS-static, so that's why this unsafe block needs to exist.
             //
-            // If `f` panics, we could end up with a dangling reference in the TLS block.
-            // A catch_unwind block could fix this problem
+            // If `f` unwinds, the SerializerGuard destructor will restore the previous value of
+            // SERIALIZER, so there's no risk of leaving a dangling reference.
             let hacked_context = unsafe {
                 core::mem::transmute::<&dyn EntitySerializer, &'static dyn EntitySerializer>(
                     context,
                 )
             };
-            std::mem::replace(&mut *existing, Some(hacked_context))
+            let prev = std::mem::replace(&mut *existing, Some(hacked_context));
+            SerializerGuard(prev)
         };
 
-        let result = (f)();
-
-        // swap context back out of LTS, putting back what we took out.
-        {
-            let mut existing = cell.borrow_mut();
-            *existing = prev;
-        }
-
-        result
+        (f)()
     })
 }
 

From f3f10e9b376fb10170b712d5f57fa8c9e952887f Mon Sep 17 00:00:00 2001
From: Autumn <autumnontape@tutanota.com>
Date: Wed, 30 Sep 2020 16:23:45 -0700
Subject: [PATCH 2/3] add test for panicking in run_as_context

---
 src/internals/serialize/mod.rs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/src/internals/serialize/mod.rs b/src/internals/serialize/mod.rs
index d4d455f1..df578879 100644
--- a/src/internals/serialize/mod.rs
+++ b/src/internals/serialize/mod.rs
@@ -565,4 +565,19 @@ mod test {
 
         assert_eq!(8, world.len());
     }
+
+    #[test]
+    fn run_as_context_panic() {
+        std::panic::catch_unwind(|| {
+            let registry = Registry::<i32>::default();
+
+            super::WorldSerializer::with_entity_serializer(&registry, &mut |canon| {
+                super::id::run_as_context(canon, || panic!());
+            });
+        })
+        .unwrap_err();
+
+        // run the serialize_bincode test again
+        serialize_bincode();
+    }
 }

From 5919bd2c2e422f27288940a1bbf34f5acd06685d Mon Sep 17 00:00:00 2001
From: Autumn <autumnontape@tutanota.com>
Date: Wed, 30 Sep 2020 18:13:35 -0700
Subject: [PATCH 3/3] avoid calling LocalKey::with unnecessarily in
 SerializerGuard::drop

---
 src/internals/serialize/id.rs | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/internals/serialize/id.rs b/src/internals/serialize/id.rs
index 54ff9cff..44493ae8 100644
--- a/src/internals/serialize/id.rs
+++ b/src/internals/serialize/id.rs
@@ -27,15 +27,16 @@ thread_local! {
 
 /// Runs the provided function with this canon as context for Entity (de)serialization.
 pub fn run_as_context<F: FnOnce() -> R, R>(context: &dyn EntitySerializer, f: F) -> R {
-    struct SerializerGuard(Option<&'static dyn EntitySerializer>);
+    struct SerializerGuard<'a> {
+        cell: &'a RefCell<Option<&'static dyn EntitySerializer>>,
+        prev: Option<&'static dyn EntitySerializer>,
+    }
 
-    impl Drop for SerializerGuard {
+    impl Drop for SerializerGuard<'_> {
         fn drop(&mut self) {
             // swap context back out of TLS, putting back what we took out.
-            SERIALIZER.with(|cell| {
-                let mut existing = cell.borrow_mut();
-                *existing = self.0.take();
-            });
+            let mut existing = self.cell.borrow_mut();
+            *existing = self.prev;
         }
     }
 
@@ -56,7 +57,7 @@ pub fn run_as_context<F: FnOnce() -> R, R>(context: &dyn EntitySerializer, f: F)
                 )
             };
             let prev = std::mem::replace(&mut *existing, Some(hacked_context));
-            SerializerGuard(prev)
+            SerializerGuard { cell, prev }
         };
 
         (f)()