vmm: add KVM SEV-SNP support to IGVM loader

Adapt the IGVM loader to work with both MSHV and KVM backends, which
differ in page type constants, CPUID page layout, and VMSA handling.

Abstract page types into a PageTypeConfig struct populated at runtime
from the detected hypervisor, replacing hardcoded mshv_bindings constants.

Apply the VMSA register state to each vCPU via setup_sev_snp_regs(),
translating SevSelector attributes to KVM segment format using a bitfield
decoder.

KVM's SNP launch path sanitizes certain CPUID bits that could lead to
an insecure guest. If the VMM sets these bits, KVM rejects the CPUID
page import on the first attempt, requiring a retry with the
firmware-corrected values.

Pre-clear the known problematic bits before import to avoid the
reject-and-retry cycle:

- Leaf 0x1, ECX bit 24: TSC_DEADLINE (filtered by KVM)
- Leaf 0x7, EBX bit 1: SGX (filtered by KVM)
- Leaf 0x7, EDX: clear entirely (contains speculative features)
- Leaf 0x80000008, EBX bit 25: filtered by KVM
- Leaf 0x80000021, ECX: clear entirely

This keeps the CPUID page stable across launch updates and avoids
noisy error logs from the retry path.

Co-authored-by: Keith Adler <kadler@cloudflare.com>
Signed-off-by: Keith Adler <kadler@cloudflare.com>
Co-authored-by: Alex Orozco <aorozco@google.com>
Signed-off-by: Alex Orozco <aorozco@google.com>
Co-authored-by: Dylan Reid <dgreid@fb.com>
Signed-off-by: Dylan Reid <dgreid@fb.com>
Signed-off-by: Ruben Hakobyan <hruben@meta.com>

hypervisor/src/cpu.rs

@@ -591,6 +591,10 @@ pub trait Vcpu: Send + Sync {
     fn set_sev_control_register(&self, _vmsa_pfn: u64) -> Result<()> {
         unimplemented!()
     }
+    #[cfg(feature = "sev_snp")]
+    fn setup_sev_snp_regs(&self, _vmsa: igvm::snp_defs::SevVmsa) -> Result<()> {
+        unimplemented!()
+    }
 
     ///
     /// Sets the value of GIC redistributor address

hypervisor/src/kvm/mod.rs

@@ -145,10 +145,61 @@ ioctl_io_nr!(KVM_NMI, kvm_bindings::KVMIO, 0x9a);
 #[cfg(feature = "sev_snp")]
 use igvm_defs::PAGE_SIZE_4K;
 #[cfg(feature = "sev_snp")]
-use kvm_bindings::{KVM_MEMORY_ATTRIBUTE_PRIVATE, KVM_X86_SNP_VM, kvm_memory_attributes};
+use kvm_bindings::{
+    KVM_MEMORY_ATTRIBUTE_PRIVATE, KVM_X86_SNP_VM, kvm_memory_attributes, kvm_segment as Segment,
+};
+use vm_memory::GuestAddress;
 #[cfg(feature = "sev_snp")]
 use x86_64::sev;
 
+// Hardcoded GPA of a bootloader and VMSA page for KVM
+// TODO: Derive these from the IGVM file's PageData/SnpVpContext directives
+// instead of using fixed constants, to support arbitrary bootloader layouts.
+pub const BOOTLOADER_START: GuestAddress = GuestAddress(0xffc0_0000);
+pub const BOOTLOADER_SIZE: usize = 0x40_0000; // 4 MiB
+pub const KVM_VMSA_PAGE_ADDRESS: GuestAddress = GuestAddress(0xffff_ffff_f000);
+pub const KVM_VMSA_PAGE_SIZE: usize = 0x1000; // 4 KiB
+
+#[cfg(feature = "sev_snp")]
+#[bitfield_struct::bitfield(u32)]
+#[derive(PartialEq, Eq)]
+/// AMD VMCB segment attributes
+/// linux/arch/x86/include/asm/svm.h
+pub struct SegAccess {
+    #[bits(4)]
+    pub seg_type: u8,
+    pub s_code_data: bool,
+    #[bits(2)]
+    pub priv_level: u8,
+    pub present: bool,
+    pub available: bool,
+    pub l_64bit: bool,
+    pub db_size_32: bool,
+    pub granularity: bool,
+    #[bits(20)]
+    _reserved: u32,
+}
+
+#[cfg(feature = "sev_snp")]
+fn make_segment(sev_selector: igvm::snp_defs::SevSelector) -> Segment {
+    let flags = SegAccess::from_bits(sev_selector.attrib.into());
+    Segment {
+        base: sev_selector.base,
+        limit: sev_selector.limit,
+        selector: sev_selector.selector,
+        type_: flags.seg_type(),
+        s: flags.s_code_data() as u8,
+        dpl: flags.priv_level(),
+        present: flags.present() as u8,
+        avl: flags.available() as u8,
+        db: flags.db_size_32() as u8,
+        g: flags.granularity() as u8,
+        l: flags.l_64bit() as u8,
+        unusable: 0,
+        ..Default::default()
+    }
+}
+
 #[cfg(feature = "tdx")]
 const KVM_EXIT_TDX: u32 = 50;
 #[cfg(feature = "tdx")]
@@ -3238,6 +3289,81 @@ impl cpu::Vcpu for KvmVcpu {
             Ok(_) => Ok(()),
         }
     }
