From 2174bc9c1f2cb0400ed9072659fc40b8d2342444 Mon Sep 17 00:00:00 2001 From: Jonas Kruckenberg Date: Sun, 29 Dec 2024 00:07:24 +0100 Subject: [PATCH 1/5] fix(kernel): prevent `setjmp`/`longjmp` miscompilations through `call_with_setjmp` --- kernel/src/arch/riscv64/mod.rs | 2 +- kernel/src/arch/riscv64/setjmp_longjmp.rs | 132 ++++++++++++++++++++-- 2 files changed, 121 insertions(+), 13 deletions(-) diff --git a/kernel/src/arch/riscv64/mod.rs b/kernel/src/arch/riscv64/mod.rs index 5cd724ef..1483786f 100644 --- a/kernel/src/arch/riscv64/mod.rs +++ b/kernel/src/arch/riscv64/mod.rs @@ -6,7 +6,7 @@ use core::arch::asm; use riscv::sstatus::FS; use riscv::{interrupt, sie, sstatus}; -pub use setjmp_longjmp::{longjmp, setjmp, JumpBuf}; +pub use setjmp_longjmp::{call_with_setjmp, longjmp, setjmp, JmpBuf, JmpBufStruct}; pub fn finish_hart_init() { unsafe { diff --git a/kernel/src/arch/riscv64/setjmp_longjmp.rs b/kernel/src/arch/riscv64/setjmp_longjmp.rs index 673de6a7..085d6a01 100644 --- a/kernel/src/arch/riscv64/setjmp_longjmp.rs +++ b/kernel/src/arch/riscv64/setjmp_longjmp.rs @@ -36,32 +36,44 @@ //! result. In a nested calls scenario (e.g. host->wasm->host->wasm) it is therefore up to each host function //! to propagate the trap and each host function therefore gets to clean up all its resources. -use core::arch::naked_asm; +use alloc::boxed::Box; +use core::any::Any; +use core::arch::{asm, naked_asm}; +use core::marker::{PhantomData, PhantomPinned}; +use core::mem::{ManuallyDrop, MaybeUninit}; +use core::panic::UnwindSafe; use core::ptr; +use core::ptr::addr_of_mut; /// A store for the register state used by `setjmp` and `longjmp`. /// /// In essence this marks a "checkpoint" in the program's execution that can be returned to later. #[repr(C)] #[derive(Clone, Debug, Default)] -pub struct JumpBuf { +pub struct JmpBufStruct { pc: usize, s: [usize; 12], sp: usize, fs: [usize; 12], + _neither_send_nor_sync: PhantomData<*const u8>, + _not_unpin: PhantomPinned, } -impl JumpBuf { +impl JmpBufStruct { pub const fn new() -> Self { Self { pc: 0, sp: 0, s: [0; 12], fs: [0; 12], + _neither_send_nor_sync: PhantomData, + _not_unpin: PhantomPinned, } } } +pub type JmpBuf = *const JmpBufStruct; + /// Helper macro for constructing the inline assembly, used below. macro_rules! define_op { ($ins:literal, $reg:ident, $ptr_width:literal, $pos:expr, $ptr:ident) => { @@ -142,8 +154,9 @@ cfg_if::cfg_if! { /// /// Due to the weird multi-return nature of `setjmp` it is very easy to make mistakes, this /// function be used with extreme care. +#[no_mangle] #[naked] -pub unsafe extern "C" fn setjmp(buf: *mut JumpBuf) -> isize { +pub unsafe extern "C" fn setjmp(env: JmpBuf) -> isize { cfg_if::cfg_if! { if #[cfg(target_feature = "d")] { naked_asm! { @@ -212,7 +225,7 @@ pub unsafe extern "C" fn setjmp(buf: *mut JumpBuf) -> isize { /// Additionally, the whole point of this function is to continue execution at a wildly different /// address, so this might cause confusing and hard-to-debug behavior. #[naked] -pub unsafe extern "C" fn longjmp(buf: *mut JumpBuf, val: isize) -> ! { +pub unsafe extern "C" fn longjmp(env: JmpBuf, val: isize) -> ! { cfg_if::cfg_if! { if #[cfg(target_feature = "d")] { naked_asm! { @@ -271,6 +284,69 @@ pub unsafe extern "C" fn longjmp(buf: *mut JumpBuf, val: isize) -> ! { } } +/// Invokes a closure, setting up the environment for contained code to safely use `longjmp`. +/// +/// # Reason +/// +/// TODO +/// +/// # Safety +/// +/// TODO +#[inline(never)] +pub fn call_with_setjmp(f: F) -> isize +where + F: for<'a> FnOnce(&'a JmpBufStruct) -> isize, +{ + let mut f = MaybeUninit::new(f); + + extern "C" fn do_call(env: JmpBuf, closure_env_ptr: *mut F) -> isize + where + F: for<'a> FnOnce(&'a JmpBufStruct) -> isize, + { + // Dereference `closure_env_ptr` with .read() to acquire ownership of + // the FnOnce object, then call it. (See also the forget note below.) + // + // Note that `closure_env_ptr` is not a raw function pointer, it's a + // pointer to a FnOnce; the code we call comes from the generic `F`. + unsafe { (closure_env_ptr.read())(&*env) } + } + + unsafe { + let mut jbuf = MaybeUninit::::zeroed().assume_init(); + let ret: isize; + let env_ptr = addr_of_mut!(jbuf); + let closure_ptr = addr_of_mut!(f); + + // The callback is now effectively owned by `closure_env_ptr` (i.e., the + // `closure_env_ptr.read()` call in `call_from_c_to_rust` will take a + // direct bitwise copy of its state, and pass that ownership into the + // FnOnce::call_once invocation.) + // + // Therefore, we need to forget about our own ownership of the callback now. + #[allow(clippy::forget_non_drop)] + core::mem::forget(f); + + asm! { + "call {setjmp}", // fills in jbuf; future longjmp calls go here. + "bnez a0, 1f", // if return value non-zero, skip the callback invocation + "mv a0, {env_ptr}", // move saved jmp buf into do_call's first arg position + "mv a1, {closure_ptr}", // move saved closure env into do_call's second arg position + "call {do_call}", // invoke the do_call and through it the callback + "1:", // at this point, a0 carries the return value (from either outcome) + + in("a0") env_ptr, + setjmp = sym setjmp, + do_call = sym do_call::, + env_ptr = in(reg) env_ptr, + closure_ptr = in(reg) closure_ptr, + lateout("a0") ret, + clobber_abi("C") + } + ret + } +} + #[cfg(test)] mod tests { use super::*; @@ -278,24 +354,56 @@ mod tests { use core::ptr::addr_of_mut; #[ktest::test] + fn _call_with_setjmp() { + unsafe { + let ret = call_with_setjmp(|env| { + log::trace!("called 1!"); + + 1234 + }); + assert_eq!(ret, 1234); + + let ret = call_with_setjmp(|env| { + log::trace!("called 2!"); + + longjmp(env, 4321); + + 1234 + }); + assert_eq!(ret, 4321); + } + } + + #[ktest::test] + #[allow(static_mut_refs)] fn setjmp_longjmp_simple() { + // The LLVM optimizer doesn't understand the "setjmp returns twice" behaviour and would + // turn the `C += 1` into a constant store instruction instead of a load-add-store sequence. + // To force this, we use a static variable here, but forcing the location of the variable + // into a different (longer-lived) stack frame would also work. + // + // Note that this only exists to test the behaviour of setjmp/longjmp, in real code you + // should use `call_with_setjmp` as it limits much of the footguns. + + static mut C: u32 = 0; + unsafe { - let mut c = 0; - let mut buf = JumpBuf::new(); + // let mut c = 0; + let mut buf = JmpBufStruct::new(); let r = setjmp(ptr::from_mut(&mut buf)); - c += 1; + C += 1; if r == 0 { - assert_eq!(c, 1); + assert_eq!(C, 1); longjmp(ptr::from_mut(&mut buf), 1234567); } - assert_eq!(c, 2); + assert_eq!(C, 2); assert_eq!(r, 1234567); } } - static mut BUFFER_A: JumpBuf = JumpBuf::new(); - static mut BUFFER_B: JumpBuf = JumpBuf::new(); + static mut BUFFER_A: JmpBufStruct = JmpBufStruct::new(); + static mut BUFFER_B: JmpBufStruct = JmpBufStruct::new(); #[ktest::test] fn setjmp_longjmp_complex() { From be665a688808de6374379443d657070d7054b5b5 Mon Sep 17 00:00:00 2001 From: Jonas Kruckenberg Date: Sun, 29 Dec 2024 08:34:27 +0100 Subject: [PATCH 2/5] remove unused --- kernel/src/arch/mod.rs | 2 -- kernel/src/arch/riscv64/setjmp_longjmp.rs | 6 +----- kernel/src/arch/riscv64/trap_handler.rs | 9 ++++----- kernel/src/arch/riscv64/vm.rs | 11 +++++------ 4 files changed, 10 insertions(+), 18 deletions(-) diff --git a/kernel/src/arch/mod.rs b/kernel/src/arch/mod.rs index d8f7d487..a29eb27b 100644 --- a/kernel/src/arch/mod.rs +++ b/kernel/src/arch/mod.rs @@ -1,5 +1,3 @@ -#![allow(unused)] - pub use pmm::arch::*; cfg_if::cfg_if! { diff --git a/kernel/src/arch/riscv64/setjmp_longjmp.rs b/kernel/src/arch/riscv64/setjmp_longjmp.rs index 085d6a01..d2b0cd3a 100644 --- a/kernel/src/arch/riscv64/setjmp_longjmp.rs +++ b/kernel/src/arch/riscv64/setjmp_longjmp.rs @@ -36,13 +36,9 @@ //! result. In a nested calls scenario (e.g. host->wasm->host->wasm) it is therefore up to each host function //! to propagate the trap and each host function therefore gets to clean up all its resources. -use alloc::boxed::Box; -use core::any::Any; use core::arch::{asm, naked_asm}; use core::marker::{PhantomData, PhantomPinned}; -use core::mem::{ManuallyDrop, MaybeUninit}; -use core::panic::UnwindSafe; -use core::ptr; +use core::mem::MaybeUninit; use core::ptr::addr_of_mut; /// A store for the register state used by `setjmp` and `longjmp`. diff --git a/kernel/src/arch/riscv64/trap_handler.rs b/kernel/src/arch/riscv64/trap_handler.rs index 8df47415..cc58b1c4 100644 --- a/kernel/src/arch/riscv64/trap_handler.rs +++ b/kernel/src/arch/riscv64/trap_handler.rs @@ -1,4 +1,3 @@ -use crate::time::Instant; use crate::{arch, TRAP_STACK_SIZE_PAGES}; use core::arch::{asm, naked_asm}; use riscv::scause::{Exception, Interrupt, Trap}; @@ -223,22 +222,22 @@ fn default_trap_handler( let tval = stval::read(); log::error!("KERNEL LOAD PAGE FAULT: epc {epc:#x?} tval {tval:#x?}"); - unsafe { sepc::write(trap_panic_trampoline as usize) } + sepc::write(trap_panic_trampoline as usize) } Trap::Exception(Exception::StorePageFault) => { let epc = sepc::read(); let tval = stval::read(); log::error!("KERNEL STORE PAGE FAULT: epc {epc:#x?} tval {tval:#x?}"); - unsafe { sepc::write(trap_panic_trampoline as usize) } + sepc::write(trap_panic_trampoline as usize) } Trap::Interrupt(Interrupt::SupervisorTimer) => { // just clear the timer interrupt when it happens for now, this is required to make the // tests in the `time` module work - riscv::sbi::time::set_timer(u64::MAX); + riscv::sbi::time::set_timer(u64::MAX).unwrap(); } _ => { - unsafe { sepc::write(trap_panic_trampoline as usize) } + sepc::write(trap_panic_trampoline as usize) // panic!("trap_handler cause {cause:?}, a1 {a1:#x} a2 {a2:#x} a3 {a3:#x} a4 {a4:#x} a5 {a5:#x} a6 {a6:#x} a7 {a7:#x}"); } } diff --git a/kernel/src/arch/riscv64/vm.rs b/kernel/src/arch/riscv64/vm.rs index 8cbdc4a2..5d132e06 100644 --- a/kernel/src/arch/riscv64/vm.rs +++ b/kernel/src/arch/riscv64/vm.rs @@ -1,4 +1,3 @@ -use crate::arch; use core::alloc::Layout; use core::num::NonZeroUsize; use loader_api::BootInfo; @@ -9,7 +8,7 @@ const KERNEL_ASID: usize = 0; pub fn init( boot_info: &BootInfo, - frame_alloc: &mut BuddyAllocator, + _frame_alloc: &mut BuddyAllocator, ) -> crate::Result { let (mut arch, mut flush) = pmm::AddressSpace::from_active(KERNEL_ASID, boot_info.physical_memory_map.start); @@ -39,17 +38,17 @@ fn unmap_loader(boot_info: &BootInfo, arch: &mut pmm::AddressSpace, flush: &mut struct IgnoreAlloc; impl FrameAllocator for IgnoreAlloc { - fn allocate_contiguous(&mut self, layout: Layout) -> Option { + fn allocate_contiguous(&mut self, _layout: Layout) -> Option { unimplemented!() } - fn deallocate_contiguous(&mut self, addr: PhysicalAddress, layout: Layout) {} + fn deallocate_contiguous(&mut self, _addr: PhysicalAddress, _layout: Layout) {} - fn allocate_contiguous_zeroed(&mut self, layout: Layout) -> Option { + fn allocate_contiguous_zeroed(&mut self, _layout: Layout) -> Option { unimplemented!() } - fn allocate_partial(&mut self, layout: Layout) -> Option<(PhysicalAddress, usize)> { + fn allocate_partial(&mut self, _layout: Layout) -> Option<(PhysicalAddress, usize)> { unimplemented!() } From d9bb00d2ff9ea546ded76c9e2ada820246c62510 Mon Sep 17 00:00:00 2001 From: Jonas Kruckenberg Date: Sun, 29 Dec 2024 08:47:36 +0100 Subject: [PATCH 3/5] docs --- kernel/src/arch/riscv64/setjmp_longjmp.rs | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/kernel/src/arch/riscv64/setjmp_longjmp.rs b/kernel/src/arch/riscv64/setjmp_longjmp.rs index d2b0cd3a..8293f028 100644 --- a/kernel/src/arch/riscv64/setjmp_longjmp.rs +++ b/kernel/src/arch/riscv64/setjmp_longjmp.rs @@ -282,13 +282,21 @@ pub unsafe extern "C" fn longjmp(env: JmpBuf, val: isize) -> ! { /// Invokes a closure, setting up the environment for contained code to safely use `longjmp`. /// -/// # Reason +/// This function acts as a somewhat-safe wrapper around `setjmp` that prevents LLVM miscompilations +/// caused by the fact that its optimization passes don't know about the *returns-twice* nature of `setjmp`. /// -/// TODO +/// Note for the pedantic: Yes LLVM *could* know about this, and does have logic to handle it, but Rust +/// has (sensibly) decided to remove the `returns_twice` attribute from the language, so instead +/// we have to rely on this wrapper. /// /// # Safety /// -/// TODO +/// While `longjmp` is still very sketchy, skipping over destructors and such, this function does +/// the necessary ceremony to ensure safe, Rust compatible usage of `setjmp`. In particular, it ensures +/// that the `JmpBuf` cannot be leaked out of the closure, and that it cannot be shared between +/// threads. +// The code below is adapted from https://github.com/pnkfelix/cee-scape/blob/d6ffeca6bd56b46b83c8c9118dbe75e38d423d28/src/asm_based.rs +// which in turn is adapted from this Zulip thread: https://rust-lang.zulipchat.com/#narrow/stream/210922-project-ffi-unwind/topic/cost.20of.20supporting.20longjmp.20without.20annotations/near/301840755 #[inline(never)] pub fn call_with_setjmp(f: F) -> isize where @@ -305,7 +313,7 @@ where // // Note that `closure_env_ptr` is not a raw function pointer, it's a // pointer to a FnOnce; the code we call comes from the generic `F`. - unsafe { (closure_env_ptr.read())(&*env) } + unsafe { closure_env_ptr.read()(&*env) } } unsafe { From 8554cc699513fe4887da4e1d33b232f70f6c9d01 Mon Sep 17 00:00:00 2001 From: Jonas Kruckenberg Date: Sun, 29 Dec 2024 08:58:22 +0100 Subject: [PATCH 4/5] Update setjmp_longjmp.rs --- kernel/src/arch/riscv64/setjmp_longjmp.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/kernel/src/arch/riscv64/setjmp_longjmp.rs b/kernel/src/arch/riscv64/setjmp_longjmp.rs index 8293f028..aa326c8f 100644 --- a/kernel/src/arch/riscv64/setjmp_longjmp.rs +++ b/kernel/src/arch/riscv64/setjmp_longjmp.rs @@ -392,7 +392,6 @@ mod tests { static mut C: u32 = 0; unsafe { - // let mut c = 0; let mut buf = JmpBufStruct::new(); let r = setjmp(ptr::from_mut(&mut buf)); From 69321e078be44d952b1eab3033a09e7badbc69b7 Mon Sep 17 00:00:00 2001 From: Jonas Kruckenberg Date: Sun, 29 Dec 2024 09:27:59 +0100 Subject: [PATCH 5/5] Update setjmp_longjmp.rs --- kernel/src/arch/riscv64/setjmp_longjmp.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/kernel/src/arch/riscv64/setjmp_longjmp.rs b/kernel/src/arch/riscv64/setjmp_longjmp.rs index aa326c8f..69ea3ac8 100644 --- a/kernel/src/arch/riscv64/setjmp_longjmp.rs +++ b/kernel/src/arch/riscv64/setjmp_longjmp.rs @@ -360,7 +360,7 @@ mod tests { #[ktest::test] fn _call_with_setjmp() { unsafe { - let ret = call_with_setjmp(|env| { + let ret = call_with_setjmp(|_env| { log::trace!("called 1!"); 1234 @@ -371,8 +371,6 @@ mod tests { log::trace!("called 2!"); longjmp(env, 4321); - - 1234 }); assert_eq!(ret, 4321); }