refactor(snp): implement memory encryption as layout change callbacks

Signed-off-by: Changyuan Lyu <changyuanl@google.com>

Author: Lencerf

alioth/src/board/board.rs

@@ -245,12 +245,12 @@ where
             coco: config.coco.clone(),
         };
         let mut vm = hv.create_vm(&vm_config)?;
-        let vm_memory = vm.create_vm_memory()?;
+        let vm_memory = Arc::new(vm.create_vm_memory()?);
         let arch = ArchBoard::new(hv, &vm, &config)?;
 
         let board = Board {
             vm,
-            memory: Memory::new(vm_memory),
+            memory: Memory::new(vm_memory.clone()),
             arch,
             config,
             payload: RwLock::new(None),
@@ -273,7 +273,7 @@ where
             cond_var: Condvar::new(),
         };
 
-        board.coco_init()?;
+        board.coco_init(vm_memory)?;
 
         Ok(board)
     }

alioth/src/board/board_aarch64.rs

@@ -154,7 +154,7 @@ where
         Ok(())
     }
 
-    pub fn coco_init(&self) -> Result<()> {
+    pub fn coco_init(&self, _: Arc<V::Memory>) -> Result<()> {
         Ok(())
     }
 

alioth/src/board/board_x86_64/board_x86_64.rs

@@ -297,34 +297,23 @@ where
             callbacks: Mutex::new(vec![]),
         };
         memory.add_region(0, Arc::new(region_low))?;
-        if let Some(coco) = &self.config.coco {
-            memory.register_encrypted_pages(&pages_low)?;
-            if let Coco::AmdSnp { .. } = coco {
-                memory.mark_private_memory(0, low_mem_size as _, true)?;
-            }
-        }
+
         if config.mem.size > RAM_32_SIZE {
             let mem_hi_size = config.mem.size - RAM_32_SIZE;
             let mem_hi = self.create_ram_pages(mem_hi_size, c"ram-high")?;
             let region_hi = MemRegion::with_ram(mem_hi.clone(), MemRegionType::Ram);
             memory.add_region(MEM_64_START, Arc::new(region_hi))?;
-            if let Some(coco) = &self.config.coco {
-                memory.register_encrypted_pages(&mem_hi)?;
-                if let Coco::AmdSnp { .. } = coco {
-                    memory.mark_private_memory(MEM_64_START as _, mem_hi_size as _, true)?;
-                }
-            }
         }
         Ok(())
     }
 
-    pub fn coco_init(&self) -> Result<()> {
+    pub fn coco_init(&self, memory: Arc<V::Memory>) -> Result<()> {
         let Some(coco) = &self.config.coco else {
             return Ok(());
         };
         match coco {
-            Coco::AmdSev { policy } => self.vm.sev_launch_start(*policy)?,
-            Coco::AmdSnp { policy } => self.vm.snp_launch_start(*policy)?,
+            Coco::AmdSev { policy } => self.sev_init(*policy, memory)?,
+            Coco::AmdSnp { policy } => self.snp_init(*policy, memory)?,
             Coco::IntelTdx { attr } => todo!("Intel TDX {attr:?}"),
         }
         Ok(())

alioth/src/board/board_x86_64/sev.rs

@@ -13,19 +13,21 @@
 // limitations under the License.
 
 use std::iter::zip;
+use std::sync::Arc;
 use std::sync::atomic::Ordering;
 
 use zerocopy::FromZeros;
 
 use crate::arch::layout::MEM_64_START;
 use crate::arch::reg::{Reg, SegAccess, SegReg, SegRegVal};
-use crate::arch::sev::{SevPolicy, SnpPageType};
+use crate::arch::sev::{SevPolicy, SnpPageType, SnpPolicy};
 use crate::board::{Board, Result, VcpuGuard};
 use crate::firmware::ovmf::sev::{
     SevDescType, SevMetadataDesc, SnpCpuidFunc, SnpCpuidInfo, parse_desc, parse_sev_ap_eip,
 };
-use crate::hv::{Vcpu, Vm};
+use crate::hv::{Vcpu, Vm, VmMemory};
 use crate::mem::mapped::ArcMemPages;
+use crate::mem::{self, LayoutChanged, MarkPrivateMemory};
 
 impl<V> Board<V>
 where
@@ -162,4 +164,38 @@ where
         )?;
         Ok(())
     }
