From b207b7b57da9e117ddb59f9b1269207a5de1e24a Mon Sep 17 00:00:00 2001 From: Jonathan Reem Date: Tue, 27 Jan 2015 16:46:46 -0500 Subject: [PATCH 1/4] Buffer: a safe utility collection for allocating memory. Buffer provides an entirely safe interface for managing memory within other collections. It provides allocate, reallocate, and deallocate functionality and obeys RAII principles (deallocate on drop). This can be used to simplify the implementation of any data structure which does allocation not through another structure like Vec. From the docs: A safe wrapper around a heap allocated buffer of Ts, tracking capacity only. Buffer makes no promises about the actual contents of this memory, that's up to the user of the structure and can be manipulated using the standard pointer utilities, accessible through the impl of `Deref` for `Buffer`. As a result of this hands-off approach, `Buffer`s destructor does not attempt to drop any of the contained elements; the destructor simply frees the contained memory. You can think of `Buffer` as an approximation for `Box<[T]>` where the elements are not guaranteed to be valid/initialized. It is meant to be used as a building block for other collections, so they do not have to concern themselves with the minutiae of allocating, reallocating, and deallocating memory. --- src/lib.rs | 2 + src/util/buffer.rs | 266 +++++++++++++++++++++++++++++++++++++++++++++ src/util/mod.rs | 3 + 3 files changed, 271 insertions(+) create mode 100644 src/util/buffer.rs create mode 100644 src/util/mod.rs diff --git a/src/lib.rs b/src/lib.rs index 2eddd1b..d9acbcd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,6 +28,7 @@ #[cfg(test)] extern crate test; extern crate core; extern crate traverse; +extern crate alloc; @@ -85,5 +86,6 @@ pub mod trie_set { +pub mod util; pub mod proto; diff --git a/src/util/buffer.rs b/src/util/buffer.rs new file mode 100644 index 0000000..0e76664 --- /dev/null +++ b/src/util/buffer.rs @@ -0,0 +1,266 @@ +//! A memory buffer with resizing. + +use core::nonzero::NonZero; +use std::num::Int; +use std::ops::Deref; +use std::rt::heap; +use std::mem; + +/// A safe wrapper around a heap allocated buffer of Ts, tracking capacity only. +/// +/// Buffer makes no promises about the actual contents of this memory, that's up +/// to the user of the structure and can be manipulated using the standard pointer +/// utilities, accessible through the impl of `Deref` for `Buffer`. +/// +/// As a result of this hands-off approach, `Buffer`s destructor does not attempt +/// to drop any of the contained elements; the destructor simply frees the contained +/// memory. +/// +/// You can think of `Buffer` as an approximation for `Box<[T]>` where the elements +/// are not guaranteed to be valid/initialized. It is meant to be used as a building +/// block for other collections, so they do not have to concern themselves with the +/// minutiae of allocating, reallocating, and deallocating memory. +pub struct Buffer { + buffer: NonZero<*mut T>, + cap: usize +} + +unsafe impl Send for Buffer { } +unsafe impl Sync for Buffer { } + +impl Buffer { + /// Create a new, empty Buffer. + /// + /// ``` + /// # use collect::util::buffer::Buffer; + /// + /// let buffer: Buffer = Buffer::new(); + /// assert_eq!(buffer.capacity(), 0); + /// ``` + pub fn new() -> Buffer { + Buffer { + buffer: unsafe { NonZero::new(heap::EMPTY as *mut T) }, + cap: 0 + } + } + + /// Create a new buffer with space for cap Ts. + /// + /// ``` + /// # use collect::util::buffer::Buffer; + /// + /// let buffer: Buffer = Buffer::allocate(128); + /// assert_eq!(buffer.capacity(), 128); + /// ``` + pub fn allocate(cap: usize) -> Buffer { + if cap == 0 { return Buffer::new() } + + Buffer { + buffer: unsafe { allocate(NonZero::new(cap)) }, + cap: cap + } + } + + /// Reallocate this buffer to fit a new number of Ts. + /// + /// ``` + /// # use collect::util::buffer::Buffer; + /// + /// let mut buffer: Buffer = Buffer::allocate(128); + /// assert_eq!(buffer.capacity(), 128); + /// + /// buffer.reallocate(1024); + /// assert_eq!(buffer.capacity(), 1024); + /// ``` + pub fn reallocate(&mut self, cap: usize) { + *self = if self.cap == 0 || cap == 0 { + Buffer::allocate(cap) + } else { + Buffer { + buffer: unsafe { + reallocate(self.buffer, + NonZero::new(self.cap), + NonZero::new(cap)) + }, + cap: cap + } + } + } + + /// Get the current capacity of the Buffer. + /// + /// ``` + /// # use collect::util::buffer::Buffer; + /// + /// let buffer: Buffer = Buffer::allocate(128); + /// assert_eq!(buffer.capacity(), 128); + /// ``` + pub fn capacity(&self) -> usize { + self.cap + } +} + +#[unsafe_destructor] +impl Drop for Buffer { + /// Safety Note: + /// + /// The Buffer will *only* deallocate the contained memory. It will + /// *not* run any destructors on data in that memory. + fn drop(&mut self) { + if self.cap == 0 { return } + unsafe { deallocate(self.buffer, NonZero::new(self.cap)) } + } +} + +impl Deref for Buffer { + type Target = *mut T; + + fn deref(&self) -> &*mut T { + &*self.buffer + } +} + +// Typed Allocation Utilities +// +// Unlike std::rt::heap these check for zero-sized types, capacity overflow, +// oom etc. and calculate the appropriate size and alignment themselves. + +unsafe fn allocate(cap: NonZero) -> NonZero<*mut T> { + if mem::size_of::() == 0 { return empty() } + + // Allocate + let ptr = heap::allocate(allocation_size::(cap), mem::align_of::()); + + // Check for allocation failure + if ptr.is_null() { ::alloc::oom() } + + NonZero::new(ptr as *mut T) +} + +unsafe fn reallocate(ptr: NonZero<*mut T>, + old_cap: NonZero, + new_cap: NonZero) -> NonZero<*mut T> { + if mem::size_of::() == 0 { return empty() } + + let old_size = allocation_size::(old_cap); + let new_size = allocation_size::(new_cap); + + // Reallocate + let ptr = heap::reallocate(*ptr as *mut u8, old_size, new_size, mem::align_of::()); + + // Check for allocation failure + if ptr.is_null() { ::alloc::oom() } + + NonZero::new(ptr as *mut T) +} + +unsafe fn deallocate(ptr: NonZero<*mut T>, cap: NonZero) { + if mem::size_of::() == 0 { return } + + let old_size = allocation_size::(cap); + + heap::deallocate(*ptr as *mut u8, old_size, mem::align_of::()) +} + +fn allocation_size(cap: NonZero) -> usize { + mem::size_of::().checked_mul(*cap).expect("Capacity overflow") +} + +fn empty() -> NonZero<*mut T> { + unsafe { NonZero::new(heap::EMPTY as *mut T) } +} + +#[cfg(test)] +mod test { + use std::ptr; + + use util::buffer::Buffer; + use super::empty; + + #[test] + fn test_empty() { + let buffer: Buffer = Buffer::new(); + assert_eq!(buffer.cap, 0); + assert_eq!(buffer.buffer, empty()); + } + + #[test] + fn test_allocate() { + let buffer: Buffer = Buffer::allocate(8); + + assert_eq!(buffer.cap, 8); + + unsafe { + ptr::write(buffer.offset(0), 8); + ptr::write(buffer.offset(1), 4); + ptr::write(buffer.offset(3), 5); + ptr::write(buffer.offset(5), 3); + ptr::write(buffer.offset(7), 6); + + assert_eq!(ptr::read(buffer.offset(0)), 8); + assert_eq!(ptr::read(buffer.offset(1)), 4); + assert_eq!(ptr::read(buffer.offset(3)), 5); + assert_eq!(ptr::read(buffer.offset(5)), 3); + assert_eq!(ptr::read(buffer.offset(7)), 6); + }; + + // Try a large buffer + let buffer: Buffer = Buffer::allocate(1024 * 1024); + + unsafe { + ptr::write(buffer.offset(1024 * 1024 - 1), 12); + assert_eq!(ptr::read(buffer.offset(1024 * 1024 - 1)), 12); + }; + } + + #[test] + fn test_reallocate() { + let mut buffer: Buffer = Buffer::allocate(8); + assert_eq!(buffer.cap, 8); + + buffer.reallocate(16); + assert_eq!(buffer.cap, 16); + + unsafe { + // Put some data in the buffer + ptr::write(buffer.offset(0), 8); + ptr::write(buffer.offset(1), 4); + ptr::write(buffer.offset(5), 3); + ptr::write(buffer.offset(7), 6); + }; + + // Allocate so in-place fails. + let _: Buffer = Buffer::allocate(128); + + buffer.reallocate(32); + + unsafe { + // Ensure the data is still there. + assert_eq!(ptr::read(buffer.offset(0)), 8); + assert_eq!(ptr::read(buffer.offset(1)), 4); + assert_eq!(ptr::read(buffer.offset(5)), 3); + assert_eq!(ptr::read(buffer.offset(7)), 6); + }; + } + + #[test] + #[should_fail = "Capacity overflow."] + fn test_allocate_capacity_overflow() { + let _: Buffer = Buffer::allocate(10_000_000_000_000_000_000); + } + + #[test] + #[should_fail = "Capacity overflow."] + fn test_fresh_reallocate_capacity_overflow() { + let mut buffer: Buffer = Buffer::new(); + buffer.reallocate(10_000_000_000_000_000_000); + } + + #[test] + #[should_fail = "Capacity overflow."] + fn test_reallocate_capacity_overflow() { + let mut buffer: Buffer = Buffer::allocate(128); + buffer.reallocate(10_000_000_000_000_000_000); + } +} + diff --git a/src/util/mod.rs b/src/util/mod.rs new file mode 100644 index 0000000..030a9be --- /dev/null +++ b/src/util/mod.rs @@ -0,0 +1,3 @@ +/// Utilities for implementing other data structures. + +pub mod buffer; From 68778f2d069be763158d0a1a15bd3c7f5b4a97ac Mon Sep 17 00:00:00 2001 From: Jonathan Reem Date: Tue, 27 Jan 2015 17:50:35 -0500 Subject: [PATCH 2/4] Buffer: Don't check for overflow for old sizes. These checks aren't needed since the old size was checked when it was originally allocated. --- src/util/buffer.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/util/buffer.rs b/src/util/buffer.rs index 0e76664..70b3f04 100644 --- a/src/util/buffer.rs +++ b/src/util/buffer.rs @@ -142,7 +142,7 @@ unsafe fn reallocate(ptr: NonZero<*mut T>, new_cap: NonZero) -> NonZero<*mut T> { if mem::size_of::() == 0 { return empty() } - let old_size = allocation_size::(old_cap); + let old_size = unchecked_allocation_size::(old_cap); let new_size = allocation_size::(new_cap); // Reallocate @@ -157,7 +157,7 @@ unsafe fn reallocate(ptr: NonZero<*mut T>, unsafe fn deallocate(ptr: NonZero<*mut T>, cap: NonZero) { if mem::size_of::() == 0 { return } - let old_size = allocation_size::(cap); + let old_size = unchecked_allocation_size::(cap); heap::deallocate(*ptr as *mut u8, old_size, mem::align_of::()) } @@ -166,6 +166,10 @@ fn allocation_size(cap: NonZero) -> usize { mem::size_of::().checked_mul(*cap).expect("Capacity overflow") } +fn unchecked_allocation_size(cap: NonZero) -> usize { + mem::size_of::() * (*cap) +} + fn empty() -> NonZero<*mut T> { unsafe { NonZero::new(heap::EMPTY as *mut T) } } From e750506603b58fbf73cb48b108678024e2b75881 Mon Sep 17 00:00:00 2001 From: Jonathan Reem Date: Wed, 28 Jan 2015 13:50:54 -0500 Subject: [PATCH 3/4] Buffer: Fix a use-after-free in reallocate. Also provide forward-compatibility for the move to `NonZero>` in the future and tweak some docs. --- src/util/buffer.rs | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/src/util/buffer.rs b/src/util/buffer.rs index 70b3f04..ebe8f19 100644 --- a/src/util/buffer.rs +++ b/src/util/buffer.rs @@ -39,13 +39,15 @@ impl Buffer { /// ``` pub fn new() -> Buffer { Buffer { - buffer: unsafe { NonZero::new(heap::EMPTY as *mut T) }, + buffer: empty() cap: 0 } } /// Create a new buffer with space for cap Ts. /// + /// Unlike `std::rt::heap::allocate`, cap == 0 is allowed. + /// /// ``` /// # use collect::util::buffer::Buffer; /// @@ -63,6 +65,8 @@ impl Buffer { /// Reallocate this buffer to fit a new number of Ts. /// + /// Unlike `std::rt::heap::reallocate`, cap == 0 is allowed. + /// /// ``` /// # use collect::util::buffer::Buffer; /// @@ -73,17 +77,17 @@ impl Buffer { /// assert_eq!(buffer.capacity(), 1024); /// ``` pub fn reallocate(&mut self, cap: usize) { - *self = if self.cap == 0 || cap == 0 { - Buffer::allocate(cap) + if self.cap == 0 || cap == 0 { + // Safe to drop the old buffer because either it never + // allocated or we're getting rid of the allocation. + *self = Buffer::allocate(cap) } else { - Buffer { - buffer: unsafe { - reallocate(self.buffer, - NonZero::new(self.cap), - NonZero::new(cap)) - }, - cap: cap - } + let buffer = mem::replace(&mut self.buffer, empty()); + self.buffer = unsafe { + reallocate(buffer, NonZero::new(self.cap), + NonZero::new(cap)) + }; + self.cap = cap; } } From cc50d2b8f9fffa45be0ce21c844df30f36acae95 Mon Sep 17 00:00:00 2001 From: Jonathan Reem Date: Wed, 28 Jan 2015 22:11:31 -0500 Subject: [PATCH 4/4] Buffer: Fix an incorrect free caused by capacity overflow during reallocation. By only swapping out the buffer and not also changing the capacity to 0 during reallocation, a capacity overflow could cause the drop impl of Buffer to run when the buffer was empty() and the capacity was non-zero, causing an incorrect call to deallocate. We now also swap the capacity to 0, so that when reallocate is called (the site of a possible panic) the buffer is in a consistent empty state. Also, in order to avoid leaking memory on oom (ya kind of useless, but still) reallocate now deallocates the old ptr if reallocation fails. --- src/util/buffer.rs | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/src/util/buffer.rs b/src/util/buffer.rs index ebe8f19..36bfd11 100644 --- a/src/util/buffer.rs +++ b/src/util/buffer.rs @@ -39,7 +39,7 @@ impl Buffer { /// ``` pub fn new() -> Buffer { Buffer { - buffer: empty() + buffer: empty(), cap: 0 } } @@ -82,9 +82,14 @@ impl Buffer { // allocated or we're getting rid of the allocation. *self = Buffer::allocate(cap) } else { + // We need to set the capacity to 0 because if the capacity + // overflows unwinding is triggered, which if we don't + // change the capacity would try to free empty(). + let old_cap = mem::replace(&mut self.cap, 0); let buffer = mem::replace(&mut self.buffer, empty()); + self.buffer = unsafe { - reallocate(buffer, NonZero::new(self.cap), + reallocate(buffer, NonZero::new(old_cap), NonZero::new(cap)) }; self.cap = cap; @@ -150,12 +155,17 @@ unsafe fn reallocate(ptr: NonZero<*mut T>, let new_size = allocation_size::(new_cap); // Reallocate - let ptr = heap::reallocate(*ptr as *mut u8, old_size, new_size, mem::align_of::()); + let new = heap::reallocate(*ptr as *mut u8, old_size, new_size, mem::align_of::()); // Check for allocation failure - if ptr.is_null() { ::alloc::oom() } + if new.is_null() { + // Cleanup old buffer. + deallocate(ptr, old_cap); - NonZero::new(ptr as *mut T) + ::alloc::oom() + } + + NonZero::new(new as *mut T) } unsafe fn deallocate(ptr: NonZero<*mut T>, cap: NonZero) {