diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs index 785c1c8026a1c..f2dd2444e2392 100644 --- a/library/alloc/src/rc.rs +++ b/library/alloc/src/rc.rs @@ -2090,8 +2090,6 @@ impl Rc { } 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 = UniqueRcUninit::new(&**this, this.alloc.clone()); unsafe { @@ -2104,10 +2102,18 @@ impl Rc { 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()); } diff --git a/library/alloc/src/sync.rs b/library/alloc/src/sync.rs index d0f59c8873300..bb46f14fb9a36 100644 --- a/library/alloc/src/sync.rs +++ b/library/alloc/src/sync.rs @@ -2539,16 +2539,27 @@ impl Arc { // 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, + } + 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 = 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 Arc { 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 { diff --git a/library/alloctests/tests/arc.rs b/library/alloctests/tests/arc.rs index 4b4d1787ace63..0a7f9e82a9c4b 100644 --- a/library/alloctests/tests/arc.rs +++ b/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>, + } + unsafe impl Allocator for PanickingCloneAlloc { + fn allocate( + &self, + layout: std::alloc::Layout, + ) -> Result, std::alloc::AllocError> { + System.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: std::ptr::NonNull, 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! +} diff --git a/library/alloctests/tests/rc.rs b/library/alloctests/tests/rc.rs index 0417206d713ef..698585e94f8a6 100644 --- a/library/alloctests/tests/rc.rs +++ b/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! +}