Fix multiple logic bugs in Arc::make_mut

library/alloc/src/rc.rs

@@ -2090,8 +2090,6 @@ impl<T: ?Sized + CloneToUninit, A: Allocator + Clone> Rc<T, A> {
         } else if Rc::weak_count(this) != 0 {
             // Can just steal the data, all that's left is Weaks
 
-            // We don't need panic-protection like the above branch does, but we might as well
-            // use the same mechanism.
             let mut in_progress: UniqueRcUninit<T, A> =
                 UniqueRcUninit::new(&**this, this.alloc.clone());
             unsafe {
@@ -2104,10 +2102,18 @@ impl<T: ?Sized + CloneToUninit, A: Allocator + Clone> Rc<T, A> {
                     size_of_val,
                 );
 
+                // This leaves us with 0 strong refs, so the data has
+                // effectively been moved to the new rc.
                 this.inner().dec_strong();
+
                 // Remove implicit strong-weak ref (no need to craft a fake
                 // Weak here -- we know other Weaks can clean up for us)
                 this.inner().dec_weak();
+
+                // Last chance to not accidentally forget the allocator.
+                // Only drop at the end of the scope to avoid panics.
+                let _alloc = ptr::read(&this.alloc);
+
                 // Replace `this` with newly constructed Rc that has the moved data.
                 ptr::write(this, in_progress.into_rc());
             }

library/alloc/src/sync.rs

@@ -2539,16 +2539,27 @@ impl<T: ?Sized + CloneToUninit, A: Allocator + Clone> Arc<T, A> {
             // usize::MAX (i.e., locked), since the weak count can only be
             // locked by a thread with a strong reference.
 
-            // Materialize our own implicit weak pointer, so that it can clean
-            // up the ArcInner as needed.
-            let _weak = Weak { ptr: this.ptr, alloc: this.alloc.clone() };
+            // Guard against panics while using the allocator.
+            // If we unwind before the Arc is overwritten, we expose a strong
+            // count of 0, resulting in a UAF (#155746, #157203).
+            // Until the new Arc is written, the old Arc must remain valid
+            struct Guard<'a, T: ?Sized> {
+                inner: &'a ArcInner<T>,
+            }
+            impl<'a, T: ?Sized> Drop for Guard<'a, T> {
+                fn drop(&mut self) {
+                    self.inner.strong.store(1, Release);
+                }
+            }
+            let guard = Guard { inner: this.inner() };
 
             // Can just steal the data, all that's left is Weaks
-            //
-            // We don't need panic-protection like the above branch does, but we might as well
-            // use the same mechanism.
+            // Note that this can panic in two ways:
+            // - The allocation can fail
+            // - The allocator clone can fail
             let mut in_progress: UniqueArcUninit<T, A> =
                 UniqueArcUninit::new(&**this, this.alloc.clone());
+
             unsafe {
                 // Initialize `in_progress` with move of **this.
                 // We have to express this in terms of bytes because `T: ?Sized`; there is no
@@ -2559,6 +2570,15 @@ impl<T: ?Sized + CloneToUninit, A: Allocator + Clone> Arc<T, A> {
                     size_of_val,
                 );
 
+                // We are now safe from panics.
+                mem::forget(guard);
+
+                // Materialize our own implicit weak pointer, so that it can clean
+                // up the ArcInner as needed.
+                // Make sure the allocator is not leaked when the Arc is overwritten.
+                // Only drop at the end of the scope to avoid panics.
+                let _weak = Weak { ptr: this.ptr, alloc: ptr::read(&this.alloc) };
+
                 ptr::write(this, in_progress.into_arc());
             }
         } else {

library/alloctests/tests/arc.rs

@@ -328,3 +328,65 @@ mod pin_coerce_unsized {
         arg
     }
 }
+
+/// Test that `Arc::make_mut` does not forget an allocator when it steals the data.
+#[test]
+fn issue_158875_make_mut_dont_leak_allocator() {
+    use std::alloc::Global;
+
+    let alloc = Rc::new(Global);
+
+    {
+        let mut arc = Arc::new_in(123, alloc.clone());
+        let weak = Arc::downgrade(&arc); // create a weak so make_mut steals the data
+        _ = Arc::make_mut(&mut arc);
+        assert_eq!(weak.upgrade(), None);
+    }
+
+    assert_eq!(Rc::strong_count(&alloc), 1); // if this is >1, we have a memory leak!
+}
+
+/// Test that `Arc::make_mut` does not cause a UAF if the allocator panics on
+/// clone when it steals the data.
+#[test]
+#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
+fn issue_155746_make_mut_panic_safety() {
+    use std::alloc::{Allocator, System};
+    use std::panic::AssertUnwindSafe;
+
+    #[derive(Default)]
+    struct PanickingCloneAlloc {
+        do_panic: Rc<Cell<bool>>,
+    }
+    unsafe impl Allocator for PanickingCloneAlloc {
+        fn allocate(
+            &self,
+            layout: std::alloc::Layout,
+        ) -> Result<std::ptr::NonNull<[u8]>, std::alloc::AllocError> {
+            System.allocate(layout)
+        }
+
+        unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) {
+            unsafe { System.deallocate(ptr, layout) }
+        }
+    }
+    impl Clone for PanickingCloneAlloc {
+        fn clone(&self) -> Self {
+            if self.do_panic.get() { panic!() } else { Self { do_panic: self.do_panic.clone() } }
+        }
+    }
+
+    let alloc = PanickingCloneAlloc::default();
+    let mut arc = Arc::new_in(vec![vec![1]], alloc.clone());
+
+    let _weak = Arc::downgrade(&arc); // create a weak so make_mut steals the data
+
+    alloc.do_panic.set(true);
+    std::panic::catch_unwind(AssertUnwindSafe(|| {
+        Arc::make_mut(&mut arc);
+    }))
+    .unwrap_err();
+
+    assert_eq!(*arc, [[1]]);
+    assert_eq!(Arc::strong_count(&arc), 1); // if this is 0, we have a UAF!
+}

library/alloctests/tests/rc.rs

@@ -950,3 +950,20 @@ fn test_unique_rc_unsizing_coercion() {
     let rc: Rc<[u8]> = UniqueRc::into_rc(rc);
     assert_eq!(*rc, [123, 0, 0]);
 }
+
+/// Test that `Rc::make_mut` does not forget an allocator when it steals the data.
+#[test]
+fn issue_158875_make_mut_dont_leak_allocator() {
+    use std::alloc::Global;
+
+    let alloc = Rc::new(Global);
+
+    {
+        let mut arc = Rc::new_in(123, alloc.clone());
+        let weak = Rc::downgrade(&arc); // create a weak so make_mut steals the data
+        _ = Rc::make_mut(&mut arc);
+        assert_eq!(weak.upgrade(), None);
+    }
+
+    assert_eq!(Rc::strong_count(&alloc), 1); // if this is >1, we have a memory leak!
+}