task: use TPR protection for scheduler locks

Protecting scheduler locks with TPR means that high-priority interrupts
can be serviced while scheduler work is underway, except during the
actual task switch, which requires interrupts to be disabled to maintain
stack consistency.

Signed-off-by: Jon Lange <[email protected]>

kernel/src/task/schedule.rs

@@ -38,11 +38,12 @@ use crate::cpu::msr::write_msr;
 use crate::cpu::percpu::{irq_nesting_count, this_cpu};
 use crate::cpu::shadow_stack::{is_cet_ss_supported, IS_CET_SUPPORTED, PL0_SSP};
 use crate::cpu::sse::{sse_restore_context, sse_save_context};
-use crate::cpu::IrqGuard;
+use crate::cpu::{IrqGuard, TprGuard};
 use crate::error::SvsmError;
 use crate::fs::Directory;
-use crate::locking::SpinLock;
+use crate::locking::SpinLockTpr;
 use crate::mm::{STACK_TOTAL_SIZE, SVSM_CONTEXT_SWITCH_SHADOW_STACK, SVSM_CONTEXT_SWITCH_STACK};
+use crate::types::TPR_LOCK;
 use alloc::sync::Arc;
 use core::arch::{asm, global_asm};
 use core::cell::OnceCell;
@@ -231,7 +232,7 @@ impl TaskList {
     }
 }
 
-pub static TASKLIST: SpinLock<TaskList> = SpinLock::new(TaskList::new());
+pub static TASKLIST: SpinLockTpr<TaskList> = SpinLockTpr::new(TaskList::new());
 
 /// Creates, initializes and starts a new kernel task. Note that the task has
 /// already started to run before this function returns.
