From 51ebc3ac92e60c1ef46fb99a028e8d175d361084 Mon Sep 17 00:00:00 2001
From: Muminul Islam <muislam@microsoft.com>
Date: Thu, 18 Jan 2024 17:36:28 -0800
Subject: [PATCH] vmm: set SEV control register for SEV-Enabled guest

Set the SEV control register so we know where to
start running.  This register configures the
SEV feature control state on a virtual processor.

Signed-off-by: Jinank Jain <jinankjain@microsoft.com>
Signed-off-by: Muminul Islam <muislam@microsoft.com>
---
 vmm/src/cpu.rs | 29 +++++++++++++++++++++++++++++
 vmm/src/vm.rs  | 19 +++++++++++++++----
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index e79ec275a..f36db223c 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -184,6 +184,10 @@ pub enum Error {
 
     #[error("Maximum number of vCPUs exceeds host limit")]
     MaximumVcpusExceeded,
+
+    #[cfg(feature = "sev_snp")]
+    #[error("Failed to set sev control register: {0}")]
+    SetSevControlRegister(#[source] hypervisor::HypervisorCpuError),
 }
 pub type Result<T> = result::Result<T, Error>;
 
@@ -434,6 +438,13 @@ impl Vcpu {
     pub fn run(&self) -> std::result::Result<VmExit, HypervisorCpuError> {
         self.vcpu.run()
     }
+
+    #[cfg(feature = "sev_snp")]
+    pub fn set_sev_control_register(&self, vmsa_pfn: u64) -> Result<()> {
+        self.vcpu
+            .set_sev_control_register(vmsa_pfn)
+            .map_err(Error::SetSevControlRegister)
+    }
 }
 
 impl Pausable for Vcpu {}
@@ -482,6 +493,8 @@ pub struct CpuManager {
     affinity: BTreeMap<u8, Vec<usize>>,
     dynamic: bool,
     hypervisor: Arc<dyn hypervisor::Hypervisor>,
+    #[cfg(feature = "sev_snp")]
+    sev_snp_enabled: bool,
 }
 
 const CPU_ENABLE_FLAG: usize = 0;
@@ -630,6 +643,7 @@ impl CpuManager {
         vm_ops: Arc<dyn VmOps>,
         #[cfg(feature = "tdx")] tdx_enabled: bool,
         numa_nodes: &NumaNodes,
+        #[cfg(feature = "sev_snp")] sev_snp_enabled: bool,
     ) -> Result<Arc<Mutex<CpuManager>>> {
         if u32::from(config.max_vcpus) > hypervisor.get_max_vcpus() {
             return Err(Error::MaximumVcpusExceeded);
@@ -721,6 +735,8 @@ impl CpuManager {
             affinity,
             dynamic,
             hypervisor: hypervisor.clone(),
+            #[cfg(feature = "sev_snp")]
+            sev_snp_enabled,
         })))
     }
 
@@ -806,6 +822,14 @@ impl CpuManager {
     ) -> Result<()> {
         let mut vcpu = vcpu.lock().unwrap();
 
+        #[cfg(feature = "sev_snp")]
+        if self.sev_snp_enabled {
+            if let Some((kernel_entry_point, _)) = boot_setup {
+                vcpu.set_sev_control_register(
+                    kernel_entry_point.entry_addr.0 / crate::igvm::HV_PAGE_SIZE,
+                )?;
+            }
+        }
         #[cfg(target_arch = "x86_64")]
         assert!(!self.cpuid.is_empty());
 
@@ -1814,6 +1838,11 @@ impl CpuManager {
             .unwrap();
         Ok(leaf_info)
     }
+
+    #[cfg(feature = "sev_snp")]
+    pub(crate) fn sev_snp_enabled(&self) -> bool {
+        self.sev_snp_enabled
+    }
 }
 
 struct Cpu {
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index b8d6599d8..ab0f41774 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -544,6 +544,8 @@ impl Vm {
             #[cfg(feature = "tdx")]
             tdx_enabled,
             &numa_nodes,
+            #[cfg(feature = "sev_snp")]
+            sev_snp_enabled,
         )
         .map_err(Error::CpuManager)?;
 
@@ -998,12 +1000,21 @@ impl Vm {
         memory_manager: Arc<Mutex<MemoryManager>>,
         cpu_manager: Arc<Mutex<cpu::CpuManager>>,
     ) -> Result<EntryPoint> {
-        let res = igvm_loader::load_igvm(&igvm, memory_manager, cpu_manager, "")
+        let res = igvm_loader::load_igvm(&igvm, memory_manager, cpu_manager.clone(), "")
             .map_err(Error::IgvmLoad)?;
 
-        Ok(EntryPoint {
-            entry_addr: vm_memory::GuestAddress(res.vmsa.rip),
-        })
+        cfg_if::cfg_if! {
+            if #[cfg(feature = "sev_snp")] {
+                let entry_point = if cpu_manager.lock().unwrap().sev_snp_enabled() {
+                    EntryPoint { entry_addr: vm_memory::GuestAddress(res.vmsa_gpa) }
+                } else {
+                    EntryPoint {entry_addr: vm_memory::GuestAddress(res.vmsa.rip) }
+                };
+            } else {
+               let entry_point = EntryPoint { entry_addr: vm_memory::GuestAddress(res.vmsa.rip) };
+            }
+        };
+        Ok(entry_point)
     }
 
     #[cfg(target_arch = "x86_64")]