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.

Sign up for GitHub

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

Jump to bottom

Fix UB in GenericMemoryAllocator::deallocate #2572

Merged
merged 2 commits into from
Oct 10, 2024
Merged

Fix UB in GenericMemoryAllocator::deallocate #2572

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
767cd27
Fix UB in `GenericMemoryAllocator::deallocate`
marc0246 Oct 7, 2024
0db670a
Use the slab allocator
marc0246 Oct 7, 2024
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
196 changes: 151 additions & 45 deletions vulkano/src/memory/allocator/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,10 +206,7 @@
//! [`Rc`]: std::rc::Rc
//! [region]: Suballocator#regions

mod layout;
pub mod suballocator;

use self::{array_vec::ArrayVec, suballocator::Region};
use self::{aliasable_box::AliasableBox, array_vec::ArrayVec, suballocator::Region};
pub use self::{
layout::DeviceLayout,
suballocator::{
Expand All @@ -229,6 +226,7 @@ use crate::{
};
use ash::vk::MAX_MEMORY_TYPES;
use parking_lot::{Mutex, MutexGuard};
use slabbin::SlabAllocator;
use std::{
error::Error,
fmt::{Debug, Display, Error as FmtError, Formatter},
Expand All @@ -239,6 +237,9 @@ use std::{
sync::Arc,
};

mod layout;
pub mod suballocator;

/// General-purpose memory allocators which allocate from any memory type dynamically as needed.
///
/// # Safety
Expand Down Expand Up @@ -952,7 +953,7 @@ impl<S> GenericMemoryAllocator<S> {
// This is a false-positive, we only use this const for static initialization.
#[allow(clippy::declare_interior_mutable_const)]
const EMPTY_POOL: DeviceMemoryPool<S> = DeviceMemoryPool {
blocks: Mutex::new(Vec::new()),
blocks: Mutex::new(DeviceMemoryBlockVec::new()),
property_flags: MemoryPropertyFlags::empty(),
atom_size: DeviceAlignment::MIN,
block_size: 0,
Expand Down Expand Up @@ -1179,13 +1180,14 @@ unsafe impl<S: Suballocator + Send + 'static> MemoryAllocator for GenericMemoryA

layout = layout.align_to(pool.atom_size).unwrap();

let mut blocks = pool.blocks.lock();
let blocks = &mut *pool.blocks.lock();
let vec = &mut blocks.vec;

// TODO: Incremental sorting
blocks.sort_by_key(|block| block.free_size());
let (Ok(idx) | Err(idx)) = blocks.binary_search_by_key(&size, |block| block.free_size());
vec.sort_by_key(|block| block.free_size());
let (Ok(idx) | Err(idx)) = vec.binary_search_by_key(&size, |block| block.free_size());

for block in &mut blocks[idx..] {
for block in &mut vec[idx..] {
if let Ok(allocation) =
block.allocate(layout, allocation_type, self.buffer_image_granularity)
{
Expand Down Expand Up @@ -1216,7 +1218,7 @@ unsafe impl<S: Suballocator + Send + 'static> MemoryAllocator for GenericMemoryA
export_handle_types,
) {
Ok(device_memory) => {
break DeviceMemoryBlock::new(device_memory);
break DeviceMemoryBlock::new(device_memory, &blocks.block_allocator);
}
// Retry up to 3 times, halving the allocation size each time so long as the
// resulting size is still large enough.
Expand All @@ -1230,8 +1232,8 @@ unsafe impl<S: Suballocator + Send + 'static> MemoryAllocator for GenericMemoryA
}
};

blocks.push(block);
let block = blocks.last_mut().unwrap();
vec.push(block);
let block = vec.last_mut().unwrap();

match block.allocate(layout, allocation_type, self.buffer_image_granularity) {
Ok(allocation) => Ok(allocation),
Expand Down Expand Up @@ -1454,15 +1456,16 @@ unsafe impl<S: Suballocator + Send + 'static> MemoryAllocator for GenericMemoryA
unsafe fn deallocate(&self, allocation: MemoryAlloc) {
if let Some(suballocation) = allocation.suballocation {
let memory_type_index = allocation.device_memory.memory_type_index();
let pool = self.pools[memory_type_index as usize].blocks.lock();
let blocks = self.pools[memory_type_index as usize].blocks.lock();
let vec = &blocks.vec;
let block_ptr = allocation
.allocation_handle
.0
.as_ptr()
.cast::<DeviceMemoryBlock<S>>();

// TODO: Maybe do a similar check for dedicated blocks.
debug_assert!(
pool.iter().any(|block| ptr::addr_of!(**block) == block_ptr),
vec.iter().any(|block| ptr::addr_of!(**block) == block_ptr),
"attempted to deallocate a memory block that does not belong to this allocator",
);

Expand All @@ -1477,8 +1480,6 @@ unsafe impl<S: Suballocator + Send + 'static> MemoryAllocator for GenericMemoryA
// SAFETY: The caller must guarantee that `allocation` refers to a currently allocated
// allocation of `self`.
block.deallocate(suballocation);

drop(pool);
}
}
}
Expand Down Expand Up @@ -1546,7 +1547,7 @@ unsafe impl<S> DeviceOwned for GenericMemoryAllocator<S> {
/// A pool of [`DeviceMemory`] blocks within [`GenericMemoryAllocator`], specific to a memory type.
#[derive(Debug)]
pub struct DeviceMemoryPool<S> {
blocks: Mutex<Vec<Box<DeviceMemoryBlock<S>>>>,
blocks: Mutex<DeviceMemoryBlockVec<S>>,
// This is cached here for faster access, so we don't need to hop through 3 pointers.
property_flags: MemoryPropertyFlags,
atom_size: DeviceAlignment,
Expand All @@ -1564,13 +1565,38 @@ impl<S> DeviceMemoryPool<S> {
#[inline]
pub fn blocks(&self) -> DeviceMemoryBlocks<'_, S> {
DeviceMemoryBlocks {
inner: MutexGuard::leak(self.blocks.lock()).iter(),
inner: MutexGuard::leak(self.blocks.lock()).vec.iter(),
// SAFETY: We have just locked the pool above.
_guard: unsafe { DeviceMemoryPoolGuard::new(self) },
}
}
}

impl<S> Drop for DeviceMemoryPool<S> {
fn drop(&mut self) {
let blocks = self.blocks.get_mut();

for block in &mut blocks.vec {
unsafe { AliasableBox::drop(block, &blocks.block_allocator) };
}
}
}

#[derive(Debug)]
struct DeviceMemoryBlockVec<S> {
vec: Vec<AliasableBox<DeviceMemoryBlock<S>>>,
block_allocator: SlabAllocator<DeviceMemoryBlock<S>>,
}

impl<S> DeviceMemoryBlockVec<S> {
const fn new() -> Self {
DeviceMemoryBlockVec {
vec: Vec::new(),
block_allocator: SlabAllocator::new(32),
}
}
}

/// A [`DeviceMemory`] block within a [`DeviceMemoryPool`].
#[derive(Debug)]
pub struct DeviceMemoryBlock<S> {
Expand All @@ -1580,36 +1606,23 @@ pub struct DeviceMemoryBlock<S> {
}

impl<S: Suballocator> DeviceMemoryBlock<S> {
fn new(device_memory: Arc<DeviceMemory>) -> Box<Self> {
fn new(
device_memory: Arc<DeviceMemory>,
block_allocator: &SlabAllocator<Self>,
) -> AliasableBox<Self> {
let suballocator = S::new(
Region::new(0, device_memory.allocation_size())
.expect("we somehow managed to allocate more than `DeviceLayout::MAX_SIZE` bytes"),
);

Box::new(DeviceMemoryBlock {
device_memory,
suballocator,
allocation_count: 0,
})
}

fn allocate(
&mut self,
layout: DeviceLayout,
allocation_type: AllocationType,
buffer_image_granularity: DeviceAlignment,
) -> Result<MemoryAlloc, SuballocatorError> {
let suballocation =
self.suballocator
.allocate(layout, allocation_type, buffer_image_granularity)?;

self.allocation_count += 1;

Ok(MemoryAlloc {
device_memory: self.device_memory.clone(),
suballocation: Some(suballocation),
allocation_handle: AllocationHandle::from_ptr(<*mut _>::cast(self)),
})
AliasableBox::new(
DeviceMemoryBlock {
device_memory,
suballocator,
allocation_count: 0,
},
block_allocator,
)
}

unsafe fn deallocate(&mut self, suballocation: Suballocation) {
Expand Down Expand Up @@ -1646,9 +1659,34 @@ impl<S: Suballocator> DeviceMemoryBlock<S> {
}
}

impl<S: Suballocator> AliasableBox<DeviceMemoryBlock<S>> {
fn allocate(
&mut self,
layout: DeviceLayout,
allocation_type: AllocationType,
buffer_image_granularity: DeviceAlignment,
) -> Result<MemoryAlloc, SuballocatorError> {
let suballocation =
self.suballocator
.allocate(layout, allocation_type, buffer_image_granularity)?;

self.allocation_count += 1;

// It is paramount to soundness that the pointer we give out doesn't go through `DerefMut`,
// as such a pointer would become invalidated when another allocation is made.
let ptr = AliasableBox::as_mut_ptr(self);

Ok(MemoryAlloc {
device_memory: self.device_memory.clone(),
suballocation: Some(suballocation),
allocation_handle: AllocationHandle::from_ptr(<*mut _>::cast(ptr)),
})
}
}

/// An iterator over the [`DeviceMemoryBlock`]s within a [`DeviceMemoryPool`].
pub struct DeviceMemoryBlocks<'a, S> {
inner: slice::Iter<'a, Box<DeviceMemoryBlock<S>>>,
inner: slice::Iter<'a, AliasableBox<DeviceMemoryBlock<S>>>,
_guard: DeviceMemoryPoolGuard<'a, S>,
}

Expand Down Expand Up @@ -1887,3 +1925,71 @@ mod array_vec {
}
}
}

mod aliasable_box {
use slabbin::SlabAllocator;
use std::{
fmt,
marker::PhantomData,
ops::{Deref, DerefMut},
panic::{RefUnwindSafe, UnwindSafe},
ptr::NonNull,
};

pub struct AliasableBox<T> {
ptr: NonNull<T>,
marker: PhantomData<T>,
}

unsafe impl<T: Send> Send for AliasableBox<T> {}
unsafe impl<T: Sync> Sync for AliasableBox<T> {}

impl<T: UnwindSafe> UnwindSafe for AliasableBox<T> {}
impl<T: RefUnwindSafe> RefUnwindSafe for AliasableBox<T> {}

impl<T> Unpin for AliasableBox<T> {}

impl<T> AliasableBox<T> {
pub fn new(value: T, allocator: &SlabAllocator<T>) -> Self {
let ptr = allocator.allocate();

unsafe { ptr.as_ptr().write(value) };

AliasableBox {
ptr,
marker: PhantomData,
}
}

pub fn as_mut_ptr(this: &mut Self) -> *mut T {
this.ptr.as_ptr()
}

pub unsafe fn drop(this: &mut Self, allocator: &SlabAllocator<T>) {
unsafe { this.ptr.as_ptr().drop_in_place() };
unsafe { allocator.deallocate(this.ptr) };
}
}

impl<T: fmt::Debug> fmt::Debug for AliasableBox<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}

impl<T> Deref for AliasableBox<T> {
type Target = T;

#[inline]
fn deref(&self) -> &Self::Target {
unsafe { self.ptr.as_ref() }
}
}

impl<T> DerefMut for AliasableBox<T> {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { self.ptr.as_mut() }
}
}
}