Merge pull request #27 from steffahn/fix_memory_leaks

Fix memory leaks on panicking destructors

src/lib.rs

@@ -362,7 +362,7 @@ macro_rules! self_cell {
                 owner_ptr.write(owner);
 
                 // Drop guard that cleans up should building the dependent panic.
-                let mut drop_guard =
+                let drop_guard =
                     $crate::unsafe_self_cell::OwnerAndCellDropGuard::new(joined_ptr);
 
                 // Initialize dependent with owner reference in final place.

src/unsafe_self_cell.rs

@@ -1,5 +1,5 @@
 use core::marker::PhantomData;
-use core::mem::transmute;
+use core::mem::{self, transmute};
 use core::ptr::{drop_in_place, read, NonNull};
 
 extern crate alloc;
@@ -82,24 +82,31 @@ impl<ContainedIn, Owner, DependentStatic> UnsafeSelfCell<ContainedIn, Owner, Dep
         let joined_ptr =
             transmute::<NonNull<u8>, NonNull<JoinedCell<Owner, Dependent>>>(self.joined_void_ptr);
 
+        // Also used in case drop_in_place(...dependent) fails
+        let _guard = OwnerAndCellDropGuard { joined_ptr };
+
         // IMPORTANT dependent must be dropped before owner.
         // We don't want to rely on an implicit order of struct fields.
         // So we drop the struct, field by field manually.
         drop_in_place(&mut (*joined_ptr.as_ptr()).dependent);
-        drop_in_place(&mut (*joined_ptr.as_ptr()).owner);
-
-        let layout = Layout::new::<JoinedCell<Owner, Dependent>>();
 
-        dealloc(self.joined_void_ptr.as_ptr(), layout);
+        // Dropping owner
+        // and deallocating
+        // due to _guard at end of scope.
     }
 
     pub unsafe fn into_owner<Dependent>(self) -> Owner {
         let joined_ptr =
             transmute::<NonNull<u8>, NonNull<JoinedCell<Owner, Dependent>>>(self.joined_void_ptr);
 
+        // In case drop_in_place(...dependent) fails
+        let drop_guard = OwnerAndCellDropGuard::new(joined_ptr);
+
         // Drop dependent
         drop_in_place(&mut (*joined_ptr.as_ptr()).dependent);
 
+        mem::forget(drop_guard);
+
         let owner_ptr: *const Owner = &(*joined_ptr.as_ptr()).owner;
 
         // Move owner out so it can be returned.
@@ -149,16 +156,31 @@ impl<Owner, Dependent> OwnerAndCellDropGuard<Owner, Dependent> {
 
 impl<Owner, Dependent> Drop for OwnerAndCellDropGuard<Owner, Dependent> {
     fn drop(&mut self) {
+        struct DeallocGuard {
+            ptr: *mut u8,
+            layout: Layout,
+        }
+        impl Drop for DeallocGuard {
+            fn drop(&mut self) {
+                unsafe { dealloc(self.ptr, self.layout) }
+            }
+        }
+
+        // Deallocate even when the drop_in_place(...owner) panics
+        let _guard = DeallocGuard {
+            ptr: unsafe {
+                transmute::<*mut JoinedCell<Owner, Dependent>, *mut u8>(self.joined_ptr.as_ptr())
+            },
+            layout: Layout::new::<JoinedCell<Owner, Dependent>>(),
+        };
+
         unsafe {
             // We must only drop owner and the struct itself,
             // The whole point of this drop guard is to clean up the partially
             // initialized struct should building the dependent fail.
             drop_in_place(&mut (*self.joined_ptr.as_ptr()).owner);
-
-            let layout = Layout::new::<JoinedCell<Owner, Dependent>>();
-            let joined_void_ptr =
-                transmute::<*mut JoinedCell<Owner, Dependent>, *mut u8>(self.joined_ptr.as_ptr());
-            dealloc(joined_void_ptr, layout);
         }
+
+        // Deallocation happens at end of scope
     }
 }

tests/self_cell.rs

@@ -388,7 +388,7 @@ fn drop_order() {
     }
 
     self_cell! {
-        struct Foo {
+        struct DropOrder {
             owner: Owner,
 
             #[covariant]
@@ -397,13 +397,13 @@ fn drop_order() {
     }
 
     let drops: Rc<RefCell<Vec<Dropped>>> = <_>::default();
-    let foo = Foo::new(Owner(drops.clone()), |o| Dependent(o, drops.clone()));
-    drop(foo);
+    let cell = DropOrder::new(Owner(drops.clone()), |o| Dependent(o, drops.clone()));
+    drop(cell);
     assert_eq!(&drops.borrow()[..], &[Dropped::Dependent, Dropped::Owner]);
 }
 
 #[test]
-fn into_owner_drop_dependent_panic() {
+fn into_owner_drop_dependent_without_panic() {
     // This test resulted in a double-free in a previous version of self-cell
     type O = Cell<Option<Box<u8>>>;
 
@@ -428,6 +428,96 @@ fn into_owner_drop_dependent_panic() {
     assert!(s.into_owner().into_inner().is_none());
 }
 
+#[test]
+#[should_panic] // but should not leak or double-free
+fn into_owner_drop_dependent_with_panic() {
+    type O = Cell<Option<Box<u8>>>;
+
+    self_cell! {
+        struct S {
+            owner: O,
+
+            #[covariant]
+            dependent: D,
+        }
+    }
+
+    struct D<'a>(&'a O);
+
+    impl Drop for D<'_> {
+        fn drop(&mut self) {
+            self.0.take();
+            panic!();
+        }
+    }
+
+    let s = S::new(Cell::new(Some(Box::new(42))), |o| D(o));
+    s.into_owner();
+}
+
+#[test]
+fn drop_panic_owner() {
+    #[derive(Clone, Debug, PartialEq)]
+    struct Owner(String);
+
+    type Dependent<'a> = &'a Owner;
+
+    self_cell! {
+        struct DropPanicOwner {
+            owner: Owner,
+
+            #[covariant]
+            dependent: Dependent,
+        }
+    }
+
+    impl Drop for Owner {
+        fn drop(&mut self) {
+            panic!()
+        }
+    }
+
+    let owner = Owner("s t e f f a h n <3 and some padding against sbo".into());
+
+    let cell = DropPanicOwner::new(owner.clone(), |o| o);
+    assert_eq!(cell.borrow_owner(), &owner);
+    assert_eq!(cell.borrow_dependent(), &&owner);
+
+    assert!(std::panic::catch_unwind(move || drop(cell)).is_err());
+    assert!(std::panic::catch_unwind(move || drop(owner)).is_err());
+}
+
+#[test]
+fn drop_panic_dependent() {
+    #[derive(Clone, Debug, PartialEq)]
+    struct Owner(String);
+
+    struct Dependent<'a>(&'a Owner);
+
+    self_cell! {
+        struct DropPanicDependent {
+            owner: Owner,
+
+            #[covariant]
+            dependent: Dependent,
+        }
+    }
+
+    impl Drop for Dependent<'_> {
+        fn drop(&mut self) {
+            panic!()
+        }
+    }
+
+    let owner = Owner("s t e f f a h n <3".into());
+
+    let cell = DropPanicDependent::new(owner.clone(), |o| Dependent(o));
+    assert_eq!(cell.borrow_owner(), &owner);
+    assert_eq!(cell.borrow_dependent().0, &owner);
+
+    assert!(std::panic::catch_unwind(move || drop(cell)).is_err());
+}
+
 #[test]
 fn dependent_mutate() {
     let mut ast_cell = PackedAstCell::new("Egal in welchen Farben ihr den ..".into(), |owner| {