kernel: add memory token mapping per-instance

Author: Qix-

oro-kernel/src/iface/kernel/mem_token_v0.rs

@@ -1,23 +1,37 @@
 //! Allows for querying information about memory tokens.
+//!
+//! # ⚠️ Stability Note ⚠️
+//! **This interface is incomplete.** Mappings that work now may not work in the future.
+//!
+//! The interface currently **does not** check for executable mapping conflicts -
+//! i.e. code, data, and rodata program segments mapped in as part of the module.
+//!
+//! Mapping operations to set the base of a token that overlaps with those regions
+//! **will not fail**, but also **will not be mapped**, as there will be no page
+//! fault for accesses to those regions.
+//!
+//! In the future, setting a `base` that *does* conflict **will** fail. Please be
+//! extra careful about your base addresses and spans when using this interface
+//! in order to be future-proof.
 
-use oro_mem::mapper::MapError;
 use oro_sysabi::{key, syscall::Error as SysError};
 
 use super::KernelInterface;
-use crate::{arch::Arch, syscall::InterfaceResponse, tab::Tab, thread::Thread, token::Token};
+use crate::{
+	arch::Arch, instance::TokenMapError, syscall::InterfaceResponse, tab::Tab, thread::Thread,
+	token::Token,
+};
 
 /// Interface specific errors.
 #[derive(Debug, Clone, Copy)]
 #[repr(u64)]
 pub enum Error {
 	/// An address conflict (existing mapping) was encountered when mapping a token.
-	Conflict    = key!("conflict"),
+	Conflict   = key!("conflict"),
 	/// The requested address is not aligned to the page size.
-	NotAligned  = key!("align"),
+	NotAligned = key!("align"),
 	/// The requested address is out of the address space range.
-	OutOfRange  = key!("range"),
-	/// The system ran out of memory trying to service the request.
-	OutOfMemory = key!("oom"),
+	OutOfRange = key!("range"),
 }
 
 /// Version 0 of the memory token query interface.
@@ -83,16 +97,16 @@ impl KernelInterface for MemTokenV0 {
 						);
 					};
 