+
+    #[cfg(feature = "sev_snp")]
+    fn set_sev_control_register(&self, _vmsa_pfn: u64) -> cpu::Result<()> {
+        Ok(())
+    }
+
+    #[cfg(feature = "sev_snp")]
+    fn setup_sev_snp_regs(&self, vmsa: igvm::snp_defs::SevVmsa) -> cpu::Result<()> {
+        let mut sregs = self
+            .fd
+            .get_sregs()
+            .map_err(|e: kvm_ioctls::Error| cpu::HypervisorCpuError::GetSpecialRegs(e.into()))?;
+        sregs.cs = make_segment(vmsa.cs);
+        sregs.ds = make_segment(vmsa.ds);
+        sregs.es = make_segment(vmsa.es);
+        sregs.fs = make_segment(vmsa.fs);
+        sregs.gs = make_segment(vmsa.gs);
+        sregs.ss = make_segment(vmsa.ss);
+        sregs.tr = make_segment(vmsa.tr);
+        sregs.ldt = make_segment(vmsa.ldtr);
+
+        sregs.cr0 = vmsa.cr0;
+        sregs.cr4 = vmsa.cr4;
+        sregs.cr3 = vmsa.cr3;
+        sregs.efer = vmsa.efer;
+
+        sregs.idt.base = vmsa.idtr.base;
+        sregs.idt.limit = vmsa
+            .idtr
+            .limit
+            .try_into()
+            .map_err(|e: std::num::TryFromIntError| {
+                cpu::HypervisorCpuError::SetSpecialRegs(anyhow!(e))
+            })?;
+        sregs.gdt.base = vmsa.gdtr.base;
+        sregs.gdt.limit = vmsa
+            .gdtr
+            .limit
+            .try_into()
+            .map_err(|e: std::num::TryFromIntError| {
+                cpu::HypervisorCpuError::SetSpecialRegs(anyhow!(e))
+            })?;
+        self.fd
+            .set_sregs(&sregs)
+            .map_err(|e: kvm_ioctls::Error| cpu::HypervisorCpuError::SetSpecialRegs(e.into()))?;
+
+        let mut regs = self
+            .fd
+            .get_regs()
+            .map_err(|e: kvm_ioctls::Error| cpu::HypervisorCpuError::GetRegister(e.into()))?;
+        regs.rip = vmsa.rip;
+        regs.rdx = vmsa.rdx;
+        regs.rflags = vmsa.rflags;
+        regs.rsp = vmsa.rsp;
+        regs.rax = vmsa.rax;
+        regs.rbx = vmsa.rbx;
+        regs.rcx = vmsa.rcx;
+        regs.rbp = vmsa.rbp;
+        regs.rsi = vmsa.rsi;
+        regs.rdi = vmsa.rdi;
+        regs.r8 = vmsa.r8;
+        regs.r9 = vmsa.r9;
+        regs.r10 = vmsa.r10;
+        regs.r11 = vmsa.r11;
+        regs.r12 = vmsa.r12;
+        regs.r13 = vmsa.r13;
+        regs.r14 = vmsa.r14;
+        regs.r15 = vmsa.r15;
+
+        self.fd
+            .set_regs(&regs)
+            .map_err(|e: kvm_ioctls::Error| cpu::HypervisorCpuError::SetRegister(e.into()))?;
+
+        Ok(())
+    }
 }
 
 impl KvmVcpu {

vmm/src/config.rs

@@ -352,6 +352,9 @@ pub enum ValidationError {
     #[cfg(feature = "sev_snp")]
     #[error("Invalid host data format")]
     InvalidHostData,
+    #[cfg(all(feature = "sev_snp", feature = "igvm"))]
+    #[error("SEV-SNP requires an IGVM payload (--payload igvm=<path>)")]
+    SevSnpRequiresIgvm,
     /// Restore expects all net ids that have fds
     #[error("Net id {0} is associated with FDs and is required")]
     RestoreMissingRequiredNetId(String),
@@ -2823,12 +2826,25 @@ impl VmConfig {
 
         #[cfg(feature = "sev_snp")]
         {
-            let host_data_opt = &self.payload.as_ref().unwrap().host_data;
-
-            if let Some(host_data) = host_data_opt
-                && host_data.len() != 64
-            {
-                return Err(ValidationError::InvalidHostData);
+            let sev_snp_enabled = self.platform.as_ref().is_some_and(|p| p.sev_snp);
+            if sev_snp_enabled {
+                let host_data_opt = &self.payload.as_ref().unwrap().host_data;
+                if let Some(host_data) = host_data_opt
+                    && host_data.len() != 64
+                {
+                    return Err(ValidationError::InvalidHostData);
+                }
+                // KVM SEV-SNP requires an IGVM payload to initialise the VMSA.
+                // Without IGVM the vCPU register state is undefined and VM entry fails.
+                #[cfg(feature = "igvm")]
+                if self
+                    .payload
+                    .as_ref()
+                    .and_then(|p| p.igvm.as_ref())
+                    .is_none()
+                {
+                    return Err(ValidationError::SevSnpRequiresIgvm);
+                }
             }
         }
         // The 'conflict' check is introduced in commit 24438e0390d3

vmm/src/cpu.rs

@@ -212,6 +212,9 @@ pub enum Error {
     #[cfg(feature = "sev_snp")]
     #[error("Failed to set sev control register")]
     SetSevControlRegister(#[source] hypervisor::HypervisorCpuError),
+    #[cfg(feature = "sev_snp")]
+    #[error("Failed to set up SEV-SNP vCPU registers")]
+    SetupSevSnpRegs(#[source] hypervisor::HypervisorCpuError),
 
     #[cfg(target_arch = "x86_64")]
     #[error("Failed to inject NMI")]
@@ -644,6 +647,13 @@ impl Vcpu {
             .map_err(Error::SetSevControlRegister)
     }
 
+    #[cfg(feature = "sev_snp")]
+    pub fn setup_sev_snp_regs(&self, vmsa: igvm::snp_defs::SevVmsa) -> Result<()> {
+        self.vcpu
+            .setup_sev_snp_regs(vmsa)
+            .map_err(Error::SetupSevSnpRegs)
+    }
+
     ///
     /// Sets the vCPU's GIC redistributor base address.
     ///
@@ -2199,7 +2209,7 @@ impl CpuManager {
         &self.vcpus_kill_signalled
     }
 
-    #[cfg(feature = "igvm")]
+    #[cfg(all(feature = "igvm", feature = "mshv"))]
     pub(crate) fn get_cpuid_leaf(
         &self,
         cpu_id: u8,
@@ -2222,6 +2232,11 @@ impl CpuManager {
         self.sev_snp_enabled
     }
 
+    #[cfg(feature = "igvm")]
+    pub(crate) fn hypervisor_type(&self) -> hypervisor::HypervisorType {
+        self.hypervisor.hypervisor_type()
+    }
+
     pub(crate) fn nmi(&mut self) -> Result<()> {
         self.vcpus_kick_signalled.store(true, Ordering::SeqCst);
         self.signal_vcpus()?;

vmm/src/igvm/igvm_loader.rs

@@ -7,6 +7,7 @@ use std::ffi::CString;
 use std::mem::size_of;
 use std::sync::{Arc, Mutex};
 
+use hypervisor::HypervisorType;
 use igvm::snp_defs::SevVmsa;
 use igvm::{IgvmDirectiveHeader, IgvmFile, IgvmPlatformHeader};
 #[cfg(feature = "sev_snp")]
@@ -15,13 +16,20 @@ use igvm_defs::{
     IGVM_VHS_PARAMETER, IGVM_VHS_PARAMETER_INSERT, IgvmPageDataType, IgvmPlatformType,
 };
 use log::debug;
+#[cfg(all(feature = "kvm", feature = "sev_snp"))]
+use log::error;
 #[cfg(feature = "sev_snp")]
 use log::info;
+#[cfg(feature = "mshv")]
 use mshv_bindings::*;
 use thiserror::Error;
 #[cfg(feature = "sev_snp")]
-use vm_memory::{GuestAddress, GuestAddressSpace, GuestMemory};
+use vm_memory::{Bytes, GuestAddress, GuestAddressSpace, GuestMemory};
+#[cfg(all(feature = "kvm", feature = "sev_snp"))]
+use vm_migration::Snapshottable;
 use zerocopy::IntoBytes;
+#[cfg(feature = "sev_snp")]
+use zerocopy::{FromBytes, FromZeros};
 
 #[cfg(feature = "sev_snp")]
 use crate::GuestMemoryMmap;
@@ -30,6 +38,36 @@ use crate::igvm::loader::Loader;
 use crate::igvm::{BootPageAcceptance, HV_PAGE_SIZE, IgvmLoadedInfo, StartupMemoryType};
 use crate::memory_manager::{Error as MemoryManagerError, MemoryManager};
 
+#[cfg(feature = "sev_snp")]
+const ISOLATED_PAGE_SHIFT: u32 = 12;
+#[cfg(feature = "sev_snp")]
+const SNP_CPUID_LIMIT: u32 = 64;
+// see section 7.1
+// https://www.amd.com/content/dam/amd/en/documents/epyc-technical-docs/specifications/56860.pdf
+#[cfg(feature = "sev_snp")]
+#[repr(C)]
+#[derive(Debug, Clone, PartialEq, Eq, IntoBytes, FromBytes)]
+pub struct SnpCpuidFunc {
+    pub eax_in: u32,
+    pub ecx_in: u32,
+    pub xcr0_in: u64,
+    pub xss_in: u64,
+    pub eax: u32,
+    pub ebx: u32,
+    pub ecx: u32,
+    pub edx: u32,
+    pub reserved: u64,
+}
+
+#[cfg(feature = "sev_snp")]
+#[repr(C)]
+#[derive(Debug, Clone, FromBytes, IntoBytes)]
+pub struct SnpCpuidInfo {
+    pub count: u32,
+    pub _reserved1: u32,
+    pub _reserved2: u64,
+    pub entries: [SnpCpuidFunc; SNP_CPUID_LIMIT as usize],
+}
 #[derive(Debug, Error)]
 pub enum Error {
     #[error("command line is not a valid C string")]
@@ -54,6 +92,30 @@ pub enum Error {
     MemoryManager(MemoryManagerError),
     #[error("IGVM file not provided")]
     MissingIgvm,
+    #[error("Error applying VMSA to vCPU registers: {0}")]
+    SetVmsa(#[source] crate::cpu::Error),
+}
+
+// KVM SNP page types — linux/arch/x86/include/uapi/asm/sev-guest.h
+#[cfg(feature = "kvm")]
+const KVM_SNP_PAGE_TYPE_NORMAL: u32 = 1;
+#[cfg(feature = "kvm")]
+const KVM_SNP_PAGE_TYPE_VMSA: u32 = 2;
+#[cfg(feature = "kvm")]
+const KVM_SNP_PAGE_TYPE_UNMEASURED: u32 = 4;
+#[cfg(feature = "kvm")]
+const KVM_SNP_PAGE_TYPE_SECRETS: u32 = 5;
+#[cfg(feature = "kvm")]
+const KVM_SNP_PAGE_TYPE_CPUID: u32 = 6;
+
+// Consolidated page type/size configuration per hypervisor.
+struct PageTypeConfig {
+    isolated_page_size_4kb: u32,
+    normal: u32,
+    unmeasured: u32,
+    cpuid: u32,
+    secrets: u32,
+    vmsa: u32,
 }
 
 #[allow(dead_code)]
@@ -151,6 +213,10 @@ pub fn extract_sev_features(igvm_file: &IgvmFile) -> u64 {
 /// Right now it only supports SNP based isolation.
 /// We can boot legacy VM with an igvm file without
 /// any isolation.
+///
+/// NOTE: KVM and MSHV have different page type values and CPUID/VMSA handling.
+/// Hypervisor-specific code paths are gated by runtime type checks. A future
+/// refactor could split these into separate KVM/MSHV loader implementations.
 #[allow(clippy::needless_pass_by_value)]
 pub fn load_igvm(
     igvm_file: IgvmFile,
@@ -159,6 +225,28 @@ pub fn load_igvm(
     cmdline: &str,
     #[cfg(feature = "sev_snp")] host_data: &Option<String>,
 ) -> Result<Box<IgvmLoadedInfo>, Error> {
+    let hypervisor_type = cpu_manager.lock().unwrap().hypervisor_type();
+    let page_types = match hypervisor_type {
+        #[cfg(feature = "mshv")]
+        HypervisorType::Mshv => PageTypeConfig {
+            isolated_page_size_4kb: mshv_bindings::hv_isolated_page_size_HV_ISOLATED_PAGE_SIZE_4KB,
+            normal: mshv_bindings::hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_NORMAL,
+            unmeasured: mshv_bindings::hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_UNMEASURED,
+            cpuid: mshv_bindings::hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_NORMAL,
+            secrets: mshv_bindings::hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_UNMEASURED,
+            vmsa: mshv_bindings::hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_VMSA,
+        },
+        #[cfg(feature = "kvm")]
+        HypervisorType::Kvm => PageTypeConfig {
+            isolated_page_size_4kb: HV_PAGE_SIZE as u32,
+            normal: KVM_SNP_PAGE_TYPE_NORMAL,
+            unmeasured: KVM_SNP_PAGE_TYPE_UNMEASURED,
+            cpuid: KVM_SNP_PAGE_TYPE_CPUID,
+            secrets: KVM_SNP_PAGE_TYPE_SECRETS,
+            vmsa: KVM_SNP_PAGE_TYPE_VMSA,
+        },
+    };
+
     let mut loaded_info: Box<IgvmLoadedInfo> = Box::default();
     let command_line = CString::new(cmdline).map_err(Error::InvalidCommandLine)?;
     let memory = memory_manager.lock().as_ref().unwrap().guest_memory();
@@ -173,6 +261,8 @@ pub fn load_igvm(
             .map_err(Error::FailedToDecodeHostData)?;
     }
 
+    #[cfg(feature = "sev_snp")]
+    let sev_snp_enabled = cpu_manager.lock().unwrap().sev_snp_enabled();
     let mask = match &igvm_file.platforms()[0] {
         IgvmPlatformHeader::SupportedPlatform(info) => {
             debug_assert!(info.platform_type == IgvmPlatformType::SEV_SNP);
@@ -205,15 +295,15 @@ pub fn load_igvm(
                         if flags.unmeasured() {
                             gpas.push(GpaPages {
                                 gpa: *gpa,
-                                page_type: hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_UNMEASURED,
-                                page_size: hv_isolated_page_size_HV_ISOLATED_PAGE_SIZE_4KB,
+                                page_type: page_types.unmeasured,
+                                page_size: page_types.isolated_page_size_4kb,
                             });
                             BootPageAcceptance::ExclusiveUnmeasured
                         } else {
                             gpas.push(GpaPages {
                                 gpa: *gpa,
-                                page_type: hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_NORMAL,
-                                page_size: hv_isolated_page_size_HV_ISOLATED_PAGE_SIZE_4KB,
+                                page_type: page_types.normal,
+                                page_size: page_types.isolated_page_size_4kb,
                             });
                             BootPageAcceptance::Exclusive
                         }
@@ -221,43 +311,46 @@ pub fn load_igvm(
                     IgvmPageDataType::SECRETS => {
                         gpas.push(GpaPages {
                             gpa: *gpa,
-                            page_type: hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_SECRETS,
-                            page_size: hv_isolated_page_size_HV_ISOLATED_PAGE_SIZE_4KB,
+                            page_type: page_types.secrets,
+                            page_size: page_types.isolated_page_size_4kb,
                         });
                         BootPageAcceptance::SecretsPage
                     }
                     IgvmPageDataType::CPUID_DATA => {
-                        // SAFETY: CPUID is readonly
-                        unsafe {
-                            let cpuid_page_p: *mut hv_psp_cpuid_page =
-                                data.as_ptr() as *mut hv_psp_cpuid_page; // as *mut hv_psp_cpuid_page;
-                            let cpuid_page: &mut hv_psp_cpuid_page = &mut *cpuid_page_p;
-                            for i in 0..cpuid_page.count {
-                                let leaf = cpuid_page.cpuid_leaf_info[i as usize];
-                                let mut in_leaf = cpu_manager
-                                    .lock()
-                                    .unwrap()
-                                    .get_cpuid_leaf(
-                                        0,
-                                        leaf.eax_in,
-                                        leaf.ecx_in,
-                                        leaf.xfem_in,
-                                        leaf.xss_in,
-                                    )
-                                    .unwrap();
-                                if leaf.eax_in == 1 {
-                                    in_leaf[2] &= 0x7FFFFFFF;
+                        #[cfg(feature = "mshv")]
+                        if hypervisor_type == HypervisorType::Mshv {
+                            // SAFETY: CPUID is readonly
+                            unsafe {
+                                let cpuid_page_p: *mut hv_psp_cpuid_page =
+                                    data.as_ptr() as *mut hv_psp_cpuid_page; // as *mut hv_psp_cpuid_page;
+                                let cpuid_page: &mut hv_psp_cpuid_page = &mut *cpuid_page_p;
+                                for i in 0..cpuid_page.count {
+                                    let leaf = cpuid_page.cpuid_leaf_info[i as usize];
+                                    let mut in_leaf = cpu_manager
+                                        .lock()
+                                        .unwrap()
+                                        .get_cpuid_leaf(
+                                            0,
+                                            leaf.eax_in,
+                                            leaf.ecx_in,
+                                            leaf.xfem_in,
+                                            leaf.xss_in,
+                                        )
+                                        .unwrap();
+                                    if leaf.eax_in == 1 {
+                                        in_leaf[2] &= 0x7FFFFFFF;
+                                    }
+                                    cpuid_page.cpuid_leaf_info[i as usize].eax_out = in_leaf[0];
+                                    cpuid_page.cpuid_leaf_info[i as usize].ebx_out = in_leaf[1];
+                                    cpuid_page.cpuid_leaf_info[i as usize].ecx_out = in_leaf[2];
+                                    cpuid_page.cpuid_leaf_info[i as usize].edx_out = in_leaf[3];
                                 }
-                                cpuid_page.cpuid_leaf_info[i as usize].eax_out = in_leaf[0];
-                                cpuid_page.cpuid_leaf_info[i as usize].ebx_out = in_leaf[1];
-                                cpuid_page.cpuid_leaf_info[i as usize].ecx_out = in_leaf[2];
-                                cpuid_page.cpuid_leaf_info[i as usize].edx_out = in_leaf[3];
                             }
                         }
                         gpas.push(GpaPages {
                             gpa: *gpa,
-                            page_type: hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_CPUID,
-                            page_size: hv_isolated_page_size_HV_ISOLATED_PAGE_SIZE_4KB,
+                            page_type: page_types.cpuid,
+                            page_size: page_types.isolated_page_size_4kb,
                         });
                         BootPageAcceptance::CpuidPage
                     }
@@ -265,9 +358,69 @@ pub fn load_igvm(
                     _ => todo!("unsupported IgvmPageDataType"),
                 };
 
-                loader
-                    .import_pages(gpa / HV_PAGE_SIZE, 1, acceptance, data)
-                    .map_err(Error::Loader)?;
+                #[allow(unused_mut)]
+                let mut imported_page = false;
+                #[cfg(all(feature = "kvm", feature = "sev_snp"))]
+                if hypervisor_type == HypervisorType::Kvm
+                    && *data_type == IgvmPageDataType::CPUID_DATA
+                {
+                    let mut new_cp = SnpCpuidInfo::new_zeroed();
+
+                    let entries = cpu_manager.lock().unwrap().common_cpuid();
+                    let cp_count = std::cmp::min(SNP_CPUID_LIMIT as usize, entries.len());
+                    // TODO: Filter cpuid rather than truncate
+                    for (i, entry) in entries.iter().enumerate().take(cp_count) {
+                        new_cp.entries[i].eax_in = entry.function;
+                        new_cp.entries[i].ecx_in = entry.index;
+                        new_cp.entries[i].eax = entry.eax;
+                        new_cp.entries[i].ebx = entry.ebx;
+                        new_cp.entries[i].ecx = entry.ecx;
+                        new_cp.entries[i].edx = entry.edx;
+                        /*
+                         * Guest kernels will calculate EBX themselves using the 0xD
+                         * subfunctions corresponding to the individual XSAVE areas, so only
+                         * encode the base XSAVE size in the initial leaves, corresponding
+                         * to the initial XCR0=1 state. (https://tinyurl.com/qemu-cpuid)
+                         */
+                        if new_cp.entries[i].eax_in == 0xd
+                            && (new_cp.entries[i].ecx_in == 0x0 || new_cp.entries[i].ecx_in == 0x1)
+                        {
+                            new_cp.entries[i].ebx = 0x240;
+                            new_cp.entries[i].xcr0_in = 1;
+                            new_cp.entries[i].xss_in = 0;
+                        }
+
+                        // KVM SNP launch may reject a CPUID page with bits it intends
+                        // to sanitize internally. Pre-clearing the known unsafe bits keeps
+                        // the CPUID page stable across launch updates.
+                        match (new_cp.entries[i].eax_in, new_cp.entries[i].ecx_in) {
+                            (0x1, 0x0) => {
+                                new_cp.entries[i].ecx &= !(1 << 24);
+                            }
+                            (0x7, 0x0) => {
+                                new_cp.entries[i].ebx &= !0x2;
+                                new_cp.entries[i].edx = 0;
+                            }
+                            (0x80000008, 0x0) => {
+                                new_cp.entries[i].ebx &= !0x0200_0000;
+                            }
+                            (0x80000021, 0x0) => {
+                                new_cp.entries[i].ecx = 0;
+                            }
+                            _ => {}
+                        }
+                    }
+                    new_cp.count = cp_count as u32;
+                    loader
+                        .import_pages(gpa / HV_PAGE_SIZE, 1, acceptance, new_cp.as_mut_bytes())
+                        .map_err(Error::Loader)?;
+                    imported_page = true;
+                }
+                if !imported_page {
+                    loader
+                        .import_pages(gpa / HV_PAGE_SIZE, 1, acceptance, data)
+                        .map_err(Error::Loader)?;
+                }
             }
             IgvmDirectiveHeader::ParameterArea {
                 number_of_bytes,
@@ -299,16 +452,16 @@ pub fn load_igvm(
             IgvmDirectiveHeader::MmioRanges(_info) => {
                 todo!("unsupported IgvmPageDataType");
             }
-            IgvmDirectiveHeader::MemoryMap(_info) => {
+            IgvmDirectiveHeader::MemoryMap(_info) =>
+            {
                 #[cfg(feature = "sev_snp")]
-                {
+                if sev_snp_enabled {
                     let guest_mem = memory_manager.lock().unwrap().boot_guest_memory();
                     let memory_map = generate_memory_map(&guest_mem)?;
                     import_parameter(&mut parameter_areas, _info, memory_map.as_bytes())?;
+                } else {
+                    todo!("Not implemented");
                 }
-
-                #[cfg(not(feature = "sev_snp"))]
-                todo!("Not implemented");
             }
             IgvmDirectiveHeader::CommandLine(info) => {
                 import_parameter(&mut parameter_areas, info, command_line.as_bytes_with_nul())?;
@@ -336,7 +489,7 @@ pub fn load_igvm(
                 vmsa,
             } => {
                 assert_eq!(gpa % HV_PAGE_SIZE, 0);
-                let mut data: [u8; 4096] = [0; 4096];
+                let mut data: [u8; HV_PAGE_SIZE as usize] = [0; HV_PAGE_SIZE as usize];
                 let len = size_of::<SevVmsa>();
                 loaded_info.vmsa_gpa = *gpa;
                 loaded_info.vmsa = **vmsa;
@@ -348,10 +501,28 @@ pub fn load_igvm(
                         .map_err(Error::Loader)?;
                 }
 
+                // Set vCPU initial register state from VMSA before SNP_LAUNCH_FINISH
+                #[cfg(all(feature = "kvm", feature = "sev_snp"))]
+                if hypervisor_type == HypervisorType::Kvm {
+                    let vcpus = cpu_manager.lock().unwrap().vcpus();
+                    for vcpu in vcpus {
+                        let vcpu_locked = vcpu.lock().unwrap();
+                        let vcpu_id: u16 = vcpu_locked.id().parse().unwrap();
+                        if vcpu_id == *vp_index {
+                            vcpu_locked
+                                .setup_sev_snp_regs(loaded_info.vmsa)
+                                .map_err(Error::SetVmsa)?;
+                            vcpu_locked
+                                .set_sev_control_register(0)
+                                .map_err(Error::SetVmsa)?;
+                        }
+                    }
+                }
+
                 gpas.push(GpaPages {
                     gpa: *gpa,
-                    page_type: hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_VMSA,
-                    page_size: hv_isolated_page_size_HV_ISOLATED_PAGE_SIZE_4KB,
+                    page_type: page_types.vmsa,
+                    page_size: page_types.isolated_page_size_4kb,
                 });
             }
             IgvmDirectiveHeader::SnpIdBlock {
@@ -419,8 +590,8 @@ pub fn load_igvm(
                 *area = ParameterAreaState::Inserted;
                 gpas.push(GpaPages {
                     gpa: *gpa,
-                    page_type: hv_isolated_page_type_HV_ISOLATED_PAGE_TYPE_UNMEASURED,
-                    page_size: hv_isolated_page_size_HV_ISOLATED_PAGE_SIZE_4KB,
+                    page_type: page_types.unmeasured,
+                    page_size: page_types.isolated_page_size_4kb,
                 });
             }
             IgvmDirectiveHeader::ErrorRange { .. } => {
@@ -433,7 +604,7 @@ pub fn load_igvm(
     }
 
     #[cfg(feature = "sev_snp")]
-    {
+    if sev_snp_enabled {
         memory_manager
             .lock()
             .unwrap()
@@ -471,7 +642,7 @@ pub fn load_igvm(
             // of PFN for importing the isolated pages
             let pfns: Vec<u64> = group
                 .iter()
-                .map(|gpa| gpa.gpa >> HV_HYP_PAGE_SHIFT)
+                .map(|gpa| gpa.gpa >> ISOLATED_PAGE_SHIFT)
                 .collect();
             let guest_memory = memory_manager.lock().unwrap().guest_memory().memory();
             let uaddrs: Vec<_> = group
@@ -483,17 +654,50 @@ pub fn load_igvm(
                     uaddr_base + uaddr_offset
                 })
                 .collect();
-            memory_manager
+            #[cfg(feature = "kvm")]
+            let page_type = group[0].page_type;
+            let mut new_cp = SnpCpuidInfo::new_zeroed();
+            let _ = guest_memory.read(new_cp.as_mut_bytes(), GuestAddress(group[0].gpa));
+            let import_result = memory_manager
                 .lock()
                 .unwrap()
                 .vm
                 .import_isolated_pages(
                     group[0].page_type,
-                    hv_isolated_page_size_HV_ISOLATED_PAGE_SIZE_4KB,
+                    page_types.isolated_page_size_4kb,
                     &pfns,
                     &uaddrs,
                 )
-                .map_err(Error::ImportIsolatedPages)?;
+                .map_err(Error::ImportIsolatedPages);
+            #[cfg(feature = "kvm")]
+            if hypervisor_type == HypervisorType::Kvm
+                && import_result.is_err()
+                && page_type == page_types.cpuid
+            {
+                // When we import the CPUID page, the firmware will change any cpuid fns that
+                // could lead to an insecure guest, we must then make sure to import the updated cpuid
+                // https://elixir.bootlin.com/linux/v6.11/source/arch/x86/kvm/svm/sev.c#L2322
+                let mut updated_cp = SnpCpuidInfo::new_zeroed();
+                let _ = guest_memory.read(updated_cp.as_mut_bytes(), GuestAddress(group[0].gpa));
+                for (set, got) in std::iter::zip(new_cp.entries.iter(), updated_cp.entries.iter()) {
+                    if set != got {
+                        error!("Set cpuid fn: {set:#x?}, but firmware expects: {got:#x?}");
+                    }
+                }
+                memory_manager
+                    .lock()
+                    .unwrap()
+                    .vm
+                    .import_isolated_pages(
+                        group[0].page_type,
+                        page_types.isolated_page_size_4kb,
+                        &pfns,
+                        &uaddrs,
+                    )
+                    .map_err(Error::ImportIsolatedPages)?;
+                continue;
+            }
+            import_result?;
         }
 
         info!(
@@ -502,13 +706,23 @@ pub fn load_igvm(
             gpas.len()
         );
 
+        let id_block_enabled = if hypervisor_type == HypervisorType::Mshv {
+            1
+        } else {
+            0
+        };
+
         now = Instant::now();
         // Call Complete Isolated Import since we are done importing isolated pages
         memory_manager
             .lock()
             .unwrap()
             .vm
-            .complete_isolated_import(loaded_info.snp_id_block, host_data_contents, 1)
+            .complete_isolated_import(
+                loaded_info.snp_id_block,
+                host_data_contents,
+                id_block_enabled,
+            )
             .map_err(Error::CompleteIsolatedImport)?;
 
         info!(