refactor(mem): decouple hypervisor and struct Memory

alioth/src/board/board.rs

@@ -43,10 +43,10 @@ use crate::device::MmioDev;
 #[cfg(target_arch = "x86_64")]
 use crate::device::fw_cfg::FwCfg;
 use crate::errors::{DebugTrace, trace_error};
-use crate::hv::{CocoSpec, Hypervisor, Vm, VmSpec};
+use crate::hv::{CocoSpec, Hypervisor, MemMapOption, Vm, VmSpec};
 use crate::loader::PayloadSpec;
 use crate::mem::mapped::ArcMemPages;
-use crate::mem::{MemBackend, MemRegion, MemRegionType, MemSpec, Memory};
+use crate::mem::{self, LayoutChanged, MemBackend, MemRegion, MemRegionType, MemSpec, Memory};
 use crate::pci::bus::PciBus;
 
 #[cfg(target_arch = "aarch64")]
@@ -162,11 +162,45 @@ impl BoardSpec {
     }
 }
 
+#[derive(Debug)]
+struct ChangeVmMemory<V> {
+    vm: Arc<V>,
+}
+
+impl<V> LayoutChanged for ChangeVmMemory<V>
+where
+    V: Vm,
+{
+    fn ram_added(&self, gpa: u64, pages: &ArcMemPages) -> mem::Result<()> {
+        let opt = MemMapOption {
+            read: true,
+            write: true,
+            exec: true,
+            log_dirty: false,
+        };
+        self.vm.map(gpa, pages.size(), pages.addr(), opt)?;
+        Ok(())
+    }
+
+    fn ram_removed(&self, gpa: u64, pages: &ArcMemPages) -> mem::Result<()> {
+        self.vm.unmap(gpa, pages.size())?;
+        Ok(())
+    }
+
+    fn dev_mem_added(&self, gpa: u64, pages: &ArcMemPages) -> mem::Result<()> {
+        self.ram_added(gpa, pages)
+    }
+
+    fn dev_mem_removed(&self, gpa: u64, pages: &ArcMemPages) -> mem::Result<()> {
+        self.ram_removed(gpa, pages)
+    }
+}
+
 pub struct Board<V>
 where
     V: Vm,
 {
-    pub vm: V,
+    pub vm: Arc<V>,
     pub memory: Memory,
     pub arch: ArchBoard<V>,
     pub spec: BoardSpec,
@@ -191,13 +225,14 @@ where
         let vm_spec = VmSpec {
             coco: spec.coco.clone(),
         };
-        let mut vm = hv.create_vm(&vm_spec)?;
-        let vm_memory = Arc::new(vm.create_vm_memory()?);
+        let vm = Arc::new(hv.create_vm(&vm_spec)?);
         let arch = ArchBoard::new(hv, &vm, &spec)?;
+        let memory = Memory::new();
+        memory.register_change_callback(Box::new(ChangeVmMemory { vm: vm.clone() }))?;
 
         let board = Board {
             vm,
-            memory: Memory::new(vm_memory.clone()),
+            memory,
             arch,
             spec,
             payload: RwLock::new(None),
@@ -208,7 +243,7 @@ where
             fw_cfg: Mutex::new(None),
         };
 
-        board.coco_init(vm_memory)?;
+        board.coco_init()?;
 
         Ok(board)
     }

alioth/src/board/board_amd64/board_amd64.rs

@@ -236,14 +236,14 @@ where
         Ok(())
     }
 
