From ef4270e473b89e3f47cd3bb4999287820ea86162 Mon Sep 17 00:00:00 2001 From: 4lDO2 <4lDO2@protonmail.com> Date: Sat, 6 Feb 2021 23:32:43 +0100 Subject: [PATCH 1/3] WIP: Attempt to rewrite switch_to in assembly. This is due to a warning in more recent compilers, which forbid anything but a single inline assembly block, in naked functions. It does unfortunately triple fault right now, but I hope I may be able to fix it soon. --- src/context/arch/x86_64.rs | 124 ++++++++++++++++++++++++------------- src/context/switch.rs | 34 +++++----- 2 files changed, 97 insertions(+), 61 deletions(-) diff --git a/src/context/arch/x86_64.rs b/src/context/arch/x86_64.rs index 64b0268..d81ad6d 100644 --- a/src/context/arch/x86_64.rs +++ b/src/context/arch/x86_64.rs @@ -1,5 +1,7 @@ use core::mem; -use core::sync::atomic::{AtomicBool, Ordering}; +use core::sync::atomic::{AtomicU8, Ordering}; + +use alloc::sync::Arc; use crate::syscall::FloatRegisters; @@ -7,14 +9,13 @@ use crate::syscall::FloatRegisters; /// Compare and exchange this to true when beginning a context switch on any CPU /// The `Context::switch_to` function will set it back to false, allowing other CPU's to switch /// This must be done, as no locks can be held on the stack during switch -pub static CONTEXT_SWITCH_LOCK: AtomicBool = AtomicBool::new(false); +pub static CONTEXT_SWITCH_LOCK: AtomicU8 = AtomicU8::new(AbiCompatBool::False as u8); const ST_RESERVED: u128 = 0xFFFF_FFFF_FFFF_0000_0000_0000_0000_0000; #[derive(Clone, Debug)] +#[repr(C)] pub struct Context { - /// FX valid? - loadable: bool, /// FX location fx: usize, /// Page table pointer @@ -34,13 +35,22 @@ pub struct Context { /// Base pointer rbp: usize, /// Stack pointer - rsp: usize + rsp: usize, + /// FX valid? + loadable: AbiCompatBool, +} + +#[repr(u8)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum AbiCompatBool { + False, + True, } impl Context { pub fn new() -> Context { Context { - loadable: false, + loadable: AbiCompatBool::False, fx: 0, cr3: 0, rflags: 0, @@ -59,7 +69,7 @@ impl Context { } pub fn get_fx_regs(&self) -> Option { - if !self.loadable { + if self.loadable == AbiCompatBool::False { return None; } let mut regs = unsafe { *(self.fx as *const FloatRegisters) }; @@ -74,7 +84,7 @@ impl Context { } pub fn set_fx_regs(&mut self, mut new: FloatRegisters) -> bool { - if !self.loadable { + if self.loadable == AbiCompatBool::False { return false; } @@ -126,53 +136,81 @@ impl Context { self.rsp += mem::size_of::(); value } +} - /// Switch to the next context by restoring its stack and registers - /// Check disassembly! - #[cold] - #[inline(never)] - #[naked] - pub unsafe fn switch_to(&mut self, next: &mut Context) { - asm!("fxsave64 [{}]", in(reg) (self.fx)); - self.loadable = true; - if next.loadable { - asm!("fxrstor64 [{}]", in(reg) (next.fx)); - }else{ - asm!("fninit"); - } +/// Switch to the next context by restoring its stack and registers +/// Check disassembly! +#[cold] +#[inline(never)] +#[naked] +pub unsafe extern "C" fn switch_to(_prev: &mut Context, _next: &mut Context) { + asm!( + // As a quick reminder for those who are unfamiliar with the System V ABI (extern "C"): + // + // - the current parameters are passed in the registers `rdi`, `rsi`, + // - we can modify scratch registers, e.g. rax + // - we cannot change callee-preserved registers arbitrarily, e.g. rbx, which is why we + // store them here in the first place. + "mov rax, [rdi + 0x00] // load `prev.fx` + fxsave64 [rax] // save processor simd state in `prev.fx` - asm!("mov {}, cr3", out(reg) (self.cr3)); - if next.cr3 != self.cr3 { - asm!("mov cr3, {}", in(reg) (next.cr3)); - } + // set `prev.loadable` to true + mov BYTE PTR [rdi + 0x50], {true} + // compare `next.loadable` with true + cmp BYTE PTR [rsi + 0x50], {true} + je switch_to.next_is_loadable - asm!("pushfq ; pop {}", out(reg) (self.rflags)); - asm!("push {} ; popfq", in(reg) (next.rflags)); + fninit + jmp switch_to.after_fx - asm!("mov {}, rbx", out(reg) (self.rbx)); - asm!("mov rbx, {}", in(reg) (next.rbx)); + switch_to.next_is_loadable: + mov rax, [rsi + 0x00] + fxrstor64 [rax] - asm!("mov {}, r12", out(reg) (self.r12)); - asm!("mov r12, {}", in(reg) (next.r12)); + switch_to.after_fx: + mov rcx, cr3 + mov [rdi + 0x08], rcx + mov rax, [rsi + 0x08] + cmp rax, rcx - asm!("mov {}, r13", out(reg) (self.r13)); - asm!("mov r13, {}", in(reg) (next.r13)); + je switch_to.same_cr3 + mov cr3, rax - asm!("mov {}, r14", out(reg) (self.r14)); - asm!("mov r14, {}", in(reg) (next.r14)); + switch_to.same_cr3: + mov [rdi + 0x18], rbx + mov rbx, [rsi + 0x18] - asm!("mov {}, r15", out(reg) (self.r15)); - asm!("mov r15, {}", in(reg) (next.r15)); + mov [rdi + 0x20], r12 + mov r12, [rsi + 0x20] - asm!("mov {}, rsp", out(reg) (self.rsp)); - asm!("mov rsp, {}", in(reg) (next.rsp)); + mov [rdi + 0x28], r13 + mov r13, [rsi + 0x28] - asm!("mov {}, rbp", out(reg) (self.rbp)); - asm!("mov rbp, {}", in(reg) (next.rbp)); + mov [rdi + 0x30], r14 + mov r14, [rsi + 0x30] + + mov [rdi + 0x38], r15 + mov r15, [rsi + 0x38] + + mov [rdi + 0x40], rsp + mov rsp, [rsi + 0x40] + + mov [rdi + 0x48], rbp + mov rbp, [rsi + 0x48] + + pushfq + pop QWORD PTR [rdi + 0x10] + + push QWORD PTR [rsi + 0x10] + popfq // Unset global lock after loading registers but before switch - CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst); - } + xor eax, eax + xchg BYTE PTR [rip+{switch_lock}], al", + + true = const(AbiCompatBool::True as u8), + switch_lock = sym CONTEXT_SWITCH_LOCK, + ); } #[allow(dead_code)] diff --git a/src/context/switch.rs b/src/context/switch.rs index cc8ac69..1cf3357 100644 --- a/src/context/switch.rs +++ b/src/context/switch.rs @@ -79,7 +79,7 @@ pub unsafe fn switch() -> bool { let ticks = PIT_TICKS.swap(0, Ordering::SeqCst); // Set the global lock to avoid the unsafe operations below from causing issues - while arch::CONTEXT_SWITCH_LOCK.compare_and_swap(false, true, Ordering::SeqCst) { + while arch::CONTEXT_SWITCH_LOCK.compare_and_swap(0_u8, 1_u8, Ordering::SeqCst) == 0_u8 { interrupt::pause(); } @@ -163,29 +163,27 @@ pub unsafe fn switch() -> bool { to_context.arch.signal_stack(signal_handler, sig); } - let from_ptr: *mut Context = &mut *from_context_guard; - let to_ptr: *mut Context = &mut *to_context; + let from_arch_ptr: *mut arch::Context = &mut from_context_guard.arch; + let to_arch: &mut arch::Context = &mut to_context.arch; - // FIXME: Ensure that this critical section is somehow still protected by the lock, and not - // just for other processes' context switching, but for other operations which do not - // require CONTEXT_SWITCH_LOCK. - // - // What I suggest, is that we wrap the Context struct (typically stored as - // `Arc>`), into a wrapper with interior locking for the inner context type - // (which could be something like `RwLock`). The wrapper would also contain - // a field of type `UnsafeCell`, which would be accessible if and only if the - // `running` field is set to false, making that field function as a lock. - drop(from_context_guard); - drop(from_context_lock); + core::mem::forget(from_context_guard); + + /* + let mut from_context_lock = Arc::clone(&from_context_lock); + let mut to_context_lock = Arc::clone(&to_context_lock); + */ + + arch::switch_to(&mut *from_arch_ptr, to_arch); + + /* to_context_lock.force_write_unlock(); - drop(to_context_lock); - - (*from_ptr).arch.switch_to(&mut (*to_ptr).arch); + from_context_lock.force_write_unlock(); + */ true } else { // No target was found, unset global lock and return - arch::CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst); + arch::CONTEXT_SWITCH_LOCK.store(0_u8, Ordering::SeqCst); false } From 47c3b2269ff17a7e7c30a570d52edbf8c4bfe1e8 Mon Sep 17 00:00:00 2001 From: 4lDO2 <4lDO2@protonmail.com> Date: Sun, 7 Feb 2021 11:59:46 +0100 Subject: [PATCH 2/3] Fix context switching. Previously there was a triple fault, due to a combination of reasons (e.g. rsp and rbp being ordered in the struct and in the assembly). Now, the locks will be held __all the way until the new context__ has been switched to, which completely eliminates any possibility that the "pcid fault" originates here. While I am unsure whether this will work, this could also be an opportunity to be able to remove CONTEXT_SWITCH_LOCK fully. --- src/context/arch/x86_64.rs | 40 ++++++++++++++++++++----------- src/context/mod.rs | 2 ++ src/context/switch.rs | 49 ++++++++++++++++++++++++++++---------- 3 files changed, 64 insertions(+), 27 deletions(-) diff --git a/src/context/arch/x86_64.rs b/src/context/arch/x86_64.rs index d81ad6d..aeb6fd9 100644 --- a/src/context/arch/x86_64.rs +++ b/src/context/arch/x86_64.rs @@ -1,7 +1,5 @@ use core::mem; -use core::sync::atomic::{AtomicU8, Ordering}; - -use alloc::sync::Arc; +use core::sync::atomic::AtomicBool; use crate::syscall::FloatRegisters; @@ -9,7 +7,7 @@ use crate::syscall::FloatRegisters; /// Compare and exchange this to true when beginning a context switch on any CPU /// The `Context::switch_to` function will set it back to false, allowing other CPU's to switch /// This must be done, as no locks can be held on the stack during switch -pub static CONTEXT_SWITCH_LOCK: AtomicU8 = AtomicU8::new(AbiCompatBool::False as u8); +pub static CONTEXT_SWITCH_LOCK: AtomicBool = AtomicBool::new(false); const ST_RESERVED: u128 = 0xFFFF_FFFF_FFFF_0000_0000_0000_0000_0000; @@ -151,8 +149,12 @@ pub unsafe extern "C" fn switch_to(_prev: &mut Context, _next: &mut Context) { // - we can modify scratch registers, e.g. rax // - we cannot change callee-preserved registers arbitrarily, e.g. rbx, which is why we // store them here in the first place. - "mov rax, [rdi + 0x00] // load `prev.fx` - fxsave64 [rax] // save processor simd state in `prev.fx` + " + // load `prev.fx` + mov rax, [rdi + 0x00] + + // save processor SSE/FPU/AVX state in `prev.fx` pointee + fxsave64 [rax] // set `prev.loadable` to true mov BYTE PTR [rdi + 0x50], {true} @@ -168,6 +170,7 @@ pub unsafe extern "C" fn switch_to(_prev: &mut Context, _next: &mut Context) { fxrstor64 [rax] switch_to.after_fx: + // Save the current CR3, and load the next CR3 if not identical mov rcx, cr3 mov [rdi + 0x08], rcx mov rax, [rsi + 0x08] @@ -177,6 +180,7 @@ pub unsafe extern "C" fn switch_to(_prev: &mut Context, _next: &mut Context) { mov cr3, rax switch_to.same_cr3: + // Save old registers, and load new ones mov [rdi + 0x18], rbx mov rbx, [rsi + 0x18] @@ -192,24 +196,32 @@ pub unsafe extern "C" fn switch_to(_prev: &mut Context, _next: &mut Context) { mov [rdi + 0x38], r15 mov r15, [rsi + 0x38] - mov [rdi + 0x40], rsp - mov rsp, [rsi + 0x40] + mov [rdi + 0x40], rbp + mov rbp, [rsi + 0x40] - mov [rdi + 0x48], rbp - mov rbp, [rsi + 0x48] + mov [rdi + 0x48], rsp + mov rsp, [rsi + 0x48] + // push RFLAGS (can only be modified via stack) pushfq + // pop RFLAGS into `self.rflags` pop QWORD PTR [rdi + 0x10] + // push `next.rflags` push QWORD PTR [rsi + 0x10] + // pop into RFLAGS popfq - // Unset global lock after loading registers but before switch - xor eax, eax - xchg BYTE PTR [rip+{switch_lock}], al", + // When we return, we cannot even guarantee that the return address on the stack, points to + // the calling function, `context::switch`. Thus, we have to execute this Rust hook by + // ourselves, which will unlock the contexts before the later switch. + + call {switch_hook} + + ", true = const(AbiCompatBool::True as u8), - switch_lock = sym CONTEXT_SWITCH_LOCK, + switch_hook = sym crate::context::switch_finish_hook, ); } diff --git a/src/context/mod.rs b/src/context/mod.rs index 3d8b7ec..efc6bcf 100644 --- a/src/context/mod.rs +++ b/src/context/mod.rs @@ -34,6 +34,8 @@ pub mod signal; /// Timeout handling pub mod timeout; +pub use self::switch::switch_finish_hook; + /// Limit on number of contexts pub const CONTEXT_MAX_CONTEXTS: usize = (isize::max_value() as usize) - 1; diff --git a/src/context/switch.rs b/src/context/switch.rs index 1cf3357..484e7a0 100644 --- a/src/context/switch.rs +++ b/src/context/switch.rs @@ -1,8 +1,11 @@ +use core::cell::Cell; use core::ops::Bound; use core::sync::atomic::Ordering; use alloc::sync::Arc; +use spin::RwLock; + use crate::context::signal::signal_handler; use crate::context::{arch, contexts, Context, Status, CONTEXT_ID}; use crate::gdt; @@ -64,6 +67,24 @@ unsafe fn update(context: &mut Context, cpu_id: usize) { } } +struct SwitchResult { + prev_lock: Arc>, + next_lock: Arc>, +} + +pub unsafe extern "C" fn switch_finish_hook() { + if let Some(SwitchResult { prev_lock, next_lock }) = SWITCH_RESULT.take() { + prev_lock.force_write_unlock(); + next_lock.force_write_unlock(); + } else { + panic!("SWITCH_RESULT was not set"); + } + arch::CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst); +} + +#[thread_local] +static SWITCH_RESULT: Cell> = Cell::new(None); + unsafe fn runnable(context: &Context, cpu_id: usize) -> bool { // Switch to context if it needs to run, is not currently running, and is owned by the current CPU !context.running && !context.ptrace_stop && context.status == Status::Runnable && context.cpu_id == Some(cpu_id) @@ -79,7 +100,7 @@ pub unsafe fn switch() -> bool { let ticks = PIT_TICKS.swap(0, Ordering::SeqCst); // Set the global lock to avoid the unsafe operations below from causing issues - while arch::CONTEXT_SWITCH_LOCK.compare_and_swap(0_u8, 1_u8, Ordering::SeqCst) == 0_u8 { + while arch::CONTEXT_SWITCH_LOCK.compare_and_swap(false, true, Ordering::SeqCst) { interrupt::pause(); } @@ -164,26 +185,28 @@ pub unsafe fn switch() -> bool { } let from_arch_ptr: *mut arch::Context = &mut from_context_guard.arch; - let to_arch: &mut arch::Context = &mut to_context.arch; - core::mem::forget(from_context_guard); - /* - let mut from_context_lock = Arc::clone(&from_context_lock); - let mut to_context_lock = Arc::clone(&to_context_lock); - */ + let prev_arch: &mut arch::Context = &mut *from_arch_ptr; + let next_arch: &mut arch::Context = &mut to_context.arch; - arch::switch_to(&mut *from_arch_ptr, to_arch); + // to_context_guard only exists as a raw pointer, but is still locked - /* - to_context_lock.force_write_unlock(); - from_context_lock.force_write_unlock(); - */ + SWITCH_RESULT.set(Some(SwitchResult { + prev_lock: from_context_lock, + next_lock: to_context_lock, + })); + + arch::switch_to(prev_arch, next_arch); + + // NOTE: After switch_to is called, the return address can even be different from the + // current return address, meaning that we cannot use local variables here, and that we + // need to use the `switch_finish_hook` to be able to release the locks. true } else { // No target was found, unset global lock and return - arch::CONTEXT_SWITCH_LOCK.store(0_u8, Ordering::SeqCst); + arch::CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst); false } From a706a0dae43b509d048deef53a6186fdeb7a9c0c Mon Sep 17 00:00:00 2001 From: 4lDO2 <4lDO2@protonmail.com> Date: Fri, 12 Feb 2021 18:10:44 +0100 Subject: [PATCH 3/3] Rewrite signal_handler_wrapper as single asm block. The reason for these types of rewrites, is that more recent Rust compilers have started to deprecate naked functions that consist of more than only a single asm block, as they can trigger all sorts of UB. --- src/context/arch/x86_64.rs | 53 +++++++++++++++++++------------------- 1 file changed, 27 insertions(+), 26 deletions(-) diff --git a/src/context/arch/x86_64.rs b/src/context/arch/x86_64.rs index aeb6fd9..ee7ce5b 100644 --- a/src/context/arch/x86_64.rs +++ b/src/context/arch/x86_64.rs @@ -245,37 +245,38 @@ pub struct SignalHandlerStack { #[naked] unsafe extern fn signal_handler_wrapper() { #[inline(never)] - unsafe fn inner(stack: &SignalHandlerStack) { + unsafe extern "C" fn inner(stack: &SignalHandlerStack) { (stack.handler)(stack.sig); } // Push scratch registers - asm!("push rax - push rcx - push rdx - push rdi - push rsi - push r8 - push r9 - push r10 - push r11"); + asm!( + " + push rax + push rcx + push rdx + push rdi + push rsi + push r8 + push r9 + push r10 + push r11 - // Get reference to stack variables - let rsp: usize; - asm!("mov {}, rsp", out(reg) rsp); + mov rdi, rsp + call {inner} - // Call inner rust function - inner(&*(rsp as *const SignalHandlerStack)); + pop r11 + pop r10 + pop r9 + pop r8 + pop rsi + pop rdi + pop rdx + pop rcx + pop rax + add rsp, 16 + ", - // Pop scratch registers, error code, and return - asm!("pop r11 - pop r10 - pop r9 - pop r8 - pop rsi - pop rdi - pop rdx - pop rcx - pop rax - add rsp, 16"); + inner = sym inner, + ); }