pw_kernel: Use MaybeUninit<u8> for stack memory

This solves two problems:
- It ensures that values which are not entirely initialized (e.g. which
  contain padding bytes) can be written to the stack before the thread
  is spawned [1]
- It ensures that pointers written to the stack retain provenance [2]

[1] After a thread is spawned, the Rust Abstract Machine (AM) sees it as
    stack memory, and not as a Rust variable. At that point, this
    concern no longer applies.
[2] rust-lang/unsafe-code-guidelines#286 (comment)

Change-Id: I5eece2c6181f6dc793815208a2ad2b3fdbfa456c
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/287296
Docs-Not-Needed: Joshua Liebow-Feeser <[email protected]>
Presubmit-Verified: CQ Bot Account <[email protected]>
Lint: Lint 🤖 <[email protected]>
Commit-Queue: Auto-Submit <[email protected]>
Reviewed-by: Erik Gilling <[email protected]>
Pigweed-Auto-Submit: Joshua Liebow-Feeser <[email protected]>

Author: joshlf

pw_kernel/kernel/arch.rs

@@ -27,6 +27,8 @@ mod host;
 #[cfg(feature = "arch_host")]
 pub use host::Arch;
 
+use core::mem::MaybeUninit;
+
 use crate::scheduler::{thread::Stack, SchedulerState};
 use crate::sync::spinlock::SpinLockGuard;
 