-    pub fn coco_init(&self, memory: Arc<V::Memory>) -> Result<()> {
+    pub fn coco_init(&self) -> Result<()> {
         let Some(coco) = &self.spec.coco else {
             return Ok(());
         };
         match coco {
-            CocoSpec::AmdSev { policy } => self.sev_init(*policy, memory)?,
-            CocoSpec::AmdSnp { policy } => self.snp_init(*policy, memory)?,
-            CocoSpec::IntelTdx { attr } => self.tdx_init(*attr, memory)?,
+            CocoSpec::AmdSev { policy } => self.sev_init(*policy)?,
+            CocoSpec::AmdSnp { policy } => self.snp_init(*policy)?,
+            CocoSpec::IntelTdx { attr } => self.tdx_init(*attr)?,
         }
         Ok(())
     }

alioth/src/board/board_amd64/sev.rs

@@ -22,7 +22,7 @@ use crate::arch::cpuid::{
 };
 use crate::arch::sev::{SevPolicy, SnpPolicy};
 use crate::board::{Board, Result, error};
-use crate::hv::{CocoSpec, Vm, VmMemory};
+use crate::hv::{CocoSpec, Vm};
 use crate::mem::mapped::ArcMemPages;
 use crate::mem::{self, LayoutChanged, MarkPrivateMemory};
 
@@ -108,33 +108,40 @@ impl<V> Board<V>
 where
     V: Vm,
 {
-    pub(crate) fn sev_init(&self, policy: SevPolicy, memory: Arc<V::Memory>) -> Result<()> {
+    pub(crate) fn sev_init(&self, policy: SevPolicy) -> Result<()> {
         self.vm.sev_launch_start(policy)?;
-        let encrypt_pages = Box::new(EncryptPages { memory });
+        let encrypt_pages = Box::new(EncryptPages {
+            vm: self.vm.clone(),
+        });
         self.memory.register_change_callback(encrypt_pages)?;
         Ok(())
     }
 
-    pub(crate) fn snp_init(&self, policy: SnpPolicy, memory: Arc<V::Memory>) -> Result<()> {
+    pub(crate) fn snp_init(&self, policy: SnpPolicy) -> Result<()> {
         self.vm.snp_launch_start(policy)?;
         let encrypt_pages = Box::new(EncryptPages {
-            memory: memory.clone(),
+            vm: self.vm.clone(),
         });
         self.memory.register_change_callback(encrypt_pages)?;
-        let mark_private_memory = Box::new(MarkPrivateMemory { memory });
+        let mark_private_memory = Box::new(MarkPrivateMemory {
+            vm: self.vm.clone(),
+        });
         self.memory.register_change_callback(mark_private_memory)?;
         Ok(())
     }
 }
 
 #[derive(Debug)]
-pub struct EncryptPages {
-    memory: Arc<dyn VmMemory>,
+pub struct EncryptPages<V> {
+    vm: Arc<V>,
 }
 
-impl LayoutChanged for EncryptPages {
+impl<V> LayoutChanged for EncryptPages<V>
+where
+    V: Vm,
+{
     fn ram_added(&self, _: u64, pages: &ArcMemPages) -> mem::Result<()> {
-        self.memory.register_encrypted_range(pages.as_slice())?;
+        self.vm.register_encrypted_range(pages.as_slice())?;
         Ok(())
     }
 

alioth/src/board/board_amd64/tdx.rs

@@ -12,20 +12,20 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::sync::Arc;
-
 use crate::arch::tdx::TdAttr;
 use crate::board::{Board, Result};
-use crate::hv::{Vm, VmMemory};
+use crate::hv::Vm;
 use crate::mem::MarkPrivateMemory;
 
 impl<V> Board<V>
 where
     V: Vm,
 {
-    pub(crate) fn tdx_init(&self, attr: TdAttr, memory: Arc<dyn VmMemory>) -> Result<()> {
+    pub(crate) fn tdx_init(&self, attr: TdAttr) -> Result<()> {
         self.vm.tdx_init_vm(attr, &self.arch.cpuids)?;
-        let mark_private_memory = Box::new(MarkPrivateMemory { memory });
+        let mark_private_memory = Box::new(MarkPrivateMemory {
+            vm: self.vm.clone(),
+        });
         self.memory.register_change_callback(mark_private_memory)?;
         Ok(())
     }

alioth/src/board/board_arm64.rs

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

alioth/src/cpu/cpu.rs

@@ -225,7 +225,7 @@ impl<V: Vm> VcpuThread<V> {
                     }
                 }
                 VmExit::ConvertMemory { gpa, size, private } => {
-                    memory.mark_private_memory(gpa, size, private)?;
+                    memory.mark_private_memory(&*self.ctx.board.vm, gpa, size, private)?;
                     VmEntry::None
                 }
             };

alioth/src/cpu/cpu_amd64/sev.rs

@@ -74,8 +74,12 @@ where
         let range_ref = ram.get_slice::<u8>(desc.base as u64, desc.len as u64)?;
         let bytes =
             unsafe { std::slice::from_raw_parts_mut(range_ref.as_ptr() as _, range_ref.len()) };
-        let memory = &self.ctx.board.memory;
-        memory.mark_private_memory(desc.base as _, desc.len as _, true)?;
+        self.ctx.board.memory.mark_private_memory(
+            &*self.ctx.board.vm,
+            desc.base as _,
+            desc.len as _,
+            true,
+        )?;
         let vm = &self.ctx.board.vm;
         let ret = vm.snp_launch_update(bytes, desc.base as _, page_type);
         if ret.is_err() && desc.type_ == SevDescType::CPUID {
@@ -93,7 +97,7 @@ where
     pub(crate) fn setup_sev(&self, fw: &mut ArcMemPages, policy: SevPolicy) -> Result<()> {
         let board = &self.ctx.board;
 
-        board.memory.register_encrypted_pages(fw)?;
+        board.vm.register_encrypted_range(fw.as_slice())?;
 
         let data = fw.as_slice_mut();
         if policy.es() {
@@ -104,8 +108,8 @@ where
     }
 
     pub(crate) fn setup_snp(&self, fw: &mut ArcMemPages) -> Result<()> {
-        let memory = &self.ctx.board.memory;
-        memory.register_encrypted_pages(fw)?;
+        let vm = &self.ctx.board.vm;
+        vm.register_encrypted_range(fw.as_slice())?;
 
         let data = fw.as_slice_mut();
         self.parse_sev_ap_eip(data)?;
@@ -114,8 +118,8 @@ where
         }
         let fw_gpa = MEM_64_START - data.len() as u64;
 
-        memory.mark_private_memory(fw_gpa, data.len() as _, true)?;
-        let vm = &self.ctx.board.vm;
+        let memory = &self.ctx.board.memory;
+        memory.mark_private_memory(&**vm, fw_gpa, data.len() as _, true)?;
         vm.snp_launch_update(data, fw_gpa, SnpPageType::NORMAL)?;
         Ok(())
     }

alioth/src/cpu/cpu_amd64/tdx.rs

@@ -37,7 +37,8 @@ where
 
         let fw_gpa = MEM_64_START - data.len() as u64;
         let memory = &self.ctx.board.memory;
-        memory.mark_private_memory(fw_gpa, data.len() as _, true)?;
+        let vm = &self.ctx.board.vm;
+        memory.mark_private_memory(&**vm, fw_gpa, data.len() as _, true)?;
 
         let mut accepted = Vec::new();
         let mut hob_ram = None;

alioth/src/hv/hv.rs

@@ -88,8 +88,6 @@ pub enum Error {
     GuestMsr { error: std::io::Error },
     #[snafu(display("Failed to configure memory encryption"))]
     MemEncrypt { error: std::io::Error },
-    #[snafu(display("Cannot create multiple VM memories"))]
-    MemoryCreated,
     #[snafu(display("Failed to configure an IrqFd"))]
     IrqFd { error: std::io::Error },
     #[snafu(display("Failed to configure an IoeventFd"))]
@@ -245,23 +243,6 @@ pub trait MsiSender: Debug + Send + Sync + 'static {
     fn create_irqfd(&self) -> Result<Self::IrqFd>;
 }
 
-pub trait VmMemory: Debug + Send + Sync + 'static {
-    fn mem_map(&self, gpa: u64, size: u64, hva: usize, option: MemMapOption) -> Result<(), Error>;
-
-    fn unmap(&self, gpa: u64, size: u64) -> Result<(), Error>;
-
-    fn reset(&self) -> Result<()>;
-
-    fn register_encrypted_range(&self, _range: &[u8]) -> Result<()> {
-        unimplemented!()
-    }
-    fn deregister_encrypted_range(&self, _range: &[u8]) -> Result<()> {
-        unimplemented!()
-    }
-
-    fn mark_private_memory(&self, gpa: u64, size: u64, private: bool) -> Result<()>;
-}
-
 pub trait IoeventFd: Debug + Send + Sync + AsFd + 'static {}
 
 pub trait IoeventFdRegistry: Debug + Send + Sync + 'static {
@@ -341,9 +322,8 @@ pub struct VmSpec {
     pub coco: Option<CocoSpec>,
 }
 
-pub trait Vm {
+pub trait Vm: Debug + Send + Sync + 'static {
     type Vcpu: Vcpu;
-    type Memory: VmMemory;
     type IrqSender: IrqSender + Send + Sync;
     type MsiSender: MsiSender;
     type IoeventFdRegistry: IoeventFdRegistry;
@@ -353,10 +333,19 @@ pub trait Vm {
         &self,
         #[cfg(target_arch = "aarch64")] devid: u32,
     ) -> Result<Self::MsiSender>;
-    fn create_vm_memory(&mut self) -> Result<Self::Memory, Error>;
     fn create_ioeventfd_registry(&self) -> Result<Self::IoeventFdRegistry>;
     fn stop_vcpu<T>(&self, identity: u64, handle: &JoinHandle<T>) -> Result<(), Error>;
 
+    fn map(&self, gpa: u64, size: u64, hva: usize, option: MemMapOption) -> Result<(), Error>;
+
+    fn unmap(&self, gpa: u64, size: u64) -> Result<(), Error>;
+
+    fn register_encrypted_range(&self, _range: &[u8]) -> Result<()>;
+
+    fn deregister_encrypted_range(&self, _range: &[u8]) -> Result<()>;
+
+    fn mark_private_memory(&self, gpa: u64, size: u64, private: bool) -> Result<()>;
+
     #[cfg(target_arch = "x86_64")]
     fn sev_launch_start(&self, policy: SevPolicy) -> Result<()>;
 

alioth/src/hv/hvf/vcpu/vcpu_test.rs

@@ -19,7 +19,7 @@ use libc::{MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, PROT_READ, PROT_WRITE, mmap};
 
 use crate::arch::reg::Reg;
 use crate::ffi;
-use crate::hv::{Hvf, Hypervisor, MemMapOption, Vcpu, Vm, VmEntry, VmExit, VmMemory, VmSpec};
+use crate::hv::{Hvf, Hypervisor, MemMapOption, Vcpu, Vm, VmEntry, VmExit, VmSpec};
 
 #[test]
 #[cfg_attr(not(feature = "test-hv"), ignore)]
@@ -75,8 +75,7 @@ fn test_vcpu_regs() {
 fn test_vcpu_run() {
     let hvf = Hvf {};
     let spec = VmSpec { coco: None };
-    let mut vm = hvf.create_vm(&spec).unwrap();
-    let memory = vm.create_vm_memory().unwrap();
+    let vm = hvf.create_vm(&spec).unwrap();
 
     let prot = PROT_WRITE | PROT_READ;
     let flag = MAP_ANONYMOUS | MAP_PRIVATE;
@@ -91,9 +90,7 @@ fn test_vcpu_run() {
         exec: true,
         ..Default::default()
     };
-    memory
-        .mem_map(0, 0x4000, user_mem as usize, mmap_option)
-        .unwrap();
+    vm.map(0, 0x4000, user_mem as usize, mmap_option).unwrap();
 
     const CODE: [u8; 20] = [
         0x00, 0x00, 0x8a, 0xd2, // mov x0, #0x5000