feedback: attempt to have a single source of truth for partition-wide objects

Author: sluck-msft

openhcl/underhill_mem/src/init.rs

@@ -6,7 +6,7 @@
 use crate::HardwareIsolatedMemoryProtector;
 use crate::MemoryAcceptor;
 use crate::mapping::GuestMemoryMapping;
-use crate::mapping::GuestValidMemory;
+use crate::mapping::GuestPartitionMemoryBuilder;
 use anyhow::Context;
 use futures::future::try_join_all;
 use guestmem::GuestMemory;
@@ -210,18 +210,19 @@ pub async fn init(params: &Init<'_>) -> anyhow::Result<MemoryMappings> {
         // kernel-registered RAM.
 
         tracing::debug!("Building valid encrypted memory view");
-        let valid_encrypted_memory = Arc::new({
-            let _span = tracing::info_span!("create encrypted memory bitmap").entered();
-            GuestValidMemory::new(params.mem_layout, true)?
-        });
+        let encrypted_memory_builder = {
+            let _span = tracing::info_span!("create encrypted memory view").entered();
+            GuestPartitionMemoryBuilder::new(params.mem_layout, Some(true))?
+        };
 
         tracing::debug!("Building VTL0 memory map");
         let vtl0_mapping = Arc::new({
             let _span = tracing::info_span!("map_vtl0_memory").entered();
-            GuestMemoryMapping::builder(0)
-                .dma_base_address(None) // prohibit direct DMA attempts until TDISP is supported
-                .use_partition_valid_memory(Some(valid_encrypted_memory.clone()))
-                .build(&gpa_fd, params.mem_layout)
+            encrypted_memory_builder
+                .build_guest_memory_mapping(
+                    &gpa_fd,
+                    GuestMemoryMapping::builder(0).dma_base_address(None),
+                )
                 .context("failed to map vtl0 memory")?
         });
 
@@ -291,27 +292,34 @@ pub async fn init(params: &Init<'_>) -> anyhow::Result<MemoryMappings> {
 
         tracing::debug!("Building shared memory map");
 
-        let valid_shared_memory = Arc::new({
-            let _span = tracing::info_span!("create shared memory bitmap").entered();
-            GuestValidMemory::new(params.complete_memory_layout, false)?
-        });
+        let shared_memory_builder = {
+            let _span = tracing::info_span!("create shared memory view").entered();
+            GuestPartitionMemoryBuilder::new(params.complete_memory_layout, Some(false))?
+        };
+
+        let valid_shared_memory = shared_memory_builder.partition_valid_memory();
 
         // Update the shared mapping bitmap for pages used by the shared
         // visibility pool to be marked as shared, since by default pages are
         // marked as no-access in the bitmap.
         tracing::debug!("Updating shared mapping bitmaps");
         for range in params.shared_pool {
-            valid_shared_memory.update_valid(range.range, true);
+            valid_shared_memory
+                .as_ref()
+                .unwrap()
+                .update_valid(range.range, true);
         }
 
         let shared_mapping = Arc::new({
             let _span = tracing::info_span!("map_shared_memory").entered();
-            GuestMemoryMapping::builder(shared_offset)
-                .shared(true)
-                .use_partition_valid_memory(Some(valid_shared_memory.clone()))
-                .ignore_registration_failure(params.boot_init.is_none())
-                .dma_base_address(Some(dma_base_address))
-                .build(&gpa_fd, params.complete_memory_layout)
+            shared_memory_builder
+                .build_guest_memory_mapping(
+                    &gpa_fd,
+                    GuestMemoryMapping::builder(shared_offset)
+                        .shared(true)
+                        .ignore_registration_failure(params.boot_init.is_none())
+                        .dma_base_address(Some(dma_base_address)),
+                )
                 .context("failed to map shared memory")?
         });
 
@@ -365,8 +373,11 @@ pub async fn init(params: &Init<'_>) -> anyhow::Result<MemoryMappings> {
         let private_vtl0_memory = GuestMemory::new("trusted", vtl0_mapping.clone());
 
         let protector = Arc::new(HardwareIsolatedMemoryProtector::new(
-            valid_encrypted_memory.clone(),
-            valid_shared_memory.clone(),
+            encrypted_memory_builder
+                .partition_valid_memory()
+                .unwrap()
+                .clone(),
+            valid_shared_memory.unwrap().clone(),
             vtl0_mapping.clone(),
             params.mem_layout.clone(),
             acceptor.as_ref().unwrap().clone(),
@@ -385,16 +396,22 @@ pub async fn init(params: &Init<'_>) -> anyhow::Result<MemoryMappings> {
             }),
         }
     } else {
+        let partition_memory_builder = GuestPartitionMemoryBuilder::new(params.mem_layout, None)?;
+
         tracing::debug!("Creating VTL0 guest memory");
         let vtl0_mapping = {
             let _span = tracing::info_span!("map_vtl0_memory").entered();
             let base_address = params.vtl0_alias_map_bit.unwrap_or(0);
+
             Arc::new(
-                GuestMemoryMapping::builder(base_address)
-                    .for_kernel_access(true)
-                    .dma_base_address(Some(base_address))
-                    .ignore_registration_failure(params.boot_init.is_none())
-                    .build(&gpa_fd, params.mem_layout)
+                partition_memory_builder
+                    .build_guest_memory_mapping(
+                        &gpa_fd,
+                        GuestMemoryMapping::builder(base_address)
+                            .for_kernel_access(true)
+                            .dma_base_address(Some(base_address))
+                            .ignore_registration_failure(params.boot_init.is_none()),
+                    )
                     .context("failed to map vtl0 memory")?,
             )
         };
@@ -424,13 +441,17 @@ pub async fn init(params: &Init<'_>) -> anyhow::Result<MemoryMappings> {
                 tracing::debug!("Creating VTL 1 memory map");
 
                 let _span = tracing::info_span!("map_vtl1_memory").entered();
-                let vtl1_mapping = GuestMemoryMapping::builder(0)
-                    .for_kernel_access(true)
-                    .dma_base_address(Some(0))
-                    .ignore_registration_failure(params.boot_init.is_none())
-                    .build(&gpa_fd, params.mem_layout)
-                    .context("failed to map vtl1 memory")?;
-                Some(Arc::new(vtl1_mapping))
+                Some(Arc::new(
+                    partition_memory_builder
+                        .build_guest_memory_mapping(
+                            &gpa_fd,
+                            GuestMemoryMapping::builder(0)
+                                .for_kernel_access(true)
+                                .dma_base_address(Some(0))
+                                .ignore_registration_failure(params.boot_init.is_none()),
+                        )
+                        .context("failed to map vtl1 memory")?,
+                ))
             }
         } else {
             None

openhcl/underhill_mem/src/mapping.rs

@@ -21,6 +21,45 @@ use std::sync::Arc;
 use thiserror::Error;
 use vm_topology::memory::MemoryLayout;
 
+pub struct GuestPartitionMemoryBuilder<'a> {
+    memory_layout: &'a MemoryLayout,
+    valid_memory: Option<Arc<GuestValidMemory>>,
+}
+
+impl<'a> GuestPartitionMemoryBuilder<'a> {
+    /// valid_bitmap_state is valid when allocating a tracking bitmap is needed
+    /// for memory access, and specifies the initial state of the bitmap.
+    ///
+    /// This is used to support tracking the shared/encrypted state of each
+    /// page.
+    pub fn new(
+        memory_layout: &'a MemoryLayout,
+        valid_bitmap_state: Option<bool>,
+    ) -> Result<Self, MappingError> {
+        let valid_memory = valid_bitmap_state
+            .map(|state| GuestValidMemory::new(memory_layout, state).map(Arc::new))
+            .transpose()?;
+        Ok(Self {
+            memory_layout,
+            valid_memory,
+        })
+    }
+
+    pub fn partition_valid_memory(&self) -> Option<Arc<GuestValidMemory>> {
+        self.valid_memory.clone()
+    }
+
+    pub fn build_guest_memory_mapping(
+        &self,
+        mshv_vtl_low: &MshvVtlLow,
+        memory_mapping_builder: &mut GuestMemoryMappingBuilder,
+    ) -> Result<GuestMemoryMapping, MappingError> {
+        memory_mapping_builder
+            .use_partition_valid_memory(self.valid_memory.clone())
+            .build(mshv_vtl_low, self.memory_layout)
+    }
+}
+
 /// Partition-wide (cross-vtl) tracking of valid memory that can be used in
 /// individual GuestMemoryMappings.
 #[derive(Debug)]
@@ -30,10 +69,7 @@ pub struct GuestValidMemory {
 }
 
 impl GuestValidMemory {
-    pub fn new(
-        memory_layout: &MemoryLayout,
-        valid_bitmap_state: bool,
-    ) -> Result<Self, MappingError> {
+    fn new(memory_layout: &MemoryLayout, valid_bitmap_state: bool) -> Result<Self, MappingError> {
         let valid_bitmap = {
             let mut bitmap = {
                 // Calculate the total size of the address space by looking at the ending region.
@@ -205,15 +241,9 @@ pub struct GuestMemoryMappingBuilder {
 }
 
 impl GuestMemoryMappingBuilder {
-    /// Set whether to allocate a tracking for memory access, and specify the
-    /// initial state of the bitmap.
-    ///
-    /// This is used to support tracking the shared/encrypted state of each
-    /// page.
-    ///
     /// FUTURE: use bitmaps to track VTL permissions as well, to support guest
     /// VSM for hardware-isolated VMs.
-    pub fn use_partition_valid_memory(
+    fn use_partition_valid_memory(
         &mut self,
         valid_memory: Option<Arc<GuestValidMemory>>,
     ) -> &mut Self {
@@ -284,7 +314,7 @@ impl GuestMemoryMappingBuilder {
     /// If `bitmap_state` is `Some`, a bitmap is created to track the
     /// accessibility state of each page in the lower VTL memory. The bitmap is
     /// initialized to the provided state.
-    pub fn build(
+    fn build(
         &self,
         mshv_vtl_low: &MshvVtlLow,
         memory_layout: &MemoryLayout,