Save/restore interrupt state during context switch (#1036)

Previously, there were no guarantees about the value of the
interrupt flag when context switching. 
If the context switch was voluntary, i.e. a task called `schedule()`
to yield the CPU, interrupts would most likely be enabled.
If a task was preempted, interrupts would be disabled.

This could result in a preempted task A switching to a
task B that voluntarily yielded, causing task B to return from
the call to `schedule()` with interrupts disabled, which 
usually isn't desirable. 

Now, this can no longer happen due to saving and restoring 
the state of interrupts for a task when context switching
away from or to it, respectively. 

* On aarch64, this corresponds to the DAIF+NZCV bitflags of `PSTATE`. 
* On x86_64, this corresponds to the IF bit flag of the `RFLAGS` register. 
 

Signed-off-by: Klimenty Tsoutsman <[email protected]>
Co-authored-by: Nathan Royer <[email protected]>

kernel/context_switch_regular/src/aarch64.rs

@@ -19,6 +19,9 @@ pub struct ContextRegular {
     // x30 stores the return address
     // it's used by the `ret` instruction
     x30_link_register: usize,
+
+    // only NZCV & DAIF bits are saved and restored
+    pstate: usize,
 }
 
 impl ContextRegular {
@@ -36,6 +39,9 @@ impl ContextRegular {
             // x30 stores the return address
             // it's used by the `ret` instruction
             x30_link_register: start_address,
+
+            // interrupts are initially unmasked/enabled
+            pstate: 0,
         }
     }
 
@@ -88,6 +94,12 @@ macro_rules! save_registers_regular {
             stp x25, x26, [sp, #8 * 2 * 3]
             stp x27, x28, [sp, #8 * 2 * 4]
             stp x29, x30, [sp, #8 * 2 * 5]
+
+            // Push an OR of DAIF and NZCV flags of PSTATE
+            mrs x29, DAIF
+            mrs x30, NZCV
+            orr x29, x29, x30
+            str x29, [sp, #8 * 2 * 6]
         "#
     );
 }
@@ -119,13 +131,19 @@ macro_rules! restore_registers_regular {
     () => (
         // Restore the next task's general purpose registers.
         r#"
+            // Pop DAIF and NZCV flags of PSTATE
+            // These MSRs discard irrelevant bits; no AND is required.
+            ldr x29, [sp, #8 * 2 * 6]
+            msr DAIF, x29
+            msr NZCV, x29
+
             // Pop registers from the stack, two at a time.
-            ldp x19, x20, [sp, #8 * 2 * 0]
-            ldp x21, x22, [sp, #8 * 2 * 1]
-            ldp x23, x24, [sp, #8 * 2 * 2]
-            ldp x25, x26, [sp, #8 * 2 * 3]
-            ldp x27, x28, [sp, #8 * 2 * 4]
             ldp x29, x30, [sp, #8 * 2 * 5]
+            ldp x27, x28, [sp, #8 * 2 * 4]
+            ldp x25, x26, [sp, #8 * 2 * 3]
+            ldp x23, x24, [sp, #8 * 2 * 2]
+            ldp x21, x22, [sp, #8 * 2 * 1]
+            ldp x19, x20, [sp, #8 * 2 * 0]
 
             // Move the stack pointer back up.
             add sp,  sp,  #8 * 2 * 6

kernel/context_switch_regular/src/x86_64.rs

@@ -10,6 +10,7 @@ use zerocopy::FromBytes;
 #[derive(FromBytes)]
 #[repr(C, packed)]
 pub struct ContextRegular {
+    rflags: usize,
     r15: usize, 
     r14: usize,
     r13: usize,
@@ -31,6 +32,9 @@ impl ContextRegular {
     /// Task containing it to begin its execution at the given `rip`.
     pub fn new(rip: usize) -> ContextRegular {
         ContextRegular {
+            // The ninth bit is the interrupt enable flag. When a task is first
+            // run, interrupts should already be enabled.
+            rflags: 1 << 9,
             r15: 0,
             r14: 0,
             r13: 0,
@@ -88,6 +92,7 @@ macro_rules! save_registers_regular {
             push r13
             push r14
             push r15
+            pushfq
         "#
     );
 }
@@ -122,6 +127,7 @@ macro_rules! restore_registers_regular {
     () => (
         // Restore the next task's general purpose registers.
         r#" 
+            popfq
             pop r15
             pop r14
             pop r13

kernel/spawn/Cargo.toml

@@ -11,8 +11,6 @@ log = "0.4.8"
 spin = "0.9.4"
 lazy_static = { features = ["spin_no_std"], version = "1.4.0" }
 
-irq_safety = { git = "https://github.com/theseus-os/irq_safety" }
-
 debugit = { path = "../../libs/debugit" }
 
 memory = { path = "../memory" }

kernel/spawn/src/lib.rs

@@ -28,7 +28,6 @@ use log::{error, info, debug};
 use cpu::CpuId;
 use debugit::debugit;
 use spin::Mutex;
-use irq_safety::enable_interrupts;
 use memory::{get_kernel_mmi_ref, MmiRef};
 use stack::Stack;
 use task::{Task, TaskRef, RestartInfo, RunState, JoinableTaskRef, ExitableTaskRef, FailureCleanupFunction};
@@ -663,6 +662,12 @@ where
     R: Send + 'static,
     F: FnOnce(A) -> R,
 {
+    // The first time a task runs, its entry function `task_wrapper()` is
+    // jumped to from the `task_switch()` function, right after the context
+    // switch occured. However, we set the context of the new task to have
+    // interrupts enabled (in `ContextRegular::new`), so interrupts are enabled
+    // as soon as the new task is switched to.
+
     let task_entry_func;
     let task_arg;
     let recovered_preemption_guard;
@@ -707,16 +712,7 @@ where
         );
     }
 
-    // The first time that a task runs, its entry function `task_wrapper()` is jumped to
-    // from the `task_switch()` function, right after the context switch occurred.
-    // Since the `task_switch()` function was originally invoked from an interrupt handler,
-    // interrupts were disabled but never had the chance to be re-enabled
-    // because we did not return from the interrupt handler as usual.
-    // Therefore, we need to re-enable interrupts and preemption here to ensure that
-    // things continue to run as normal, with interrupts and preemption enabled
-    // such that we can properly preempt this task.
     drop(recovered_preemption_guard);
-    enable_interrupts();
 
     // This synchronizes with the acquire fence in `JoinableTaskRef::join()`.
     fence(Ordering::Release);