+
+    pub(crate) fn sev_init(&self, policy: SevPolicy, memory: Arc<V::Memory>) -> Result<()> {
+        self.vm.sev_launch_start(policy)?;
+        let encrypt_pages = Box::new(EncryptPages { memory });
+        self.memory.register_change_callback(encrypt_pages)?;
+        Ok(())
+    }
+
+    pub(crate) fn snp_init(&self, policy: SnpPolicy, memory: Arc<V::Memory>) -> Result<()> {
+        self.vm.snp_launch_start(policy)?;
+        let encrypt_pages = Box::new(EncryptPages {
+            memory: memory.clone(),
+        });
+        self.memory.register_change_callback(encrypt_pages)?;
+        let mark_private_memory = Box::new(MarkPrivateMemory { memory });
+        self.memory.register_change_callback(mark_private_memory)?;
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+pub struct EncryptPages {
+    memory: Arc<dyn VmMemory>,
+}
+
+impl LayoutChanged for EncryptPages {
+    fn ram_added(&self, _: u64, pages: &ArcMemPages) -> mem::Result<()> {
+        self.memory.register_encrypted_range(pages.as_slice())?;
+        Ok(())
+    }
+
+    fn ram_removed(&self, _: u64, _: &ArcMemPages) -> mem::Result<()> {
+        Ok(())
+    }
 }

alioth/src/mem/mem.rs

@@ -264,21 +264,21 @@ pub struct Memory {
     callbacks: Mutex<LayoutCallbacks>,
     ram_bus: Arc<RamBus>,
     mmio_bus: RwLock<MmioBus>,
-    vm_memory: Box<dyn VmMemory>,
+    vm_memory: Arc<dyn VmMemory>,
 
     #[cfg(target_arch = "x86_64")]
     io_bus: RwLock<MmioBus>,
     io_regions: Mutex<Addressable<Arc<IoRegion>>>,
 }
 
 impl Memory {
-    pub fn new(vm_memory: impl VmMemory) -> Self {
+    pub fn new(vm_memory: Arc<dyn VmMemory>) -> Self {
         Memory {
             regions: Mutex::new(Addressable::new()),
             callbacks: Mutex::new(LayoutCallbacks::default()),
             ram_bus: Arc::new(RamBus::new()),
             mmio_bus: RwLock::new(MmioBus::new()),
-            vm_memory: Box::new(vm_memory),
+            vm_memory,
             #[cfg(target_arch = "x86_64")]
             io_bus: RwLock::new(MmioBus::new()),
             io_regions: Mutex::new(Addressable::new()),
@@ -397,7 +397,7 @@ impl Memory {
                 }
                 MemRange::Ram(r) => {
                     self.ram_bus.remove(gpa)?;
-                    for callback in &callbacks.changed {
+                    for callback in callbacks.changed.iter().rev() {
                         callback.ram_removed(gpa, r)?;
                     }
                     self.unmap_from_vm(gpa, r)?;
@@ -605,3 +605,19 @@ impl Memory {
         }
     }
 }
+
+#[derive(Debug)]
+pub struct MarkPrivateMemory {
+    pub memory: Arc<dyn VmMemory>,
+}
+
+impl LayoutChanged for MarkPrivateMemory {
+    fn ram_added(&self, gpa: u64, pages: &ArcMemPages) -> Result<()> {
+        self.memory.mark_private_memory(gpa, pages.size(), true)?;
+        Ok(())
+    }
+
+    fn ram_removed(&self, _: u64, _: &ArcMemPages) -> Result<()> {
+        Ok(())
+    }
+}

alioth/src/pci/config_test.rs

@@ -291,7 +291,7 @@ impl VmMemory for FakeVmMemory {
 
 #[test]
 fn test_mem_bar_layout_change() {
-    let memory = Memory::new(FakeVmMemory);
+    let memory = Memory::new(Arc::new(FakeVmMemory));
     let header = fixture_emulated_header();
 
     let callback = assert_matches!(
@@ -351,7 +351,7 @@ fn test_mem_bar_layout_change() {
 
 #[test]
 fn test_io_bar_layout_change() {
-    let memory = Memory::new(FakeVmMemory);
+    let memory = Memory::new(Arc::new(FakeVmMemory));
     let header = fixture_emulated_header();
 
     let callback = assert_matches!(