Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Check for overflow in DroplessArena and align returned memory #73237

Merged
merged 2 commits into from
Jun 16, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 41 additions & 31 deletions src/librustc_arena/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -333,13 +333,6 @@ impl Default for DroplessArena {
}

impl DroplessArena {
#[inline]
fn align(&self, align: usize) {
let final_address = ((self.ptr.get() as usize) + align - 1) & !(align - 1);
self.ptr.set(final_address as *mut u8);
assert!(self.ptr <= self.end);
}

#[inline(never)]
#[cold]
fn grow(&self, additional: usize) {
Expand Down Expand Up @@ -370,30 +363,50 @@ impl DroplessArena {
}
}

/// Allocates a byte slice with specified size and alignment from the
/// current memory chunk. Returns `None` if there is no free space left to
/// satisfy the request.
#[inline]
pub fn alloc_raw(&self, bytes: usize, align: usize) -> &mut [u8] {
unsafe {
assert!(bytes != 0);

self.align(align);
fn alloc_raw_without_grow(&self, bytes: usize, align: usize) -> Option<*mut u8> {
let ptr = self.ptr.get() as usize;
let end = self.end.get() as usize;
// The allocation request fits into the current chunk iff:
//
// let aligned = align_to(ptr, align);
// ptr <= aligned && aligned + bytes <= end
//
// Except that we work with fixed width integers and need to be careful
// about potential overflow in the calcuation. If the overflow does
// happen, then we definitely don't have enough free and need to grow
// the arena.
let aligned = ptr.checked_add(align - 1)? & !(align - 1);
let new_ptr = aligned.checked_add(bytes)?;
if new_ptr <= end {
self.ptr.set(new_ptr as *mut u8);
Some(aligned as *mut u8)
} else {
None
}
}

let future_end = intrinsics::arith_offset(self.ptr.get(), bytes as isize);
if (future_end as *mut u8) > self.end.get() {
self.grow(bytes);
#[inline]
pub fn alloc_raw(&self, bytes: usize, align: usize) -> *mut u8 {
assert!(bytes != 0);
loop {
if let Some(a) = self.alloc_raw_without_grow(bytes, align) {
break a;
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens here on overflow? It'll call self.grow(bytes). Then what will happen -- panic, success? I think the possibilities need to be clearer here. In particular, alloc_raw_fast() returns None in two cases: "overflow", and "need to grow". I feel like those two cases need to be distinguished.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From perspective of code that tried to use the current chunk the situation is the same: there is no capacity left and arena needs to grow. Overflow does not necessarily preclude later success if say the chunk had merely a high address, requested capacity was large, and the calculation overflowed.

The code then delegates the responsibility to handle the situation to grow and subsequently to RawVec. Under assumption that grow doesn't hit checked_mul(2).unwrap() (this could be saturated to isize::MAX to be a little bit friendlier to 32-bit systems), the allocation is performed by RawVec which panics when requested capacity exceeds isize::MAX (on 32-bit systems) and aborts on OOM.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the explanation. Could you put that information into a comment? :)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now with additional comments.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for that. Sorry to nitpick this further, but you've added a comment explaining the high-level behaviour, which is fairly obvious, while not adding a comment distinguishing the overflow vs. need-to-grow cases, which is less obvious.

Can you add a comment inside the function explaining those two cases, much like your comment above? I specifically would like to see explanation of the "high address + large capacity" case, in particular. Thanks.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added a different comment explaining the abstract check we want to perform.
In my view the overflow is the same as the need-to-grow case, which is why I
didn't want to comment about the distinction. I just don't see any.


let ptr = self.ptr.get();
// Set the pointer past ourselves
self.ptr.set(intrinsics::arith_offset(self.ptr.get(), bytes as isize) as *mut u8);
slice::from_raw_parts_mut(ptr, bytes)
// No free space left. Allocate a new chunk to satisfy the request.
// On failure the grow will panic or abort.
self.grow(bytes);
}
}

#[inline]
pub fn alloc<T>(&self, object: T) -> &mut T {
assert!(!mem::needs_drop::<T>());

let mem = self.alloc_raw(mem::size_of::<T>(), mem::align_of::<T>()) as *mut _ as *mut T;
let mem = self.alloc_raw(mem::size_of::<T>(), mem::align_of::<T>()) as *mut T;

unsafe {
// Write into uninitialized memory.
Expand All @@ -418,13 +431,11 @@ impl DroplessArena {
assert!(mem::size_of::<T>() != 0);
assert!(!slice.is_empty());

let mem = self.alloc_raw(slice.len() * mem::size_of::<T>(), mem::align_of::<T>()) as *mut _
as *mut T;
let mem = self.alloc_raw(slice.len() * mem::size_of::<T>(), mem::align_of::<T>()) as *mut T;

unsafe {
let arena_slice = slice::from_raw_parts_mut(mem, slice.len());
arena_slice.copy_from_slice(slice);
arena_slice
mem.copy_from_nonoverlapping(slice.as_ptr(), slice.len());
slice::from_raw_parts_mut(mem, slice.len())
}
}

Expand Down Expand Up @@ -467,7 +478,7 @@ impl DroplessArena {
return &mut [];
}
let size = len.checked_mul(mem::size_of::<T>()).unwrap();
let mem = self.alloc_raw(size, mem::align_of::<T>()) as *mut _ as *mut T;
let mem = self.alloc_raw(size, mem::align_of::<T>()) as *mut T;
unsafe { self.write_from_iter(iter, len, mem) }
}
(_, _) => {
Expand All @@ -482,7 +493,7 @@ impl DroplessArena {
let len = vec.len();
let start_ptr = self
.alloc_raw(len * mem::size_of::<T>(), mem::align_of::<T>())
as *mut _ as *mut T;
as *mut T;
vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
vec.set_len(0);
slice::from_raw_parts_mut(start_ptr, len)
Expand Down Expand Up @@ -526,8 +537,7 @@ pub struct DropArena {
impl DropArena {
#[inline]
pub unsafe fn alloc<T>(&self, object: T) -> &mut T {
let mem =
self.arena.alloc_raw(mem::size_of::<T>(), mem::align_of::<T>()) as *mut _ as *mut T;
let mem = self.arena.alloc_raw(mem::size_of::<T>(), mem::align_of::<T>()) as *mut T;
// Write into uninitialized memory.
ptr::write(mem, object);
let result = &mut *mem;
Expand All @@ -550,7 +560,7 @@ impl DropArena {
let start_ptr = self
.arena
.alloc_raw(len.checked_mul(mem::size_of::<T>()).unwrap(), mem::align_of::<T>())
as *mut _ as *mut T;
as *mut T;

let mut destructors = self.destructors.borrow_mut();
// Reserve space for the destructors so we can't panic while adding them
Expand Down
2 changes: 1 addition & 1 deletion src/librustc_middle/ty/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ impl<T: Copy> List<T> {
.dropless
.alloc_raw(size, cmp::max(mem::align_of::<T>(), mem::align_of::<usize>()));
unsafe {
let result = &mut *(mem.as_mut_ptr() as *mut List<T>);
let result = &mut *(mem as *mut List<T>);
// Write the length
result.len = slice.len();

Expand Down