From fa61678a7d471cfc8bf166979674d4574a5d3bec Mon Sep 17 00:00:00 2001 From: Giacomo Stevanato Date: Sat, 10 Sep 2022 11:26:24 +0200 Subject: [PATCH 1/5] Fix leaking in inplace collection when destructor panics --- library/alloc/src/vec/in_place_collect.rs | 8 ++++++-- library/alloc/src/vec/in_place_drop.rs | 15 +++++++++++++++ library/alloc/src/vec/mod.rs | 2 +- 3 files changed, 22 insertions(+), 3 deletions(-) diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs index 55dcb84ad16f9..359f1a9cec251 100644 --- a/library/alloc/src/vec/in_place_collect.rs +++ b/library/alloc/src/vec/in_place_collect.rs @@ -138,7 +138,7 @@ use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce}; use core::mem::{self, ManuallyDrop}; use core::ptr::{self}; -use super::{InPlaceDrop, SpecFromIter, SpecFromIterNested, Vec}; +use super::{InPlaceDrop, InPlaceDstBufDrop, SpecFromIter, SpecFromIterNested, Vec}; /// Specialization marker for collecting an iterator pipeline into a Vec while reusing the /// source allocation, i.e. executing the pipeline in place. @@ -193,12 +193,16 @@ where // Drop any remaining values at the tail of the source but prevent drop of the allocation // itself once IntoIter goes out of scope. - // If the drop panics then we also leak any elements collected into dst_buf. + // If the drop panics then we also try to drop the destination buffer and its elements. + // This is safe because `forget_allocation_drop_remaining` forgets the allocation *before* + // trying to drop the remaining elements. // // Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce // contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the // module documenttation why this is ok anyway. + let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap }; src.forget_allocation_drop_remaining(); + mem::forget(dst_guard); let vec = unsafe { Vec::from_raw_parts(dst_buf, len, cap) }; diff --git a/library/alloc/src/vec/in_place_drop.rs b/library/alloc/src/vec/in_place_drop.rs index 1b1ef9130face..25ca33c6a7bf0 100644 --- a/library/alloc/src/vec/in_place_drop.rs +++ b/library/alloc/src/vec/in_place_drop.rs @@ -22,3 +22,18 @@ impl Drop for InPlaceDrop { } } } + +// A helper struct for in-place collection that drops the destination allocation and elements, +// to avoid leaking them if some other destructor panics. +pub(super) struct InPlaceDstBufDrop { + pub(super) ptr: *mut T, + pub(super) len: usize, + pub(super) cap: usize, +} + +impl Drop for InPlaceDstBufDrop { + #[inline] + fn drop(&mut self) { + unsafe { super::Vec::from_raw_parts(self.ptr, self.len, self.cap) }; + } +} diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index fa9f2131c0c1d..acfbb98272dd5 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -125,7 +125,7 @@ use self::set_len_on_drop::SetLenOnDrop; mod set_len_on_drop; #[cfg(not(no_global_oom_handling))] -use self::in_place_drop::InPlaceDrop; +use self::in_place_drop::{InPlaceDrop, InPlaceDstBufDrop}; #[cfg(not(no_global_oom_handling))] mod in_place_drop; From f81b07e9470728a311a5d816616762e37b00f29f Mon Sep 17 00:00:00 2001 From: Giacomo Stevanato Date: Sat, 10 Sep 2022 11:27:21 +0200 Subject: [PATCH 2/5] Adapt inplace collection leak test to check for no leaks --- library/alloc/tests/vec.rs | 36 ++++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index d94da8f5f5a0e..ebf5cbd0e60e0 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -1110,11 +1110,13 @@ fn test_from_iter_specialization_panic_during_iteration_drops() { } #[test] -fn test_from_iter_specialization_panic_during_drop_leaks() { - static mut DROP_COUNTER: usize = 0; +fn test_from_iter_specialization_panic_during_drop_doesnt_leak() { + static mut DROP_COUNTER_SHOULD_BE_DROPPED: usize = 0; + static mut DROP_COUNTER_DROPPED_TWICE: usize = 0; #[derive(Debug)] enum Droppable { + ShouldBeDropped, DroppedTwice(Box), PanicOnDrop, } @@ -1122,11 +1124,17 @@ fn test_from_iter_specialization_panic_during_drop_leaks() { impl Drop for Droppable { fn drop(&mut self) { match self { + Droppable::ShouldBeDropped => { + unsafe { + DROP_COUNTER_SHOULD_BE_DROPPED += 1; + } + println!("Dropping ShouldBeDropped!") + } Droppable::DroppedTwice(_) => { unsafe { - DROP_COUNTER += 1; + DROP_COUNTER_DROPPED_TWICE += 1; } - println!("Dropping!") + println!("Dropping DroppedTwice!") } Droppable::PanicOnDrop => { if !std::thread::panicking() { @@ -1137,21 +1145,17 @@ fn test_from_iter_specialization_panic_during_drop_leaks() { } } - let mut to_free: *mut Droppable = core::ptr::null_mut(); - let mut cap = 0; - let _ = std::panic::catch_unwind(AssertUnwindSafe(|| { - let mut v = vec![Droppable::DroppedTwice(Box::new(123)), Droppable::PanicOnDrop]; - to_free = v.as_mut_ptr(); - cap = v.capacity(); - let _ = v.into_iter().take(0).collect::>(); + let v = vec![ + Droppable::ShouldBeDropped, + Droppable::DroppedTwice(Box::new(123)), + Droppable::PanicOnDrop, + ]; + let _ = v.into_iter().take(1).collect::>(); })); - assert_eq!(unsafe { DROP_COUNTER }, 1); - // clean up the leak to keep miri happy - unsafe { - drop(Vec::from_raw_parts(to_free, 0, cap)); - } + assert_eq!(unsafe { DROP_COUNTER_SHOULD_BE_DROPPED }, 1); + assert_eq!(unsafe { DROP_COUNTER_DROPPED_TWICE }, 1); } // regression test for issue #85322. Peekable previously implemented InPlaceIterable, From f52082f54321e56fa4dbd9194c1cfd61089e2729 Mon Sep 17 00:00:00 2001 From: Giacomo Stevanato Date: Sat, 10 Sep 2022 12:29:16 +0200 Subject: [PATCH 3/5] Update documentation --- library/alloc/src/vec/in_place_collect.rs | 7 +++++-- library/alloc/src/vec/into_iter.rs | 7 ++++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs index 359f1a9cec251..83ed3d915a291 100644 --- a/library/alloc/src/vec/in_place_collect.rs +++ b/library/alloc/src/vec/in_place_collect.rs @@ -55,6 +55,9 @@ //! This is handled by the [`InPlaceDrop`] guard for sink items (`U`) and by //! [`vec::IntoIter::forget_allocation_drop_remaining()`] for remaining source items (`T`). //! +//! If dropping any remaining source item (`T`) panics then [`InPlaceDstBufDrop`] will handle dropping +//! the already collected sink items (`U`) and freeing the allocation. +//! //! [`vec::IntoIter::forget_allocation_drop_remaining()`]: super::IntoIter::forget_allocation_drop_remaining() //! //! # O(1) collect @@ -194,8 +197,8 @@ where // Drop any remaining values at the tail of the source but prevent drop of the allocation // itself once IntoIter goes out of scope. // If the drop panics then we also try to drop the destination buffer and its elements. - // This is safe because `forget_allocation_drop_remaining` forgets the allocation *before* - // trying to drop the remaining elements. + // This is safe because `forget_allocation_drop_remaining` immediately forgets the allocation + // and won't panic before that. // // Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce // contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index 1b483e3fc7793..b5a03392a1797 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/library/alloc/src/vec/into_iter.rs @@ -96,13 +96,16 @@ impl IntoIter { } /// Drops remaining elements and relinquishes the backing allocation. + /// This method guarantees it won't panic before relinquishing + /// the backing allocation. /// /// This is roughly equivalent to the following, but more efficient /// /// ``` /// # let mut into_iter = Vec::::with_capacity(10).into_iter(); + /// let mut into_iter = std::mem::replace(&mut into_iter, Vec::new().into_iter()); /// (&mut into_iter).for_each(core::mem::drop); - /// unsafe { core::ptr::write(&mut into_iter, Vec::new().into_iter()); } + /// std::mem::forget(into_iter); /// ``` /// /// This method is used by in-place iteration, refer to the vec::in_place_collect @@ -119,6 +122,8 @@ impl IntoIter { self.ptr = self.buf.as_ptr(); self.end = self.buf.as_ptr(); + // Dropping the remaining elements can panic, so this needs to be + // done only after updating the other fields. unsafe { ptr::drop_in_place(remaining); } From dad049cb5cb3d259836cfe6a9160521d9d4809ca Mon Sep 17 00:00:00 2001 From: Giacomo Stevanato Date: Sat, 10 Sep 2022 12:29:25 +0200 Subject: [PATCH 4/5] Update test --- library/alloc/tests/vec.rs | 67 +++++++++++++++++++------------------- 1 file changed, 34 insertions(+), 33 deletions(-) diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index ebf5cbd0e60e0..990ce38e22edd 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -1111,51 +1111,52 @@ fn test_from_iter_specialization_panic_during_iteration_drops() { #[test] fn test_from_iter_specialization_panic_during_drop_doesnt_leak() { - static mut DROP_COUNTER_SHOULD_BE_DROPPED: usize = 0; - static mut DROP_COUNTER_DROPPED_TWICE: usize = 0; + static mut DROP_COUNTER_OLD: [usize; 5] = [0; 5]; + static mut DROP_COUNTER_NEW: [usize; 2] = [0; 2]; #[derive(Debug)] - enum Droppable { - ShouldBeDropped, - DroppedTwice(Box), - PanicOnDrop, + struct Old(usize); + + impl Drop for Old { + fn drop(&mut self) { + unsafe { + DROP_COUNTER_OLD[self.0] += 1; + } + + if self.0 == 3 { + panic!(); + } + + println!("Dropped Old: {}", self.0); + } } - impl Drop for Droppable { + #[derive(Debug)] + struct New(usize); + + impl Drop for New { fn drop(&mut self) { - match self { - Droppable::ShouldBeDropped => { - unsafe { - DROP_COUNTER_SHOULD_BE_DROPPED += 1; - } - println!("Dropping ShouldBeDropped!") - } - Droppable::DroppedTwice(_) => { - unsafe { - DROP_COUNTER_DROPPED_TWICE += 1; - } - println!("Dropping DroppedTwice!") - } - Droppable::PanicOnDrop => { - if !std::thread::panicking() { - panic!(); - } - } + unsafe { + DROP_COUNTER_NEW[self.0] += 1; } + + println!("Dropped New: {}", self.0); } } let _ = std::panic::catch_unwind(AssertUnwindSafe(|| { - let v = vec![ - Droppable::ShouldBeDropped, - Droppable::DroppedTwice(Box::new(123)), - Droppable::PanicOnDrop, - ]; - let _ = v.into_iter().take(1).collect::>(); + let v = vec![Old(0), Old(1), Old(2), Old(3), Old(4)]; + let _ = v.into_iter().map(|x| New(x.0)).take(2).collect::>(); })); - assert_eq!(unsafe { DROP_COUNTER_SHOULD_BE_DROPPED }, 1); - assert_eq!(unsafe { DROP_COUNTER_DROPPED_TWICE }, 1); + assert_eq!(unsafe { DROP_COUNTER_OLD[0] }, 1); + assert_eq!(unsafe { DROP_COUNTER_OLD[1] }, 1); + assert_eq!(unsafe { DROP_COUNTER_OLD[2] }, 1); + assert_eq!(unsafe { DROP_COUNTER_OLD[3] }, 1); + assert_eq!(unsafe { DROP_COUNTER_OLD[4] }, 1); + + assert_eq!(unsafe { DROP_COUNTER_NEW[0] }, 1); + assert_eq!(unsafe { DROP_COUNTER_NEW[1] }, 1); } // regression test for issue #85322. Peekable previously implemented InPlaceIterable, From 1750c7bdd36ec18324423bd30867e39d787d5977 Mon Sep 17 00:00:00 2001 From: Giacomo Stevanato Date: Sat, 10 Sep 2022 15:00:37 +0200 Subject: [PATCH 5/5] Clarify documentation --- library/alloc/src/vec/in_place_collect.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs index 83ed3d915a291..9f2555ee1611a 100644 --- a/library/alloc/src/vec/in_place_collect.rs +++ b/library/alloc/src/vec/in_place_collect.rs @@ -194,11 +194,10 @@ where ); } - // Drop any remaining values at the tail of the source but prevent drop of the allocation - // itself once IntoIter goes out of scope. - // If the drop panics then we also try to drop the destination buffer and its elements. + // The ownership of the allocation and the new `T` values is temporarily moved into `dst_guard`. // This is safe because `forget_allocation_drop_remaining` immediately forgets the allocation - // and won't panic before that. + // before any panic can occur in order to avoid any double free, and then proceeds to drop + // any remaining values at the tail of the source. // // Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce // contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the