-					i.map_token(&token, virt).map_or_else(
+					i.try_map_token_at(&token, virt).map_or_else(
 						|err| {
 							InterfaceResponse::immediate(
 								SysError::InterfaceError,
 								match err {
-									MapError::Exists => Error::Conflict as u64,
-									MapError::OutOfMemory => Error::OutOfMemory as u64,
-									MapError::VirtNotAligned => Error::NotAligned as u64,
-									MapError::VirtOutOfRange
-									| MapError::VirtOutOfAddressSpaceRange => Error::OutOfRange as u64,
+									TokenMapError::Conflict => Error::Conflict as u64,
+									TokenMapError::VirtNotAligned => Error::NotAligned as u64,
+									TokenMapError::VirtOutOfRange => Error::OutOfRange as u64,
+									// NOTE(qix-): We already handled this at the beginning of the match.
+									TokenMapError::BadToken => unreachable!(),
 								},
 							)
 						},

oro-kernel/src/instance.rs

@@ -1,6 +1,9 @@
 //! Module instance types and functionality.
 
-use oro_mem::mapper::{AddressSegment, AddressSpace as _, MapError};
+use oro_mem::{
+	mapper::{AddressSegment, AddressSpace as _, MapError},
+	phys::PhysAddr,
+};
 
 use crate::{
 	AddressSpace, Kernel, UserHandle,
@@ -63,6 +66,13 @@ pub struct Instance<A: Arch> {
 	/// If a token is here, the instance is allowed to map it
 	/// into its address space.
 	tokens: Table<Tab<Token>>,
+	/// The virtual memory mappings of virtual addresses to tokens.
+	/// Values are `(token, page_id)` pairs.
+	// TODO(qix-): This assumes a 4096 byte page size, and is also
+	// TODO(qix-): not a very efficient data structure. It will be
+	// TODO(qix-): replaced with a more efficient data structure
+	// TODO(qix-): in the future.
+	token_vmap: Table<(Tab<Token>, usize)>,
 }
 
 impl<A: Arch> Instance<A> {
@@ -103,6 +113,7 @@ impl<A: Arch> Instance<A> {
 				handle,
 				data: TypeTable::new(),
 				tokens: Table::new(),
+				token_vmap: Table::new(),
 			})
 			.ok_or(MapError::OutOfMemory)?;
 
@@ -151,11 +162,110 @@ impl<A: Arch> Instance<A> {
 		self.tokens.insert_tab(token)
 	}
 
-	/// Maps a [`Token`] into the instance's address space.
+	/// Maps (sets the base of and reserves) a [`Token`] into the instance's address space.
 	///
-	/// Returns the address segment of the mapping.
-	pub fn map_token(&self, token: &Tab<Token>, virt: usize) -> Result<(), MapError> {
-		todo!("map token: {:016X} -> {virt:016X}", token.id());
+	/// **This does not immediately map in any memory.** It only marks the range
+	/// in the _internal kernel_ address space as reserved for the token, to be
+	/// committed later (typically via a page fault calling [`Instance::try_commit_token_at`]).
+	pub fn try_map_token_at(
+		&mut self,
+		token: &Tab<Token>,
+		virt: usize,
+	) -> Result<(), TokenMapError> {
+		if !self.tokens.contains(token.id()) {
+			return Err(TokenMapError::BadToken);
+		}
+
+		token.with(|t| {
+			match t {
+				Token::Normal(t) => {
+					debug_assert!(
+						t.page_size() == 4096,
+						"page size != 4096 is not implemented"
+					);
+					debug_assert!(
+						t.page_size().is_power_of_two(),
+						"page size is not a power of 2"
+					);
+
+					if (virt & (t.page_size() - 1)) != 0 {
+						return Err(TokenMapError::VirtNotAligned);
+					}
+
+					let segment = AddressSpace::<A>::user_data();
+
+					// Make sure that none of the tokens exist in the vmap.
+					for page_idx in 0..t.page_count() {
+						let page_base = virt + (page_idx * t.page_size());
+
+						if !virt_resides_within::<A>(&segment, page_base)
+							|| !virt_resides_within::<A>(&segment, page_base + t.page_size() - 1)
+						{
+							return Err(TokenMapError::VirtOutOfRange);
+						}
+
+						if self.token_vmap.contains(
+							u64::try_from(page_base).map_err(|_| TokenMapError::VirtOutOfRange)?,
+						) {
+							return Err(TokenMapError::Conflict);
+						}
+					}
+
+					// Everything's okay, map them into the vmap now.
+					for page_idx in 0..t.page_count() {
+						let page_base = virt + (page_idx * t.page_size());
+						// NOTE(qix-): We can use `as` here since we already check the page base above.
+						self.token_vmap
+							.insert(page_base as u64, (token.clone(), page_idx));
+					}
+
+					Ok(())
+				}
+			}
+		})
+	}
+
+	/// Commits a [`Token`] into the instance's address space at
+	/// the specified virtual address. Returns a mapping error
+	/// if the mapping could not be completed.
+	///
+	/// **The `maybe_unaligned_virt` parameter is not guaranteed to be aligned
+	/// to any page boundary.** In most cases, it is coming directly from a
+	/// userspace application (typically via a fault).
+	///
+	/// The `Token` must have been previously mapped via [`Instance::try_map_token_at`],
+	/// or else this method will fail.
+	pub fn try_commit_token_at(&self, maybe_unaligned_virt: usize) -> Result<(), TryCommitError> {
+		// TODO(qix-): We always assume a 4096 page boundary. This will change in the future.
+		let virt = maybe_unaligned_virt & !0xFFF;
+
+		if let Some((token, page_idx)) = u64::try_from(virt)
+			.ok()
+			.and_then(|virt| self.token_vmap.get(virt))
+		{
+			token.with_mut(|t| {
+				match t {
+					Token::Normal(t) => {
+						debug_assert!(*page_idx < t.page_count());
+						debug_assert!(
+							t.page_size() == 4096,
+							"page size != 4096 is not implemented"
+						);
+						let page_base = virt + (*page_idx * t.page_size());
+						let segment = AddressSpace::<A>::user_data();
+						let phys = t
+							.get_or_allocate(*page_idx)
+							.ok_or(TryCommitError::MapError(MapError::OutOfMemory))?;
+						segment
+							.map(self.handle.mapper(), page_base, phys.address_u64())
+							.map_err(TryCommitError::MapError)?;
+						Ok(())
+					}
+				}
+			})
+		} else {
+			Err(TryCommitError::BadVirt)
+		}
 	}
 
 	/// Returns the instance's address space handle.
@@ -177,3 +287,37 @@ impl<A: Arch> Instance<A> {
 		&mut self.data
 	}
 }
+
+/// An error returned by [`Instance::try_map_token_at`].
+#[derive(Debug, Clone, Copy)]
+pub enum TokenMapError {
+	/// The virtual address is not aligned.
+	VirtNotAligned,
+	/// The virtual address is out of range for the
+	/// thread's address space.
+	VirtOutOfRange,
+	/// The token was not found for the instance.
+	BadToken,
+	/// The mapping conflicts (overlaps) with another mapping.
+	Conflict,
+}
+
+/// Checks if the given virtual address resides within the given address segment.
+#[inline]
+fn virt_resides_within<A: Arch>(
+	segment: &<AddressSpace<A> as ::oro_mem::mapper::AddressSpace>::UserSegment,
+	virt: usize,
+) -> bool {
+	// NOTE(qix-): Range is *inclusive*.
+	let (first, last) = segment.range();
+	virt >= first && virt <= last
+}
+
+/// An error returned by [`Instance::try_commit_token_at`].
+#[derive(Debug, Clone, Copy)]
+pub enum TryCommitError {
+	/// The virtual address was not found in the virtual map.
+	BadVirt,
+	/// Mapping the token failed.
+	MapError(MapError),
+}

oro-kernel/src/scheduler.rs

@@ -261,15 +261,25 @@ impl<A: Arch> Scheduler<A> {
 		let Some(thread) = self.current.take() else {
 			panic!("event_page_fault() called with no current thread");
 		};
+
 		let id = thread.id();
-		unsafe {
-			thread.with_mut(|t| t.terminate());
-		}
-		dbg_warn!(
-			"thread {:#016X} terminated due to page fault: {fault_type:?} at {vaddr:016X}",
-			id
-		);
-		let switch = Switch::from_schedule_action(self.pick_user_thread(), Some(id));
+
+		let instance = thread.with(|t| t.instance().clone());
+
+		let switch = if instance.with_mut(|i| i.try_commit_token_at(vaddr)).is_err() {
+			unsafe {
+				thread.with_mut(|t| t.terminate());
+			}
+			dbg_warn!(
+				"thread {:#016X} terminated due to page fault: {fault_type:?} at {vaddr:016X}",
+				id
+			);
+			Switch::from_schedule_action(self.pick_user_thread(), Some(id))
+		} else {
+			self.current = Some(thread.clone());
+			Switch::UserResume(thread, None)
+		};
+
 		self.kernel.handle().schedule_timer(1000);
 		switch
 	}

oro-kernel/src/table.rs

@@ -97,6 +97,12 @@ impl<T: Sized, Alloc: Allocator + Default> Table<T, Alloc> {
 	pub fn remove(&mut self, id: u64) -> Option<T> {
 		self.0.remove(&id)
 	}
+
+	/// Returns whether or not the given key exists in the table.
+	#[inline]
+	pub fn contains(&self, key: u64) -> bool {
+		self.0.contains_key(&key)
+	}
 }
 
 /// A [`Table`] wrapper that allows for artibtrary singleton values

oro-kernel/src/token.rs

@@ -75,7 +75,11 @@ impl NormalToken {
 	pub(crate) fn new_4kib(pages: usize) -> Self {
 		Self {
 			page_size:  4096,
-			phys_addrs: Vec::with_capacity(pages),
+			phys_addrs: {
+				let mut v = Vec::with_capacity(pages);
+				v.resize_with(pages, || None);
+				v
+			},
 			committed:  0,
 		}
 	}