@@ -354,10 +355,12 @@ unsafe fn switch_to(prev: *const Task, next: *const Task) {
 /// function has ran it is safe to call [`schedule()`] on the current CPU.
 pub fn schedule_init() {
     unsafe {
-        let guard = IrqGuard::new();
+        let tpr_guard = TprGuard::raise(TPR_LOCK);
+        let irq_guard = IrqGuard::new();
         let next = task_pointer(this_cpu().schedule_init());
         switch_to(null_mut(), next);
-        drop(guard);
+        drop(irq_guard);
+        drop(tpr_guard);
     }
 }
 
@@ -373,7 +376,7 @@ pub fn schedule() {
     // check if preemption is safe
     preemption_checks();
 
-    let guard = IrqGuard::new();
+    let tpr_guard = TprGuard::raise(TPR_LOCK);
 
     let work = this_cpu().schedule_prepare();
 
@@ -402,16 +405,20 @@ pub fn schedule() {
             let b = task_pointer(next);
             sse_save_context(u64::from((*a).xsa.vaddr()));
 
-            // Switch tasks
+            // Switch tasks.  This requires disabling interrupts so no
+            // interrupt can be taken while tke stack is in an inconsistent
+            // state.
+            let guard = IrqGuard::new();
             switch_to(a, b);
+            drop(guard);
 
             // We're now in the context of task pointed to by 'a'
             // which was previously scheduled out.
             sse_restore_context(u64::from((*a).xsa.vaddr()));
         }
     }
 
-    drop(guard);
+    drop(tpr_guard);
 
     // If the previous task had terminated then we can release
     // it's reference here.

kernel/src/task/tasks.rs

@@ -16,14 +16,15 @@ use core::sync::atomic::{AtomicU32, Ordering};
 use crate::address::{Address, VirtAddr};
 use crate::cpu::idt::svsm::return_new_task;
 use crate::cpu::irq_state::EFLAGS_IF;
+use crate::cpu::lower_tpr;
 use crate::cpu::percpu::PerCpu;
 use crate::cpu::shadow_stack::is_cet_ss_supported;
 use crate::cpu::sse::{get_xsave_area_size, sse_restore_context};
 use crate::cpu::X86ExceptionContext;
 use crate::cpu::{irqs_enable, X86GeneralRegs};
 use crate::error::SvsmError;
 use crate::fs::{opendir, Directory, FileHandle};
-use crate::locking::{RWLock, SpinLock};
+use crate::locking::{RWLockTpr, SpinLockTpr};
 use crate::mm::pagetable::{PTEntryFlags, PageTable};
 use crate::mm::vm::{
     Mapping, ShadowStackInit, VMFileMappingFlags, VMKernelShadowStack, VMKernelStack, VMR,
@@ -35,7 +36,7 @@ use crate::mm::{
 };
 use crate::platform::SVSM_PLATFORM;
 use crate::syscall::{Obj, ObjError, ObjHandle};
-use crate::types::{SVSM_USER_CS, SVSM_USER_DS};
+use crate::types::{SVSM_USER_CS, SVSM_USER_DS, TPR_LOCK};
 use crate::utils::MemoryRegion;
 use intrusive_collections::{intrusive_adapter, LinkedListAtomicLink};
 
@@ -134,7 +135,7 @@ pub struct Task {
     pub exception_shadow_stack: VirtAddr,
 
     /// Page table that is loaded when the task is scheduled
-    pub page_table: SpinLock<PageBox<PageTable>>,
+    pub page_table: SpinLockTpr<PageBox<PageTable>>,
 
     /// Task virtual memory range for use at CPL 0
     vm_kernel_range: VMR,
@@ -143,7 +144,7 @@ pub struct Task {
     vm_user_range: Option<VMR>,
 
     /// State relevant for scheduler
-    sched_state: RWLock<TaskSchedState>,
+    sched_state: RWLockTpr<TaskSchedState>,
 
     /// ID of the task
     id: u32,
@@ -158,7 +159,7 @@ pub struct Task {
     runlist_link: LinkedListAtomicLink,
 
     /// Objects shared among threads within the same process
-    objs: Arc<RWLock<BTreeMap<ObjHandle, Arc<dyn Obj>>>>,
+    objs: Arc<RWLockTpr<BTreeMap<ObjHandle, Arc<dyn Obj>>>>,
 }
 
 // SAFETY: Send + Sync is required for Arc<Task> to implement Send. All members
@@ -248,10 +249,10 @@ impl Task {
             xsa,
             stack_bounds: bounds,
             exception_shadow_stack,
-            page_table: SpinLock::new(pgtable),
+            page_table: SpinLockTpr::new(pgtable),
             vm_kernel_range,
             vm_user_range: None,
-            sched_state: RWLock::new(TaskSchedState {
+            sched_state: RWLockTpr::new(TaskSchedState {
                 idle_task: false,
                 state: TaskState::RUNNING,
                 cpu: cpu.get_apic_id(),
@@ -260,7 +261,7 @@ impl Task {
             rootdir: opendir("/")?,
             list_link: LinkedListAtomicLink::default(),
             runlist_link: LinkedListAtomicLink::default(),
-            objs: Arc::new(RWLock::new(BTreeMap::new())),
+            objs: Arc::new(RWLockTpr::new(BTreeMap::new())),
         }))
     }
 
@@ -331,10 +332,10 @@ impl Task {
             xsa,
             stack_bounds: bounds,
             exception_shadow_stack,
-            page_table: SpinLock::new(pgtable),
+            page_table: SpinLockTpr::new(pgtable),
             vm_kernel_range,
             vm_user_range: Some(vm_user_range),
-            sched_state: RWLock::new(TaskSchedState {
+            sched_state: RWLockTpr::new(TaskSchedState {
                 idle_task: false,
                 state: TaskState::RUNNING,
                 cpu: cpu.get_apic_id(),
@@ -343,7 +344,7 @@ impl Task {
             rootdir: root,
             list_link: LinkedListAtomicLink::default(),
             runlist_link: LinkedListAtomicLink::default(),
-            objs: Arc::new(RWLock::new(BTreeMap::new())),
+            objs: Arc::new(RWLockTpr::new(BTreeMap::new())),
         }))
     }
 
@@ -698,12 +699,13 @@ unsafe fn setup_new_task(xsa_addr: u64) {
     // Re-enable IRQs here, as they are still disabled from the
     // schedule()/sched_init() functions. After the context switch the IrqGuard
     // from the previous task is not dropped, which causes IRQs to stay
-    // disabled in the new task.
+    // disabled in the new task.  TPR needs to be lowered for the same reason.
     // This only needs to be done for the first time a task runs. Any
-    // subsequent task switches will go through schedule() and there the guard
-    // is dropped, re-enabling IRQs.
+    // subsequent task switches will go through schedule() and there the guards
+    // are dropped, re-enabling IRQs and restoring TPR.
 
     irqs_enable();
+    lower_tpr(TPR_LOCK);
 
     // SAFETY: The caller takes responsibility for the correctness of the save
     // area address.