From 2fefd9ad07122a4cf1329d52f7c5fe7ccd0893c9 Mon Sep 17 00:00:00 2001 From: Yechan Bae Date: Wed, 3 Feb 2021 16:36:33 -0500 Subject: [PATCH 1/4] Fixes #80335 (cherry picked from commit 6d43225bfb08ec91f7476b76c7fec632c4a096ef) --- library/alloc/src/str.rs | 42 ++++++++++++++++++++++---------------- library/alloc/tests/str.rs | 30 +++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 18 deletions(-) diff --git a/library/alloc/src/str.rs b/library/alloc/src/str.rs index 8c64bc0801259..89e6eee925e10 100644 --- a/library/alloc/src/str.rs +++ b/library/alloc/src/str.rs @@ -90,8 +90,8 @@ impl> Join<&str> for [S] { } } -macro_rules! spezialize_for_lengths { - ($separator:expr, $target:expr, $iter:expr; $($num:expr),*) => { +macro_rules! specialize_for_lengths { + ($separator:expr, $target:expr, $iter:expr; $($num:expr),*) => {{ let mut target = $target; let iter = $iter; let sep_bytes = $separator; @@ -102,7 +102,8 @@ macro_rules! spezialize_for_lengths { $num => { for s in iter { copy_slice_and_advance!(target, sep_bytes); - copy_slice_and_advance!(target, s.borrow().as_ref()); + let content_bytes = s.borrow().as_ref(); + copy_slice_and_advance!(target, content_bytes); } }, )* @@ -110,11 +111,13 @@ macro_rules! spezialize_for_lengths { // arbitrary non-zero size fallback for s in iter { copy_slice_and_advance!(target, sep_bytes); - copy_slice_and_advance!(target, s.borrow().as_ref()); + let content_bytes = s.borrow().as_ref(); + copy_slice_and_advance!(target, content_bytes); } } } - }; + target + }} } macro_rules! copy_slice_and_advance { @@ -153,7 +156,7 @@ where // if the `len` calculation overflows, we'll panic // we would have run out of memory anyway and the rest of the function requires // the entire Vec pre-allocated for safety - let len = sep_len + let reserved_len = sep_len .checked_mul(iter.len()) .and_then(|n| { slice.iter().map(|s| s.borrow().as_ref().len()).try_fold(n, usize::checked_add) @@ -161,22 +164,25 @@ where .expect("attempt to join into collection with len > usize::MAX"); // crucial for safety - let mut result = Vec::with_capacity(len); - assert!(result.capacity() >= len); + let mut result = Vec::with_capacity(reserved_len); + debug_assert!(result.capacity() >= reserved_len); result.extend_from_slice(first.borrow().as_ref()); unsafe { - { - let pos = result.len(); - let target = result.get_unchecked_mut(pos..len); - - // copy separator and slices over without bounds checks - // generate loops with hardcoded offsets for small separators - // massive improvements possible (~ x2) - spezialize_for_lengths!(sep, target, iter; 0, 1, 2, 3, 4); - } - result.set_len(len); + let pos = result.len(); + let target = result.get_unchecked_mut(pos..reserved_len); + + // copy separator and slices over without bounds checks + // generate loops with hardcoded offsets for small separators + // massive improvements possible (~ x2) + let remain = specialize_for_lengths!(sep, target, iter; 0, 1, 2, 3, 4); + + // issue #80335: A weird borrow implementation can return different + // slices for the length calculation and the actual copy, so + // `remain.len()` might be non-zero. + let result_len = reserved_len - remain.len(); + result.set_len(result_len); } result } diff --git a/library/alloc/tests/str.rs b/library/alloc/tests/str.rs index 604835e6cc4a6..6df8d8c2f354f 100644 --- a/library/alloc/tests/str.rs +++ b/library/alloc/tests/str.rs @@ -160,6 +160,36 @@ fn test_join_for_different_lengths_with_long_separator() { test_join!("～～～～～a～～～～～bc", ["", "a", "bc"], "～～～～～"); } +#[test] +fn test_join_isue_80335() { + use core::{borrow::Borrow, cell::Cell}; + + struct WeirdBorrow { + state: Cell, + } + + impl Default for WeirdBorrow { + fn default() -> Self { + WeirdBorrow { state: Cell::new(false) } + } + } + + impl Borrow for WeirdBorrow { + fn borrow(&self) -> &str { + let state = self.state.get(); + if state { + "0" + } else { + self.state.set(true); + "123456" + } + } + } + + let arr: [WeirdBorrow; 3] = Default::default(); + test_join!("0-0-0", arr, "-"); +} + #[test] #[cfg_attr(miri, ignore)] // Miri is too slow fn test_unsafe_slice() { From cc4d6d644b6eb9840abf13f8abe882b6e6ee1960 Mon Sep 17 00:00:00 2001 From: Yechan Bae Date: Sat, 20 Mar 2021 13:42:54 -0400 Subject: [PATCH 2/4] Update the comment (cherry picked from commit 26a62701e42d10c03ce5f2f911e7d5edeefa2f0f) --- library/alloc/src/str.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/library/alloc/src/str.rs b/library/alloc/src/str.rs index 89e6eee925e10..73401fcc8a9dc 100644 --- a/library/alloc/src/str.rs +++ b/library/alloc/src/str.rs @@ -163,7 +163,7 @@ where }) .expect("attempt to join into collection with len > usize::MAX"); - // crucial for safety + // prepare an uninitialized buffer let mut result = Vec::with_capacity(reserved_len); debug_assert!(result.capacity() >= reserved_len); @@ -178,9 +178,9 @@ where // massive improvements possible (~ x2) let remain = specialize_for_lengths!(sep, target, iter; 0, 1, 2, 3, 4); - // issue #80335: A weird borrow implementation can return different - // slices for the length calculation and the actual copy, so - // `remain.len()` might be non-zero. + // A weird borrow implementation may return different + // slices for the length calculation and the actual copy. + // Make sure we don't expose uninitialized bytes to the caller. let result_len = reserved_len - remain.len(); result.set_len(result_len); } From 6ed66ebdded175cc8739cd19459554a6562f660e Mon Sep 17 00:00:00 2001 From: The8472 Date: Mon, 29 Mar 2021 04:22:34 +0200 Subject: [PATCH 3/4] add testcase for double-drop during Vec in-place collection (cherry picked from commit fa89c0fbcfa8f4d44f153b1195ec5a305540ffc4) --- library/alloc/tests/vec.rs | 38 +++++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index c142536cd2dfb..b926c697d58ab 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -1027,7 +1027,7 @@ fn test_from_iter_specialization_head_tail_drop() { } #[test] -fn test_from_iter_specialization_panic_drop() { +fn test_from_iter_specialization_panic_during_iteration_drops() { let drop_count: Vec<_> = (0..=2).map(|_| Rc::new(())).collect(); let src: Vec<_> = drop_count.iter().cloned().collect(); let iter = src.into_iter(); @@ -1050,6 +1050,42 @@ fn test_from_iter_specialization_panic_drop() { ); } +#[test] +fn test_from_iter_specialization_panic_during_drop_leaks() { + static mut DROP_COUNTER: usize = 0; + + #[derive(Debug)] + enum Droppable { + DroppedTwice(Box), + PanicOnDrop, + } + + impl Drop for Droppable { + fn drop(&mut self) { + match self { + Droppable::DroppedTwice(_) => { + unsafe { + DROP_COUNTER += 1; + } + println!("Dropping!") + } + Droppable::PanicOnDrop => { + if !std::thread::panicking() { + panic!(); + } + } + } + } + } + + let _ = std::panic::catch_unwind(AssertUnwindSafe(|| { + let v = vec![Droppable::DroppedTwice(Box::new(123)), Droppable::PanicOnDrop]; + let _ = v.into_iter().take(0).collect::>(); + })); + + assert_eq!(unsafe { DROP_COUNTER }, 1); +} + #[test] fn test_cow_from() { let borrowed: &[_] = &["borrowed", "(slice)"]; From 836fef007ea502fb37a9308be394ad5a4a889d8f Mon Sep 17 00:00:00 2001 From: The8472 Date: Mon, 29 Mar 2021 04:22:48 +0200 Subject: [PATCH 4/4] fix double-drop in in-place collect specialization (cherry picked from commit 421f5d282a51e130d3ca7c4524d8ad6753437da9) --- library/alloc/src/vec/into_iter.rs | 27 ++++++++++++++------- library/alloc/src/vec/source_iter_marker.rs | 4 +-- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index bcbdffabc7fbe..324e894bafd23 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/library/alloc/src/vec/into_iter.rs @@ -85,20 +85,29 @@ impl IntoIter { ptr::slice_from_raw_parts_mut(self.ptr as *mut T, self.len()) } - pub(super) fn drop_remaining(&mut self) { - unsafe { - ptr::drop_in_place(self.as_mut_slice()); - } - self.ptr = self.end; - } + /// Drops remaining elements and relinquishes the backing allocation. + /// + /// This is roughly equivalent to the following, but more efficient + /// + /// ``` + /// # let mut into_iter = Vec::::with_capacity(10).into_iter(); + /// (&mut into_iter).for_each(core::mem::drop); + /// unsafe { core::ptr::write(&mut into_iter, Vec::new().into_iter()); } + /// ``` + pub(super) fn forget_allocation_drop_remaining(&mut self) { + let remaining = self.as_raw_mut_slice(); - /// Relinquishes the backing allocation, equivalent to - /// `ptr::write(&mut self, Vec::new().into_iter())` - pub(super) fn forget_allocation(&mut self) { + // overwrite the individual fields instead of creating a new + // struct and then overwriting &mut self. + // this creates less assembly self.cap = 0; self.buf = unsafe { NonNull::new_unchecked(RawVec::NEW.ptr()) }; self.ptr = self.buf.as_ptr(); self.end = self.buf.as_ptr(); + + unsafe { + ptr::drop_in_place(remaining); + } } } diff --git a/library/alloc/src/vec/source_iter_marker.rs b/library/alloc/src/vec/source_iter_marker.rs index 8c0e95559fa15..9301f7a5184e1 100644 --- a/library/alloc/src/vec/source_iter_marker.rs +++ b/library/alloc/src/vec/source_iter_marker.rs @@ -78,9 +78,9 @@ where } // drop any remaining values at the tail of the source - src.drop_remaining(); // but prevent drop of the allocation itself once IntoIter goes out of scope - src.forget_allocation(); + // if the drop panics then we also leak any elements collected into dst_buf + src.forget_allocation_drop_remaining(); let vec = unsafe { let len = dst.offset_from(dst_buf) as usize;