X86_64 msr cpuid refactor

src/cpu-template-helper/Cargo.toml

@@ -15,6 +15,7 @@ tracing = ["log-instrument", "vmm/tracing"]
 [dependencies]
 clap = { version = "4.6.1", features = ["derive", "string"] }
 displaydoc = "0.2.6"
+kvm-bindings = "0.14.0"
 libc = "0.2.186"
 log-instrument = { path = "../log-instrument", optional = true }
 serde = { version = "1.0.228", features = ["derive"] }

src/cpu-template-helper/src/template/dump/x86_64.rs

@@ -1,14 +1,13 @@
 // Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::collections::BTreeMap;
-
+use kvm_bindings::Msrs;
 use vmm::MSR_RANGE;
 use vmm::arch::x86_64::generated::msr_index::*;
 use vmm::arch::x86_64::msr::MsrRange;
 use vmm::cpu_config::templates::{CpuConfiguration, CustomCpuTemplate, RegisterValueFilter};
 use vmm::cpu_config::x86_64::cpuid::common::get_vendor_id_from_host;
-use vmm::cpu_config::x86_64::cpuid::{Cpuid, VENDOR_ID_AMD};
+use vmm::cpu_config::x86_64::cpuid::{Cpuid, KvmCpuidFlags, VENDOR_ID_AMD};
 use vmm::cpu_config::x86_64::custom_cpu_template::{
     CpuidLeafModifier, CpuidRegister, CpuidRegisterModifier, RegisterModifier,
 };
