Add stage-2 pagetables

This makes the RITM memory invisible from the guest.

Author: m4tx

.github/workflows/rust.yml

@@ -31,3 +31,26 @@ jobs:
       - uses: actions/checkout@v6
       - name: Format Rust code
         run: cargo fmt --all -- --check
+
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - name: Install aarch64 toolchain
+        uses: dtolnay/rust-toolchain@v1
+        with:
+          toolchain: stable
+          targets: aarch64-unknown-none
+          components: llvm-tools
+      - name: Install QEMU
+        run: |
+          sudo apt-get update && \
+            sudo apt-get install -y --no-install-recommends qemu-system-arm ipxe-qemu seabios
+      - name: Install cargo-binutils
+        uses: taiki-e/install-action@v2
+        with:
+          tool: cargo-binutils
+      - name: Prepare QEMU
+        run: sudo chown $(whoami):$(whoami) /dev/vhost-vsock && sudo chmod g+rw /dev/vhost-vsock
+      - name: Run the tests
+        run: make test

Cargo.lock

@@ -1,3 +1,9 @@
+[workspace]
+members = [
+    ".",
+    "tests/isolation_test",
+]
+
 [package]
 name = "ritm"
 version = "0.1.0"
@@ -10,7 +16,7 @@ keywords = ["arm", "aarch64", "cortex-a", "osdev"]
 categories = ["embedded", "no-std"]
 
 [dependencies]
-aarch64-paging = { version = "0.11", default-features = false }
+aarch64-paging = "0.11"
 aarch64-rt = { version = "0.4", features = ["el2", "exceptions", "initial-pagetable", "psci"], default-features = false }
 arm-pl011-uart = "0.4"
 arm-psci = "0.2"
