From 535914fd50656cc0dc3d5c1559a887acf92495d5 Mon Sep 17 00:00:00 2001
From: Tyler Fanelli <tfanelli@redhat.com>
Date: Thu, 12 Dec 2024 01:23:17 -0500
Subject: [PATCH] amd-sev: Set guest_memfd memory regions in memory_init

Signed-off-by: Tyler Fanelli <tfanelli@redhat.com>
---
 src/vmm/src/linux/tee/amdsnp.rs | 40 ++--------------------------
 src/vmm/src/linux/vstate.rs     | 46 ++++++++++++++++++++++++++++++++-
 2 files changed, 47 insertions(+), 39 deletions(-)

diff --git a/src/vmm/src/linux/tee/amdsnp.rs b/src/vmm/src/linux/tee/amdsnp.rs
index e0d99e92..05e0633b 100644
--- a/src/vmm/src/linux/tee/amdsnp.rs
+++ b/src/vmm/src/linux/tee/amdsnp.rs
@@ -12,10 +12,7 @@ use sev::{
     firmware::{guest::GuestPolicy, host::Firmware},
 };
 
-use kvm_bindings::{
-    kvm_create_guest_memfd, kvm_userspace_memory_region2, CpuId, KVM_CPUID_FLAG_SIGNIFCANT_INDEX,
-    KVM_MEM_GUEST_MEMFD,
-};
+use kvm_bindings::{CpuId, KVM_CPUID_FLAG_SIGNIFCANT_INDEX};
 use kvm_ioctls::VmFd;
 use vm_memory::{
     Bytes, GuestAddress, GuestMemory, GuestMemoryMmap, GuestMemoryRegion, GuestRegionMmap,
@@ -90,41 +87,8 @@ impl AmdSnp {
     pub fn vm_prepare(
         &self,
         vm_fd: &VmFd,
-        guest_mem: &GuestMemoryMmap,
+        _guest_mem: &GuestMemoryMmap,
     ) -> Result<Launcher<Started, RawFd, RawFd>, Error> {
-        for region in guest_mem.iter() {
-            // It's safe to unwrap because the guest address is valid.
-            let host_addr = guest_mem.get_host_address(region.start_addr()).unwrap();
-
-            // Create guest_memfd struct.
-            let gmem = kvm_create_guest_memfd {
-                size: region.len(),
-                flags: 0, //Unused.
-                reserved: [0; 6],
-            };
-
-            // Create KVM guest_memfd.
-            let fd = vm_fd.create_guest_memfd(gmem).unwrap();
-
-            // Create memory region.
-            let mem_region = kvm_userspace_memory_region2 {
-                slot: 0,
-                flags: KVM_MEM_GUEST_MEMFD,
-                guest_phys_addr: region.start_addr().0,
-                memory_size: region.len(),
-                userspace_addr: host_addr as u64,
-                guest_memfd_offset: 0,
-                guest_memfd: fd as u32,
-                pad1: 0,
-                pad2: [0; 14],
-            };
-
-            // Set the memory region.
-            unsafe {
-                vm_fd.set_user_memory_region2(mem_region).unwrap();
-            }
-        }
-
         let vm_rfd = vm_fd.as_raw_fd();
         let fw_rfd = self.fw.as_raw_fd();
 
diff --git a/src/vmm/src/linux/vstate.rs b/src/vmm/src/linux/vstate.rs
index e755fc56..2d026c02 100644
--- a/src/vmm/src/linux/vstate.rs
+++ b/src/vmm/src/linux/vstate.rs
@@ -38,6 +38,8 @@ use crate::vmm_config::machine_config::CpuFeaturesTemplate;
 use arch::aarch64::gic::GICDevice;
 #[cfg(target_arch = "x86_64")]
 use cpuid::{c3, filter_cpuid, t2, VmSpec};
+#[cfg(not(feature = "amd-sev"))]
+use kvm_bindings::kvm_userspace_memory_region;
 #[cfg(target_arch = "x86_64")]
 use kvm_bindings::{
     kvm_clock_data, kvm_debugregs, kvm_irqchip, kvm_lapic_state, kvm_mp_state, kvm_pit_config,
@@ -45,7 +47,11 @@ use kvm_bindings::{
     Msrs, KVM_CLOCK_TSC_STABLE, KVM_IRQCHIP_IOAPIC, KVM_IRQCHIP_PIC_MASTER, KVM_IRQCHIP_PIC_SLAVE,
     KVM_MAX_CPUID_ENTRIES, KVM_PIT_SPEAKER_DUMMY,
 };
-use kvm_bindings::{kvm_userspace_memory_region, KVM_API_VERSION};
+
+#[cfg(feature = "amd-sev")]
+use kvm_bindings::{kvm_create_guest_memfd, kvm_userspace_memory_region2, KVM_MEM_GUEST_MEMFD};
+
+use kvm_bindings::KVM_API_VERSION;
 use kvm_ioctls::*;
 use utils::eventfd::EventFd;
 use utils::signal::{register_signal_handler, sigrtmin, Killable};
@@ -562,6 +568,8 @@ impl Vm {
         if guest_mem.num_regions() > kvm_max_memslots {
             return Err(Error::NotEnoughMemorySlots);
         }
+
+        #[cfg(not(feature = "amd-sev"))]
         for region in guest_mem.iter() {
             // It's safe to unwrap because the guest address is valid.
             let host_addr = guest_mem.get_host_address(region.start_addr()).unwrap();
@@ -583,6 +591,42 @@ impl Vm {
             self.next_mem_slot += 1;
         }
 
+        #[cfg(feature = "amd-sev")]
+        for region in guest_mem.iter() {
+            // It's safe to unwrap because the guest address is valid.
+            let host_addr = guest_mem.get_host_address(region.start_addr()).unwrap();
+
+            // Create guest_memfd struct.
+            let gmem = kvm_create_guest_memfd {
+                size: region.len(),
+                flags: 0, // Unused.
+                reserved: [0; 6],
+            };
+
+            // Create KVM guest_memfd.
+            let guest_memfd = self.fd.create_guest_memfd(gmem).unwrap();
+
+            // Create memory region.
+            let memory_region = kvm_userspace_memory_region2 {
+                slot: self.next_mem_slot,
+                flags: KVM_MEM_GUEST_MEMFD,
+                guest_phys_addr: region.start_addr().raw_value(),
+                memory_size: region.len(),
+                userspace_addr: host_addr as u64,
+                guest_memfd_offset: 0,
+                guest_memfd: guest_memfd as u32,
+                pad1: 0,
+                pad2: [0; 14],
+            };
+
+            // Set the memory region.
+            unsafe {
+                self.fd.set_user_memory_region2(memory_region).unwrap();
+            }
+
+            self.next_mem_slot += 1;
+        }
+
         #[cfg(target_arch = "x86_64")]
         self.fd
             .set_tss_address(arch::x86_64::layout::KVM_TSS_ADDRESS as usize)