@@ -68,7 +70,7 @@ pub trait ThreadState {
     fn initialize_user_frame(
         &mut self,
         kernel_stack: Stack,
-        initial_sp: *mut u8,
+        initial_sp: *mut MaybeUninit<u8>,
         initial_function: extern "C" fn(usize, usize),
         args: (usize, usize),
     );

pw_kernel/kernel/arch/arm_cortex_m/threads.rs

@@ -13,7 +13,7 @@
 // the License.
 
 use core::arch::asm;
-use core::mem;
+use core::mem::{self, MaybeUninit};
 
 use cortex_m::peripheral::SCB;
 
@@ -169,7 +169,7 @@ impl super::super::ThreadState for ArchThreadState {
     fn initialize_user_frame(
         &mut self,
         kernel_stack: Stack,
-        initial_sp: *mut u8,
+        initial_sp: *mut MaybeUninit<u8>,
         initial_function: extern "C" fn(usize, usize),
         args: (usize, usize),
     ) {

pw_kernel/kernel/arch/host.rs

@@ -12,6 +12,8 @@
 // License for the specific language governing permissions and limitations under
 // the License.
 
+use core::mem::MaybeUninit;
+
 use pw_log::info;
 
 use crate::arch::ArchInterface;
@@ -48,7 +50,7 @@ impl super::ThreadState for ThreadState {
     fn initialize_user_frame(
         &mut self,
         _kernel_stack: Stack,
-        _initial_sp: *mut u8,
+        _initial_sp: *mut MaybeUninit<u8>,
         _initial_function: extern "C" fn(usize, usize),
         _args: (usize, usize),
     ) {

pw_kernel/kernel/arch/riscv/threads.rs

@@ -16,7 +16,7 @@ use crate::arch::{Arch, ArchInterface};
 use crate::scheduler::{self, thread::Stack, SchedulerState};
 use crate::sync::spinlock::SpinLockGuard;
 use core::arch::naked_asm;
-use core::mem;
+use core::mem::{self, MaybeUninit};
 use log_if::debug_if;
 
 const LOG_CONTEXT_SWITCH: bool = false;
@@ -114,7 +114,7 @@ impl super::super::ThreadState for ArchThreadState {
     fn initialize_user_frame(
         &mut self,
         kernel_stack: Stack,
-        initial_sp: *mut u8,
+        initial_sp: *mut MaybeUninit<u8>,
         initial_function: extern "C" fn(usize, usize),
         args: (usize, usize),
     ) {

pw_kernel/kernel/lib.rs

@@ -37,6 +37,10 @@ pub use scheduler::{
 };
 pub use timer::{Clock, Duration};
 
+// Used by the `init_thread!` macro.
+#[doc(hidden)]
+pub use scheduler::thread::{StackStorage, StackStorageExt};
+
 #[no_mangle]
 #[allow(non_snake_case)]
 pub extern "C" fn pw_assert_HandleFailure() -> ! {
@@ -99,10 +103,11 @@ macro_rules! init_thread {
         info!("initializing thread: {}", $name as &'static str);
         thread.initialize_kernel_thread(
             {
-                static mut STACK: [u8; $stack_size] = [0; $stack_size];
+                static mut STACK_STORAGE: $crate::StackStorage<{ $stack_size }> =
+                    $crate::StackStorageExt::ZEROED;
                 #[allow(static_mut_refs)]
                 unsafe {
-                    Stack::from_slice(&STACK)
+                    Stack::from_slice(&STACK_STORAGE)
                 }
             },
             $entry,
@@ -136,17 +141,19 @@ macro_rules! init_non_priv_thread {
         );
         thread.initialize_non_priv_thread(
             {
-                static mut STACK: [u8; $stack_size] = [0; $stack_size];
+                static mut STACK_STORAGE: $crate::StackStorage<{ $stack_size }> =
+                    $crate::StackStorageExt::ZEROED;
                 #[allow(static_mut_refs)]
                 unsafe {
-                    Stack::from_slice(&STACK)
+                    Stack::from_slice(&STACK_STORAGE)
                 }
             },
             {
-                static mut STACK: [u8; $stack_size] = [0; $stack_size];
+                static mut STACK_STORAGE: $crate::StackStorage<{ $stack_size }> =
+                    $crate::StackStorageExt::ZEROED;
                 #[allow(static_mut_refs)]
                 unsafe {
-                    Stack::from_slice(&STACK)
+                    Stack::from_slice(&STACK_STORAGE)
                 }
             },
             $entry,

pw_kernel/kernel/scheduler/thread.rs

@@ -13,6 +13,7 @@
 // the License.
 
 use core::cell::UnsafeCell;
+use core::mem::MaybeUninit;
 
 use list::*;
 use pw_log::info;
@@ -21,19 +22,46 @@ use crate::arch::ThreadState;
 
 use super::SCHEDULER_STATE;
 
+/// The memory backing a thread's stack before it has been started.
+///
+/// After a thread has been started, ownership of its stack's memory is (from
+/// the Rust Abstract Machine (AM) perspective) relinquished, and so the type we
+/// use to represent that memory is irrelevant.
+///
+/// However, while we are initializing a thread in preparation for starting it,
+/// we are operating on that memory as a normal Rust variable, and so its type
+/// is important.
+///
+/// Using `MaybeUninit<u8>` instead of `u8` is important for two reasons:
+/// - It ensures that it is sound to write values which are not entirely
+///   initialized (e.g., which contain padding bytes).
+/// - It ensures that pointers written to the stack [retain
+///   provenance][provenance].
+///
+/// [provenance]: https://github.com/rust-lang/unsafe-code-guidelines/issues/286#issuecomment-2837585644
+pub type StackStorage<const N: usize> = [MaybeUninit<u8>; N];
+
+pub trait StackStorageExt {
+    const ZEROED: Self;
+}
+
+impl<const N: usize> StackStorageExt for StackStorage<N> {
+    const ZEROED: StackStorage<N> = [MaybeUninit::new(0); N];
+}
+
 #[derive(Clone, Copy)]
 pub struct Stack {
     // Starting (lowest) address of the stack.  Inclusive.
-    start: *const u8,
+    start: *const MaybeUninit<u8>,
 
     // Ending (highest) address of the stack.  Exclusive.
-    end: *const u8,
+    end: *const MaybeUninit<u8>,
 }
 
 #[allow(dead_code)]
 impl Stack {
-    pub const fn from_slice(slice: &[u8]) -> Self {
-        let start: *const u8 = slice.as_ptr();
+    pub const fn from_slice(slice: &[MaybeUninit<u8>]) -> Self {
+        let start: *const MaybeUninit<u8> = slice.as_ptr();
         // Safety: offset based on known size of slice.
         let end = unsafe { start.add(slice.len()) };
         Self { start, end }
@@ -46,23 +74,27 @@ impl Stack {
         }
     }
 
-    pub fn start(&self) -> *const u8 {
+    pub fn start(&self) -> *const MaybeUninit<u8> {
         self.start
     }
-    pub fn end(&self) -> *const u8 {
+    pub fn end(&self) -> *const MaybeUninit<u8> {
         self.end
     }
 
-    pub fn initial_sp(&self, alignment: usize) -> *const u8 {
+    pub fn initial_sp(&self, alignment: usize) -> *const MaybeUninit<u8> {
         // Use a zero sized allocation to align the initial stack pointer.
         Self::aligned_stack_allocation(self.end, 0, alignment)
     }
 
-    pub fn contains(&self, ptr: *const u8) -> bool {
+    pub fn contains(&self, ptr: *const MaybeUninit<u8>) -> bool {
         ptr >= self.start && ptr < self.end
     }
 
-    pub fn aligned_stack_allocation(sp: *const u8, size: usize, alignment: usize) -> *const u8 {
+    pub fn aligned_stack_allocation(
+        sp: *const MaybeUninit<u8>,
+        size: usize,
+        alignment: usize,
+    ) -> *const MaybeUninit<u8> {
         let sp = sp.wrapping_byte_sub(size);
         let offset = sp.align_offset(alignment);
         if offset > 0 {
@@ -239,7 +271,7 @@ impl Thread {
                 // the RISC-V calling convention.   Ideally this would be
                 // architecture dependant.  However, this value will eventually
                 // be passed in from user space.
-                main_stack.initial_sp(16) as *mut u8,
+                main_stack.initial_sp(16) as *mut MaybeUninit<u8>,
                 Self::trampoline,
                 args,
             );