Fix panic-safety in AlignedBox::realloc (double-free on panicking element Drop)

src/lib.rs

@@ -274,22 +274,85 @@ impl<T> AlignedBox<[T]> {
         let new_layout = alloc::alloc::Layout::from_size_align(memsize, self.layout.align())
             .map_err(|_| AlignedBoxError::InvalidAlign)?;
 
-        // SAFETY:
-        // * self.container is not used afterwards but re-assigned with a new
-        //   instance of std::mem::ManuallyDrop<alloc::boxed::Box> in all possible cases.
         let b = unsafe { std::mem::ManuallyDrop::take(&mut self.container) };
         let ptr = alloc::boxed::Box::into_raw(b);
+        let elem_ptr = ptr as *mut T;
+
+        // Drop the tail elements, back to front, under a guard. Each element's
+        // Drop is user-controlled and may panic. If it does, `self.container`
+        // still holds the old pointer while ownership has moved to `ptr`, so
+        // `AlignedBox`'s own `Drop` would free those elements a second time.
+        // The guard restores a consistent state on unwind, so every element is
+        // freed exactly once.
+        //
+        // Restoring means shrinking the allocation to the still-live prefix,
+        // not just the box: `container.len() * size_of::<T>()` has to keep
+        // matching `layout.size()`, otherwise the later `dealloc` gets either
+        // the wrong size or a pointer whose provenance is too narrow. If that
+        // realloc fails the original allocation is untouched, so the dropped
+        // slots are reinitialized and the full-length box is restored, the same
+        // way the realloc-failure path below does it.
+        struct ShrinkGuard<'a, T> {
+            container: *mut std::mem::ManuallyDrop<alloc::boxed::Box<[T]>>,
+            layout: *mut alloc::alloc::Layout,
+            elem_ptr: *mut T,
+            valid: usize,
+            old_nelems: usize,
+            initializer: &'a dyn Fn(*mut T),
+        }
+        impl<T> Drop for ShrinkGuard<'_, T> {
+            fn drop(&mut self) {
+                // SAFETY:
+                // * elem_ptr was allocated by the global allocator with *self.layout
+                // * valid >= nelems >= 1, so memsize is never zero
+                // * the caller must ensure initializer does not read or drop the
+                //   old value behind the pointer
+                unsafe {
+                    let old_layout = *self.layout;
+                    let memsize = std::mem::size_of::<T>() * self.valid;
+                    let new_ptr =
+                        alloc::alloc::realloc(self.elem_ptr as *mut u8, old_layout, memsize)
+                            as *mut T;
+
+                    if new_ptr.is_null() {
+                        for i in self.valid..self.old_nelems {
+                            (self.initializer)(self.elem_ptr.add(i));
+                        }
+                        let slice = std::slice::from_raw_parts_mut(self.elem_ptr, self.old_nelems);
+                        *self.container =
+                            std::mem::ManuallyDrop::new(alloc::boxed::Box::from_raw(slice));
+                    } else {
+                        // Shrinking a valid layout keeps it valid.
+                        let layout = alloc::alloc::Layout::from_size_align_unchecked(
+                            memsize,
+                            old_layout.align(),
+                        );
+                        let slice = std::slice::from_raw_parts_mut(new_ptr, self.valid);
+                        *self.container =
+                            std::mem::ManuallyDrop::new(alloc::boxed::Box::from_raw(slice));
+                        *self.layout = layout;
+                    }
+                }
+            }
+        }
 
-        // Drop any values that will be deallocated by realloc
-        for i in nelems..old_nelems {
-            // SAFETY:
-            // * (*ptr)[i] is valid for R/W, properly aligned and valid to drop
-            // * the element will not be read as it will either be re-initialized using
-            //   std::ptr::write or the slice behind ptr will not be used anymore.
+        let mut guard = ShrinkGuard::<T> {
+            container: &mut self.container as *mut _,
+            layout: &mut self.layout as *mut _,
+            elem_ptr,
+            valid: old_nelems,
+            old_nelems,
+            initializer: &initializer,
+        };
+        for i in (nelems..old_nelems).rev() {
+            // Set before the drop: element `i` is the one that may unwind, so
+            // it stays out of the restored prefix and is never dropped twice.
+            guard.valid = i;
             unsafe {
-                std::ptr::drop_in_place(&mut (*ptr)[i]);
+                std::ptr::drop_in_place(elem_ptr.add(i));
             }
         }
+        std::mem::forget(guard);
 
         // SAFETY:
         // * ptr has been assigned by the same (that is: global) allocator
@@ -810,4 +873,46 @@ mod tests {
         println!("100"); // Stack of this function
         std::thread::sleep(std::time::Duration::from_secs(DELAY));
     }
+
+    #[test]
+    fn realloc_shrink_panicking_drop_is_sound() {
+        use std::panic::{AssertUnwindSafe, catch_unwind};
+        use std::sync::atomic::{AtomicBool, Ordering};
+
+        let _m = SEQ_TEST_MUTEX.read().unwrap();
+
+        static FIRED: AtomicBool = AtomicBool::new(false);
+
+        #[derive(Clone)]
+        struct PanicOnDrop {
+            heap: String,
+        }
+        impl Default for PanicOnDrop {
+            fn default() -> Self {
+                PanicOnDrop {
+                    heap: "x".repeat(64),
+                }
+            }
+        }
+        impl Drop for PanicOnDrop {
+            fn drop(&mut self) {
+                let _ = self.heap.len();
+                if !FIRED.swap(true, Ordering::SeqCst) {
+                    panic!("drop panic");
+                }
+            }
+        }
+
+        // Shrink the box; realloc drops the tail elements. A panic in an
+        // element's Drop leaves self.container pointing at the half-processed
+        // buffer, and AlignedBox's Drop then frees it again -- a double-free.
+        let mut b: AlignedBox<[PanicOnDrop]> = AlignedBox::slice_from_default(64, 8).unwrap();
+
+        let result = catch_unwind(AssertUnwindSafe(|| {
+            let _ = b.realloc_with_default(2);
+        }));
+        assert!(result.is_err(), "expected an element Drop to unwind");
+
+        drop(b);
+    }
 }