hypervisor: x86_64: represent collection of MSRs as a map

Replace the Vec<Register> with a simpler Map<u32, u64>.

This changes the snapshot JSON schema - the "msrs" field will now be a
JSON dictionary (object) with the MSR index as key rather than a list of
objects.

Change-Id: I71a26dec6bcdaae0b66d497818a65b8c143eea8b
Reviewed-on: https://chromium-review.googlesource.com/c/crosvm/crosvm/+/5436912
Reviewed-by: Frederick Mayle <fmayle@google.com>
Commit-Queue: Daniel Verkamp <dverkamp@chromium.org>

8 files changed, 81 insertions, 127 deletions

diff --git a/devices/tests/irqchip/userspace.rs b/devices/tests/irqchip/userspace.rs
index d799ddd3d..d39ca0ac3 100644
--- a/devices/tests/irqchip/userspace.rs
+++ b/devices/tests/irqchip/userspace.rs
@@ -4,6 +4,7 @@
 
 #![cfg(target_arch = "x86_64")]
 
+use std::collections::BTreeMap;
 use std::sync::Arc;
 use std::thread;
 use std::time::Duration;
@@ -761,7 +762,7 @@ impl VcpuX86_64 for FakeVcpu {
     fn get_msr(&self, _msr_index: u32) -> Result<u64> {
         unimplemented!()
     }
-    fn get_all_msrs(&self) -> Result<Vec<Register>> {
+    fn get_all_msrs(&self) -> Result<BTreeMap<u32, u64>> {
         unimplemented!()
     }
     fn set_msr(&self, _msr_index: u32, _value: u64) -> Result<()> {
diff --git a/hypervisor/src/haxm/vcpu.rs b/hypervisor/src/haxm/vcpu.rs
index 07785d294..63752d1c8 100644
--- a/hypervisor/src/haxm/vcpu.rs
+++ b/hypervisor/src/haxm/vcpu.rs
@@ -5,6 +5,7 @@
 use core::ffi::c_void;
 use std::arch::x86_64::CpuidResult;
 use std::cmp::min;
+use std::collections::BTreeMap;
 use std::intrinsics::copy_nonoverlapping;
 use std::mem::size_of;
 
@@ -496,7 +497,7 @@ impl VcpuX86_64 for HaxmVcpu {
         Ok(msr_data.entries[0].value)
     }
 
-    fn get_all_msrs(&self) -> Result<Vec<Register>> {
+    fn get_all_msrs(&self) -> Result<BTreeMap<u32, u64>> {
         Err(Error::new(EOPNOTSUPP))
     }
 
diff --git a/hypervisor/src/kvm/x86_64.rs b/hypervisor/src/kvm/x86_64.rs
index 404b80966..2b37380db 100644
--- a/hypervisor/src/kvm/x86_64.rs
+++ b/hypervisor/src/kvm/x86_64.rs
@@ -3,6 +3,7 @@
 // found in the LICENSE file.
 
 use std::arch::x86_64::CpuidResult;
+use std::collections::BTreeMap;
 
 use base::errno_result;
 use base::error;
@@ -941,7 +942,7 @@ impl VcpuX86_64 for KvmVcpu {
         Ok(value)
     }
 
-    fn get_all_msrs(&self) -> Result<Vec<Register>> {
+    fn get_all_msrs(&self) -> Result<BTreeMap<u32, u64>> {
         let msr_index_list = self.kvm.get_msr_index_list()?;
         let mut kvm_msrs = vec_with_array_field::<kvm_msrs, kvm_msr_entry>(msr_index_list.len());
         kvm_msrs[0].nmsrs = msr_index_list.len() as u32;
@@ -980,15 +981,13 @@ impl VcpuX86_64 for KvmVcpu {
         // SAFETY:
         // Safe because we trust the kernel to return the correct array length on success.
         let msrs = unsafe {
-            kvm_msrs[0]
-                .entries
-                .as_slice(count)
-                .iter()
-                .map(|kvm_msr| Register {
-                    id: kvm_msr.index,
-                    value: kvm_msr.data,
-                })
-                .collect()
+            BTreeMap::from_iter(
+                kvm_msrs[0]
+                    .entries
+                    .as_slice(count)
+                    .iter()
+                    .map(|kvm_msr| (kvm_msr.index, kvm_msr.data)),
+            )
         };
 
         Ok(msrs)
diff --git a/hypervisor/src/whpx/vcpu.rs b/hypervisor/src/whpx/vcpu.rs
index ae7d0f812..5c9afc724 100644
--- a/hypervisor/src/whpx/vcpu.rs
+++ b/hypervisor/src/whpx/vcpu.rs
@@ -4,6 +4,7 @@
 
 use core::ffi::c_void;
 use std::arch::x86_64::CpuidResult;
+use std::collections::BTreeMap;
 use std::convert::TryInto;
 use std::mem::size_of;
 use std::sync::Arc;
@@ -1149,7 +1150,7 @@ impl VcpuX86_64 for WhpxVcpu {
         Ok(value)
     }
 
-    fn get_all_msrs(&self) -> Result<Vec<Register>> {
+    fn get_all_msrs(&self) -> Result<BTreeMap<u32, u64>> {
         // Note that some members of VALID_MSRS cannot be fetched from WHPX with
         // WHvGetVirtualProcessorRegisters per the HTLFS, so we enumerate all of
         // permitted MSRs here.
@@ -1179,12 +1180,9 @@ impl VcpuX86_64 for WhpxVcpu {
             .iter()
             .map(|msr_index| {
                 let value = self.get_msr(*msr_index)?;
-                Ok(Register {
-                    id: *msr_index,
-                    value,
-                })
+                Ok((*msr_index, value))
             })
-            .collect::<Result<Vec<Register>>>()?;
+            .collect::<Result<BTreeMap<u32, u64>>>()?;
 
         Ok(registers)
     }
@@ -1588,7 +1586,7 @@ mod tests {
 
         // Our MSR buffer is init'ed to zeros in the registers. The APIC base will be non-zero, so
         // by asserting that we know the MSR fetch actually did get us data.
-        let apic_base = all_msrs.iter().find(|reg| reg.id == MSR_APIC_BASE).unwrap();
-        assert_ne!(apic_base.value, 0);
+        let apic_base = all_msrs.get(&MSR_APIC_BASE).unwrap();
+        assert_ne!(*apic_base, 0);
     }
 }
diff --git a/hypervisor/src/x86_64.rs b/hypervisor/src/x86_64.rs
index 64714e8cc..eec0e2de5 100644
--- a/hypervisor/src/x86_64.rs
+++ b/hypervisor/src/x86_64.rs
@@ -6,7 +6,7 @@ use std::arch::x86_64::CpuidResult;
 #[cfg(any(unix, feature = "haxm", feature = "whpx"))]
 use std::arch::x86_64::__cpuid;
 use std::arch::x86_64::_rdtsc;
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::collections::HashSet;
 
 use anyhow::Context;
@@ -133,7 +133,7 @@ pub trait VcpuX86_64: Vcpu {
     fn get_msr(&self, msr_index: u32) -> Result<u64>;
 
     /// Gets the model-specific registers. Returns all the MSRs for the VCPU.
-    fn get_all_msrs(&self) -> Result<Vec<Register>>;
+    fn get_all_msrs(&self) -> Result<BTreeMap<u32, u64>>;
 
     /// Sets a single model-specific register's value.
     fn set_msr(&self, msr_index: u32, value: u64) -> Result<()>;
@@ -260,17 +260,12 @@ pub trait VcpuX86_64: Vcpu {
         self.set_debugregs(&snapshot.debug_regs)?;
         self.set_xcrs(&snapshot.xcrs)?;
 
-        let mut msrs = HashMap::new();
-        for reg in self.get_all_msrs()? {
-            msrs.insert(reg.id, reg.value);
-        }
-
-        for &msr in snapshot.msrs.iter() {
-            if Some(&msr.value) == msrs.get(&msr.id) {
+        for (msr_index, value) in snapshot.msrs.iter() {
+            if self.get_msr(*msr_index) == Ok(*value) {
                 continue; // no need to set MSR since the values are the same.
             }
-            if let Err(e) = self.set_msr(msr.id, msr.value) {
-                if msr_allowlist.contains(&msr.id) {
+            if let Err(e) = self.set_msr(*msr_index, *value) {
+                if msr_allowlist.contains(msr_index) {
                     warn!(
                         "Failed to set MSR. MSR might not be supported in this kernel. Err: {}",
                         e
@@ -298,7 +293,7 @@ pub struct VcpuSnapshot {
     sregs: Sregs,
     debug_regs: DebugRegs,
     xcrs: Vec<Register>,
-    msrs: Vec<Register>,
+    msrs: BTreeMap<u32, u64>,
     xsave: Xsave,
     hypervisor_data: serde_json::Value,
     tsc_offset: u64,
@@ -337,7 +332,7 @@ pub struct VcpuInitX86_64 {
     pub fpu: Fpu,
 
     /// Machine-specific registers.
-    pub msrs: Vec<Register>,
+    pub msrs: BTreeMap<u32, u64>,
 }
 
 /// Hold the CPU feature configurations that are needed to setup a vCPU.
diff --git a/x86_64/src/lib.rs b/x86_64/src/lib.rs
index 351400671..6202e2400 100644
--- a/x86_64/src/lib.rs
+++ b/x86_64/src/lib.rs
@@ -1004,8 +1004,8 @@ impl arch::LinuxArch for X8664arch {
         let pci_start = read_pci_mmio_before_32bit().start;
 
         let mut vcpu_init = vec![VcpuInitX86_64::default(); vcpu_count];
+        let mut msrs = BTreeMap::new();
 
-        let mut msrs;
         match components.vm_image {
             VmImage::Bios(ref mut bios) => {
                 // Allow a bios to hardcode CMDLINE_OFFSET and read the kernel command line from it.
@@ -1016,7 +1016,7 @@ impl arch::LinuxArch for X8664arch {
                 )
                 .map_err(Error::LoadCmdline)?;
                 Self::load_bios(&mem, bios)?;
-                msrs = regs::default_msrs();
+                regs::set_default_msrs(&mut msrs);
                 // The default values for `Regs` and `Sregs` already set up the reset vector.
             }
             VmImage::Kernel(ref mut kernel_image) => {
@@ -1039,8 +1039,8 @@ impl arch::LinuxArch for X8664arch {
                 vcpu_init[0].regs.rsp = BOOT_STACK_POINTER;
                 vcpu_init[0].regs.rsi = ZERO_PAGE_OFFSET;
 
-                msrs = regs::long_mode_msrs();
-                msrs.append(&mut regs::mtrr_msrs(&vm, pci_start));
+                regs::set_long_mode_msrs(&mut msrs);
+                regs::set_mtrr_msrs(&mut msrs, &vm, pci_start);
 
                 // Set up long mode and enable paging.
                 regs::configure_segments_and_sregs(&mem, &mut vcpu_init[0].sregs)
@@ -1111,24 +1111,24 @@ impl arch::LinuxArch for X8664arch {
 
         let vcpu_supported_var_mtrrs = regs::vcpu_supported_variable_mtrrs(vcpu);
         let num_var_mtrrs = regs::count_variable_mtrrs(&vcpu_init.msrs);
-        let msrs = if num_var_mtrrs > vcpu_supported_var_mtrrs {
+        let skip_mtrr_msrs = if num_var_mtrrs > vcpu_supported_var_mtrrs {
             warn!(
                 "Too many variable MTRR entries ({} required, {} supported),
                 please check pci_start addr, guest with pass through device may be very slow",
                 num_var_mtrrs, vcpu_supported_var_mtrrs,
             );
             // Filter out the MTRR entries from the MSR list.
-            vcpu_init
-                .msrs
-                .into_iter()
-                .filter(|&msr| !regs::is_mtrr_msr(msr.id))
-                .collect()
+            true
         } else {
-            vcpu_init.msrs
+            false
         };
 
-        for msr in msrs {
-            vcpu.set_msr(msr.id, msr.value).map_err(Error::SetupMsrs)?;
+        for (msr_index, value) in vcpu_init.msrs.into_iter() {
+            if skip_mtrr_msrs && regs::is_mtrr_msr(msr_index) {
+                continue;
+            }
+
+            vcpu.set_msr(msr_index, value).map_err(Error::SetupMsrs)?;
         }
 
         interrupts::set_lint(vcpu_id, irq_chip).map_err(Error::SetLint)?;
diff --git a/x86_64/src/regs.rs b/x86_64/src/regs.rs
index ac4c87d97..cbdd189ec 100644
--- a/x86_64/src/regs.rs
+++ b/x86_64/src/regs.rs
@@ -2,11 +2,11 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+use std::collections::BTreeMap;
 use std::mem;
 use std::result;
 
 use base::warn;
-use hypervisor::Register;
 use hypervisor::Sregs;
 use hypervisor::VcpuX86_64;
 use hypervisor::Vm;
@@ -116,105 +116,64 @@ pub fn is_mtrr_msr(id: u32) -> bool {
 }
 
 /// Returns the count of variable MTRR entries specified by the list of `msrs`.
-pub fn count_variable_mtrrs(msrs: &[Register]) -> usize {
+pub fn count_variable_mtrrs(msrs: &BTreeMap<u32, u64>) -> usize {
     // Each variable MTRR takes up two MSRs (base + mask), so divide by 2. This will also count the
     // MTRRdefType entry, but that is only one extra and the division truncates, so it won't affect
     // the final count.
-    msrs.iter().filter(|msr| is_mtrr_msr(msr.id)).count() / 2
+    msrs.keys().filter(|&msr| is_mtrr_msr(*msr)).count() / 2
 }
 
 /// Returns a set of MSRs containing the MTRR configuration.
-pub fn mtrr_msrs(vm: &dyn Vm, pci_start: u64) -> Vec<Register> {
+pub fn set_mtrr_msrs(msrs: &mut BTreeMap<u32, u64>, vm: &dyn Vm, pci_start: u64) {
     // Set pci_start .. 4G as UC
     // all others are set to default WB
     let pci_len = (1 << 32) - pci_start;
     let vecs = get_mtrr_pairs(pci_start, pci_len);
 
-    let mut entries = Vec::new();
-
     let phys_mask: u64 = (1 << vm.get_guest_phys_addr_bits()) - 1;
     for (idx, (base, len)) in vecs.iter().enumerate() {
         let reg_idx = idx as u32 * 2;
-        entries.push(Register {
-            id: MTRR_PHYS_BASE_MSR + reg_idx,
-            value: base | MTRR_MEMTYPE_UC as u64,
-        });
+        msrs.insert(MTRR_PHYS_BASE_MSR + reg_idx, base | MTRR_MEMTYPE_UC as u64);
         let mask: u64 = len.wrapping_neg() & phys_mask | MTRR_VAR_VALID;
-        entries.push(Register {
-            id: MTRR_PHYS_MASK_MSR + reg_idx,
-            value: mask,
-        });
+        msrs.insert(MTRR_PHYS_MASK_MSR + reg_idx, mask);
     }
     // Disable fixed MTRRs and enable variable MTRRs, set default type as WB
-    entries.push(Register {
-        id: crate::msr_index::MSR_MTRRdefType,
-        value: MTRR_ENABLE | MTRR_MEMTYPE_WB as u64,
-    });
-    entries
+    msrs.insert(
+        crate::msr_index::MSR_MTRRdefType,
+        MTRR_ENABLE | MTRR_MEMTYPE_WB as u64,
+    );
 }
 
 /// Returns the default value of MSRs at reset.
 ///
 /// Currently only sets IA32_TSC to 0.
-pub fn default_msrs() -> Vec<Register> {
-    vec![
-        Register {
-            id: crate::msr_index::MSR_IA32_TSC,
-            value: 0x0,
-        },
-        Register {
-            id: crate::msr_index::MSR_IA32_MISC_ENABLE,
-            value: crate::msr_index::MSR_IA32_MISC_ENABLE_FAST_STRING as u64,
-        },
-    ]
+pub fn set_default_msrs(msrs: &mut BTreeMap<u32, u64>) {
+    msrs.insert(crate::msr_index::MSR_IA32_TSC, 0x0);
+    msrs.insert(
+        crate::msr_index::MSR_IA32_MISC_ENABLE,
+        crate::msr_index::MSR_IA32_MISC_ENABLE_FAST_STRING as u64,
+    );
 }
 
 /// Configure Model specific registers for long (64-bit) mode.
-pub fn long_mode_msrs() -> Vec<Register> {
-    vec![
-        Register {
-            id: crate::msr_index::MSR_IA32_SYSENTER_CS,
-            value: 0x0,
-        },
-        Register {
-            id: crate::msr_index::MSR_IA32_SYSENTER_ESP,
-            value: 0x0,
-        },
-        Register {
-            id: crate::msr_index::MSR_IA32_SYSENTER_EIP,
-            value: 0x0,
-        },
-        // x86_64 specific msrs, we only run on x86_64 not x86
-        Register {
-            id: crate::msr_index::MSR_STAR,
-            value: 0x0,
-        },
-        Register {
-            id: crate::msr_index::MSR_CSTAR,
-            value: 0x0,
-        },
-        Register {
-            id: crate::msr_index::MSR_KERNEL_GS_BASE,
-            value: 0x0,
-        },
-        Register {
-            id: crate::msr_index::MSR_SYSCALL_MASK,
-            value: 0x0,
-        },
-        Register {
-            id: crate::msr_index::MSR_LSTAR,
-            value: 0x0,
-        },
-        // end of x86_64 specific code
-        Register {
-            id: crate::msr_index::MSR_IA32_TSC,
-            value: 0x0,
-        },
-        Register {
-            id: crate::msr_index::MSR_IA32_MISC_ENABLE,
-            value: crate::msr_index::MSR_IA32_MISC_ENABLE_FAST_STRING as u64,
-        },
-    ]
+pub fn set_long_mode_msrs(msrs: &mut BTreeMap<u32, u64>) {
+    msrs.insert(crate::msr_index::MSR_IA32_SYSENTER_CS, 0x0);
+    msrs.insert(crate::msr_index::MSR_IA32_SYSENTER_ESP, 0x0);
+    msrs.insert(crate::msr_index::MSR_IA32_SYSENTER_EIP, 0x0);
+
+    // x86_64 specific msrs, we only run on x86_64 not x86
+    msrs.insert(crate::msr_index::MSR_STAR, 0x0);
+    msrs.insert(crate::msr_index::MSR_CSTAR, 0x0);
+    msrs.insert(crate::msr_index::MSR_KERNEL_GS_BASE, 0x0);
+    msrs.insert(crate::msr_index::MSR_SYSCALL_MASK, 0x0);
+    msrs.insert(crate::msr_index::MSR_LSTAR, 0x0);
+    // end of x86_64 specific code
+
+    msrs.insert(crate::msr_index::MSR_IA32_TSC, 0x0);
+    msrs.insert(
+        crate::msr_index::MSR_IA32_MISC_ENABLE,
+        crate::msr_index::MSR_IA32_MISC_ENABLE_FAST_STRING as u64,
+    );
 }
 
 const X86_CR0_PE: u64 = 0x1;
diff --git a/x86_64/tests/integration/main.rs b/x86_64/tests/integration/main.rs
index f8a3d1ac0..9c617b27f 100644
--- a/x86_64/tests/integration/main.rs
+++ b/x86_64/tests/integration/main.rs
@@ -45,8 +45,8 @@ use x86_64::mptable;
 use x86_64::read_pci_mmio_before_32bit;
 use x86_64::read_pcie_cfg_mmio;
 use x86_64::regs::configure_segments_and_sregs;
-use x86_64::regs::long_mode_msrs;
-use x86_64::regs::mtrr_msrs;
+use x86_64::regs::set_long_mode_msrs;
+use x86_64::regs::set_mtrr_msrs;
 use x86_64::regs::setup_page_tables;
 use x86_64::smbios;
 use x86_64::X8664arch;
@@ -281,10 +281,11 @@ where
                 setup_cpuid(&hyp, &irq_chip, &vcpu, 0, 1, cpu_config).unwrap();
             }
 
-            let mut msrs = long_mode_msrs();
-            msrs.append(&mut mtrr_msrs(&vm, read_pci_mmio_before_32bit().start));
-            for msr in msrs {
-                vcpu.set_msr(msr.id, msr.value).unwrap();
+            let mut msrs = BTreeMap::new();
+            set_long_mode_msrs(&mut msrs);
+            set_mtrr_msrs(&mut msrs, &vm, read_pci_mmio_before_32bit().start);
+            for (msr_index, value) in msrs {
+                vcpu.set_msr(msr_index, value).unwrap();
             }
 
             let mut vcpu_regs = Regs {