remove dead code, add asserts, add builder pattern

Author: babayet2

vm/loader/page_table/src/x64.rs

@@ -280,7 +280,6 @@ pub struct PageTableBuilder {
     confidential_bit: Option<u32>,
     map_reset_vector: bool,
     read_only: bool,
-    debug: bool,
 }
 
 impl PteOps for PageTableBuilder {
@@ -307,7 +306,6 @@ impl PageTableBuilder {
             local_map: None,
             confidential_bit: None,
             read_only: false,
-            debug: false,
             map_reset_vector: false,
         }
     }
@@ -328,11 +326,6 @@ impl PageTableBuilder {
         self
     }
 
-    pub fn with_debug(mut self, debug: bool) -> Self {
-        self.debug = debug;
-        self
-    }
-
     /// Map the reset vector at page 0xFFFFF with a single page.
     pub fn with_reset_vector(mut self, map_reset_vector: bool) -> Self {
         self.map_reset_vector = map_reset_vector;
@@ -353,6 +346,8 @@ impl PageTableBuilder {
         page_table: &mut [PageTable],
         flattened_page_table: &'a mut [u8],
     ) -> &'a [u8] {
+        assert!(flattened_page_table.len() == (page_table.len() * PAGE_TABLE_SIZE));
+
         const SIZE_512_GB: u64 = 0x8000000000;
 
         if self.size == 0 {
@@ -497,97 +492,131 @@ impl PageTableBuilder {
     }
 }
 
-/// Build a set of X64 page tables identity mapping the bottom address
-/// space with an optional address bias.
-///
-/// An optional PML4E entry may be linked, with arguments being (link_target_gpa, linkage_gpa).
-/// link_target_gpa represents the GPA of the PML4E to link into the built page table.
-/// linkage_gpa represents the GPA at which the linked PML4E should be linked.
-pub fn build_page_tables_64<'a>(
+#[derive(Debug, Clone)]
+pub struct IdentityMapBuilder {
     page_table_gpa: u64,
-    address_bias: u64,
     identity_map_size: IdentityMapSize,
+    address_bias: u64,
     pml4e_link: Option<(u64, u64)>,
     read_only: bool,
-    page_table: &mut [PageTable],
-    flattened_page_table: &'a mut [u8],
-) -> &'a [u8] {
-    // 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)
-    //      4 or 8 PDE tables (Level 2)
-    // Note that there are no level 1 page tables, as 2MB pages are used.
-    let leaf_page_table_count = match identity_map_size {
-        IdentityMapSize::Size4Gb => 4,
-        IdentityMapSize::Size8Gb => 8,
-    };
-    let page_table_count = leaf_page_table_count + if address_bias == 0 { 2 } else { 1 };
-    let mut page_table_allocator = page_table.iter_mut().enumerate();
-
-    // Allocate single PDPTE table.
-    let pdpte_table = if address_bias == 0 {
-        // Allocate single PML4E page table.
-        let (_, pml4e_table) = page_table_allocator
-            .next()
-            .expect("pagetable should always be available, code bug if not");
-
-        // PDPTE table is the next pagetable.
-        let (pdpte_table_index, pdpte_table) = page_table_allocator
-            .next()
-            .expect("pagetable should always be available, code bug if not");
-
-        // Set PML4E entry linking PML4E to PDPTE.
-        let output_address = page_table_gpa + pdpte_table_index as u64 * X64_PAGE_SIZE;
-        pml4e_table.entries[0].set_entry(PageTableEntryType::Pde(output_address), read_only);
-
-        // Set PML4E entry to link the additional entry if specified.
-        if let Some((link_target_gpa, linkage_gpa)) = pml4e_link {
-            assert!((linkage_gpa & 0x7FFFFFFFFF) == 0);
-            pml4e_table.entries[linkage_gpa as usize >> 39]
-                .set_entry(PageTableEntryType::Pde(link_target_gpa), read_only);
-        }
+}
 
-        pdpte_table
-    } else {
-        // PDPTE table is the first table, if no PML4E.
-        page_table_allocator
-            .next()
-            .expect("pagetable should always be available, code bug if not")
-            .1
-    };
-
-    // Build PDEs that point to 2 MB pages.
-    let top_address = match identity_map_size {
-        IdentityMapSize::Size4Gb => 0x100000000u64,
-        IdentityMapSize::Size8Gb => 0x200000000u64,
-    };
-    let mut current_va = 0;
-
-    while current_va < top_address {
-        // Allocate a new PDE table
-        let (pde_table_index, pde_table) = page_table_allocator
-            .next()
-            .expect("pagetable should always be available, code bug if not");
-
-        // Link PDPTE table to PDE table (L3 to L2)
-        let pdpte_index = get_amd64_pte_index(current_va, 2);
-        let output_address = page_table_gpa + pde_table_index as u64 * X64_PAGE_SIZE;
-        let pdpte_entry = &mut pdpte_table.entries[pdpte_index as usize];
-        assert!(!pdpte_entry.is_present());
-        pdpte_entry.set_entry(PageTableEntryType::Pde(output_address), read_only);
-
-        // Set all 2MB entries in this PDE table.
-        for entry in pde_table.iter_mut() {
-            entry.set_entry(
-                PageTableEntryType::Leaf2MbPage(current_va + address_bias),
-                read_only,
-            );
-            current_va += X64_LARGE_PAGE_SIZE;
+impl IdentityMapBuilder {
+    pub fn new(page_table_gpa: u64, identity_map_size: IdentityMapSize) -> Self {
+        IdentityMapBuilder {
+            page_table_gpa,
+            identity_map_size,
+            address_bias: 0,
+            pml4e_link: None,
+            read_only: false,
         }
     }
 
-    // Flatten page table vec into u8 vec
-    flatten_page_table(page_table, flattened_page_table, page_table_count)
+    pub fn with_address_bias(mut self, address_bias: u64) -> Self {
+        self.address_bias = address_bias;
+        self
+    }
+
+    pub fn with_pml4e_link(mut self, pml4e_link: (u64, u64)) -> Self {
+        self.pml4e_link = Some(pml4e_link);
+        self
+    }
+
+    pub fn with_read_only(mut self, read_only: bool) -> Self {
+        self.read_only = read_only;
+        self
+    }
+
+    /// Build a set of X64 page tables identity mapping the bottom address
+    /// space with an optional address bias.
+    ///
+    /// An optional PML4E entry may be linked, with arguments being (link_target_gpa, linkage_gpa).
+    /// link_target_gpa represents the GPA of the PML4E to link into the built page table.
+    /// linkage_gpa represents the GPA at which the linked PML4E should be linked.
+    pub fn build<'a>(
+        self,
+        page_table: &mut [PageTable],
+        flattened_page_table: &'a mut [u8],
+    ) -> &'a [u8] {
+        assert!(flattened_page_table.len() == (page_table.len() * PAGE_TABLE_SIZE));
+        // 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)
+        //      4 or 8 PDE tables (Level 2)
+        // Note that there are no level 1 page tables, as 2MB pages are used.
+        let leaf_page_table_count = match self.identity_map_size {
+            IdentityMapSize::Size4Gb => 4,
+            IdentityMapSize::Size8Gb => 8,
+        };
+        let page_table_count = leaf_page_table_count + if self.address_bias == 0 { 2 } else { 1 };
+        let mut page_table_allocator = page_table.iter_mut().enumerate();
+
+        // Allocate single PDPTE table.
+        let pdpte_table = if self.address_bias == 0 {
+            // Allocate single PML4E page table.
+            let (_, pml4e_table) = page_table_allocator
+                .next()
+                .expect("pagetable should always be available, code bug if not");
+
+            // PDPTE table is the next pagetable.
+            let (pdpte_table_index, pdpte_table) = page_table_allocator
+                .next()
+                .expect("pagetable should always be available, code bug if not");
+
+            // Set PML4E entry linking PML4E to PDPTE.
+            let output_address = self.page_table_gpa + pdpte_table_index as u64 * X64_PAGE_SIZE;
+            pml4e_table.entries[0]
+                .set_entry(PageTableEntryType::Pde(output_address), self.read_only);
+
+            // Set PML4E entry to link the additional entry if specified.
+            if let Some((link_target_gpa, linkage_gpa)) = self.pml4e_link {
+                assert!((linkage_gpa & 0x7FFFFFFFFF) == 0);
+                pml4e_table.entries[linkage_gpa as usize >> 39]
+                    .set_entry(PageTableEntryType::Pde(link_target_gpa), self.read_only);
+            }
+
+            pdpte_table
+        } else {
+            // PDPTE table is the first table, if no PML4E.
+            page_table_allocator
+                .next()
+                .expect("pagetable should always be available, code bug if not")
+                .1
+        };
+
+        // Build PDEs that point to 2 MB pages.
+        let top_address = match self.identity_map_size {
+            IdentityMapSize::Size4Gb => 0x100000000u64,
+            IdentityMapSize::Size8Gb => 0x200000000u64,
+        };
+        let mut current_va = 0;
+
+        while current_va < top_address {
+            // Allocate a new PDE table
+            let (pde_table_index, pde_table) = page_table_allocator
+                .next()
+                .expect("pagetable should always be available, code bug if not");
+
+            // Link PDPTE table to PDE table (L3 to L2)
+            let pdpte_index = get_amd64_pte_index(current_va, 2);
+            let output_address = self.page_table_gpa + pde_table_index as u64 * X64_PAGE_SIZE;
+            let pdpte_entry = &mut pdpte_table.entries[pdpte_index as usize];
+            assert!(!pdpte_entry.is_present());
+            pdpte_entry.set_entry(PageTableEntryType::Pde(output_address), self.read_only);
+
+            // Set all 2MB entries in this PDE table.
+            for entry in pde_table.iter_mut() {
+                entry.set_entry(
+                    PageTableEntryType::Leaf2MbPage(current_va + self.address_bias),
+                    self.read_only,
+                );
+                current_va += X64_LARGE_PAGE_SIZE;
+            }
+        }
+
+        // Flatten page table vec into u8 vec
+        flatten_page_table(page_table, flattened_page_table, page_table_count)
+    }
 }
 
 /// Align an address up to the start of the next page.

