aarch64: reclaim TT0 physical pages after transfer

oro-arch-aarch64/src/arch.rs

@@ -2,7 +2,10 @@
 
 #![allow(clippy::inline_always, clippy::verbose_bit_mask)]
 
-use crate::{mem::address_space::AddressSpaceLayout, xfer::TransferToken};
+use crate::{
+	mem::{address_space::AddressSpaceLayout, paging::PageTable},
+	xfer::TransferToken,
+};
 use core::{
 	arch::asm,
 	fmt::{self, Write},
@@ -176,14 +179,92 @@ unsafe impl Arch for Aarch64 {
 
 	unsafe fn after_transfer<A, P>(
 		_mapper: &<<Self as Arch>::AddressSpace as AddressSpace>::SupervisorHandle,
-		_translator: &P,
-		_alloc: &mut A,
-		_is_primary: bool,
+		translator: &P,
+		alloc: &mut A,
+		is_primary: bool,
 	) where
 		A: PageFrameAllocate + PageFrameFree,
 		P: PhysicalAddressTranslator,
 	{
-		// TODO(qix-)
+		// NOTE(qix-): `_mapper` isn't useful to use because it points to TT1.
+		// NOTE(qix-): We're unmapping all of TT0.
+		let tt0_phys = crate::asm::load_ttbr0();
+		let l0_virt = translator.to_virtual_addr(tt0_phys);
+		let l0 = &mut *(l0_virt as *mut PageTable);
+
+		if is_primary {
+			// TODO(qix-): This will absolutely need to be updated when different addressing
+			// TODO(qix-): types or more than 4 page table levels are supported. Even though
+			// TODO(qix-): the 'official' init routine has this tightly controlled, we can't
+			// TODO(qix-): guarantee that it'll never change.
+			for l0_idx in 0..512 {
+				let l0_entry = &mut l0[l0_idx];
+				if l0_entry.valid() {
+					// SAFETY(qix-): Guaranteed to be valid by the above checks.
+					let l1_phys = l0_entry.address(0).unwrap();
+					let l1_virt = translator.to_virtual_addr(l1_phys);
+					let l1 = &mut *(l1_virt as *mut PageTable);
+
+					for l1_idx in 0..512 {
+						let l1_entry = &mut l1[l1_idx];
+						if l1_entry.valid() {
+							// SAFETY(qix-): Guaranteed to be valid by the above checks.
+							let l2_phys = l1_entry.address(1).unwrap();
+							let l2_virt = translator.to_virtual_addr(l2_phys);
+							let l2 = &mut *(l2_virt as *mut PageTable);
+
+							for l2_idx in 0..512 {
+								let l2_entry = &mut l2[l2_idx];
+								if l2_entry.valid() {
+									// SAFETY(qix-): Guaranteed to be valid by the above checks.
+									let l3_phys = l2_entry.address(2).unwrap();
+									let l3_virt = translator.to_virtual_addr(l3_phys);
+									let l3 = &mut *(l3_virt as *mut PageTable);
+
+									for l3_idx in 0..512 {
+										let l3_entry = &mut l3[l3_idx];
+										if l3_entry.valid() {
+											// SAFETY(qix-): Guaranteed to be valid by the above checks.
+											let page_phys = l3_entry.address(3).unwrap();
+											alloc.free(page_phys);
+										}
+									}
+
+									alloc.free(l3_phys);
+								}
+							}
+
+							alloc.free(l2_phys);
+						}
+					}
+
+					alloc.free(l1_phys);
+				}
+
+				// Make sure other cores see the writes.
+				Self::strong_memory_barrier();
+			}
+		} else {
+			// We simply need to reset the L4 entries in the TT0 table.
+			// All of the addresses they have pointed to have been freed
+			// by the primary.
+			//
+			// SAFETY(qix-): The specification of this method guarantees that
+			// SAFETY(qix-): this method is called on the primary core first.
+			// SAFETY(qix-): This means that the primary core has already freed
+			// SAFETY(qix-): all of the pages that the secondary core's L0
+			// SAFETY(qix-): entries point to, and those entries are now zombies.
+			// SAFETY(qix-): We can further guarantee this is the case since
+			// SAFETY(qix-): the secondary cores shallow clone the L0 table when
+			// SAFETY(qix-): bootstrapping.
+
+			for l0_idx in 0..512 {
+				l0[l0_idx].reset();
+			}
+		}
+
+		alloc.free(tt0_phys);
+		crate::asm::store_ttbr0(0);
 	}
 }
 

oro-arch-aarch64/src/asm.rs

@@ -34,7 +34,6 @@ pub unsafe fn invalidate_tlb_el1(virt: usize) {
 /// all cores in all execution contexts. Be VERY CAREFUL
 /// that you control the value in the `TTBR0_EL1` register
 /// prior to calling this again.
-#[allow(dead_code)]
 pub unsafe fn load_ttbr0() -> u64 {
 	let ttbr0_el1: u64;
 	unsafe {
@@ -62,7 +61,6 @@ pub unsafe fn load_ttbr0() -> u64 {
 /// all cores in all execution contexts. Be VERY CAREFUL
 /// that you control the value in the `TTBR1_EL1` register
 /// prior to calling this again.
-#[allow(dead_code)]
 pub unsafe fn load_ttbr1() -> u64 {
 	let ttbr1_el1: u64;
 	unsafe {
@@ -80,7 +78,6 @@ pub unsafe fn load_ttbr1() -> u64 {
 /// Caller must ensure that the physical address is valid
 /// memory, and that page tables written to that
 /// address are valid.
-#[allow(dead_code)]
 pub unsafe fn store_ttbr0(phys: u64) {
 	unsafe {
 		asm!(
@@ -89,19 +86,3 @@ pub unsafe fn store_ttbr0(phys: u64) {
 		);
 	}
 }
-
-/// Stores a new physical address into the `TTBR1_EL1` register.
-///
-/// # Safety
-/// Caller must ensure that the physical address is valid
-/// memory, and that page tables written to that
-/// address are valid.
-#[allow(dead_code)]
-pub unsafe fn store_ttbr1(phys: u64) {
-	unsafe {
-		asm!(
-			"msr TTBR1_EL1, {0:x}",
-			in(reg) phys,
-		);
-	}
-}

oro-arch-aarch64/src/lib.rs

@@ -16,6 +16,11 @@
 //!   to set it to 16.
 //! - `TTBR0_EL1` must be left undefined or set to 0 and not be relied upon for any execution,
 //!   as the initialization subroutine will overwrite it.
+//!
+//! ### After Transfer Behavior
+//! All of TT0 is unmapped and TTBR0 is set to 0. This, however, means nothing to the
+//! preboot environment, as the preboot environment MUST NOT rely on TTBR0 for any resource
+//! allocation or mapping.
 #![no_std]
 #![deny(
 	missing_docs,