vmem: Add support for lazily zero-initialized memory regions

This is useful for hyperlight-wasm, although it can't fully provide
for all use cases in a uniform way until we have better support for
hugepages.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/hyperlight_common/src/arch/aarch64/vmem.rs

@@ -26,6 +26,13 @@ const VA_BITS: usize = 48;
 pub const ATTR_INDEX_NORMAL: u8 = 0;
 const SOFTWARE_USE_COW: u8 = 0b1;
 
+const INVALID_KIND_MIN_BIT: u8 = 2;
+const INVALID_KIND_MAX_BIT: u8 = 2;
+const INVALID_KIND_ZERO_INIT: u8 = 0b1;
+const INVALID_KIND_ZERO_INIT_PERM_READABLE_BIT: u8 = 8;
+const INVALID_KIND_ZERO_INIT_PERM_WRITABLE_BIT: u8 = 9;
+const INVALID_KIND_ZERO_INIT_PERM_EXECUTABLE_BIT: u8 = 10;
+
 // Utility structures
 impl<
     Op: TableOps<TableMovability = crate::vmem::MayMoveTable>,
@@ -172,6 +179,19 @@ unsafe fn map_page<
             false,
         ),
         MappingKind::Unmapped => 0,
+        MappingKind::ZeroInit(bm) => {
+            let mut desc = (INVALID_KIND_ZERO_INIT as u64) << INVALID_KIND_MIN_BIT;
+            if bm.readable {
+                desc |= 1 << INVALID_KIND_ZERO_INIT_PERM_READABLE_BIT
+            }
+            if bm.writable {
+                desc |= 1 << INVALID_KIND_ZERO_INIT_PERM_WRITABLE_BIT
+            }
+            if bm.executable {
+                desc |= 1 << INVALID_KIND_ZERO_INIT_PERM_EXECUTABLE_BIT
+            }
+            desc
+        }
     };
     unsafe {
         write_entry_updating(op, r.update_parent, r.entry_ptr, desc);
@@ -497,7 +517,24 @@ unsafe fn internal_walk_va_spaces<'a, Op: TableReadOps + 'a>(
                         kind,
                     }))
                 } else {
-                    None // do nothing - there is no mapping to record here
+                    let ik = bits::<INVALID_KIND_MAX_BIT, INVALID_KIND_MIN_BIT>(desc);
+                    match ik as u8 {
+                        INVALID_KIND_ZERO_INIT => Some(SpaceAwareMapping::ThisSpace(Mapping {
+                            phys_base: 0,
+                            virt_base: rq.vmin,
+                            len: PAGE_SIZE as u64,
+                            kind: MappingKind::ZeroInit(BasicMapping {
+                                readable: desc & (1 << INVALID_KIND_ZERO_INIT_PERM_READABLE_BIT)
+                                    != 0,
+                                writable: desc & (1 << INVALID_KIND_ZERO_INIT_PERM_WRITABLE_BIT)
+                                    != 0,
+                                executable: desc
+                                    & (1 << INVALID_KIND_ZERO_INIT_PERM_EXECUTABLE_BIT)
+                                    != 0,
+                            }),
+                        })),
+                        _ => None, // do nothing - there is no mapping to record here
+                    }
                 }
             });
         (root_id, iter)

src/hyperlight_common/src/arch/amd64/vmem.rs

@@ -102,6 +102,13 @@ const PAGE_PAT_WB: u64 = 0 << 7; // PAT - page attribute table index bit (0 for
 const PTE_AVL_MASK: u64 = 0x0000_0000_0000_0E00;
 const PAGE_AVL_COW: u64 = 1 << 9;
 
+const INVALID_KIND_MIN_BIT: u8 = 2;
+const INVALID_KIND_MAX_BIT: u8 = 2;
+const INVALID_KIND_ZERO_INIT: u8 = 0b1;
+const INVALID_KIND_ZERO_INIT_PERM_READABLE_BIT: u8 = 8;
+const INVALID_KIND_ZERO_INIT_PERM_WRITABLE_BIT: u8 = 9;
+const INVALID_KIND_ZERO_INIT_PERM_EXECUTABLE_BIT: u8 = 10;
+
 /// Returns PAGE_RW if writable is true, 0 otherwise
 #[inline(always)]
 const fn page_rw_flag(writable: bool) -> u64 {
@@ -275,6 +282,19 @@ unsafe fn map_page<
                 PAGE_PRESENT // P   - this entry is present
         }
         MappingKind::Unmapped => 0,
