From 84ae9699c6bb60b07329b1cf5e644d3d30eb24f3 Mon Sep 17 00:00:00 2001 From: Aaron Loucks Date: Sun, 26 May 2019 22:41:34 -0400 Subject: [PATCH 1/7] Prevent Vec::drain_filter from double dropping on panic Fixes: #60977 --- src/liballoc/tests/vec.rs | 99 +++++++++++++++++++++++++++++++++++++++ src/liballoc/vec.rs | 73 +++++++++++++++++++++++++---- 2 files changed, 162 insertions(+), 10 deletions(-) diff --git a/src/liballoc/tests/vec.rs b/src/liballoc/tests/vec.rs index 5ddac673c9ff1..c0967cd374d5e 100644 --- a/src/liballoc/tests/vec.rs +++ b/src/liballoc/tests/vec.rs @@ -945,6 +945,105 @@ fn drain_filter_complex() { } } +#[test] +fn drain_filter_consumed_panic() { + use std::rc::Rc; + use std::sync::Mutex; + + struct Check { + index: usize, + drop_counts: Rc>>, + }; + + impl Drop for Check { + fn drop(&mut self) { + self.drop_counts.lock().unwrap()[self.index] += 1; + println!("drop: {}", self.index); + } + } + + let check_count = 10; + let drop_counts = Rc::new(Mutex::new(vec![0_usize; check_count])); + let mut data: Vec = (0..check_count) + .map(|index| Check { index, drop_counts: Rc::clone(&drop_counts) }) + .collect(); + + let _ = std::panic::catch_unwind(move || { + let filter = |c: &mut Check| { + if c.index == 2 { + panic!("panic at index: {}", c.index); + } + // Verify that if the filter could panic again on another element + // that it would not cause a double panic and all elements of the + // vec would still be dropped exactly once. + if c.index == 4 { + panic!("panic at index: {}", c.index); + } + c.index < 6 + }; + let drain = data.drain_filter(filter); + + // NOTE: The DrainFilter is explictly consumed + drain.for_each(drop); + }); + + let drop_counts = drop_counts.lock().unwrap(); + assert_eq!(check_count, drop_counts.len()); + + for (index, count) in drop_counts.iter().cloned().enumerate() { + assert_eq!(1, count, "unexpected drop count at index: {} (count: {})", index, count); + } +} + +#[test] +fn drain_filter_unconsumed_panic() { + use std::rc::Rc; + use std::sync::Mutex; + + struct Check { + index: usize, + drop_counts: Rc>>, + }; + + impl Drop for Check { + fn drop(&mut self) { + self.drop_counts.lock().unwrap()[self.index] += 1; + println!("drop: {}", self.index); + } + } + + let check_count = 10; + let drop_counts = Rc::new(Mutex::new(vec![0_usize; check_count])); + let mut data: Vec = (0..check_count) + .map(|index| Check { index, drop_counts: Rc::clone(&drop_counts) }) + .collect(); + + let _ = std::panic::catch_unwind(move || { + let filter = |c: &mut Check| { + if c.index == 2 { + panic!("panic at index: {}", c.index); + } + // Verify that if the filter could panic again on another element + // that it would not cause a double panic and all elements of the + // vec would still be dropped exactly once. + if c.index == 4 { + panic!("panic at index: {}", c.index); + } + c.index < 6 + }; + let _drain = data.drain_filter(filter); + + // NOTE: The DrainFilter is dropped without being consumed + }); + + let drop_counts = drop_counts.lock().unwrap(); + assert_eq!(check_count, drop_counts.len()); + + for (index, count) in drop_counts.iter().cloned().enumerate() { + assert_eq!(1, count, "unexpected drop count at index: {} (count: {})", index, count); + } +} + #[test] fn test_reserve_exact() { // This is all the same as test_reserve diff --git a/src/liballoc/vec.rs b/src/liballoc/vec.rs index 5cb91395b7bf7..adc0929fdbed4 100644 --- a/src/liballoc/vec.rs +++ b/src/liballoc/vec.rs @@ -2120,6 +2120,7 @@ impl Vec { del: 0, old_len, pred: filter, + panic_flag: false, } } } @@ -2751,6 +2752,7 @@ pub struct DrainFilter<'a, T, F> del: usize, old_len: usize, pred: F, + panic_flag: bool, } #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")] @@ -2760,21 +2762,34 @@ impl Iterator for DrainFilter<'_, T, F> type Item = T; fn next(&mut self) -> Option { + struct SetIdxOnDrop<'a> { + idx: &'a mut usize, + new_idx: usize, + } + + impl<'a> Drop for SetIdxOnDrop<'a> { + fn drop(&mut self) { + *self.idx = self.new_idx; + } + } + unsafe { - while self.idx != self.old_len { + while self.idx < self.old_len { let i = self.idx; - self.idx += 1; let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len); - if (self.pred)(&mut v[i]) { + let mut set_idx = SetIdxOnDrop { new_idx: self.idx, idx: &mut self.idx }; + self.panic_flag = true; + let drained = (self.pred)(&mut v[i]); + self.panic_flag = false; + set_idx.new_idx += 1; + if drained { self.del += 1; return Some(ptr::read(&v[i])); - } else if self.del > 0 { + } + else if self.del > 0 { let del = self.del; let src: *const T = &v[i]; let dst: *mut T = &mut v[i - del]; - // This is safe because self.vec has length 0 - // thus its elements will not have Drop::drop - // called on them in the event of a panic. ptr::copy_nonoverlapping(src, dst, 1); } } @@ -2792,9 +2807,47 @@ impl Drop for DrainFilter<'_, T, F> where F: FnMut(&mut T) -> bool, { fn drop(&mut self) { - self.for_each(drop); - unsafe { - self.vec.set_len(self.old_len - self.del); + // If the predicate panics, we still need to backshift everything + // down after the last successfully drained element, but no additional + // elements are drained or checked. + struct BackshiftOnDrop<'a, 'b, T, F> + where + F: FnMut(&mut T) -> bool, + { + drain: &'b mut DrainFilter<'a, T, F>, + } + + impl<'a, 'b, T, F> Drop for BackshiftOnDrop<'a, 'b, T, F> + where + F: FnMut(&mut T) -> bool + { + fn drop(&mut self) { + unsafe { + while self.drain.idx < self.drain.old_len { + let i = self.drain.idx; + self.drain.idx += 1; + let v = slice::from_raw_parts_mut( + self.drain.vec.as_mut_ptr(), + self.drain.old_len, + ); + if self.drain.del > 0 { + let del = self.drain.del; + let src: *const T = &v[i]; + let dst: *mut T = &mut v[i - del]; + ptr::copy_nonoverlapping(src, dst, 1); + } + } + self.drain.vec.set_len(self.drain.old_len - self.drain.del); + } + } + } + + let backshift = BackshiftOnDrop { + drain: self + }; + + if !backshift.drain.panic_flag { + backshift.drain.for_each(drop); } } } From 17a517a42a403ec0275f1d5d286038b0acc758b6 Mon Sep 17 00:00:00 2001 From: Aaron Loucks Date: Mon, 27 May 2019 11:30:23 -0400 Subject: [PATCH 2/7] Fix formatting nit --- src/liballoc/vec.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/liballoc/vec.rs b/src/liballoc/vec.rs index adc0929fdbed4..e6ef5a8ebc993 100644 --- a/src/liballoc/vec.rs +++ b/src/liballoc/vec.rs @@ -2785,8 +2785,7 @@ impl Iterator for DrainFilter<'_, T, F> if drained { self.del += 1; return Some(ptr::read(&v[i])); - } - else if self.del > 0 { + } else if self.del > 0 { let del = self.del; let src: *const T = &v[i]; let dst: *mut T = &mut v[i - del]; From 53d46ae96e2e8ac153680cb6711c9935d6dfe0f6 Mon Sep 17 00:00:00 2001 From: Aaron Loucks Date: Mon, 27 May 2019 11:47:14 -0400 Subject: [PATCH 3/7] Add drain_filter_unconsumed test --- src/liballoc/tests/vec.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/liballoc/tests/vec.rs b/src/liballoc/tests/vec.rs index c0967cd374d5e..a89b8aa4a717b 100644 --- a/src/liballoc/tests/vec.rs +++ b/src/liballoc/tests/vec.rs @@ -1044,6 +1044,14 @@ fn drain_filter_unconsumed_panic() { } } +#[test] +fn drain_filter_unconsumed() { + let mut vec = vec![1, 2, 3, 4]; + let drain = vec.drain_filter(|&mut x| x % 2 != 0); + drop(drain); + assert_eq!(vec, [2, 4]); +} + #[test] fn test_reserve_exact() { // This is all the same as test_reserve From b13ae65b8b995bc851708380d80eae301c39f1c6 Mon Sep 17 00:00:00 2001 From: Aaron Loucks Date: Mon, 27 May 2019 12:47:47 -0400 Subject: [PATCH 4/7] Refactor DrainFilter::next and update comments --- src/liballoc/vec.rs | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/src/liballoc/vec.rs b/src/liballoc/vec.rs index e6ef5a8ebc993..99116ac2d2f13 100644 --- a/src/liballoc/vec.rs +++ b/src/liballoc/vec.rs @@ -2762,26 +2762,17 @@ impl Iterator for DrainFilter<'_, T, F> type Item = T; fn next(&mut self) -> Option { - struct SetIdxOnDrop<'a> { - idx: &'a mut usize, - new_idx: usize, - } - - impl<'a> Drop for SetIdxOnDrop<'a> { - fn drop(&mut self) { - *self.idx = self.new_idx; - } - } - unsafe { while self.idx < self.old_len { let i = self.idx; let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len); - let mut set_idx = SetIdxOnDrop { new_idx: self.idx, idx: &mut self.idx }; self.panic_flag = true; let drained = (self.pred)(&mut v[i]); self.panic_flag = false; - set_idx.new_idx += 1; + // Update the index *after* the predicate is called. If the index + // is updated prior and the predicate panics, the element at this + // index would be leaked. + self.idx += 1; if drained { self.del += 1; return Some(ptr::read(&v[i])); @@ -2806,9 +2797,6 @@ impl Drop for DrainFilter<'_, T, F> where F: FnMut(&mut T) -> bool, { fn drop(&mut self) { - // If the predicate panics, we still need to backshift everything - // down after the last successfully drained element, but no additional - // elements are drained or checked. struct BackshiftOnDrop<'a, 'b, T, F> where F: FnMut(&mut T) -> bool, @@ -2822,6 +2810,12 @@ impl Drop for DrainFilter<'_, T, F> { fn drop(&mut self) { unsafe { + // Backshift any unprocessed elements, preventing double-drop + // of any element that *should* have been previously overwritten + // but was not due to a panic in the filter predicate. This is + // implemented via drop so that it's guaranteed to run even in + // the event of a panic while consuming the remainder of the + // DrainFilter. while self.drain.idx < self.drain.old_len { let i = self.drain.idx; self.drain.idx += 1; @@ -2845,6 +2839,9 @@ impl Drop for DrainFilter<'_, T, F> drain: self }; + // Attempt to consume any remaining elements if the filter predicate + // has not yet panicked. We'll backshift any remaining elements + // whether we've already panicked or if the consumption here panics. if !backshift.drain.panic_flag { backshift.drain.for_each(drop); } From f5ab0313fe9330717e6697b148bcdd0bb78f3508 Mon Sep 17 00:00:00 2001 From: Aaron Loucks Date: Mon, 27 May 2019 12:55:59 -0400 Subject: [PATCH 5/7] Disable drain_filter tests that require catch_unwind on miri --- src/liballoc/tests/vec.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/liballoc/tests/vec.rs b/src/liballoc/tests/vec.rs index a89b8aa4a717b..07ee37412680d 100644 --- a/src/liballoc/tests/vec.rs +++ b/src/liballoc/tests/vec.rs @@ -946,6 +946,7 @@ fn drain_filter_complex() { } #[test] +#[cfg(not(miri))] // Miri does not support catching panics fn drain_filter_consumed_panic() { use std::rc::Rc; use std::sync::Mutex; @@ -996,6 +997,7 @@ fn drain_filter_consumed_panic() { } #[test] +#[cfg(not(miri))] // Miri does not support catching panics fn drain_filter_unconsumed_panic() { use std::rc::Rc; use std::sync::Mutex; From a4a6a67a0142d46205f0aa662dc29c1f37aca86d Mon Sep 17 00:00:00 2001 From: Aaron Loucks Date: Sun, 7 Jul 2019 13:26:06 -0400 Subject: [PATCH 6/7] Remove while loop in DrainFilter::drop and add additional docs --- src/liballoc/vec.rs | 40 +++++++++++++++++++++------------------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/src/liballoc/vec.rs b/src/liballoc/vec.rs index 99116ac2d2f13..4adbe6467a1e8 100644 --- a/src/liballoc/vec.rs +++ b/src/liballoc/vec.rs @@ -2748,10 +2748,19 @@ pub struct DrainFilter<'a, T, F> where F: FnMut(&mut T) -> bool, { vec: &'a mut Vec, + /// The index of the item that will be inspected by the next call to `next`. idx: usize, + /// The number of items that have been drained (removed) thus far. del: usize, + /// The original length of `vec` prior to draining. old_len: usize, + /// The filter test predicate. pred: F, + /// A flag that indicates a panic has occured in the filter test prodicate. + /// This is used as a hint in the drop implmentation to prevent consumption + /// of the remainder of the `DrainFilter`. Any unprocessed items will be + /// backshifted in the `vec`, but no further items will be dropped or + /// tested by the filter predicate. panic_flag: bool, } @@ -2810,25 +2819,18 @@ impl Drop for DrainFilter<'_, T, F> { fn drop(&mut self) { unsafe { - // Backshift any unprocessed elements, preventing double-drop - // of any element that *should* have been previously overwritten - // but was not due to a panic in the filter predicate. This is - // implemented via drop so that it's guaranteed to run even in - // the event of a panic while consuming the remainder of the - // DrainFilter. - while self.drain.idx < self.drain.old_len { - let i = self.drain.idx; - self.drain.idx += 1; - let v = slice::from_raw_parts_mut( - self.drain.vec.as_mut_ptr(), - self.drain.old_len, - ); - if self.drain.del > 0 { - let del = self.drain.del; - let src: *const T = &v[i]; - let dst: *mut T = &mut v[i - del]; - ptr::copy_nonoverlapping(src, dst, 1); - } + if self.drain.idx < self.drain.old_len && self.drain.del > 0 { + // This is a pretty messed up state, and there isn't really an + // obviously right thing to do. We don't want to keep trying + // to execute `pred`, so we just backshift all the unprocessed + // elements and tell the vec that they still exist. The backshift + // is required to prevent a double-drop of the last successfully + // drained item following a panic in the predicate. + let ptr = self.drain.vec.as_mut_ptr(); + let src = ptr.add(self.drain.idx); + let dst = src.sub(self.drain.del); + let tail_len = self.drain.old_len - self.drain.idx; + src.copy_to(dst, tail_len); } self.drain.vec.set_len(self.drain.old_len - self.drain.del); } From df5b32ee9b98838c1cccc93f75f941ecafc9d086 Mon Sep 17 00:00:00 2001 From: Aaron Loucks Date: Sun, 7 Jul 2019 16:36:19 -0400 Subject: [PATCH 7/7] Clarify double-drop comment --- src/liballoc/vec.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/liballoc/vec.rs b/src/liballoc/vec.rs index 4adbe6467a1e8..8939248adab0d 100644 --- a/src/liballoc/vec.rs +++ b/src/liballoc/vec.rs @@ -2825,7 +2825,7 @@ impl Drop for DrainFilter<'_, T, F> // to execute `pred`, so we just backshift all the unprocessed // elements and tell the vec that they still exist. The backshift // is required to prevent a double-drop of the last successfully - // drained item following a panic in the predicate. + // drained item prior to a panic in the predicate. let ptr = self.drain.vec.as_mut_ptr(); let src = ptr.add(self.drain.idx); let dst = src.sub(self.drain.del);