From c909802331a7fbe13e2a33978ad882b3b24dd440 Mon Sep 17 00:00:00 2001
From: Putta Khunchalee <ultimaweapon@outlook.com>
Date: Thu, 22 Aug 2024 23:04:48 +0700
Subject: [PATCH] Sets TTBR1_EL1 for Mac M1 (#942)

---
 src/core/src/vmm/hv/macos/cpu.rs | 33 +++++++++++++++++++++++++++-----
 src/core/src/vmm/hv/mod.rs       | 10 ++++++++--
 src/core/src/vmm/mod.rs          |  9 +++++++--
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/src/core/src/vmm/hv/macos/cpu.rs b/src/core/src/vmm/hv/macos/cpu.rs
index 711a78c3..6690398f 100644
--- a/src/core/src/vmm/hv/macos/cpu.rs
+++ b/src/core/src/vmm/hv/macos/cpu.rs
@@ -132,6 +132,7 @@ impl<'a> Cpu for HfCpu<'a> {
         Ok(HfStates {
             cpu: self,
             tcr_el1: State::None,
+            ttbr1_el1: State::None,
             sp_el1: State::None,
         })
     }
@@ -206,6 +207,8 @@ pub struct HfStates<'a, 'b> {
     #[cfg(target_arch = "aarch64")]
     tcr_el1: State<u64>,
     #[cfg(target_arch = "aarch64")]
+    ttbr1_el1: State<u64>,
+    #[cfg(target_arch = "aarch64")]
     sp_el1: State<u64>,
 }
 
@@ -280,14 +283,24 @@ impl<'a, 'b> CpuStates for HfStates<'a, 'b> {
     }
 
     #[cfg(target_arch = "aarch64")]
-    fn set_tcr_el1(&mut self, t0sz: u8, t1sz: u8) {
+    fn set_tcr_el1(&mut self, ips: u8, a1: bool, t0sz: u8, t1sz: u8) {
+        let ips: u64 = ips.into();
+        let a1: u64 = a1.into();
         let t0sz: u64 = t0sz.into();
         let t1sz: u64 = t1sz.into();
 
+        assert_eq!(ips & 0b11111000, 0);
         assert_eq!(t0sz & 0b11000000, 0);
         assert_eq!(t1sz & 0b11000000, 0);
 
-        self.tcr_el1 = State::Dirty(t1sz << 16 | t0sz);
+        self.tcr_el1 = State::Dirty(ips << 32 | a1 << 22 | t1sz << 16 | t0sz);
+    }
+
+    #[cfg(target_arch = "aarch64")]
+    fn set_ttbr1_el1(&mut self, baddr: usize) {
+        assert_eq!(baddr & 0xFFFF000000000001, 0);
+
+        self.ttbr1_el1 = State::Dirty(baddr.try_into().unwrap());
     }
 
     #[cfg(target_arch = "aarch64")]
@@ -338,7 +351,9 @@ impl<'a, 'b> CpuStates for HfStates<'a, 'b> {
     #[cfg(target_arch = "aarch64")]
     fn commit(self) -> Result<(), Self::Err> {
         use hv_sys::{
-            hv_sys_reg_t_HV_SYS_REG_SP_EL1, hv_sys_reg_t_HV_SYS_REG_TCR_EL1, hv_vcpu_set_sys_reg,
+            hv_sys_reg_t_HV_SYS_REG_SP_EL1 as HV_SYS_REG_SP_EL1,
+            hv_sys_reg_t_HV_SYS_REG_TCR_EL1 as HV_SYS_REG_TCR_EL1,
+            hv_sys_reg_t_HV_SYS_REG_TTBR1_EL1 as HV_SYS_REG_TTBR1_EL1, hv_vcpu_set_sys_reg,
         };
 
         // Set system registers.
@@ -349,11 +364,15 @@ impl<'a, 'b> CpuStates for HfStates<'a, 'b> {
         };
 
         if let State::Dirty(v) = self.tcr_el1 {
-            set_sys(hv_sys_reg_t_HV_SYS_REG_TCR_EL1, v).map_err(StatesError::SetTcrEl1Failed)?;
+            set_sys(HV_SYS_REG_TCR_EL1, v).map_err(StatesError::SetTcrEl1Failed)?;
+        }
+
+        if let State::Dirty(v) = self.ttbr1_el1 {
+            set_sys(HV_SYS_REG_TTBR1_EL1, v).map_err(StatesError::SetTtbr1El1Failed)?;
         }
 
         if let State::Dirty(v) = self.sp_el1 {
-            set_sys(hv_sys_reg_t_HV_SYS_REG_SP_EL1, v).map_err(StatesError::SetSpEl1Failed)?;
+            set_sys(HV_SYS_REG_SP_EL1, v).map_err(StatesError::SetSpEl1Failed)?;
         }
 
         Ok(())
@@ -443,6 +462,10 @@ pub enum StatesError {
     #[error("couldn't set TCR_EL1")]
     SetTcrEl1Failed(NonZero<hv_sys::hv_return_t>),
 
+    #[cfg(target_arch = "aarch64")]
+    #[error("couldn't set TTBR1_EL1")]
+    SetTtbr1El1Failed(NonZero<hv_sys::hv_return_t>),
+
     #[cfg(target_arch = "aarch64")]
     #[error("couldn't set SP_EL1")]
     SetSpEl1Failed(NonZero<hv_sys::hv_return_t>),
diff --git a/src/core/src/vmm/hv/mod.rs b/src/core/src/vmm/hv/mod.rs
index bf53e14f..49497712 100644
--- a/src/core/src/vmm/hv/mod.rs
+++ b/src/core/src/vmm/hv/mod.rs
@@ -88,9 +88,15 @@ pub trait CpuStates {
     fn set_ss(&mut self, p: bool);
 
     /// # Panics
-    /// If `t0sz` or `t1sz` larger than 6 bits.
+    /// - If `ips` greater than 7.
+    /// - If `t0sz` or `t1sz` larger than 6 bits.
     #[cfg(target_arch = "aarch64")]
-    fn set_tcr_el1(&mut self, t0sz: u8, t1sz: u8);
+    fn set_tcr_el1(&mut self, ips: u8, a1: bool, t0sz: u8, t1sz: u8);
+
+    /// # Panics
+    /// If `baddr` has non-zero on bit 0 or 48:64.
+    #[cfg(target_arch = "aarch64")]
+    fn set_ttbr1_el1(&mut self, baddr: usize);
 
     #[cfg(target_arch = "aarch64")]
     fn set_sp_el1(&mut self, v: usize);
diff --git a/src/core/src/vmm/mod.rs b/src/core/src/vmm/mod.rs
index 0ac8b48e..a37d0061 100644
--- a/src/core/src/vmm/mod.rs
+++ b/src/core/src/vmm/mod.rs
@@ -553,8 +553,13 @@ fn setup_main_cpu(cpu: &mut impl Cpu, entry: usize, map: RamMap) -> Result<(), M
         .states()
         .map_err(|e| MainCpuError::GetCpuStatesFailed(Box::new(e)))?;
 
-    // Uses 48-bit virtual addresses (64 - 16 = 48) for both kernel space and user space.
-    states.set_tcr_el1(16, 16);
+    // Uses 48-bit Intermediate Physical Address (ips = 0b101) and 48-bit virtual addresses
+    // (64 - 16 = 48) for both kernel space and user space. Use ASID from user space (TTBR0_EL1 AKA
+    // lower VA range).
+    states.set_tcr_el1(0b101, false, 16, 16);
+
+    // Set page table.
+    states.set_ttbr1_el1(map.page_table);
 
     // Set stack pointer to the kernel.
     states.set_sp_el1(map.stack_vaddr + map.stack_len); // Top-down.