Skip to content

Commit

Permalink
Resolve deadlock in RustBootServicesAllocatorDxe (#358)
Browse files Browse the repository at this point in the history 
## Description

Fix an issue where rust boot services allocator implementation could
deadlock if memory allocations were attempted at different TPLs.
Deadlock occurs in following scenario:
1. Task running at low TPL initiates an allocation request.
2. While that allocation is in progress and the lock is held, an
interrupt occurs, and a new task executes at a higher TPL
3. new task attempts an allocation. 

Deadlock occurs because the higher TPL task cannot acquire the lock
because it is held by the lower TPL task, and the lower TPL task cannot
make forward progress because it has been interrupted by the higher TPL
task.

To resolve this, this PR updates the allocation implementation to remove
the spinlock. An AtomicPtr is used to give well-ordered access to the
bootservices pointer used as the basis for the allocation
implementation. Other aspects of the implementation (i.e. creation of
the allocation tracker) are already thread-safe or are the
responsibility of the boot services implementation.

PR also contains some minor style cleanup. 

- [ ] Impacts functionality?
- [ ] Impacts security?
- [ ] Breaking change?
- [x] Includes tests?
- RustBootServicesAllocatorDxe unit tests updated to accommodate new
implementation.
- [ ] Includes documentation?

## How This Was Tested

Verified boot on QEMU.

## Integration Instructions

N/A
  • Loading branch information
@joschock @kenlautner
joschock authored and kenlautner committed Dec 14, 2023
1 parent 71fdab6 commit 0aa6322
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 76 deletions.
2 changes: 2 additions & 0 deletions MsCorePkg/Crates/RustBootServicesAllocatorDxe/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,4 +9,6 @@ path = "src/lib.rs"

[dependencies]
r-efi = {workspace=true}

[dev-dependencies]
spin = {workspace=true}
134 changes: 58 additions & 76 deletions MsCorePkg/Crates/RustBootServicesAllocatorDxe/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,19 +34,14 @@
use core::{
alloc::{GlobalAlloc, Layout},
ffi::c_void,
sync::atomic::AtomicPtr,
};

use r_efi::{
efi::{BootServices, Status},
system::BOOT_SERVICES_DATA,
};
use r_efi::efi;

/// Static GLOBAL_ALLOCATOR instance that is marked with the `#[global_allocator]` attribute.
///
/// See [`core::alloc::GlobalAlloc`] for more details on rust global allocators. This allocator must be initialized
/// before use, see [`SpinLockedAllocator::init()`].
#[cfg_attr(not(test), global_allocator)]
pub static GLOBAL_ALLOCATOR: SpinLockedAllocator = SpinLockedAllocator::new();
pub static GLOBAL_ALLOCATOR: BootServicesAllocator = BootServicesAllocator::new();

const ALLOC_TRACKER_SIG: u32 = 0x706F6F6C; //arbitrary sig

Expand All @@ -56,35 +51,26 @@ struct AllocationTracker {
orig_ptr: *mut c_void,
}

// Private unlocked allocator implementation. The public locked allocator delegates to this implementation.
struct BootServicesAllocator {
boot_services: Option<*mut BootServices>,
/// Boot services allocator implementation. Must be initialized with a boot_services pointer before use,
/// see [`BootServicesAllocator::init()`].
pub struct BootServicesAllocator {
boot_services: AtomicPtr<efi::BootServices>,
}

impl BootServicesAllocator {
// Create a new instance. const fn to allow static initialization.
const fn new() -> Self {
BootServicesAllocator { boot_services: None }
}

// initialize the allocator by providing a pointer to the global boot services table.
fn init(&mut self, boot_services: *mut BootServices) {
self.boot_services = Some(boot_services);
BootServicesAllocator { boot_services: AtomicPtr::new(core::ptr::null_mut()) }
}

// implement allocation using EFI boot services AllocatePool() call.
fn boot_services_alloc(&self, layout: Layout) -> *mut u8 {
//bail early if not initialized.
let Some(bs_ptr) = self.boot_services else { return core::ptr::null_mut() };

let bs = unsafe { bs_ptr.as_mut().expect("Boot Services pointer is null.") };

fn boot_services_alloc(&self, layout: Layout, boot_services: &efi::BootServices) -> *mut u8 {
match layout.align() {
0..=8 => {
//allocate the pointer directly since UEFI pool allocations are 8-byte aligned already.
let mut ptr: *mut c_void = core::ptr::null_mut();
match (bs.allocate_pool)(BOOT_SERVICES_DATA, layout.size(), core::ptr::addr_of_mut!(ptr)) {
Status::SUCCESS => ptr as *mut u8,
match (boot_services.allocate_pool)(efi::BOOT_SERVICES_DATA, layout.size(), core::ptr::addr_of_mut!(ptr)) {
efi::Status::SUCCESS => ptr as *mut u8,
_ => core::ptr::null_mut(),
}
}
Expand All @@ -98,8 +84,12 @@ impl BootServicesAllocator {
let expanded_size = expanded_layout.size() + expanded_layout.align();

let mut orig_ptr: *mut c_void = core::ptr::null_mut();
let final_ptr = match (bs.allocate_pool)(BOOT_SERVICES_DATA, expanded_size, core::ptr::addr_of_mut!(orig_ptr)) {
Status::SUCCESS => orig_ptr as *mut u8,
let final_ptr = match (boot_services.allocate_pool)(
efi::BOOT_SERVICES_DATA,
expanded_size,
core::ptr::addr_of_mut!(orig_ptr),
) {
efi::Status::SUCCESS => orig_ptr as *mut u8,
_ => return core::ptr::null_mut(),
};

Expand All @@ -120,16 +110,11 @@ impl BootServicesAllocator {
}

// implement dealloc (free) using EFI boot services FreePool() call.
fn boot_services_dealloc(&self, ptr: *mut u8, layout: Layout) {
//bail early if not initialized.
let Some(bs_ptr) = self.boot_services else { return };

let bs = unsafe { bs_ptr.as_mut().expect("Boot Services pointer is null.") };

fn boot_services_dealloc(&self, boot_services: &efi::BootServices, ptr: *mut u8, layout: Layout) {
match layout.align() {
0..=8 => {
//pointer was allocated directly, so free it directly.
let _ = (bs.free_pool)(ptr as *mut c_void);
let _ = (boot_services.free_pool)(ptr as *mut c_void);
}
_ => {
//pointer was potentially adjusted for alignment. Recover tracking structure to retrieve the original
Expand All @@ -142,47 +127,39 @@ impl BootServicesAllocator {
ptr.add(tracking_offset).cast::<AllocationTracker>().as_mut().expect("tracking pointer is invalid")
};
debug_assert_eq!(tracker.signature, ALLOC_TRACKER_SIG);
let _ = (bs.free_pool)(tracker.orig_ptr);
let _ = (boot_services.free_pool)(tracker.orig_ptr);
}
}
}
}

/// A spin-locked allocator implementation.
///
/// Provides a locked allocator instance (using [`spin::Mutex`]) that is suitable for use as a
/// [`core::alloc::GlobalAlloc`].
pub struct SpinLockedAllocator {
inner: spin::Mutex<BootServicesAllocator>,
}

impl SpinLockedAllocator {
// Create a new instance. const fn to allow static initialization.
const fn new() -> Self {
SpinLockedAllocator { inner: spin::Mutex::new(BootServicesAllocator::new()) }
}

/// Initialize the allocator.
///
/// This routine initializes the allocator by providing a pointer to the global EFI Boot Services table that will
/// supply pointers to the AllocatePool() and FreePool() primitives that implement memory allocation.
pub fn init(&self, boot_services: *mut BootServices) {
self.inner.lock().init(boot_services);
/// initializes the allocator instance with a pointer to the UEFI Boot Services table.
pub fn init(&self, boot_services: *mut efi::BootServices) {
self.boot_services.store(boot_services, core::sync::atomic::Ordering::SeqCst);
}
}

unsafe impl GlobalAlloc for SpinLockedAllocator {
unsafe impl GlobalAlloc for BootServicesAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
self.inner.lock().boot_services_alloc(layout)
let bs_ptr = self.boot_services.load(core::sync::atomic::Ordering::SeqCst);
if let Some(boot_services) = unsafe { bs_ptr.as_ref() } {
self.boot_services_alloc(layout, boot_services)
} else {
panic!("Attempted allocation on uninitialized allocator")
}
}

unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
self.inner.lock().boot_services_dealloc(ptr, layout)
let bs_ptr = self.boot_services.load(core::sync::atomic::Ordering::SeqCst);
if let Some(boot_services) = unsafe { bs_ptr.as_ref() } {
self.boot_services_dealloc(boot_services, ptr, layout)
} else {
panic!("Attempted deallocation on uninitialized allocator")
}
}
}

unsafe impl Sync for SpinLockedAllocator {}
unsafe impl Send for SpinLockedAllocator {}
unsafe impl Sync for BootServicesAllocator {}
unsafe impl Send for BootServicesAllocator {}

#[cfg(test)]
mod tests {
Expand All @@ -195,22 +172,19 @@ mod tests {
};
use std::alloc::System;

use r_efi::{
efi::Status,
system::{BootServices, BOOT_SERVICES_DATA},
};
use r_efi::efi;
use std::collections::BTreeMap;

use crate::{AllocationTracker, SpinLockedAllocator, ALLOC_TRACKER_SIG};
use crate::{AllocationTracker, BootServicesAllocator, ALLOC_TRACKER_SIG};

static ALLOCATION_TRACKER: spin::Mutex<BTreeMap<usize, Layout>> = spin::Mutex::new(BTreeMap::new());

extern "efiapi" fn mock_allocate_pool(
pool_type: r_efi::system::MemoryType,
pool_type: efi::MemoryType,
size: usize,
buffer: *mut *mut c_void,
) -> Status {
assert_eq!(pool_type, BOOT_SERVICES_DATA);
) -> efi::Status {
assert_eq!(pool_type, efi::BOOT_SERVICES_DATA);

unsafe {
let layout = Layout::from_size_align(size, 8).unwrap();
Expand All @@ -220,29 +194,37 @@ mod tests {
assert!(existing_key.is_none());
}

Status::SUCCESS
efi::Status::SUCCESS
}

extern "efiapi" fn mock_free_pool(buffer: *mut c_void) -> Status {
extern "efiapi" fn mock_free_pool(buffer: *mut c_void) -> efi::Status {
let layout = ALLOCATION_TRACKER.lock().remove(&(buffer as usize)).expect("freeing an un-allocated pointer");
unsafe {
System.dealloc(buffer as *mut u8, layout);
}

Status::SUCCESS
efi::Status::SUCCESS
}

fn mock_boot_services() -> BootServices {
extern "efiapi" fn mock_raise_tpl(_new_tpl: efi::Tpl) -> efi::Tpl {
efi::TPL_APPLICATION
}

extern "efiapi" fn mock_restore_tpl(_new_tpl: efi::Tpl) {}

fn mock_boot_services() -> efi::BootServices {
let boot_services = MaybeUninit::zeroed();
let mut boot_services: BootServices = unsafe { boot_services.assume_init() };
let mut boot_services: efi::BootServices = unsafe { boot_services.assume_init() };
boot_services.allocate_pool = mock_allocate_pool;
boot_services.free_pool = mock_free_pool;
boot_services.raise_tpl = mock_raise_tpl;
boot_services.restore_tpl = mock_restore_tpl;
boot_services
}

#[test]
fn basic_alloc_and_dealloc() {
static ALLOCATOR: SpinLockedAllocator = SpinLockedAllocator::new();
static ALLOCATOR: BootServicesAllocator = BootServicesAllocator::new();
ALLOCATOR.init(&mut mock_boot_services());

let layout = Layout::from_size_align(0x40, 0x8).unwrap();
Expand All @@ -256,7 +238,7 @@ mod tests {

#[test]
fn big_alignment_should_allocate_tracking_structure() {
static ALLOCATOR: SpinLockedAllocator = SpinLockedAllocator::new();
static ALLOCATOR: BootServicesAllocator = BootServicesAllocator::new();
ALLOCATOR.init(&mut mock_boot_services());

let layout = Layout::from_size_align(0x40, 0x1000).unwrap();
Expand Down

0 comments on commit 0aa6322

Please sign in to comment.