@@ -25,29 +24,33 @@ pub fn config_to_template(cpu_config: &CpuConfiguration) -> CustomCpuTemplate {
 }
 
 fn cpuid_to_modifiers(cpuid: &Cpuid) -> Vec<CpuidLeafModifier> {
-    cpuid
+    let mut result: Vec<_> = cpuid
         .inner()
+        .as_slice()
         .iter()
-        .map(|(key, entry)| {
+        .map(|entry| {
             cpuid_leaf_modifier!(
-                key.leaf,
-                key.subleaf,
-                entry.flags,
+                entry.function,
+                entry.index,
+                KvmCpuidFlags(entry.flags),
                 vec![
-                    cpuid_reg_modifier!(CpuidRegister::Eax, entry.result.eax),
-                    cpuid_reg_modifier!(CpuidRegister::Ebx, entry.result.ebx),
-                    cpuid_reg_modifier!(CpuidRegister::Ecx, entry.result.ecx),
-                    cpuid_reg_modifier!(CpuidRegister::Edx, entry.result.edx),
+                    cpuid_reg_modifier!(CpuidRegister::Eax, entry.eax),
+                    cpuid_reg_modifier!(CpuidRegister::Ebx, entry.ebx),
+                    cpuid_reg_modifier!(CpuidRegister::Ecx, entry.ecx),
+                    cpuid_reg_modifier!(CpuidRegister::Edx, entry.edx),
                 ]
             )
         })
-        .collect()
+        .collect();
+    result.sort_by_key(|m| (m.leaf, m.subleaf));
+    result
 }
 
-fn msrs_to_modifier(msrs: &BTreeMap<u32, u64>) -> Vec<RegisterModifier> {
+fn msrs_to_modifier(msrs: &Msrs) -> Vec<RegisterModifier> {
     let mut msrs: Vec<RegisterModifier> = msrs
+        .as_slice()
         .iter()
-        .map(|(index, value)| msr_modifier!(*index, *value))
+        .map(|entry| msr_modifier!(entry.index, entry.data))
         .collect();
 
     msrs.retain(|modifier| !should_exclude_msr(modifier.addr));
@@ -119,47 +122,37 @@ fn should_exclude_msr_amd(index: u32) -> bool {
 
 #[cfg(test)]
 mod tests {
-    use std::collections::BTreeMap;
-
-    use vmm::cpu_config::x86_64::cpuid::{
-        CpuidEntry, CpuidKey, CpuidRegisters, IntelCpuid, KvmCpuidFlags,
-    };
+    use kvm_bindings::kvm_msr_entry;
+    use vmm::cpu_config::x86_64::cpuid::IntelCpuid;
 
     use super::*;
 
     fn build_sample_cpuid() -> Cpuid {
-        Cpuid::Intel(IntelCpuid(BTreeMap::from([
-            (
-                CpuidKey {
-                    leaf: 0x0,
-                    subleaf: 0x0,
-                },
-                CpuidEntry {
-                    flags: KvmCpuidFlags::EMPTY,
-                    result: CpuidRegisters {
-                        eax: 0xffff_ffff,
-                        ebx: 0x0000_ffff,
-                        ecx: 0xffff_0000,
-                        edx: 0x0000_0000,
-                    },
-                },
-            ),
-            (
-                CpuidKey {
-                    leaf: 0x1,
-                    subleaf: 0x1,
+        Cpuid::Intel(IntelCpuid(
+            kvm_bindings::CpuId::from_entries(&[
+                kvm_bindings::kvm_cpuid_entry2 {
+                    function: 0x0,
+                    index: 0x0,
+                    flags: KvmCpuidFlags::EMPTY.0,
+                    eax: 0xffff_ffff,
+                    ebx: 0x0000_ffff,
+                    ecx: 0xffff_0000,
+                    edx: 0x0000_0000,
+                    ..Default::default()
                 },
-                CpuidEntry {
-                    flags: KvmCpuidFlags::SIGNIFICANT_INDEX,
-                    result: CpuidRegisters {
-                        eax: 0xaaaa_aaaa,
-                        ebx: 0xaaaa_5555,
-                        ecx: 0x5555_aaaa,
-                        edx: 0x5555_5555,
-                    },
+                kvm_bindings::kvm_cpuid_entry2 {
+                    function: 0x1,
+                    index: 0x1,
+                    flags: KvmCpuidFlags::SIGNIFICANT_INDEX.0,
+                    eax: 0xaaaa_aaaa,
+                    ebx: 0xaaaa_5555,
+                    ecx: 0x5555_aaaa,
+                    edx: 0x5555_5555,
+                    ..Default::default()
                 },
-            ),
-        ])))
+            ])
+            .unwrap(),
+        ))
     }
 
     fn build_expected_cpuid_modifiers() -> Vec<CpuidLeafModifier> {
@@ -189,24 +182,29 @@ mod tests {
         ]
     }
 
-    fn build_sample_msrs() -> BTreeMap<u32, u64> {
-        let mut map = BTreeMap::from([
+    fn build_sample_msrs() -> Msrs {
+        let entry = |index, data| kvm_msr_entry {
+            index,
+            data,
+            ..Default::default()
+        };
+        let mut entries = vec![
             // should be sorted in the result.
-            (0x1, 0xffff_ffff_ffff_ffff),
-            (0x5, 0xffff_ffff_0000_0000),
-            (0x3, 0x0000_0000_ffff_ffff),
-            (0x2, 0x0000_0000_0000_0000),
-        ]);
+            entry(0x1, 0xffff_ffff_ffff_ffff),
+            entry(0x5, 0xffff_ffff_0000_0000),
+            entry(0x3, 0x0000_0000_ffff_ffff),
+            entry(0x2, 0x0000_0000_0000_0000),
+        ];
         // should be excluded from the result.
         MSR_EXCLUSION_LIST
             .iter()
             .chain(MSR_EXCLUSION_LIST_AMD.iter())
             .for_each(|range| {
                 (range.base..(range.base + range.nmsrs)).for_each(|id| {
-                    map.insert(id, 0);
+                    entries.push(entry(id, 0));
                 })
             });
-        map
+        Msrs::from_entries(&entries).unwrap()
     }
 
     fn build_expected_msr_modifiers() -> Vec<RegisterModifier> {

src/vmm/src/arch/x86_64/msr.rs

@@ -423,21 +423,43 @@ pub fn create_boot_msr_entries() -> Vec<kvm_msr_entry> {
     ]
 }
 
+/// Insert or update an MSR entry in `msrs`, mirroring `BTreeMap::insert` semantics: when an entry
+/// with the given index already exists its `data` is overwritten, otherwise a new entry is pushed.
+///
+/// # Errors
+///
+/// Propagates [`vmm_sys_util::fam::Error`] from [`Msrs::push`] when the FAM struct is at capacity.
+pub fn msrs_insert(msrs: &mut Msrs, index: u32, data: u64) {
+    if let Some(entry) = msrs
+        .as_mut_slice()
+        .iter_mut()
+        .find(|entry| entry.index == index)
+    {
+        entry.data = data;
+    } else {
+        msrs.push(kvm_msr_entry {
+            index,
+            data,
+            ..Default::default()
+        })
+        .expect("MSRS entry count exceeded KVM_MAX_MSR_ENTRIES");
+    }
+}
+
 /// Configure Model Specific Registers (MSRs) required to boot Linux for a given x86_64 vCPU.
 ///
 /// # Arguments
 ///
 /// * `vcpu` - Structure for the VCPU that holds the VCPU's fd.
+/// * `msrs` - The MSRs to write to KVM.
 ///
 /// # Errors
 ///
 /// When:
-/// - Failed to create [`vmm_sys_util::fam::FamStructWrapper`] for MSRs.
 /// - [`kvm_ioctls::ioctls::vcpu::VcpuFd::set_msrs`] errors.
 /// - [`kvm_ioctls::ioctls::vcpu::VcpuFd::set_msrs`] fails to write all given MSRs entries.
-pub fn set_msrs(vcpu: &VcpuFd, msr_entries: &[kvm_msr_entry]) -> Result<(), MsrError> {
-    let msrs = Msrs::from_entries(msr_entries)?;
-    vcpu.set_msrs(&msrs)
+pub fn set_msrs(vcpu: &VcpuFd, msrs: &Msrs) -> Result<(), MsrError> {
+    vcpu.set_msrs(msrs)
         .map_err(MsrError::SetMsrs)
         .and_then(|msrs_written| {
             if msrs_written == msrs.as_fam_struct_ref().nmsrs as usize {
@@ -484,7 +506,8 @@ mod tests {
     fn test_setup_msrs() {
         let vcpu = create_vcpu();
         let msr_boot_entries = create_boot_msr_entries();
-        set_msrs(&vcpu, &msr_boot_entries).unwrap();
+        let msrs = Msrs::from_entries(&msr_boot_entries).unwrap();
+        set_msrs(&vcpu, &msrs).unwrap();
 
         // This test will check against the last MSR entry configured (the tenth one).
         // See create_msr_entries() for details.
@@ -512,12 +535,31 @@ mod tests {
         // Test `set_msrs()` with a valid MSR entry. It should succeed, as IA32_TSC MSR is listed
         // in supported MSRs as of now.
         let vcpu = create_vcpu();
-        let msr_entries = vec![kvm_msr_entry {
+        let msrs = Msrs::from_entries(&[kvm_msr_entry {
             index: MSR_IA32_TSC,
             data: 0,
             ..Default::default()
-        }];
-        set_msrs(&vcpu, &msr_entries).unwrap();
+        }])
+        .unwrap();
+        set_msrs(&vcpu, &msrs).unwrap();
+    }
+
+    #[test]
+    fn test_msrs_insert_overwrites_and_pushes() {
+        let mut msrs = Msrs::new(0).unwrap();
+
+        msrs_insert(&mut msrs, 0x10, 0xaa);
+        msrs_insert(&mut msrs, 0x20, 0xbb);
+        assert_eq!(msrs.as_slice().len(), 2);
+        assert_eq!(msrs.as_slice()[0].data, 0xaa);
+        assert_eq!(msrs.as_slice()[1].data, 0xbb);
+
+        msrs_insert(&mut msrs, 0x10, 0xcc);
+        assert_eq!(msrs.as_slice().len(), 2);
+        assert_eq!(msrs.as_slice()[0].index, 0x10);
+        assert_eq!(msrs.as_slice()[0].data, 0xcc);
+        assert_eq!(msrs.as_slice()[1].index, 0x20);
+        assert_eq!(msrs.as_slice()[1].data, 0xbb);
     }
 
     #[test]
@@ -526,12 +568,13 @@ mod tests {
         // listed in supported MSRs as of now. If hardware vendor adds this MSR index and KVM
         // supports this MSR, we need to change the index as needed.
         let vcpu = create_vcpu();
-        let msr_entries = vec![kvm_msr_entry {
+        let msrs = Msrs::from_entries(&[kvm_msr_entry {
             index: 2,
             ..Default::default()
-        }];
+        }])
+        .unwrap();
         assert_eq!(
-            set_msrs(&vcpu, &msr_entries).unwrap_err(),
+            set_msrs(&vcpu, &msrs).unwrap_err(),
             MsrError::SetMsrsIncomplete
         );
     }