A safer more flexible idiom for guarding boot service lifetimes

[dfoxfranke]

This issue is a tour of some scaffolding I wrote for being able to safely access a SystemTable<Boot> as a global variable. Incorporating this would make it possible to turn Logger::new into a safe call and eliminate any requirement to disable it before exiting boot services. Give this a look and let me know if you'd like me to turn it into a PR.

This code uses a couple nightly features but they're not especially needed:

#![feature(error_in_core)]
#![feature(core_intrinsics)]

use core::cell::UnsafeCell;
use core::ops::{Deref, DerefMut};
use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use uefi::prelude::*;

GlobalSystemTable stores a SystemTable<Boot> plus some bookkeeping entries for reference counting:

struct GlobalSystemTable {
    cell: UnsafeCell<Option<SystemTable<Boot>>>,
    refcnt: AtomicUsize,
    stdin_borrowed: AtomicBool,
    stdout_borrowed: AtomicBool,
    stderr_borrowed: AtomicBool,
}

We exploit the knowledge that the boot-services environment is single-threaded in order to implement Sync so we can keep it as a global.

unsafe impl Sync for GlobalSystemTable {}

static GLOBAL_SYSTEM_TABLE: GlobalSystemTable = GlobalSystemTable::new();

impl GlobalSystemTable {
    const fn new() -> GlobalSystemTable {
        GlobalSystemTable {
            cell: UnsafeCell::new(None),
            refcnt: AtomicUsize::new(0),
            stdin_borrowed: AtomicBool::new(false),
            stdout_borrowed: AtomicBool::new(false),
            stderr_borrowed: AtomicBool::new(false),
        }
    }

The first thing the EFI entry point should do is initialize it:

    unsafe fn init(&self, ptr: *mut c_void) {
        *self.cell.get() = SystemTable::from_ptr(ptr);
    }

Just like the Rc and Arc types in the standard library, the inc_refcnt helper method checks for overflow in order to be safe in the ridiculous case that code calls mem::forget a whole bunch in order to make it overflow without first running out of memory. This is the only use feature(core_intrinsics), so if this ever becomes the last blocker to being able to compile uefi on stable then just replace it with an infinite loop or whatnot, because it's just a stupid pathological case that never happens and nobody cares about.

We use relaxed memory ordering, because since we're single-threaded we don't really need an atomic at all and a Cell<usize> would be fine, but AtomicUsize provides more convenient methods.

    fn inc_refcnt(&self) {
        let old = self.refcnt.fetch_add(1, Ordering::Relaxed);
        if old == usize::MAX {
            core::intrinsics::abort()
        }
    }

The borrow method returns a Guard object which implements Deref and will decrement the reference count when dropped.

    fn borrow(&self) -> Result<Guard<'_>, GlobalSystemTableError> {
        unsafe {
            if let Some(table) = (*self.cell.get()).as_ref() {
                self.inc_refcnt();
                Ok(Guard {
                    target: table,
                    refcnt: &self.refcnt,
                })
            } else {
                Err(GlobalSystemTableError::NotInitialized)
            }
        }
    }

The stdin/stdout/stderr methods similarly return guard objects which implement DerefMut and whose Drop implementation clears the stdfoo_borrowed bit as well as decrementing the reference count. You can one stdin borrow, one stdout borrow, one stderr borrow, and N immutable borrows out all at once. This should be safe because these all operate on pairwise-disjoint sets of fields of the underlying SystemTable structure.

    fn stdin(&self) -> Result<InputGuard<'_>, GlobalSystemTableError> {
        unsafe {
            if let Some(table) = (*self.cell.get()).as_mut() {
                if self.stdin_borrowed.swap(true, Ordering::Relaxed) {
                    Err(GlobalSystemTableError::InUse)
                } else {
                    self.inc_refcnt();
                    Ok(IoGuard {
                        target: table.stdin(),
                        refcnt: &self.refcnt,
                        borrow: &self.stdin_borrowed,
                    })
                }
            } else {
                Err(GlobalSystemTableError::NotInitialized)
            }
        }
    }

    fn stdout(&self) -> Result<OutputGuard<'_>, GlobalSystemTableError> {
        unsafe {
            if let Some(table) = (*self.cell.get()).as_mut() {
                if self.stdout_borrowed.swap(true, Ordering::Relaxed) {
                    Err(GlobalSystemTableError::InUse)
                } else {
                    self.inc_refcnt();
                    Ok(IoGuard {
                        target: table.stdout(),
                        refcnt: &self.refcnt,
                        borrow: &self.stdout_borrowed,
                    })
                }
            } else {
                Err(GlobalSystemTableError::NotInitialized)
            }
        }
    }

    fn stderr(&self) -> Result<OutputGuard<'_>, GlobalSystemTableError> {
        unsafe {
            if let Some(table) = (*self.cell.get()).as_mut() {
                if self.stderr_borrowed.swap(true, Ordering::Relaxed) {
                    Err(GlobalSystemTableError::InUse)
                } else {
                    self.inc_refcnt();
                    Ok(IoGuard {
                        target: table.stdout(),
                        refcnt: &self.refcnt,
                        borrow: &self.stdout_borrowed,
                    })
                }
            } else {
                Err(GlobalSystemTableError::NotInitialized)
            }
        }
    }

