feat: implement stack protection

Author: zyuiop

src/arch/aarch64/kernel/mod.rs

@@ -14,17 +14,16 @@ pub mod serial;
 mod start;
 pub mod systemtime;
 
-use alloc::alloc::{Layout, alloc};
 use core::arch::global_asm;
 use core::sync::atomic::{AtomicPtr, AtomicU32, Ordering};
 use core::{ptr, str};
-
+use hermit_sync::InterruptTicketMutex;
 use memory_addresses::PhysAddr;
 
 use crate::arch::aarch64::kernel::core_local::*;
-use crate::arch::aarch64::mm::paging::{BasePageSize, PageSize};
 use crate::config::*;
 use crate::env;
+use crate::mm::stack_alloc::{allocate_stack, StackAllocation};
 
 #[repr(align(8))]
 pub(crate) struct AlignedAtomicU32(AtomicU32);
@@ -34,6 +33,7 @@ pub(crate) struct AlignedAtomicU32(AtomicU32);
 /// It also synchronizes initialization of CPU cores.
 pub(crate) static CPU_ONLINE: AlignedAtomicU32 = AlignedAtomicU32(AtomicU32::new(0));
 
+pub static CURRENT_STACK: InterruptTicketMutex<Option<StackAllocation>> = InterruptTicketMutex::new(None);
 pub(crate) static CURRENT_STACK_ADDRESS: AtomicPtr<u8> = AtomicPtr::new(ptr::null_mut());
 
 #[cfg(target_os = "none")]
