move MMU disable code to assembly

Author: m4tx

src/arch.rs

@@ -8,7 +8,7 @@
 
 //! Architecture-specific code.
 
-use core::arch::asm;
+use core::arch::{asm, naked_asm};
 
 /// Data Synchronization Barrier.
 pub fn dsb() {
@@ -35,6 +35,7 @@ pub fn isb() {
 }
 
 /// Invalidate all instruction caches.
+#[allow(dead_code)]
 pub fn ic_iallu() {
     // SAFETY: `ic iallu` is always safe.
     unsafe {
@@ -43,6 +44,7 @@ pub fn ic_iallu() {
 }
 
 /// Invalidate all TLB entries for EL2.
+#[allow(dead_code)]
 pub fn tlbi_alle2is() {
     // SAFETY: `tlbi alle2is` is always safe.
     unsafe {
@@ -153,7 +155,33 @@ sys_reg!(elr_el2);
 sys_reg!(sp_el1);
 sys_reg!(mpidr_el1);
 
-pub(super) fn disable_mmu_and_caches() {
+/// Disable MMU and caches.
+/// 
+/// # Safety
+/// 
+/// It is not sound to execute arbitrary Rust code after disabling the data cache, so this
+/// should only ever be called from assembly.
+/// This disables MMU, so the caller must ensure that the code will be executable at the
+/// same address.
+#[unsafe(naked)]
+pub unsafe extern "C" fn disable_mmu_and_caches() {
+    naked_asm!(
+        "mov x28, x30",
+        "bl {invalidate_caches}",
+        "mov x30, x28",
+        "msr sctlr_el2, x0",
+        "dsb sy",
+        "isb",
+        "ic iallu",
+        "tlbi alle2is",
+        "dsb sy",
+        "isb",
+        "ret",
+        invalidate_caches = sym invalidate_caches,
+    );
+}
+
+unsafe extern "C" fn invalidate_caches() -> u64 {
     invalidate_dcache();
 
     // Disable MMU and caches
@@ -165,21 +193,8 @@ pub(super) fn disable_mmu_and_caches() {
     sctlr &= !sctlr_el2::M; // MMU Enable
     sctlr &= !sctlr_el2::C; // Data Cache Enable
     sctlr &= !sctlr_el2::I; // Instruction Cache Enable
-    // SAFETY: We assume we have an identity mapped pagetables for the currently running
-    // code, so disabling MMU is safe.
-    unsafe {
-        sctlr_el2::write(sctlr);
-    }
-    dsb();
-    isb();
-
-    // Invalidate I-cache
-    ic_iallu();
-    tlbi_alle2is();
 
-    // Final synchronization
-    dsb();
-    isb();
+    sctlr
 }
 
 /// Invalidate D-cache by set/way to the point of coherency.

src/hypervisor.rs

@@ -9,7 +9,7 @@
 use core::arch::naked_asm;
 
 use aarch64_rt::{RegisterStateRef, Stack};
-use alloc::{boxed::Box, collections::btree_map::BTreeMap};
+use alloc::boxed::Box;
 use log::debug;
 use spin::mutex::SpinMutex;
 
@@ -26,8 +26,6 @@ use crate::{arch::{self, esr, far}, platform::{Platform, PlatformImpl}, simple_m
 /// address for EL1 execution that never returns.
 /// This function must be called in EL2.
 pub unsafe fn entry_point_el1(arg0: u64, arg1: u64, arg2: u64, arg3: u64, entry_point: u64) -> ! {
-    arch::disable_mmu_and_caches();
-
     // Setup EL1
     // SAFETY: We are configuring HCR_EL2 to allow EL1 execution.
     unsafe {
@@ -98,7 +96,13 @@ pub unsafe fn entry_point_el1(arg0: u64, arg1: u64, arg2: u64, arg3: u64, entry_
 /// This function must be called in EL2.
 #[unsafe(naked)]
 pub unsafe extern "C" fn eret_to_el1(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
-    naked_asm!("eret");
+    naked_asm!(
+        "mov x19, x0",
+        "bl {disable_mmu_and_caches}",
+        "mov x0, x19",
+        "eret",
+        disable_mmu_and_caches = sym arch::disable_mmu_and_caches,
+    );
 }
 
 pub fn handle_sync_lower(mut register_state: RegisterStateRef) {
@@ -349,6 +353,6 @@ fn get_secondary_stack(mpidr: u64) -> *mut Stack<SECONDARY_STACK_PAGE_COUNT> {
     if let Some(stack) = stack_map.get_mut(&mpidr) {
         &raw mut **stack
     } else {
-        &raw mut **stack_map.insert(mpidr, Default::default())
+        &raw mut **stack_map.insert(mpidr, Box::default())
     }
 }

src/main.rs

@@ -30,7 +30,6 @@ use log::{LevelFilter, info};
 use spin::mutex::{SpinMutex, SpinMutexGuard};
 
 use crate::{
-    arch::disable_mmu_and_caches,
     exceptions::Exceptions,
     platform::{BootMode, Platform, PlatformImpl},
 };
@@ -108,23 +107,14 @@ fn add_to_heap<const ORDER: usize>(heap: &mut Heap<ORDER>, range: &'static mut [
 /// # Safety
 ///
 /// `NEXT_IMAGE` must point to a valid executable piece of code which never returns.
-unsafe fn run_payload_el2(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
-    disable_mmu_and_caches();
-    // SAFETY: The caller guarantees that `NEXT_IMAGE` points to a valid executable piece of code which never returns.
-    unsafe {
-        jump_to_payload(x0, x1, x2, x3);
-    }
-}
-
-/// Jumps to the payload.
-///
-/// # Safety
-///
-/// `NEXT_IMAGE` must point to a valid executable piece of code which never returns.
 #[unsafe(naked)]
-unsafe extern "C" fn jump_to_payload(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
+unsafe extern "C" fn run_payload_el2(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
     naked_asm!(
+        "mov x19, x0",
+        "bl {disable_mmu_and_caches}",
+        "mov x0, x19",
         "b {next_image}",
+        disable_mmu_and_caches = sym arch::disable_mmu_and_caches,
         next_image = sym NEXT_IMAGE,
     );
 }