page table builder errors, arm fixes

Author: babayet2

Cargo.lock

@@ -5287,6 +5287,7 @@ name = "page_table"
 version = "0.0.0"
 dependencies = [
  "bitfield-struct 0.11.0",
+ "thiserror 2.0.16",
  "zerocopy 0.8.25",
 ]
 

tmk/tmk_vmm/src/load.rs

@@ -22,6 +22,7 @@ use vm_topology::processor::ProcessorTopology;
 use vm_topology::processor::aarch64::Aarch64Topology;
 use vm_topology::processor::x86::X86Topology;
 use zerocopy::FromBytes as _;
+use zerocopy::FromZeros;
 use zerocopy::IntoBytes;
 
 /// Loads a TMK, returning the initial registers for the BSP.
@@ -38,12 +39,17 @@ pub fn load_x86(
     let load_info = load_common(&mut loader, tmk, test)?;
 
     let page_table_base = load_info.next_available_address;
-    let page_tables = page_table::x64::build_page_tables_64(
+    let mut page_table_work_buffer: Vec<page_table::x64::PageTable> =
+        vec![page_table::x64::PageTable::new_zeroed(); page_table::x64::PAGE_TABLE_MAX_COUNT];
+    let mut page_tables: Vec<u8> = vec![0 as u8; page_table::x64::PAGE_TABLE_MAX_BYTES];
+    let page_table_builder = page_table::x64::IdentityMapBuilder::new(
         page_table_base,
-        0,
         page_table::IdentityMapSize::Size4Gb,
-        None,
     );
+    let page_tables = page_table_builder.build(
+        page_table_work_buffer.as_mut_slice(),
+        page_tables.as_mut_slice(),
+    )?;
     loader
         .import_pages(
             page_table_base >> 12,

vm/loader/page_table/Cargo.toml

@@ -8,6 +8,7 @@ rust-version.workspace = true
 
 [dependencies]
 bitfield-struct.workspace = true
+thiserror.workspace = true
 zerocopy.workspace = true
 [lints]
 workspace = true

vm/loader/page_table/src/aarch64.rs

@@ -673,13 +673,13 @@ impl<'a> Arm64PageTableSpace<'a> {
 }
 
 /// Build a set of Aarch64 page tables identity mapping the given region.
-pub fn build_identity_page_tables_aarch64<'a>(
+pub fn build_identity_page_tables_aarch64(
     page_table_gpa: u64,
     start_gpa: u64,
     size: u64,
     memory_attribute_indirection: MemoryAttributeIndirectionEl1,
-    page_table_space: &'a mut [u8],
-) -> &'a [u8] {
+    page_table_space: &mut [u8],
+) -> &[u8] {
     // start_gpa and size must be 2MB aligned.
     if !aligned(start_gpa, Arm64PageSize::Large) {
         panic!("start_gpa not 2mb aligned");
@@ -710,6 +710,8 @@ pub fn build_identity_page_tables_aarch64<'a>(
 
 #[cfg(test)]
 mod tests {
+    extern crate std;
+
     use super::*;
 
     const DUMP_PAGE_TABLES: bool = false;

vm/loader/page_table/src/lib.rs

@@ -10,6 +10,16 @@
 pub mod aarch64;
 pub mod x64;
 
+use thiserror::Error;
+
+/// Errors returned by the Page Table Builder
+#[derive(Debug, PartialEq, Eq, Error)]
+pub enum Error {
+    /// The page table [u8] buffer does not match the size of the [PageTable] buffer
+    #[error("bad buffer size")]
+    BadBufferSize,
+}
+
 /// Size of the initial identity map
 #[derive(Debug, Copy, Clone)]
 pub enum IdentityMapSize {

vm/loader/page_table/src/x64.rs

@@ -3,6 +3,7 @@
 
 //! Methods to construct page tables on x64.
 
+use crate::Error;
 use crate::IdentityMapSize;
 use zerocopy::FromBytes;
 use zerocopy::Immutable;
@@ -345,8 +346,10 @@ impl PageTableBuilder {
         self,
         page_table: &mut [PageTable],
         flattened_page_table: &'a mut [u8],
-    ) -> &'a [u8] {
-        assert!(flattened_page_table.len() == (page_table.len() * PAGE_TABLE_SIZE));
+    ) -> Result<&'a [u8], Error> {
+        if flattened_page_table.len() != (page_table.len() * PAGE_TABLE_SIZE) {
+            return Err(Error::BadBufferSize);
+        }
 
         const SIZE_512_GB: u64 = 0x8000000000;
 
@@ -488,7 +491,11 @@ impl PageTableBuilder {
         }
 
         // Flatten page table vec into u8 vec
-        flatten_page_table(page_table, flattened_page_table, page_table_index + 1)
+        Ok(flatten_page_table(
+            page_table,
+            flattened_page_table,
+            page_table_index + 1,
+        ))
     }
 }
 
@@ -537,8 +544,10 @@ impl IdentityMapBuilder {
         self,
         page_table: &mut [PageTable],
         flattened_page_table: &'a mut [u8],
-    ) -> &'a [u8] {
-        assert!(flattened_page_table.len() == (page_table.len() * PAGE_TABLE_SIZE));
+    ) -> Result<&'a [u8], Error> {
+        if flattened_page_table.len() != (page_table.len() * PAGE_TABLE_SIZE) {
+            return Err(Error::BadBufferSize);
+        }
         // Allocate page tables. There are up to 6 total page tables:
         //      1 PML4E (Level 4) (omitted if the address bias is non-zero)
         //      1 PDPTE (Level 3)
@@ -615,7 +624,11 @@ impl IdentityMapBuilder {
         }
 
         // Flatten page table vec into u8 vec
-        flatten_page_table(page_table, flattened_page_table, page_table_count)
+        Ok(flatten_page_table(
+            page_table,
+            flattened_page_table,
+            page_table_count,
+        ))
     }
 }
 

