Skip to content

Commit

Permalink
Fix UB in GenericMemoryAllocator::deallocate (#2572)
Browse files Browse the repository at this point in the history 
* Fix UB in `GenericMemoryAllocator::deallocate`

* Use the slab allocator
  • Loading branch information
@marc0246
marc0246 authored Oct 10, 2024
1 parent f6bc05d commit 3189169
Showing 1 changed file with 151 additions and 45 deletions.
196 changes: 151 additions & 45 deletions vulkano/src/memory/allocator/mod.rs
View file
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() }
}
}
}

0 comments on commit 3189169

Please sign in to comment.