@@ -25,6 +31,7 @@ spin = { version = "0.10", features = ["lazy", "once", "spin_mutex"], default-fe
 
 [patch.crates-io]
 aarch64-rt = { git = "https://github.com/m4tx/aarch64-rt.git", rev = "b5783edd5e5d5555105946c462646572d3454d99" }
+aarch64-paging = { git = "https://github.com/m4tx/aarch64-paging.git", rev = "6e28cdee945701821b7b5721b500e8e93f1a99b5" }
 
 [lints.rust]
 deprecated-safe = "warn"

Cargo.toml

@@ -13,7 +13,7 @@ PAYLOAD ?= payload.bin
 QEMU_BIN := target/ritm.qemu.bin
 QEMU_RUSTFLAGS := "--cfg platform=\"qemu\""
 
-.PHONY: all build.qemu clean clippy qemu
+.PHONY: all build.qemu clean clippy qemu test
 
 all: $(QEMU_BIN)
 
@@ -37,6 +37,9 @@ qemu: $(QEMU_BIN)
 	  -device virtconsole,chardev=char0 \
 	  -device vhost-vsock-device,id=virtiosocket0,guest-cid=102
 
+test:
+	tests/isolation_test.py
+
 clean:
 	cargo clean
 	rm -f target/*.bin

Makefile

@@ -34,6 +34,14 @@ pub fn isb() {
     }
 }
 
+/// TLBI VMALLS12E1 - VMID-based Stage-1/Stage-2 combined invalidation for the EL1&0 regime.
+pub fn tlbi_vmalls12e1() {
+    // SAFETY: TLBI VMALLS12E1 is always safe.
+    unsafe {
+        asm!("tlbi vmalls12e1", options(nostack, preserves_flags));
+    }
+}
+
 pub fn esr() -> u64 {
     let mut esr: u64;
     // SAFETY: Reading esr is always safe.
@@ -106,7 +114,8 @@ sys_reg!(ccsidr_el1);
 sys_reg!(hcr_el2, {
     RW: 1 << 31,
     TSC: 1 << 19,
-    IMO: 1 << 4
+    IMO: 1 << 4,
+    VM: 1 << 0
 });
 sys_reg!(cntvoff_el2);
 sys_reg!(cnthctl_el2, {
@@ -120,6 +129,20 @@ sys_reg!(spsr_el2, {
 sys_reg!(elr_el2);
 sys_reg!(sp_el1);
 sys_reg!(mpidr_el1);
+sys_reg!(vtcr_el2, {
+    PS_40BIT: 2 << 16,
+    TG0_4KB: 0 << 14,
+    SH0_INNER: 3 << 12,
+    ORGN0_WB_RA_WA: 1 << 10,
+    IRGN0_WB_RA_WA: 1 << 8,
+    SL0_L0: 2 << 6,
+    T0SZ_40BIT: 24
+});
+sys_reg!(vbar_el1);
+sys_reg!(elr_el1);
+sys_reg!(spsr_el1);
+sys_reg!(esr_el1);
+sys_reg!(far_el1);
 
 /// Disables MMU and caches.
 ///

src/arch.rs

@@ -33,9 +33,12 @@ pub unsafe fn entry_point_el1(arg0: u64, arg1: u64, arg2: u64, arg3: u64, entry_
     // Setup EL1
     // SAFETY: We are configuring HCR_EL2 to allow EL1 execution.
     unsafe {
+        setup_stage2();
+
         let mut hcr = arch::hcr_el2::read();
         hcr |= arch::hcr_el2::RW;
         hcr |= arch::hcr_el2::TSC;
+        hcr |= arch::hcr_el2::VM;
         hcr &= !arch::hcr_el2::IMO;
         arch::hcr_el2::write(hcr);
     }
@@ -81,6 +84,37 @@ pub unsafe fn entry_point_el1(arg0: u64, arg1: u64, arg2: u64, arg3: u64, entry_
     }
 }
 
+fn setup_stage2() {
+    debug!("Setting up stage 2 page table");
+    let idmap = Box::new(PlatformImpl::make_stage2_pagetable());
+
+    let root_pa = idmap.root_address().0;
+    debug!("Root PA: {root_pa:#x}");
+    let idmap = Box::leak(idmap);
+
+    // Activate the page table
+    // SAFETY: We are initializing the Stage 2 translation. The guest is not running yet.
+    unsafe {
+        let ttbr = idmap.activate();
+        debug!("idmap.activate() returned ttbr={ttbr:#x}");
+
+        let vtcr = arch::vtcr_el2::PS_40BIT
+            | arch::vtcr_el2::TG0_4KB
+            | arch::vtcr_el2::SH0_INNER
+            | arch::vtcr_el2::ORGN0_WB_RA_WA
+            | arch::vtcr_el2::IRGN0_WB_RA_WA
+            | arch::vtcr_el2::SL0_L0
+            | arch::vtcr_el2::T0SZ_40BIT;
+        debug!("Writing VTCR_EL2={vtcr:#x}...");
+        arch::vtcr_el2::write(vtcr);
+
+        arch::tlbi_vmalls12e1();
+        arch::dsb();
+        arch::isb();
+        debug!("Stage 2 activation complete.");
+    }
+}
+
 /// Returns to EL1.
 ///
 /// This function executes the `eret` instruction to return to EL1 with the provided arguments.
@@ -122,13 +156,53 @@ pub fn handle_sync_lower(mut register_state: RegisterStateRef) {
                 }
             }
         }
-        ExceptionClass::Unknown(_) => {
+        ExceptionClass::Unknown(val) => {
             panic!(
-                "Unexpected sync_lower, far={:#x}, register_state={register_state:?}",
+                "Unexpected sync_lower, esr={esr:#x}, ec={val:#x}, far={:#x}, register_state={register_state:?}",
                 far(),
             );
         }
+        ExceptionClass::DataAbortLowerEL => {
+            inject_data_abort(&mut register_state);
+        }
+    }
+}
+
+fn inject_data_abort(register_state: &mut RegisterStateRef) {
+    // SAFETY: We are modifying the saved register state to redirect execution.
+    let regs = unsafe { register_state.get_mut() };
+    let fault_addr = far();
+    let syndrome = esr();
+
+    debug!("Injecting data abort to guest: fault_addr={fault_addr:#x}, syndrome={syndrome:#x}");
+
+    // Read guest VBAR
+    let vbar = arch::vbar_el1::read();
+    assert!(
+        vbar != 0,
+        "Guest VBAR_EL1 is 0, cannot inject data abort. Fault addr: {fault_addr:#x}"
+    );
+    let handler = vbar + 0x200; // Current EL with SPx Sync
+
+    // Save current context to guest EL1 regs
+    // SAFETY: We are accessing EL1 system registers to inject exception.
+    unsafe {
+        arch::elr_el1::write(regs.elr as u64);
+        arch::spsr_el1::write(regs.spsr);
+        arch::esr_el1::write(syndrome);
+        arch::far_el1::write(fault_addr);
+    }
+
+    // Redirect execution
+    #[expect(
+        clippy::cast_possible_truncation,
+        reason = "only 64-bit target is supported"
+    )]
+    {
+        regs.elr = handler as usize;
     }
+    // Mask all interrupts (DAIF) and set mode to EL1h (0x5)
+    regs.spsr = 0x3C5;
 }
 
 const AARCH64_INSTRUCTION_LENGTH: usize = 4;
@@ -311,6 +385,8 @@ enum ExceptionClass {
     HvcTrappedInAArch64,
     /// SMC instruction execution in `AArch64` state.
     SmcTrappedInAArch64,
+    /// Data Abort taken without a change in Exception Level.
+    DataAbortLowerEL,
     #[allow(unused)]
     /// Unknown exception class.
     Unknown(u8),
@@ -321,6 +397,7 @@ impl ExceptionClass {
         match value {
             0x16 => Self::HvcTrappedInAArch64,
             0x17 => Self::SmcTrappedInAArch64,
+            0x24 => Self::DataAbortLowerEL,
             _ => Self::Unknown(value),
         }
     }

src/hypervisor.rs

@@ -43,9 +43,15 @@ const LOG_LEVEL: LevelFilter = LevelFilter::Info;
 const HEAP_SIZE: usize = 40 * PAGE_SIZE;
 static HEAP: SpinMutex<[u8; HEAP_SIZE]> = SpinMutex::new([0; HEAP_SIZE]);
 
+const SHARED_HEAP_SIZE: usize = 16 * PAGE_SIZE;
+static SHARED_HEAP: SpinMutex<[u8; SHARED_HEAP_SIZE]> = SpinMutex::new([0; SHARED_HEAP_SIZE]);
+
 #[global_allocator]
 static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::new();
 
+/// Heap allocator for data that needs to be shared between RITM and the guest running in EL1.
+pub static SHARED_HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::new();
+
 #[repr(align(0x200000))] // Linux requires 2MB alignment
 struct AlignImage<T>(T);
 
@@ -73,6 +79,12 @@ fn main(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
         SpinMutexGuard::leak(HEAP.try_lock().expect("failed to lock heap")).as_mut_slice(),
     );
 
+    add_to_heap(
+        SHARED_HEAP_ALLOCATOR.lock().deref_mut(),
+        SpinMutexGuard::leak(SHARED_HEAP.try_lock().expect("failed to lock shared heap"))
+            .as_mut_slice(),
+    );
+
     let fdt_address = x0 as *const u8;
     // SAFETY: We trust that the FDT pointer we were given is valid, and this is the only time we
     // use it.
@@ -138,3 +150,19 @@ unsafe fn run_payload_el1(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
         hypervisor::entry_point_el1(x0, x1, x2, x3, &raw const NEXT_IMAGE.0 as u64);
     }
 }
+
+/// Allocates a buffer from the shared heap.
+///
+/// # Panics
+///
+/// Panics if the requested size is invalid or if the allocation fails.
+pub fn shared_alloc(size: usize) -> &'static mut [u8] {
+    use core::alloc::Layout;
+    let layout = Layout::from_size_align(size, PAGE_SIZE).expect("invalid layout");
+    let ptr = SHARED_HEAP_ALLOCATOR
+        .lock()
+        .alloc(layout)
+        .expect("failed to allocate from shared heap");
+    // SAFETY: The pointer is valid and represents the requested size.
+    unsafe { core::slice::from_raw_parts_mut(ptr.as_ptr(), size) }
+}

src/main.rs

@@ -7,21 +7,21 @@
 // except according to those terms.
 
 use crate::platform::PlatformImpl;
-use aarch64_paging::descriptor::Attributes;
+use aarch64_paging::descriptor::Stage1Attributes;
 use aarch64_rt::initial_pagetable;
 
 /// Attributes to use for device memory in the initial identity map.
-pub const DEVICE_ATTRIBUTES: Attributes = Attributes::VALID
-    .union(Attributes::ATTRIBUTE_INDEX_0)
-    .union(Attributes::ACCESSED)
-    .union(Attributes::UXN);
+pub const DEVICE_ATTRIBUTES: Stage1Attributes = Stage1Attributes::VALID
+    .union(Stage1Attributes::ATTRIBUTE_INDEX_0)
+    .union(Stage1Attributes::ACCESSED)
+    .union(Stage1Attributes::UXN);
 
 /// Attributes to use for normal memory in the initial identity map.
-pub const MEMORY_ATTRIBUTES: Attributes = Attributes::VALID
-    .union(Attributes::ATTRIBUTE_INDEX_1)
-    .union(Attributes::INNER_SHAREABLE)
-    .union(Attributes::ACCESSED)
-    .union(Attributes::NON_GLOBAL);
+pub const MEMORY_ATTRIBUTES: Stage1Attributes = Stage1Attributes::VALID
+    .union(Stage1Attributes::ATTRIBUTE_INDEX_1)
+    .union(Stage1Attributes::INNER_SHAREABLE)
+    .union(Stage1Attributes::ACCESSED)
+    .union(Stage1Attributes::NON_GLOBAL);
 
 // The initial hardcoded page table used before the Rust code starts and activates the main page
 // table.

src/pagetable.rs

@@ -9,6 +9,8 @@
 #[cfg(platform = "qemu")]
 mod qemu;
 
+use aarch64_paging::descriptor::Stage2Attributes;
+use aarch64_paging::idmap::IdMap;
 use dtoolkit::fdt::Fdt;
 use embedded_io::{Write, WriteReady};
 #[cfg(platform = "qemu")]
@@ -45,6 +47,12 @@ pub trait Platform {
     fn modify_dt(&self, fdt: Fdt<'static>) -> Fdt<'static> {
         fdt
     }
+
+    /// Create stage-2 page table for use by the guest for use when booting the payload at EL1.
+    ///
+    /// The page table should typically unmap the part of the memory where RITM resides, so that
+    /// the guest cannot interact with it in any way.
+    fn make_stage2_pagetable() -> IdMap<Stage2Attributes>;
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]