From 0438ef0f1e28a528eadd88c4658622d718a67aa8 Mon Sep 17 00:00:00 2001
From: Jon Lange <jlange@microsoft.com>
Date: Wed, 12 Feb 2025 11:22:58 -0800
Subject: [PATCH] task: fix some safety issues

Starting the scheduler is an unsafe operation because it relies on there
not being a current task, which is a guarantee the compiler cannot make
on its own.

Also, a function that returns a pointer is not necessarily unsafe.  The
operation of using the pointer is where the code becomes unsafe.

Signed-off-by: Jon Lange <jlange@microsoft.com>
---
 kernel/src/cpu/smp.rs       |  7 ++++++-
 kernel/src/svsm.rs          |  7 ++++++-
 kernel/src/task/schedule.rs | 27 ++++++++++++++++++---------
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/kernel/src/cpu/smp.rs b/kernel/src/cpu/smp.rs
index c5aae5675..e3f8ed8c3 100644
--- a/kernel/src/cpu/smp.rs
+++ b/kernel/src/cpu/smp.rs
@@ -149,7 +149,12 @@ extern "C" fn start_ap() -> ! {
     this_cpu_shared().set_online();
 
     sse_init();
-    schedule_init();
+
+    // SAFETY: there is no current task running on this processor yet, so
+    // initializing the scheduler is safe.
+    unsafe {
+        schedule_init();
+    }
 
     unreachable!("Road ends here!");
 }
diff --git a/kernel/src/svsm.rs b/kernel/src/svsm.rs
index 0dccd7777..fb9c7bfec 100755
--- a/kernel/src/svsm.rs
+++ b/kernel/src/svsm.rs
@@ -262,7 +262,12 @@ extern "C" fn svsm_start(li: &KernelLaunchInfo, vb_addr: usize) -> ! {
         .expect("Failed to initialize SVSM platform object");
 
     sse_init();
-    schedule_init();
+
+    // SAFETY: there is no current task running on this processor yet, so
+    // initializing the scheduler is safe.
+    unsafe {
+        schedule_init();
+    }
 
     unreachable!("SVSM entry point terminated unexpectedly");
 }
diff --git a/kernel/src/task/schedule.rs b/kernel/src/task/schedule.rs
index a1f8d2f27..f60b0ac19 100644
--- a/kernel/src/task/schedule.rs
+++ b/kernel/src/task/schedule.rs
@@ -326,17 +326,17 @@ pub fn terminate() {
     schedule();
 }
 
-// SAFETY: This function returns a raw pointer to a task. It is safe
-// because this function is only used in the task switch code, which also only
-// takes a single reference to the next and previous tasks. Also, this
-// function works on an Arc, which ensures that only a single mutable reference
-// can exist.
-unsafe fn task_pointer(taskptr: TaskPointer) -> *const Task {
+fn task_pointer(taskptr: TaskPointer) -> *const Task {
     Arc::as_ptr(&taskptr)
 }
 
+// SAFETY: The caller is required to provide correct pointers for the previous
+// and current tasks.
 #[inline(always)]
 unsafe fn switch_to(prev: *const Task, next: *const Task) {
+    // SAFETY: Assuming the caller has provided the correct task pointers,
+    // the page table and stack information in those tasks are correct and
+    // can be used to switch to the correct page table and execution stack.
     unsafe {
         let cr3 = (*next).page_table.lock().cr3_value().bits() as u64;
 
@@ -366,13 +366,22 @@ unsafe fn switch_to(prev: *const Task, next: *const Task) {
 /// Initializes the [RunQueue] on the current CPU. It will switch to the idle
 /// task and initialize the current_task field of the RunQueue. After this
 /// function has ran it is safe to call [`schedule()`] on the current CPU.
-pub fn schedule_init() {
+///
+/// # Safety
+///
+/// This function can only be called when it is known that there is no current
+/// task.  Otherwise, the run state can become corrupted, and thus future
+/// calculation of task pointers can be incorrect.
+pub unsafe fn schedule_init() {
+    let guard = IrqGuard::new();
+    // SAFETY: The caller guarantees that there is no current task, and the
+    // pointer obtained for the next task will always be correct, thus
+    // providing a guarantee that the task switch will be safe.
     unsafe {
-        let guard = IrqGuard::new();
         let next = task_pointer(this_cpu().schedule_init());
         switch_to(null_mut(), next);
-        drop(guard);
     }
+    drop(guard);
 }
 
 fn preemption_checks() {