pvmfw: Manage dirty state in hardware if possible

Manage dirty state in hardware on devices that support `FEAT_HAFDBS`.
Check presence of feature and enable dirty state management in TCR_EL1.

Bug: 269738062
Test: atest MicrodroidTestApp

Change-Id: If9044702a02fa2f29a747d9d80cc958773bdfe17
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index c79605f..c329d26 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -12,6 +12,7 @@
     flags: ["-Dunsafe_op_in_unsafe_fn"],
     features: [
         "legacy",
+        "cpu_feat_hafdbs",
     ],
     rustlibs: [
         "libaarch64_paging",
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index ca74740..7655f40 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -24,7 +24,7 @@
 use crate::mmu;
 use crate::rand;
 use core::arch::asm;
-use core::mem::size_of;
+use core::mem::{drop, size_of};
 use core::num::NonZeroUsize;
 use core::ops::Range;
 use core::slice;
@@ -208,8 +208,6 @@
     // script prevents it from overlapping with other objects.
     let appended_data = unsafe { get_appended_data_slice() };
 
-    // Up to this point, we were using the built-in static (from .rodata) page tables.
-
     let mut page_table = mmu::PageTable::from_static_layout().map_err(|e| {
         error!("Failed to set up the dynamic page tables: {e}");
         RebootReason::InternalError
@@ -231,13 +229,9 @@
 
     let (bcc_slice, debug_policy) = appended.get_entries();
 
-    debug!("Activating dynamic page table...");
-    // SAFETY - page_table duplicates the static mappings for everything that the Rust code is
-    // aware of so activating it shouldn't have any visible effect.
-    unsafe { page_table.activate() };
-    debug!("... Success!");
-
+    // Up to this point, we were using the built-in static (from .rodata) page tables.
     MEMORY.lock().replace(MemoryTracker::new(page_table));
+
     let slices = MemorySlices::new(fdt, payload, payload_size)?;
 
     rand::init().map_err(|e| {
@@ -262,7 +256,9 @@
         error!("Failed to unshare the UART: {e}");
         RebootReason::InternalError
     })?;
-    MEMORY.lock().take().unwrap();
+
+    // Drop MemoryTracker and deactivate page table.
+    drop(MEMORY.lock().take());
 
     Ok((slices.kernel.as_ptr() as usize, next_bcc))
 }
diff --git a/pvmfw/src/helpers.rs b/pvmfw/src/helpers.rs
index 4b669d7..9c4cb1b 100644
--- a/pvmfw/src/helpers.rs
+++ b/pvmfw/src/helpers.rs
@@ -187,6 +187,16 @@
     }};
 }
 
+/// Returns `true` if hardware dirty state management is available.
+pub fn dbm_available() -> bool {
+    if !cfg!(feature = "cpu_feat_hafdbs") {
+        return false;
+    }
+    // Hardware dirty bit management available flag (ID_AA64MMFR1_EL1.HAFDBS[1])
+    const DBM_AVAILABLE: usize = 1 << 1;
+    read_sysreg!("id_aa64mmfr1_el1") & DBM_AVAILABLE != 0
+}
+
 /// Executes a data synchronization barrier.
 #[macro_export]
 macro_rules! dsb {
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index 4ed3072..2f6e7a4 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -16,9 +16,9 @@
 
 #![deny(unsafe_op_in_unsafe_fn)]
 
-use crate::helpers::{self, page_4kb_of, RangeExt, PVMFW_PAGE_SIZE, SIZE_4MB};
+use crate::helpers::{self, dbm_available, page_4kb_of, RangeExt, PVMFW_PAGE_SIZE, SIZE_4MB};
 use crate::mmu;
-use crate::{dsb, isb, tlbi};
+use crate::{dsb, isb, read_sysreg, tlbi, write_sysreg};
 use aarch64_paging::paging::{Attributes, Descriptor, MemoryRegion as VaRange};
 use alloc::alloc::alloc_zeroed;
 use alloc::alloc::dealloc;
@@ -36,8 +36,8 @@
 use core::ptr::NonNull;
 use core::result;
 use hyp::get_hypervisor;
-use log::error;
 use log::trace;
+use log::{debug, error};
 use once_cell::race::OnceBox;
 use spin::mutex::SpinMutex;
 use tinyvec::ArrayVec;
@@ -219,9 +219,22 @@
     const CAPACITY: usize = 5;
     const MMIO_CAPACITY: usize = 5;
     const PVMFW_RANGE: MemoryRange = (BASE_ADDR - SIZE_4MB)..BASE_ADDR;
+    // TCR_EL1.{HA,HD} bits controlling hardware management of access and dirty state
+    const TCR_EL1_HA_HD_BITS: usize = 3 << 39;
 
     /// Create a new instance from an active page table, covering the maximum RAM size.
-    pub fn new(page_table: mmu::PageTable) -> Self {
+    pub fn new(mut page_table: mmu::PageTable) -> Self {
+        // Activate dirty state management first, otherwise we may get permission faults immediately
+        // after activating the new page table. This has no effect before the new page table is
+        // activated because none of the entries in the initial idmap have the DBM flag.
+        Self::set_dbm_enabled(true);
+
+        debug!("Activating dynamic page table...");
+        // SAFETY - page_table duplicates the static mappings for everything that the Rust code is
+        // aware of so activating it shouldn't have any visible effect.
+        unsafe { page_table.activate() };
+        debug!("... Success!");
+
         Self {
             total: BASE_ADDR..MAX_ADDR,
             page_table,
@@ -421,12 +434,27 @@
             .modify_range(&(addr..addr + 1), &mark_dirty_block)
             .map_err(|_| MemoryTrackerError::SetPteDirtyFailed)
     }
+
+    fn set_dbm_enabled(enabled: bool) {
+        if dbm_available() {
+            let mut tcr = read_sysreg!("tcr_el1");
+            if enabled {
+                tcr |= Self::TCR_EL1_HA_HD_BITS
+            } else {
+                tcr &= !Self::TCR_EL1_HA_HD_BITS
+            };
+            // Safe because it writes to a system register and does not affect Rust.
+            unsafe { write_sysreg!("tcr_el1", tcr) }
+            isb!();
+        }
+    }
 }
 
 impl Drop for MemoryTracker {
     fn drop(&mut self) {
+        Self::set_dbm_enabled(false);
         self.flush_dirty_pages().unwrap();
-        self.unshare_all_memory()
+        self.unshare_all_memory();
     }
 }