vm/loader/src/linux.rs

@@ -129,6 +129,8 @@ pub enum Error {
     UnalignedAddress(u64),
     #[error("importer error")]
     Importer(#[source] anyhow::Error),
+    #[error("PageTableBuilder: {0}")]
+    PageTableBuilder(#[from] page_table::Error),
 }
 
 pub struct AcpiConfig<'a> {
@@ -338,11 +340,11 @@ pub fn load_config(
         IdentityMapBuilder::new(registers.page_table_address, IdentityMapSize::Size4Gb);
     let mut page_table_work_buffer: Vec<PageTable> =
         vec![PageTable::new_zeroed(); PAGE_TABLE_MAX_COUNT];
-    let mut page_table: Vec<u8> = vec![0 as u8; PAGE_TABLE_MAX_BYTES];
+    let mut page_table: Vec<u8> = vec![0; PAGE_TABLE_MAX_BYTES];
     let page_table = page_table_builder.build(
         page_table_work_buffer.as_mut_slice(),
         page_table.as_mut_slice(),
-    );
+    )?;
     assert!((page_table.len() as u64).is_multiple_of(HV_PAGE_SIZE));
     importer
         .import_pages(

vm/loader/src/paravisor.rs

@@ -86,6 +86,8 @@ pub enum Error {
     NotEnoughMemory(u64),
     #[error("importer error")]
     Importer(#[from] anyhow::Error),
+    #[error("PageTableBuilder: {0}")]
+    PageTableBuilder(#[from] page_table::Error),
 }
 
 /// Kernel Command line type.
@@ -458,12 +460,12 @@ where
 
     let mut page_table_work_buffer: Vec<PageTable> =
         vec![PageTable::new_zeroed(); PAGE_TABLE_MAX_COUNT];
-    let mut page_table: Vec<u8> = vec![0 as u8; PAGE_TABLE_MAX_BYTES];
+    let mut page_table: Vec<u8> = vec![0; PAGE_TABLE_MAX_BYTES];
 
     let page_table = page_table_builder.build(
         page_table_work_buffer.as_mut_slice(),
         page_table.as_mut_slice(),
-    );
+    )?;
 
     assert!((page_table.len() as u64).is_multiple_of(HV_PAGE_SIZE));
     let page_table_page_base = page_table_region_start / HV_PAGE_SIZE;

vm/loader/src/uefi/mod.rs

@@ -334,6 +334,8 @@ pub enum Error {
     InvalidConfigType(String),
     #[error("Importer error")]
     Importer(#[source] anyhow::Error),
+    #[error("PageTableBuilder: {0}")]
+    PageTableBuilder(#[from] page_table::Error),
 }
 
 #[derive(Debug)]
@@ -410,7 +412,7 @@ pub mod x86_64 {
 
         let mut page_table_work_buffer: Vec<PageTable> =
             vec![PageTable::new_zeroed(); PAGE_TABLE_MAX_COUNT];
-        let mut page_tables: Vec<u8> = vec![0 as u8; PAGE_TABLE_MAX_BYTES];
+        let mut page_tables: Vec<u8> = vec![0; PAGE_TABLE_MAX_BYTES];
         let page_table_builder =
             IdentityMapBuilder::new(PAGE_TABLE_GPA_BASE, IdentityMapSize::Size4Gb);
         let mut shared_vis_page_table_work_buffer: Vec<PageTable> = Vec::new();
@@ -442,21 +444,21 @@ pub mod x86_64 {
                 let shared_vis_page_tables = shared_vis_builder.build(
                     shared_vis_page_table_work_buffer.as_mut_slice(),
                     shared_vis_page_tables.as_mut_slice(),
-                );
+                )?;
 
                 let page_tables = page_table_builder
                     .with_pml4e_link((shared_vis_page_table_gpa, shared_gpa_boundary))
                     .build(
                         page_table_work_buffer.as_mut_slice(),
                         page_tables.as_mut_slice(),
-                    );
+                    )?;
 
                 (page_tables, Some(shared_vis_page_tables))
             } else {
                 let page_tables = page_table_builder.build(
                     page_table_work_buffer.as_mut_slice(),
                     page_tables.as_mut_slice(),
-                );
+                )?;
                 (page_tables, None)
             };