+        MappingKind::ZeroInit(bm) => {
+            let mut pte = (INVALID_KIND_ZERO_INIT as u64) << INVALID_KIND_MIN_BIT;
+            if bm.readable {
+                pte |= 1 << INVALID_KIND_ZERO_INIT_PERM_READABLE_BIT
+            }
+            if bm.writable {
+                pte |= 1 << INVALID_KIND_ZERO_INIT_PERM_WRITABLE_BIT
+            }
+            if bm.executable {
+                pte |= 1 << INVALID_KIND_ZERO_INIT_PERM_EXECUTABLE_BIT
+            }
+            pte
+        }
     };
     unsafe {
         write_entry_updating(op, r.update_parent, r.entry_ptr, pte);
@@ -339,34 +359,9 @@ pub unsafe fn walk_va_spaces<Op: TableReadOps>(
         .flat_map(modify_ptes::<20, 12, Op, _>);
 
         for r in iter {
-            let Some(pte) = (unsafe { read_pte_if_present(op, r.entry_ptr) }) else {
-                continue;
-            };
-            let phys_addr = pte & PTE_ADDR_MASK;
-            let sgn_bit = r.vmin >> (VA_BITS - 1);
-            let sgn_bits = 0u64.wrapping_sub(sgn_bit) << VA_BITS;
-            let virt_addr = sgn_bits | r.vmin;
-
-            let executable = (pte & PAGE_NX) == 0;
-            let avl = pte & PTE_AVL_MASK;
-            let kind = if avl == PAGE_AVL_COW {
-                MappingKind::Cow(CowMapping {
-                    readable: true,
-                    executable,
-                })
-            } else {
-                MappingKind::Basic(BasicMapping {
-                    readable: true,
-                    writable: (pte & PAGE_RW) != 0,
-                    executable,
-                })
-            };
-            mappings.push(crate::vmem::SpaceAwareMapping::ThisSpace(Mapping {
-                phys_base: phys_addr,
-                virt_base: virt_addr,
-                len: PAGE_SIZE as u64,
-                kind,
-            }));
+            if let Some(mapping) = unsafe { read_and_decode_pte(op, &r) } {
+                mappings.push(crate::vmem::SpaceAwareMapping::ThisSpace(mapping));
+            }
         }
 
         out.push((root_id, mappings));
@@ -404,6 +399,56 @@ pub unsafe fn map<Op: TableOps>(op: &Op, mapping: Mapping) {
     .for_each(drop);
 }
 
+unsafe fn read_and_decode_pte<Op: TableReadOps, P: UpdateParent<Op>>(
+    op: &Op,
+    r: &MapResponse<Op, P>,
+) -> Option<Mapping> {
+    // Re-do the sign extension
+    let sgn_bit = r.vmin >> (VA_BITS - 1);
+    let sgn_bits = 0u64.wrapping_sub(sgn_bit) << VA_BITS;
+    let virt_addr = sgn_bits | r.vmin;
+
+    let Some(pte) = (unsafe { read_pte_if_present(op, r.entry_ptr) }) else {
+        let pte: u64 = unsafe { op.read_entry(r.entry_ptr) };
+        let ik = crate::vmem::bits::<INVALID_KIND_MAX_BIT, INVALID_KIND_MIN_BIT>(pte);
+        return match ik as u8 {
+            INVALID_KIND_ZERO_INIT => Some(Mapping {
+                phys_base: 0,
+                virt_base: virt_addr,
+                len: PAGE_SIZE as u64,
+                kind: MappingKind::ZeroInit(BasicMapping {
+                    readable: pte & (1 << INVALID_KIND_ZERO_INIT_PERM_READABLE_BIT) != 0,
+                    writable: pte & (1 << INVALID_KIND_ZERO_INIT_PERM_WRITABLE_BIT) != 0,
+                    executable: pte & (1 << INVALID_KIND_ZERO_INIT_PERM_EXECUTABLE_BIT) != 0,
+                }),
+            }),
+            _ => None,
+        };
+    };
+    let phys_addr = pte & PTE_ADDR_MASK;
+
+    let executable = (pte & PAGE_NX) == 0;
+    let avl = pte & PTE_AVL_MASK;
+    let kind = if avl == PAGE_AVL_COW {
+        MappingKind::Cow(CowMapping {
+            readable: true,
+            executable,
+        })
+    } else {
+        MappingKind::Basic(BasicMapping {
+            readable: true,
+            writable: (pte & PAGE_RW) != 0,
+            executable,
+        })
+    };
+    Some(Mapping {
+        phys_base: phys_addr,
+        virt_base: virt_addr,
+        len: PAGE_SIZE as u64,
+        kind,
+    })
+}
+
 // There are no notable architecture-specific safety considerations
 // here, and the general conditions are documented in the
 // architecture-independent re-export in vmem.rs
@@ -446,35 +491,7 @@ pub unsafe fn virt_to_phys<'a, Op: TableReadOps + 'a>(
     .flat_map(modify_ptes::<29, 21, Op, _>)
     .filter_map(move |r| unsafe { require_pte_exist(op.as_ref(), r) })
     .flat_map(modify_ptes::<20, 12, Op, _>)
-    .filter_map(move |r| {
-        let pte = unsafe { read_pte_if_present(op.as_ref(), r.entry_ptr) }?;
-        let phys_addr = pte & PTE_ADDR_MASK;
-        // Re-do the sign extension
-        let sgn_bit = r.vmin >> (VA_BITS - 1);
-        let sgn_bits = 0u64.wrapping_sub(sgn_bit) << VA_BITS;
-        let virt_addr = sgn_bits | r.vmin;
-
-        let executable = (pte & PAGE_NX) == 0;
-        let avl = pte & PTE_AVL_MASK;
-        let kind = if avl == PAGE_AVL_COW {
-            MappingKind::Cow(CowMapping {
-                readable: true,
-                executable,
-            })
-        } else {
-            MappingKind::Basic(BasicMapping {
-                readable: true,
-                writable: (pte & PAGE_RW) != 0,
-                executable,
-            })
-        };
-        Some(Mapping {
-            phys_base: phys_addr,
-            virt_base: virt_addr,
-            len: PAGE_SIZE as u64,
-            kind,
-        })
-    })
+    .filter_map(move |r| unsafe { read_and_decode_pte(op.as_ref(), &r) })
 }
 
 const VA_BITS: usize = 48; // We use 48-bit virtual addresses at the moment.

src/hyperlight_common/src/vmem.rs

@@ -413,13 +413,14 @@ pub struct CowMapping {
 #[derive(Debug, PartialEq, Clone, Copy)]
 pub enum MappingKind {
     Unmapped,
+    ZeroInit(BasicMapping),
     Basic(BasicMapping),
     Cow(CowMapping),
     /* TODO: What useful things other than basic mappings actually
      * require touching the tables? */
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 pub struct Mapping {
     pub phys_base: u64,
     pub virt_base: u64,

src/hyperlight_guest_bin/src/arch/aarch64/exception/handle.rs

@@ -3,9 +3,7 @@
 use core::fmt::Write;
 
 use hyperlight_common::arch::exn::{DataFault, DataFaultKind, Exception, decode_syndrome};
-use hyperlight_common::vmem::{
-    BasicMapping, CowMapping, MappingKind, PAGE_SIZE, PhysAddr, VirtAddr,
-};
+use hyperlight_common::vmem::{BasicMapping, Mapping, MappingKind, PAGE_SIZE, PhysAddr, VirtAddr};
 use hyperlight_guest::error::ErrorCode;
 use hyperlight_guest::exit::write_abort;
 use hyperlight_guest::layout::{MAIN_STACK_LIMIT_GVA, MAIN_STACK_TOP_GVA};
@@ -37,7 +35,7 @@ fn handle_stack_fault(far: u64) {
     }
 }
 
-fn handle_cow_fault(_orig_phys: PhysAddr, virt: VirtAddr, perms: CowMapping) {
+fn handle_cow_fault(orig_phys: Option<PhysAddr>, virt: VirtAddr, perms: BasicMapping) {
     unsafe {
         let new_page = hyperlight_guest::prim_alloc::alloc_phys_pages(1);
         let target_virt = virt as *mut u8;
@@ -50,22 +48,17 @@ fn handle_cow_fault(_orig_phys: PhysAddr, virt: VirtAddr, perms: CowMapping) {
             // should never reach beyond this call.
             unreachable!();
         };
-        core::ptr::copy(target_virt, scratch_mapping_access, PAGE_SIZE);
+        // Only copy if there is an original page from which to copy
+        if orig_phys.is_some() {
+            core::ptr::copy(target_virt, scratch_mapping_access, PAGE_SIZE);
+        }
         // todo(multithreading): this will definitely require a
         // break-before-make sequence
         crate::paging::map_region(
             new_page,
             target_virt,
             PAGE_SIZE as u64,
-            MappingKind::Basic(BasicMapping {
-                // Inherit R bit from the original mapping (always 1 at the moment)
-                readable: perms.readable,
-                // If we got here, the original marking was marked
-                // CoW, so the copied mapping should always be
-                // writable
-                writable: true,
-                executable: perms.executable,
-            }),
+            MappingKind::Basic(perms),
         );
         // This is updating an entry that was already valid, changing
         // its OA, so we need to actually invalidate the TLB for it.
@@ -86,7 +79,7 @@ pub extern "Rust" fn _debug_print(x: &str) {
     hyperlight_guest::exit::debug_print(x);
 }
 
-fn handle_internal_fault(exn: Exception, far: u64) -> bool {
+fn handle_internal_fault(exn: Exception, far: u64, orig_mapping: Option<Mapping>) -> bool {
     match exn {
         Exception::DataFault(DataFault {
             from_lower_el: false,
@@ -96,6 +89,15 @@ fn handle_internal_fault(exn: Exception, far: u64) -> bool {
             if (MAIN_STACK_LIMIT_GVA..MAIN_STACK_TOP_GVA).contains(&far) {
                 handle_stack_fault(far);
                 true
+            } else if let Some(
+                m @ Mapping {
+                    kind: MappingKind::ZeroInit(bm),
+                    ..
+                },
+            ) = orig_mapping
+            {
+                handle_cow_fault(None, m.virt_base, bm);
+                true
             } else {
                 false
             }
@@ -106,12 +108,23 @@ fn handle_internal_fault(exn: Exception, far: u64) -> bool {
             kind: DataFaultKind::PermissionFault(_),
             ..
         }) => {
-            let mut orig_mappings = crate::paging::virt_to_phys(far);
-            if let Some(mapping) = orig_mappings.next()
-                && let None = orig_mappings.next()
-                && let MappingKind::Cow(cm) = mapping.kind
+            if let Some(
+                m @ Mapping {
+                    kind: MappingKind::Cow(cm),
+                    ..
+                },
+            ) = orig_mapping
             {
-                handle_cow_fault(mapping.phys_base, mapping.virt_base, cm);
+                let bm = BasicMapping {
+                    // Inherit R bit from the original mapping (always 1 at the moment)
+                    readable: cm.readable,
+                    // If we got here, the original marking was marked
+                    // CoW, so the copied mapping should always be
+                    // writable
+                    writable: true,
+                    executable: cm.executable,
+                };
+                handle_cow_fault(Some(m.phys_base), m.virt_base, bm);
                 true
             } else {
                 false
@@ -131,7 +144,9 @@ pub(super) extern "C" fn handle_exception(
 
     if typ == ExceptionType::Synchronous && from == ExceptionFrom::CurrentSP0 {
         let exn = decode_syndrome(esr);
-        if handle_internal_fault(exn, far) {
+        let mut orig_mappings = crate::paging::virt_to_phys(far);
+        let orig_mapping = orig_mappings.next();
+        if handle_internal_fault(exn, far, orig_mapping) {
             return;
         }
     }