@@ -105,10 +105,9 @@ fn finish_processor_init() {
 	debug!("Initialized processor {}", core_id());
 
 	// Allocate stack for the CPU and pass the addresses.
-	let layout = Layout::from_size_align(KERNEL_STACK_SIZE, BasePageSize::SIZE as usize).unwrap();
-	let stack = unsafe { alloc(layout) };
-	assert!(!stack.is_null());
-	CURRENT_STACK_ADDRESS.store(stack, Ordering::Relaxed);
+	let stack = allocate_stack(KERNEL_STACK_SIZE);
+	CURRENT_STACK_ADDRESS.store(stack.stack_start().as_mut_ptr(), Ordering::Relaxed);
+	let _ = CURRENT_STACK.lock().insert(stack);
 }
 
 pub fn boot_next_processor() {

src/arch/aarch64/kernel/scheduler.rs

@@ -6,16 +6,16 @@ use core::sync::atomic::Ordering;
 use aarch64_cpu::asm::barrier::{SY, isb};
 use aarch64_cpu::registers::*;
 use align_address::Align;
-use free_list::{PageLayout, PageRange};
-use memory_addresses::{PhysAddr, VirtAddr};
+use memory_addresses::VirtAddr;
 
 use crate::arch::aarch64::kernel::CURRENT_STACK_ADDRESS;
 use crate::arch::aarch64::kernel::core_local::core_scheduler;
-use crate::arch::aarch64::mm::paging::{BasePageSize, PageSize, PageTableEntryFlags};
-use crate::mm::{FrameAlloc, PageAlloc, PageRangeAllocator};
+use crate::arch::aarch64::mm::paging::{BasePageSize, PageSize};
 use crate::scheduler::PerCoreSchedulerExt;
 use crate::scheduler::task::{Task, TaskFrame};
 use crate::{DEFAULT_STACK_SIZE, KERNEL_STACK_SIZE};
+use crate::arch::kernel::CURRENT_STACK;
+use crate::mm::stack_alloc::{allocate_stack, StackAllocation};
 
 #[derive(Debug)]
 #[repr(C, packed)]
@@ -94,156 +94,52 @@ pub(crate) struct State {
 	pub x30: u64,
 }
 
-pub struct BootStack {
-	/// Stack for kernel tasks
-	stack: VirtAddr,
-}
-
-pub struct CommonStack {
-	/// Start address of allocated virtual memory region
-	virt_addr: VirtAddr,
-	/// Start address of allocated virtual memory region
-	phys_addr: PhysAddr,
-	/// Total size of all stacks
-	total_size: usize,
-}
-
-pub enum TaskStacks {
-	Boot(BootStack),
-	Common(CommonStack),
+pub struct TaskStacks {
+	kernel_stack: StackAllocation,
+	user_stack: Option<StackAllocation>,
 }
 
 impl TaskStacks {
-	/// Size of the debug marker at the very top of each stack.
-	///
-	/// We have a marker at the very top of the stack for debugging (`0xdeadbeef`), which should not be overridden.
-	pub const MARKER_SIZE: usize = 0x10;
-
 	pub fn new(size: usize) -> Self {
 		let user_stack_size = if size < KERNEL_STACK_SIZE {
 			KERNEL_STACK_SIZE
 		} else {
 			size.align_up(BasePageSize::SIZE as usize)
 		};
-		let total_size = user_stack_size + DEFAULT_STACK_SIZE;
-		let layout = PageLayout::from_size(total_size + 3 * BasePageSize::SIZE as usize).unwrap();
-		let page_range = PageAlloc::allocate(layout).unwrap();
-		let virt_addr = VirtAddr::from(page_range.start());
-		let frame_layout = PageLayout::from_size(total_size).unwrap();
-		let frame_range = FrameAlloc::allocate(frame_layout)
-			.expect("Failed to allocate Physical Memory for TaskStacks");
-		let phys_addr = PhysAddr::from(frame_range.start());
-
-		debug!(
-			"Create stacks at {:p} with a size of {} KB",
-			virt_addr,
-			total_size >> 10
-		);
-
-		let mut flags = PageTableEntryFlags::empty();
-		flags.normal().writable().execute_disable();
 
-		// map kernel stack into the address space
-		crate::arch::mm::paging::map::<BasePageSize>(
-			virt_addr + BasePageSize::SIZE,
-			phys_addr,
-			DEFAULT_STACK_SIZE / BasePageSize::SIZE as usize,
-			flags,
-		);
+		let kernel_stack = allocate_stack(DEFAULT_STACK_SIZE);
+		let user_stack = allocate_stack(user_stack_size);
 
-		// map user stack into the address space
-		crate::arch::mm::paging::map::<BasePageSize>(
-			virt_addr + DEFAULT_STACK_SIZE + 2 * BasePageSize::SIZE,
-			phys_addr + DEFAULT_STACK_SIZE,
-			user_stack_size / BasePageSize::SIZE as usize,
-			flags,
-		);
-
-		// clear user stack
-		unsafe {
-			(virt_addr + DEFAULT_STACK_SIZE + 2 * BasePageSize::SIZE)
-				.as_mut_ptr::<u8>()
-				.write_bytes(0, user_stack_size);
+		TaskStacks {
+			kernel_stack, user_stack: Some(user_stack)
 		}
-
-		TaskStacks::Common(CommonStack {
-			virt_addr,
-			phys_addr,
-			total_size,
-		})
 	}
 
 	pub fn from_boot_stacks() -> TaskStacks {
-		let stack = VirtAddr::from_ptr(CURRENT_STACK_ADDRESS.load(Ordering::Relaxed));
-		debug!("Using boot stack {stack:p}");
-
-		TaskStacks::Boot(BootStack { stack })
-	}
-
-	pub fn get_user_stack_size(&self) -> usize {
-		match self {
-			TaskStacks::Boot(_) => 0,
-			TaskStacks::Common(stacks) => stacks.total_size - DEFAULT_STACK_SIZE,
-		}
-	}
+		let kernel_stack = if let Some(stack) = CURRENT_STACK.lock().as_ref() {
+			stack.weak()
+		} else {
+			let stack = VirtAddr::from_ptr(CURRENT_STACK_ADDRESS.load(Ordering::Relaxed));
+			debug!("Using boot stack {stack:p}");
 
-	pub fn get_user_stack(&self) -> VirtAddr {
-		match self {
-			TaskStacks::Boot(_) => VirtAddr::zero(),
-			TaskStacks::Common(stacks) => {
-				stacks.virt_addr + DEFAULT_STACK_SIZE + 2 * BasePageSize::SIZE
+			unsafe {
+				StackAllocation::new_external(stack, KERNEL_STACK_SIZE)
 			}
-		}
-	}
+		};
 
-	pub fn get_kernel_stack(&self) -> VirtAddr {
-		match self {
-			TaskStacks::Boot(stacks) => stacks.stack,
-			TaskStacks::Common(stacks) => stacks.virt_addr + BasePageSize::SIZE,
+		Self {
+			kernel_stack, user_stack: None
 		}
 	}
 
-	pub fn get_kernel_stack_size(&self) -> usize {
-		match self {
-			TaskStacks::Boot(_) => KERNEL_STACK_SIZE,
-			TaskStacks::Common(_) => DEFAULT_STACK_SIZE,
-		}
+	#[inline(always)]
+	pub fn get_user_stack(&self) -> Option<&StackAllocation> {
+		self.user_stack.as_ref()
 	}
-}
-
-impl Drop for TaskStacks {
-	fn drop(&mut self) {
-		// we should never deallocate a boot stack
-		match self {
-			TaskStacks::Boot(_) => {}
-			TaskStacks::Common(stacks) => {
-				debug!(
-					"Deallocating stacks at {:p} with a size of {} KB",
-					stacks.virt_addr,
-					stacks.total_size >> 10,
-				);
 
-				crate::arch::mm::paging::unmap::<BasePageSize>(
-					stacks.virt_addr,
-					stacks.total_size / BasePageSize::SIZE as usize + 3,
-				);
-				let range = PageRange::from_start_len(
-					stacks.virt_addr.as_usize(),
-					stacks.total_size + 3 * BasePageSize::SIZE as usize,
-				)
-				.unwrap();
-				unsafe {
-					PageAlloc::deallocate(range);
-				}
-
-				let range =
-					PageRange::from_start_len(stacks.phys_addr.as_usize(), stacks.total_size)
-						.unwrap();
-				unsafe {
-					FrameAlloc::deallocate(range);
-				}
-			}
-		}
+	#[inline(always)]
+	pub fn get_kernel_stack(&self) -> &StackAllocation {
+		&self.kernel_stack
 	}
 }
 
@@ -288,10 +184,7 @@ impl TaskFrame for Task {
 		}
 
 		unsafe {
-			// Set a marker for debugging at the very top.
-			let mut stack = self.stacks.get_kernel_stack() + self.stacks.get_kernel_stack_size()
-				- TaskStacks::MARKER_SIZE;
-			*stack.as_mut_ptr::<u64>() = 0xdead_beefu64;
+			let mut stack = self.stacks.get_kernel_stack().top_of_stack();
 
 			// Put the State structure expected by the ASM switch() function on the stack.
 			stack -= size_of::<State>();
@@ -314,12 +207,8 @@ impl TaskFrame for Task {
 
 			// Set the task's stack pointer entry to the stack we have just crafted.
 			self.last_stack_pointer = stack;
+			self.user_stack_pointer = self.stacks.get_user_stack().unwrap().top_of_stack();
 
-			// initialize user-level stack
-			self.user_stack_pointer = self.stacks.get_user_stack()
-				+ self.stacks.get_user_stack_size()
-				- TaskStacks::MARKER_SIZE;
-			*self.user_stack_pointer.as_mut_ptr::<u64>() = 0xdead_beefu64;
 			(*state).sp_el0 = self.user_stack_pointer.as_u64();
 		}
 	}

src/arch/aarch64/kernel/start.rs

@@ -10,8 +10,8 @@ use aarch64_cpu::asm::barrier::{SY, dsb};
 use hermit_entry::Entry;
 use hermit_entry::boot_info::RawBootInfo;
 
-use crate::arch::aarch64::kernel::scheduler::TaskStacks;
 use crate::{KERNEL_STACK_SIZE, env};
+use crate::mm::stack_alloc;
 
 /*
  * Memory types available.
@@ -102,7 +102,7 @@ pub unsafe extern "C" fn _start(boot_info: Option<&'static RawBootInfo>, cpu_id:
 		"b {pre_init}",
 
 		cpu_online = sym super::CPU_ONLINE,
-		stack_top_offset = const KERNEL_STACK_SIZE - TaskStacks::MARKER_SIZE,
+		stack_top_offset = const KERNEL_STACK_SIZE - stack_alloc::MARKER_SIZE,
 		current_stack_address = sym super::CURRENT_STACK_ADDRESS,
 		pre_init = sym pre_init,
 	)
@@ -239,7 +239,7 @@ pub(crate) unsafe extern "C" fn smp_start() -> ! {
 
 		mair_el1 = const mair(0x00, MT_DEVICE_nGnRnE) | mair(0x04, MT_DEVICE_nGnRE) | mair(0x0c, MT_DEVICE_GRE) | mair(0x44, MT_NORMAL_NC) | mair(0xff, MT_NORMAL),
 		tcr_bits = const tcr_size(VA_BITS) | TCR_TG1_4K | TCR_FLAGS,
-		stack_top_offset = const KERNEL_STACK_SIZE - TaskStacks::MARKER_SIZE,
+		stack_top_offset = const KERNEL_STACK_SIZE - stack_alloc::MARKER_SIZE,
 		current_stack_address = sym super::CURRENT_STACK_ADDRESS,
 		sctlr_el1 = const SCTLR_EL1,
 		ttbr0 = sym TTBR0,

src/arch/riscv64/kernel/core_local.rs

@@ -19,7 +19,7 @@ pub struct CoreLocal {
 	core_id: CoreId,
 	/// Scheduler of the current Core.
 	scheduler: Cell<*mut PerCoreScheduler>,
-	/// Start address of the kernel stack
+	/// Start address of the kernel stack (appears unused)
 	pub kernel_stack: Cell<u64>,
 	/// The core-local async executor.
 	ex: StaticLocalExecutor<RawSpinMutex, RawRwSpinLock>,

src/arch/riscv64/kernel/mod.rs

@@ -26,6 +26,7 @@ use crate::config::KERNEL_STACK_SIZE;
 use crate::env;
 use crate::init_cell::InitCell;
 use crate::mm::{FrameAlloc, PageRangeAllocator};
+use crate::mm::stack_alloc::allocate_stack;
 
 // Used to store information about available harts. The index of the hart in the vector
 // represents its CpuId and does not need to match its hart_id
@@ -148,10 +149,9 @@ pub fn boot_next_processor() {
 
 	{
 		debug!("Allocating stack for hard_id {next_hart_id}");
-		let frame_layout = PageLayout::from_size(KERNEL_STACK_SIZE).unwrap();
-		let frame_range =
-			FrameAlloc::allocate(frame_layout).expect("Failed to allocate boot stack for new core");
-		let stack = ptr::with_exposed_provenance_mut(frame_range.start());
+
+        let stack = allocate_stack(KERNEL_STACK_SIZE).leak();
+		let stack = ptr::with_exposed_provenance_mut(stack.stack_start().as_usize());
 		CURRENT_STACK_ADDRESS.store(stack, Ordering::Relaxed);
 	}