From 8ddf61330d73bd1b33ed01e03f2bf0b8aaba8d11 Mon Sep 17 00:00:00 2001 From: Matt Brubeck Date: Wed, 22 Apr 2020 13:34:45 -0700 Subject: [PATCH] Fix leak on panic in `insert_many`. Fixes #208. --- src/lib.rs | 58 ++++++++++++++++++++++++++---- src/tests.rs | 99 ++++++++++++++++++++++++++++++++++++++-------------- 2 files changed, 125 insertions(+), 32 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 203a154..cf79ad8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -98,7 +98,7 @@ use core::hint::unreachable_unchecked; use core::iter::{repeat, FromIterator, FusedIterator, IntoIterator}; use core::mem; use core::mem::MaybeUninit; -use core::ops::{self, RangeBounds}; +use core::ops::{self, Range, RangeBounds}; use core::ptr::{self, NonNull}; use core::slice::{self, SliceIndex}; @@ -975,8 +975,6 @@ impl SmallVec { /// Insert multiple elements at position `index`, shifting all following elements toward the /// back. - /// - /// Note: when the iterator panics, this can leak memory. pub fn insert_many>(&mut self, index: usize, iterable: I) { let iter = iterable.into_iter(); if index == self.len() { @@ -991,13 +989,19 @@ impl SmallVec { unsafe { let old_len = self.len(); assert!(index <= old_len); - let mut ptr = self.as_mut_ptr().add(index); + let start = self.as_mut_ptr(); + let mut ptr = start.add(index); // Move the trailing elements. ptr::copy(ptr, ptr.add(lower_size_bound), old_len - index); // In case the iterator panics, don't double-drop the items we just copied above. - self.set_len(index); + self.set_len(0); + let mut guard = DropOnPanic { + start, + skip: index..(index + lower_size_bound), + len: old_len + lower_size_bound, + }; let mut num_added = 0; for element in iter { @@ -1005,13 +1009,21 @@ impl SmallVec { if num_added >= lower_size_bound { // Iterator provided more elements than the hint. Move trailing items again. self.reserve(1); - ptr = self.as_mut_ptr().add(index); + let start = self.as_mut_ptr(); + ptr = start.add(index); cur = ptr.add(num_added); ptr::copy(cur, cur.add(1), old_len - index); + + guard.start = start; + guard.len += 1; + guard.skip.end += 1; } ptr::write(cur, element); + guard.skip.start += 1; num_added += 1; } + mem::forget(guard); + if num_added < lower_size_bound { // Iterator provided fewer elements than the hint ptr::copy( @@ -1023,6 +1035,24 @@ impl SmallVec { self.set_len(old_len + num_added); } + + struct DropOnPanic { + start: *mut T, + skip: Range, + len: usize, + } + + impl Drop for DropOnPanic { + fn drop(&mut self) { + for i in 0..self.len { + if !self.skip.contains(&i) { + unsafe { + ptr::drop_in_place(self.start.add(i)); + } + } + } + } + } } /// Convert a SmallVec to a Vec, without reallocating if the SmallVec has already spilled onto @@ -1249,6 +1279,22 @@ impl SmallVec { data: SmallVecData::from_heap(ptr, length), } } + + /// Returns a raw pointer to the vector's buffer. + pub fn as_ptr(&self) -> *const A::Item { + // We shadow the slice method of the same name to avoid going through + // `deref`, which creates an intermediate reference that may place + // additional safety constraints on the contents of the slice. + self.triple().0 + } + + /// Returns a raw mutable pointer to the vector's buffer. + pub fn as_mut_ptr(&mut self) -> *mut A::Item { + // We shadow the slice method of the same name to avoid going through + // `deref_mut`, which creates an intermediate reference that may place + // additional safety constraints on the contents of the slice. + self.triple_mut().0 + } } impl SmallVec diff --git a/src/tests.rs b/src/tests.rs index f88e341..0452ae8 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1,4 +1,4 @@ -use crate::{SmallVec, smallvec}; +use crate::{smallvec, SmallVec}; use std::iter::FromIterator; @@ -320,13 +320,23 @@ fn test_insert_many_long_hint() { ); } -#[test] // https://github.com/servo/rust-smallvec/issues/96 -fn test_insert_many_panic() { +mod insert_many_panic { + use crate::{smallvec, SmallVec}; + use alloc::boxed::Box; + struct PanicOnDoubleDrop { dropped: Box, } + impl PanicOnDoubleDrop { + fn new() -> Self { + Self { + dropped: Box::new(false), + } + } + } + impl Drop for PanicOnDoubleDrop { fn drop(&mut self) { assert!(!*self.dropped, "already dropped"); @@ -334,38 +344,75 @@ fn test_insert_many_panic() { } } - struct BadIter; + /// Claims to yield `hint` items, but actually yields `count`, then panics. + struct BadIter { + hint: usize, + count: usize, + } + impl Iterator for BadIter { type Item = PanicOnDoubleDrop; fn size_hint(&self) -> (usize, Option) { - (1, None) + (self.hint, None) } fn next(&mut self) -> Option { - panic!() + if self.count == 0 { + panic!() + } + self.count -= 1; + Some(PanicOnDoubleDrop::new()) } } - // These boxes are leaked on purpose by panicking `insert_many`, - // so we clean them up manually to appease Miri's leak checker. - let mut box1 = Box::new(false); - let mut box2 = Box::new(false); + #[test] + fn panic_early_at_start() { + let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = + smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),]; + let result = ::std::panic::catch_unwind(move || { + vec.insert_many(0, BadIter { hint: 1, count: 0 }); + }); + assert!(result.is_err()); + } - let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = vec![ - PanicOnDoubleDrop { - dropped: unsafe { Box::from_raw(&mut *box1) }, - }, - PanicOnDoubleDrop { - dropped: unsafe { Box::from_raw(&mut *box2) }, - }, - ] - .into(); - let result = ::std::panic::catch_unwind(move || { - vec.insert_many(0, BadIter); - }); - assert!(result.is_err()); - - drop(box1); - drop(box2); + #[test] + fn panic_early_in_middle() { + let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = + smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),]; + let result = ::std::panic::catch_unwind(move || { + vec.insert_many(1, BadIter { hint: 4, count: 2 }); + }); + assert!(result.is_err()); + } + + #[test] + fn panic_early_at_end() { + let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = + smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),]; + let result = ::std::panic::catch_unwind(move || { + vec.insert_many(2, BadIter { hint: 3, count: 1 }); + }); + assert!(result.is_err()); + } + + #[test] + fn panic_late_at_start() { + let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = + smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),]; + let result = ::std::panic::catch_unwind(move || { + vec.insert_many(0, BadIter { hint: 3, count: 5 }); + }); + assert!(result.is_err()); + } + + #[test] + fn panic_late_at_end() { + let mut vec: SmallVec<[PanicOnDoubleDrop; 0]> = + smallvec![PanicOnDoubleDrop::new(), PanicOnDoubleDrop::new(),]; + let result = ::std::panic::catch_unwind(move || { + vec.insert_many(2, BadIter { hint: 3, count: 5 }); + }); + assert!(result.is_err()); + } } #[test]