schedule: switch to special stack during context switches

We need to guard against IRQs coming in after switching to the new page tables
and before switching to the new stack.

Signed-off-by: Tom Dohrmann <[email protected]>

Author: Freax13

Diff for: kernel/src/cpu/percpu.rs

@@ -25,8 +25,8 @@ use crate::mm::vm::{
     VMReserved, VMR,
 };
 use crate::mm::{
-    virt_to_phys, PageBox, SVSM_PERCPU_BASE, SVSM_PERCPU_CAA_BASE, SVSM_PERCPU_END,
-    SVSM_PERCPU_TEMP_BASE_2M, SVSM_PERCPU_TEMP_BASE_4K, SVSM_PERCPU_TEMP_END_2M,
+    virt_to_phys, PageBox, SVSM_CONTEXT_SWITCH_STACK, SVSM_PERCPU_BASE, SVSM_PERCPU_CAA_BASE,
+    SVSM_PERCPU_END, SVSM_PERCPU_TEMP_BASE_2M, SVSM_PERCPU_TEMP_BASE_4K, SVSM_PERCPU_TEMP_END_2M,
     SVSM_PERCPU_TEMP_END_4K, SVSM_PERCPU_VMSA_BASE, SVSM_SHADOW_STACKS_INIT_TASK,
     SVSM_SHADOW_STACK_ISST_DF_BASE, SVSM_STACKS_INIT_TASK, SVSM_STACK_IST_DF_BASE,
 };
@@ -499,6 +499,11 @@ impl PerCpu {
         Ok(())
     }
 
+    fn allocate_context_switch_stack(&self) -> Result<(), SvsmError> {
+        self.allocate_stack(SVSM_CONTEXT_SWITCH_STACK)?;
+        Ok(())
+    }
+
     fn allocate_ist_stacks(&self) -> Result<(), SvsmError> {
         let double_fault_stack = self.allocate_stack(SVSM_STACK_IST_DF_BASE)?;
         self.ist.double_fault_stack.set(Some(double_fault_stack));
@@ -629,6 +634,9 @@ impl PerCpu {
             self.allocate_init_shadow_stack()?;
         }
 
+        // Allocate per-cpu context switch stack
+        self.allocate_context_switch_stack()?;
+
         // Allocate IST stacks
         self.allocate_ist_stacks()?;
 

Diff for: kernel/src/mm/address_space.rs

@@ -172,8 +172,11 @@ pub const SVSM_STACKS_INIT_TASK: VirtAddr = SVSM_PERCPU_STACKS_BASE;
 pub const SVSM_SHADOW_STACKS_INIT_TASK: VirtAddr =
     SVSM_STACKS_INIT_TASK.const_add(STACK_TOTAL_SIZE);
 
+/// Stack address to use during context switches
+pub const SVSM_CONTEXT_SWITCH_STACK: VirtAddr = SVSM_SHADOW_STACKS_INIT_TASK.const_add(PAGE_SIZE);
+
 ///  IST Stacks base address
-pub const SVSM_STACKS_IST_BASE: VirtAddr = SVSM_SHADOW_STACKS_INIT_TASK.const_add(PAGE_SIZE);
+pub const SVSM_STACKS_IST_BASE: VirtAddr = SVSM_CONTEXT_SWITCH_STACK.const_add(STACK_TOTAL_SIZE);
 
 /// DoubleFault IST stack base address
 pub const SVSM_STACK_IST_DF_BASE: VirtAddr = SVSM_STACKS_IST_BASE;

Diff for: kernel/src/task/schedule.rs

@@ -32,7 +32,7 @@ extern crate alloc;
 
 use super::INITIAL_TASK_ID;
 use super::{Task, TaskListAdapter, TaskPointer, TaskRunListAdapter};
-use crate::address::Address;
+use crate::address::{Address, VirtAddr};
 use crate::cpu::msr::write_msr;
 use crate::cpu::percpu::{irq_nesting_count, this_cpu};
 use crate::cpu::shadow_stack::PL0_SSP;
@@ -41,6 +41,7 @@ use crate::cpu::sse::sse_save_context;
 use crate::cpu::IrqGuard;
 use crate::error::SvsmError;
 use crate::locking::SpinLock;
+use crate::mm::{STACK_TOTAL_SIZE, SVSM_CONTEXT_SWITCH_STACK};
 use alloc::sync::Arc;
 use core::arch::{asm, global_asm};
 use core::cell::OnceCell;
@@ -414,6 +415,9 @@ global_asm!(
         jz      1f
         movq    %rsp, {TASK_RSP_OFFSET}(%rsi)
 
+        // Switch to a stack pointer that's valid in both the old and new page tables.
+        mov ${CONTEXT_SWITCH_STACK}, %rsp
+
     1:
         // Switch to the new task state
         mov     %rdx, %cr3
@@ -445,5 +449,16 @@ global_asm!(
         ret
     "#,
     TASK_RSP_OFFSET = const offset_of!(Task, rsp),
+    CONTEXT_SWITCH_STACK = const CONTEXT_SWITCH_STACK.as_usize(),
     options(att_syntax)
 );
+
+/// The location of a cpu-local stack that's mapped into every set of page
+/// tables for use during context switches.
+///
+/// If an IRQ is raised after switching the page tables but before switching
+/// to the new stack, the CPU will try to access the old stack in the new page
+/// tables. To protect against this, we switch to another stack that's mapped
+/// into both the old and the new set of page tables. That way we always have a
+/// valid stack to handle exceptions on.
+const CONTEXT_SWITCH_STACK: VirtAddr = SVSM_CONTEXT_SWITCH_STACK.const_add(STACK_TOTAL_SIZE);