The deinit method checks that the reference count is zero, and then deinitializes the global and hands back an owned SystemTable<Boot>.

    fn deinit(&self) -> Result<SystemTable<Boot>, GlobalSystemTableError> {
        unsafe {
            if self.refcnt.load(Ordering::Relaxed) != 0 {
                Err(GlobalSystemTableError::InUse)
            } else {
                match core::ptr::replace(self.cell.get(), None) {
                    None => Err(GlobalSystemTableError::NotInitialized),
                    Some(table) => Ok(table),
                }
            }
        }
    }
} // impl GlobalSystemTable

For completeness, here are the (unremarkable) implementations of Guard, IoGuard, and GlobalSystemTableError:

#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum GlobalSystemTableError {
    NotInitialized,
    InUse,
}

impl core::fmt::Display for GlobalSystemTableError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::NotInitialized => "Not initialized".fmt(f),
            Self::InUse => "In use".fmt(f),
        }
    }
}

impl core::error::Error for GlobalSystemTableError {}

struct Guard<'a> {
    target: &'a SystemTable<Boot>,
    refcnt: &'a AtomicUsize,
}

impl Deref for Guard<'_> {
    type Target = SystemTable<Boot>;

    fn deref(&self) -> &Self::Target {
        self.target
    }
}

impl Drop for Guard<'_> {
    fn drop(&mut self) {
        self.refcnt.fetch_sub(1, Ordering::Relaxed);
    }
}
struct IoGuard<'a, Target> {
    target: &'a mut Target,
    refcnt: &'a AtomicUsize,
    borrow: &'a AtomicBool,
}

type InputGuard<'a> = IoGuard<'a, Input>;
type OutputGuard<'a> = IoGuard<'a, Output<'a>>;

impl<Target> Deref for IoGuard<'_, Target> {
    type Target = Target;

    fn deref(&self) -> &Self::Target {
        self.target
    }
}

impl<Target> DerefMut for IoGuard<'_, Target> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.target
    }
}

impl<Target> Drop for IoGuard<'_, Target> {
    fn drop(&mut self) {
        self.refcnt.fetch_sub(1, Ordering::Relaxed);
        self.borrow.store(false, Ordering::Relaxed);
    }
}

Finally, as an example, here's how I used this code to implement a completely safe panic handler. Logger could be written the same way.

#[panic_handler]
fn panic_handler(info: &core::panic::PanicInfo) -> ! {
    if let Ok(mut stderr) = GLOBAL_SYSTEM_TABLE.stderr() {
        let _ = match (info.location(), info.message()) {
            (None, None) => writeln!(stderr, "PANIC"),
            (None, Some(message)) => writeln!(stderr, "PANIC: {}", message),
            (Some(location), None) => writeln!(stderr, "PANIC at {}", location),
            (Some(location), Some(message)) => {
                writeln!(stderr, "PANIC at {}: {}", location, message)
            }
        };
    }
    loop {
        x86_64::instructions::hlt();
    }
}

Another good thing to add if you merge this code would be a #[gentry] attribute macro, which you can put on a main function which takes no arguments. It would initialize GLOBAL_SYSTEM_TABLE as well as another global GLOBAL_IMAGE_HANDLE which would be a OnceCell<Handle>.

[nicholasbishop]

Thanks for writing this up! If you have time, I think turning this into a PR would be helpful to continue discussion.