vm/loader/src/linux.rs

@@ -21,12 +21,12 @@ use bitfield_struct::bitfield;
 use hvdef::HV_PAGE_SIZE;
 use loader_defs::linux as defs;
 use page_table::IdentityMapSize;
+use page_table::x64::IdentityMapBuilder;
 use page_table::x64::PAGE_TABLE_MAX_BYTES;
 use page_table::x64::PAGE_TABLE_MAX_COUNT;
 use page_table::x64::PageTable;
 use page_table::x64::align_up_to_large_page_size;
 use page_table::x64::align_up_to_page_size;
-use page_table::x64::build_page_tables_64;
 use std::ffi::CString;
 use thiserror::Error;
 use vm_topology::memory::MemoryLayout;
@@ -334,15 +334,12 @@ pub fn load_config(
     check_address_alignment(registers.gdt_address)?;
     import_default_gdt(importer, registers.gdt_address / HV_PAGE_SIZE).map_err(Error::Importer)?;
     check_address_alignment(registers.page_table_address)?;
+    let page_table_builder =
+        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];
-    build_page_tables_64(
-        registers.page_table_address,
-        0,
-        IdentityMapSize::Size4Gb,
-        None,
-        false,
+    let page_table = page_table_builder.build(
         page_table_work_buffer.as_mut_slice(),
         page_table.as_mut_slice(),
     );
@@ -353,7 +350,7 @@ pub fn load_config(
             page_table.len() as u64 / HV_PAGE_SIZE,
             "linux-pagetables",
             BootPageAcceptance::Exclusive,
-            page_table.as_slice(),
+            page_table,
         )
         .map_err(Error::Importer)?;
 

vm/loader/src/uefi/mod.rs

@@ -371,11 +371,11 @@ pub mod x86_64 {
     use crate::uefi::get_sec_entry_point_offset;
     use hvdef::HV_PAGE_SIZE;
     use page_table::IdentityMapSize;
+    use page_table::x64::IdentityMapBuilder;
     use page_table::x64::PAGE_TABLE_MAX_BYTES;
     use page_table::x64::PAGE_TABLE_MAX_COUNT;
     use page_table::x64::PageTable;
     use page_table::x64::align_up_to_page_size;
-    use page_table::x64::build_page_tables_64;
     use zerocopy::FromZeros;
     use zerocopy::IntoBytes;
 
@@ -411,6 +411,8 @@ 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 page_table_builder =
+            IdentityMapBuilder::new(PAGE_TABLE_GPA_BASE, IdentityMapSize::Size4Gb);
         let mut shared_vis_page_table_work_buffer: Vec<PageTable> = Vec::new();
         let mut shared_vis_page_tables: Vec<u8> = Vec::new();
         let (page_tables, shared_vis_page_tables) =
@@ -430,38 +432,28 @@ pub mod x86_64 {
                 shared_vis_page_table_work_buffer
                     .resize(PAGE_TABLE_MAX_COUNT, PageTable::new_zeroed());
                 shared_vis_page_tables.resize(PAGE_TABLE_MAX_BYTES, 0);
+                let shared_vis_builder =
+                    IdentityMapBuilder::new(shared_vis_page_table_gpa, IdentityMapSize::Size4Gb)
+                        .with_address_bias(shared_gpa_boundary);
 
                 // The extra page tables are placed after the first config blob
                 // page.  They will be accounted for when the IGVM parameters are
                 // built.
-                let shared_vis_page_tables = build_page_tables_64(
-                    shared_vis_page_table_gpa,
-                    shared_gpa_boundary,
-                    IdentityMapSize::Size4Gb,
-                    None,
-                    false,
+                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 = build_page_tables_64(
-                    PAGE_TABLE_GPA_BASE,
-                    0,
-                    IdentityMapSize::Size4Gb,
-                    Some((shared_vis_page_table_gpa, shared_gpa_boundary)),
-                    false,
-                    page_table_work_buffer.as_mut_slice(),
-                    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 = build_page_tables_64(
-                    PAGE_TABLE_GPA_BASE,
-                    0,
-                    IdentityMapSize::Size4Gb,
-                    None,
-                    false,
+                let page_tables = page_table_builder.build(
                     page_table_work_buffer.as_mut_slice(),
                     page_tables.as_mut_slice(),
                 );