Use BTreeMap/Vec instead of HashMap/HashSet when dealing with MSRs

As the previous commit shows, the order of restoration of MSRs can be
important. Thus, we should not effectively randomly shuffle them
whenever we construct a VM. With this commit, we preserve the order in
which KVM tells us about MSRs in `KVM_GET_MSR_INDEX_LIST`, under the
assumption that if any dependencies exist, KVM will account for them as
part of the above IOCTL.

Signed-off-by: Patrick Roy <[email protected]>

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

@@ -1,7 +1,7 @@
 // Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 
 use vmm::arch::x86_64::msr::MsrRange;
 use vmm::arch_gen::x86::msr_index::*;
@@ -44,7 +44,7 @@ fn cpuid_to_modifiers(cpuid: &Cpuid) -> Vec<CpuidLeafModifier> {
         .collect()
 }
 
-fn msrs_to_modifier(msrs: &HashMap<u32, u64>) -> Vec<RegisterModifier> {
+fn msrs_to_modifier(msrs: &BTreeMap<u32, u64>) -> Vec<RegisterModifier> {
     let mut msrs: Vec<RegisterModifier> = msrs
         .iter()
         .map(|(index, value)| msr_modifier!(*index, *value))
@@ -189,8 +189,8 @@ mod tests {
         ]
     }
 
-    fn build_sample_msrs() -> HashMap<u32, u64> {
-        let mut map = HashMap::from([
+    fn build_sample_msrs() -> BTreeMap<u32, u64> {
+        let mut map = BTreeMap::from([
             // should be sorted in the result.
             (0x1, 0xffff_ffff_ffff_ffff),
             (0x5, 0xffff_ffff_0000_0000),

src/vmm/src/cpu_config/x86_64/cpuid/common.rs

@@ -53,7 +53,7 @@ pub fn get_vendor_id_from_host() -> Result<[u8; 12], GetCpuidError> {
 }
 
 /// Returns MSRs to be saved based on CPUID features that are enabled.
-pub(crate) fn msrs_to_save_by_cpuid(cpuid: &kvm_bindings::CpuId) -> std::collections::HashSet<u32> {
+pub(crate) fn msrs_to_save_by_cpuid(cpuid: &kvm_bindings::CpuId) -> Vec<u32> {
     /// Memory Protection Extensions
     const MPX_BITINDEX: u32 = 14;
 
@@ -81,7 +81,7 @@ pub(crate) fn msrs_to_save_by_cpuid(cpuid: &kvm_bindings::CpuId) -> std::collect
         }};
     }
 
-    let mut msrs = std::collections::HashSet::new();
+    let mut msrs = Vec::new();
 
     // Macro used for easy definition of CPUID-MSR dependencies.
     macro_rules! cpuid_msr_dep {

src/vmm/src/cpu_config/x86_64/mod.rs

@@ -10,7 +10,7 @@ pub mod static_cpu_templates;
 /// Module with test utils for custom CPU templates
 pub mod test_utils;
 
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 
 use self::custom_cpu_template::CpuidRegister;
 use super::templates::CustomCpuTemplate;
@@ -37,7 +37,7 @@ pub struct CpuConfiguration {
     /// Register values as a key pair for model specific registers
     /// Key: MSR address
     /// Value: MSR value
-    pub msrs: HashMap<u32, u64>,
+    pub msrs: BTreeMap<u32, u64>,
 }
 
 impl CpuConfiguration {
@@ -187,14 +187,14 @@ mod tests {
     fn supported_cpu_config() -> CpuConfiguration {
         CpuConfiguration {
             cpuid: build_supported_cpuid(),
-            msrs: HashMap::from([(0x8000, 0b1000), (0x9999, 0b1010)]),
+            msrs: BTreeMap::from([(0x8000, 0b1000), (0x9999, 0b1010)]),
         }
     }
 
     fn unsupported_cpu_config() -> CpuConfiguration {
         CpuConfiguration {
             cpuid: build_supported_cpuid(),
-            msrs: HashMap::from([(0x8000, 0b1000), (0x8001, 0b1010)]),
+            msrs: BTreeMap::from([(0x8000, 0b1000), (0x8001, 0b1010)]),
         }
     }
 

src/vmm/src/vstate/vcpu/mod.rs

@@ -694,6 +694,8 @@ pub enum VcpuEmulation {
 pub mod tests {
     #![allow(clippy::undocumented_unsafe_blocks)]
 
+    #[cfg(target_arch = "x86_64")]
+    use std::collections::BTreeMap;
     use std::sync::{Arc, Barrier, Mutex};
 
     use linux_loader::loader::KernelLoader;
@@ -919,7 +921,7 @@ pub mod tests {
                         smt: false,
                         cpu_config: CpuConfiguration {
                             cpuid: Cpuid::try_from(_vm.supported_cpuid().clone()).unwrap(),
-                            msrs: std::collections::HashMap::new(),
+                            msrs: BTreeMap::new(),
                         },
                     },
                 )

src/vmm/src/vstate/vcpu/x86_64.rs

@@ -5,7 +5,7 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the THIRD-PARTY file.
 
-use std::collections::{HashMap, HashSet};
+use std::collections::BTreeMap;
 use std::fmt::Debug;
 
 use kvm_bindings::{
@@ -143,7 +143,9 @@ pub struct KvmVcpu {
     pub fd: VcpuFd,
     /// Vcpu peripherals, such as buses
     pub(super) peripherals: Peripherals,
-    msrs_to_save: HashSet<u32>,
+    /// The list of MSRs to include in a VM snapshot, in the same order as KVM returned them
+    /// from KVM_GET_MSR_INDEX_LIST
+    msrs_to_save: Vec<u32>,
 }
 
 /// Vcpu peripherals
@@ -172,7 +174,7 @@ impl KvmVcpu {
             index,
             fd: kvm_vcpu,
             peripherals: Default::default(),
-            msrs_to_save: vm.msrs_to_save().as_slice().iter().copied().collect(),
+            msrs_to_save: vm.msrs_to_save().as_slice().to_vec(),
         })
     }
 
@@ -455,8 +457,8 @@ impl KvmVcpu {
     pub fn get_msrs(
         &self,
         msr_index_iter: impl ExactSizeIterator<Item = u32>,
-    ) -> Result<HashMap<u32, u64>, KvmVcpuError> {
-        let mut msrs: HashMap<u32, u64> = HashMap::new();
+    ) -> Result<BTreeMap<u32, u64>, KvmVcpuError> {
+        let mut msrs = BTreeMap::new();
         self.get_msr_chunks(msr_index_iter)?
             .iter()
             .for_each(|msr_chunk| {
@@ -716,7 +718,7 @@ mod tests {
     use std::os::unix::io::AsRawFd;
 
     use kvm_bindings::kvm_msr_entry;
-    use kvm_ioctls::Cap;
+    use kvm_ioctls::{Cap, Kvm};
 
     use super::*;
     use crate::arch::x86_64::cpu_model::CpuModel;
@@ -901,7 +903,7 @@ mod tests {
             smt: false,
             cpu_config: CpuConfiguration {
                 cpuid: Cpuid::try_from(vm.supported_cpuid().clone()).unwrap(),
-                msrs: HashMap::new(),
+                msrs: BTreeMap::new(),
             },
         };
         vcpu.configure(&vm_mem, GuestAddress(0), &vcpu_config)
@@ -963,7 +965,7 @@ mod tests {
             smt: false,
             cpu_config: CpuConfiguration {
                 cpuid: Cpuid::try_from(vm.supported_cpuid().clone()).unwrap(),
-                msrs: HashMap::new(),
+                msrs: BTreeMap::new(),
             },
         };
         vcpu.configure(&vm_mem, GuestAddress(0), &vcpu_config)
@@ -1163,4 +1165,34 @@ mod tests {
             &[(MSR_IA32_TSC_DEADLINE, 1), (MSR_IA32_TSC, 2)],
         );
     }
+
+    #[test]
+    fn test_get_msr_chunks_preserved_order() {
+        // Regression test for #4666
+
+        let kvm = Kvm::new().unwrap();
+        let vm = Vm::new(Vec::new()).unwrap();
+        let vcpu = KvmVcpu::new(0, &vm).unwrap();
+
+        // The list of supported MSR indices, in the order they were returned by KVM
+        let msrs_to_save = crate::arch::x86_64::msr::get_msrs_to_save(&kvm).unwrap();
+        // The MSRs after processing. The order should be identical to the one returned by KVM, with
+        // the exception of deferred MSRs, which should be moved to the end (but show up in the same
+        // order as they are listed in [`DEFERRED_MSRS`].
+        let msr_chunks = vcpu
+            .get_msr_chunks(vcpu.msrs_to_save.iter().copied())
+            .unwrap();
+
+        msr_chunks
+            .iter()
+            .flat_map(|chunk| chunk.as_slice().iter())
+            .zip(
+                msrs_to_save
+                    .as_slice()
+                    .iter()
+                    .filter(|&idx| !DEFERRED_MSRS.contains(idx))
+                    .chain(DEFERRED_MSRS.iter()),
+            )
+            .for_each(|(left, &right)| assert_eq!(left.index, right));
+    }
 }