Fix UB in the locking code of context switching

The spin crate considers it UB to call force_write_unlock while there is
a threads trying to obtain a read lock on the same rwlock.

This also switches to the spinning_lot crate for the context rwlock as
it has support for a write guard keeping a reference to the rwlock using
an Arc instead of a reference.

src/context/list.rs

@@ -1,7 +1,7 @@
 use alloc::{collections::BTreeMap, sync::Arc};
 use core::{iter, mem};
 
-use spin::RwLock;
+use spinning_top::RwSpinlock;
 
 use super::context::{Context, ContextId};
 use crate::syscall::error::{Error, Result, EAGAIN};
@@ -9,7 +9,7 @@ use crate::syscall::error::{Error, Result, EAGAIN};
 /// Context list type
 pub struct ContextList {
     // Using a BTreeMap for it's range method
-    map: BTreeMap<ContextId, Arc<RwLock<Context>>>,
+    map: BTreeMap<ContextId, Arc<RwSpinlock<Context>>>,
     next_id: usize,
 }
 
@@ -23,7 +23,7 @@ impl ContextList {
     }
 
     /// Get the nth context.
-    pub fn get(&self, id: ContextId) -> Option<&Arc<RwLock<Context>>> {
+    pub fn get(&self, id: ContextId) -> Option<&Arc<RwSpinlock<Context>>> {
         self.map.get(&id)
     }
 
@@ -31,7 +31,7 @@ impl ContextList {
     pub fn ancestors(
         &'_ self,
         id: ContextId,
-    ) -> impl Iterator<Item = (ContextId, &Arc<RwLock<Context>>)> + '_ {
+    ) -> impl Iterator<Item = (ContextId, &Arc<RwSpinlock<Context>>)> + '_ {
         iter::successors(
             self.get(id).map(|context| (id, context)),
             move |(_id, context)| {
@@ -43,25 +43,30 @@ impl ContextList {
     }
 
     /// Get the current context.
-    pub fn current(&self) -> Option<&Arc<RwLock<Context>>> {
+    pub fn current(&self) -> Option<&Arc<RwSpinlock<Context>>> {
         self.map.get(&super::context_id())
     }
 
-    pub fn iter(&self) -> ::alloc::collections::btree_map::Iter<ContextId, Arc<RwLock<Context>>> {
+    pub fn iter(
+        &self,
+    ) -> ::alloc::collections::btree_map::Iter<ContextId, Arc<RwSpinlock<Context>>> {
         self.map.iter()
     }
 
     pub fn range(
         &self,
         range: impl core::ops::RangeBounds<ContextId>,
-    ) -> ::alloc::collections::btree_map::Range<'_, ContextId, Arc<RwLock<Context>>> {
+    ) -> ::alloc::collections::btree_map::Range<'_, ContextId, Arc<RwSpinlock<Context>>> {
         self.map.range(range)
     }
 
-    pub(crate) fn insert_context_raw(&mut self, id: ContextId) -> Result<&Arc<RwLock<Context>>> {
+    pub(crate) fn insert_context_raw(
+        &mut self,
+        id: ContextId,
+    ) -> Result<&Arc<RwSpinlock<Context>>> {
         assert!(self
             .map
-            .insert(id, Arc::new(RwLock::new(Context::new(id)?)))
+            .insert(id, Arc::new(RwSpinlock::new(Context::new(id)?)))
             .is_none());
 
         Ok(self
@@ -71,7 +76,7 @@ impl ContextList {
     }
 
     /// Create a new context.
-    pub fn new_context(&mut self) -> Result<&Arc<RwLock<Context>>> {
+    pub fn new_context(&mut self) -> Result<&Arc<RwSpinlock<Context>>> {
         // Zero is not a valid context ID, therefore add 1.
         //
         // FIXME: Ensure the number of CPUs can't switch between new_context calls.
@@ -98,7 +103,7 @@ impl ContextList {
     }
 
     /// Spawn a context from a function.
-    pub fn spawn(&mut self, func: extern "C" fn()) -> Result<&Arc<RwLock<Context>>> {
+    pub fn spawn(&mut self, func: extern "C" fn()) -> Result<&Arc<RwSpinlock<Context>>> {
         let context_lock = self.new_context()?;
         {
             let mut context = context_lock.write();
@@ -129,7 +134,7 @@ impl ContextList {
         Ok(context_lock)
     }
 
-    pub fn remove(&mut self, id: ContextId) -> Option<Arc<RwLock<Context>>> {
+    pub fn remove(&mut self, id: ContextId) -> Option<Arc<RwSpinlock<Context>>> {
         self.map.remove(&id)
     }
 }

src/context/mod.rs

@@ -5,6 +5,7 @@
 use alloc::{borrow::Cow, sync::Arc};
 
 use spin::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+use spinning_top::RwSpinlock;
 
 use crate::{
     paging::{RmmA, RmmArch, TableKind},
@@ -102,7 +103,7 @@ pub fn context_id() -> ContextId {
     PercpuBlock::current().switch_internals.context_id()
 }
 
-pub fn current() -> Result<Arc<RwLock<Context>>> {
+pub fn current() -> Result<Arc<RwSpinlock<Context>>> {
     contexts()
         .current()
         .ok_or(Error::new(ESRCH))

src/context/switch.rs

@@ -1,8 +1,6 @@
-use core::{cell::Cell, ops::Bound, sync::atomic::Ordering};
+use core::{cell::Cell, mem, ops::Bound, sync::atomic::Ordering};
 
-use alloc::sync::Arc;
-
-use spin::{RwLock, RwLockWriteGuard};
+use spinning_top::guard::ArcRwSpinlockWriteGuard;
 
 use crate::{
     context::{arch, contexts, signal::signal_handler, Context},
@@ -89,8 +87,8 @@ unsafe fn update_runnable(context: &mut Context, cpu_id: LogicalCpuId) -> bool {
 }
 
 struct SwitchResult {
-    prev_lock: Arc<RwLock<Context>>,
-    next_lock: Arc<RwLock<Context>>,
+    _prev_guard: ArcRwSpinlockWriteGuard<Context>,
+    _next_guard: ArcRwSpinlockWriteGuard<Context>,
 }
 
 pub fn tick() {
@@ -106,13 +104,8 @@ pub fn tick() {
 }
 
 pub unsafe extern "C" fn switch_finish_hook() {
-    if let Some(SwitchResult {
-        prev_lock,
-        next_lock,
-    }) = PercpuBlock::current().switch_internals.switch_result.take()
-    {
-        prev_lock.force_write_unlock();
-        next_lock.force_write_unlock();
+    if let Some(switch_result) = PercpuBlock::current().switch_internals.switch_result.take() {
+        drop(switch_result);
     } else {
         // TODO: unreachable_unchecked()?
         crate::arch::stop::emergency_reset();
@@ -151,7 +144,7 @@ pub unsafe fn switch() -> bool {
         let prev_context_lock = contexts
             .current()
             .expect("context::switch: not inside of context");
-        let prev_context_guard = prev_context_lock.write();
+        let prev_context_guard = prev_context_lock.write_arc();
 
         let idle_id = percpu.switch_internals.idle_id();
         let mut skip_idle = true;
@@ -179,17 +172,12 @@ pub unsafe fn switch() -> bool {
             }
 
             // Lock next context
-            let mut next_context_guard = next_context_lock.write();
+            let mut next_context_guard = next_context_lock.write_arc();
 
             // Update state of next context and check if runnable
             if update_runnable(&mut *next_context_guard, cpu_id) {
                 // Store locks for previous and next context
-                switch_context_opt = Some((
-                    Arc::clone(prev_context_lock),
-                    RwLockWriteGuard::leak(prev_context_guard) as *mut Context,
-                    Arc::clone(next_context_lock),
-                    RwLockWriteGuard::leak(next_context_guard) as *mut Context,
-                ));
+                switch_context_opt = Some((prev_context_guard, next_context_guard));
                 break;
             } else {
                 continue;
@@ -198,16 +186,14 @@ pub unsafe fn switch() -> bool {
     };
 
     // Switch process states, TSS stack pointer, and store new context ID
-    if let Some((prev_context_lock, prev_context_ptr, next_context_lock, next_context_ptr)) =
-        switch_context_opt
-    {
+    if let Some((mut prev_context_guard, mut next_context_guard)) = switch_context_opt {
         // Set old context as not running and update CPU time
-        let prev_context = &mut *prev_context_ptr;
+        let prev_context = &mut *prev_context_guard;
         prev_context.running = false;
         prev_context.cpu_time += switch_time.saturating_sub(prev_context.switch_time);
 
         // Set new context as running and set switch time
-        let next_context = &mut *next_context_ptr;
+        let next_context = &mut *next_context_guard;
         next_context.running = true;
         next_context.cpu_id = Some(cpu_id);
         next_context.switch_time = switch_time;
@@ -233,12 +219,18 @@ pub unsafe fn switch() -> bool {
             }
         }
 
+        // FIXME set th switch result in arch::switch_to instead
+        let prev_context =
+            mem::transmute::<&'_ mut Context, &'_ mut Context>(&mut *prev_context_guard);
+        let next_context =
+            mem::transmute::<&'_ mut Context, &'_ mut Context>(&mut *next_context_guard);
+
         percpu
             .switch_internals
             .switch_result
             .set(Some(SwitchResult {
-                prev_lock: prev_context_lock,
-                next_lock: next_context_lock,
+                _prev_guard: prev_context_guard,
+                _next_guard: next_context_guard,
             }));
 
         arch::switch_to(prev_context, next_context);