From b6944e48bd00ac6f78ec2acd2f5e6fde8d01982a Mon Sep 17 00:00:00 2001
From: Ruihan Li <lrh2000@pku.edu.cn>
Date: Wed, 20 Nov 2024 09:43:15 +0800
Subject: [PATCH] `CpuLocalCell::as_ptr_mut` should be safe

---
 ostd/src/cpu/local/cell.rs | 21 +++++++++++++--------
 ostd/src/task/processor.rs |  4 ++--
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/ostd/src/cpu/local/cell.rs b/ostd/src/cpu/local/cell.rs
index 057ae148a..466e24c25 100644
--- a/ostd/src/cpu/local/cell.rs
+++ b/ostd/src/cpu/local/cell.rs
@@ -79,7 +79,7 @@ macro_rules! cpu_local_cell {
 pub struct CpuLocalCell<T: 'static>(UnsafeCell<T>);
 
 impl<T: 'static> CpuLocalCell<T> {
-    /// Initialize a CPU-local object.
+    /// Initializes a CPU-local object.
     ///
     /// Please do not call this function directly. Instead, use the
     /// `cpu_local!` macro.
@@ -94,17 +94,22 @@ impl<T: 'static> CpuLocalCell<T> {
         Self(UnsafeCell::new(val))
     }
 
-    /// Get access to the underlying value through a raw pointer.
+    /// Gets access to the underlying value through a raw pointer.
     ///
     /// This function calculates the virtual address of the CPU-local object
     /// based on the CPU-local base address and the offset in the BSP.
     ///
-    /// # Safety
-    ///
-    /// The caller should ensure that within the entire execution of this
-    /// function, no interrupt or preemption can occur. Otherwise, the
-    /// returned pointer may points to the variable in another CPU.
-    pub unsafe fn as_ptr_mut(&'static self) -> *mut T {
+    /// This method is safe, but using the returned pointer will be unsafe.
+    /// Specifically,
+    /// - Preemption should be disabled from the time this method is called
+    ///   to the time the pointer is used. Otherwise, the pointer may point
+    ///   to the variable on another CPU, making it difficult or impossible
+    ///   to determine if the data can be borrowed.
+    /// - If the variable can be used in interrupt handlers, borrowing the
+    ///   data should be done with interrupts disabled. Otherwise, more care
+    ///   must be taken to ensure that the borrowing rules are correctly
+    ///   enforced, since the interrupts may come asynchronously.
+    pub fn as_ptr_mut(&'static self) -> *mut T {
         super::has_init::assert_true();
 
         let offset = {
diff --git a/ostd/src/task/processor.rs b/ostd/src/task/processor.rs
index 9fb495052..b1650ff8b 100644
--- a/ostd/src/task/processor.rs
+++ b/ostd/src/task/processor.rs
@@ -45,8 +45,8 @@ pub(super) fn switch_to_task(next_task: Arc<Task>) {
         // Throughout this method, the task's context is alive and can be exclusively used.
         current_task.ctx.get()
     } else {
-        // SAFETY: Interrupts are disabled, so the pointer is safe to be fetched.
-        unsafe { BOOTSTRAP_CONTEXT.as_ptr_mut() }
+        // Throughout this method, interrupts are disabled and the context can be exclusively used.
+        BOOTSTRAP_CONTEXT.as_ptr_mut()
     };
 
     let next_task_ctx_ptr = next_task.ctx().get().cast_const();