From c0b6d85673560f9890dfd050f85f72e9bf636f0f Mon Sep 17 00:00:00 2001
From: Liu Jiang <gerry@linux.alibaba.com>
Date: Sat, 5 Jun 2021 17:08:00 +0800
Subject: [PATCH]

---
 src/ioctls/vcpu.rs | 49 +++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 46 insertions(+), 3 deletions(-)

diff --git a/src/ioctls/vcpu.rs b/src/ioctls/vcpu.rs
index f9e11426..3bb9d5bd 100644
--- a/src/ioctls/vcpu.rs
+++ b/src/ioctls/vcpu.rs
@@ -1417,6 +1417,8 @@ mod tests {
 
     use super::*;
     use ioctls::system::Kvm;
+    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+    use vmm_sys_util::fam::FamStruct;
     #[cfg(any(
         target_arch = "x86",
         target_arch = "x86_64",
@@ -1512,8 +1514,10 @@ mod tests {
         let kvm = Kvm::new().unwrap();
         if kvm.check_extension(Cap::ExtCpuid) {
             let vm = kvm.create_vm().unwrap();
-            let mut cpuid = kvm.get_supported_cpuid(KVM_MAX_CPUID_ENTRIES).unwrap();
+            //let mut cpuid = kvm.get_supported_cpuid(KVM_MAX_CPUID_ENTRIES).unwrap();
+            let mut cpuid = kvm.get_supported_cpuid(48).unwrap();
             let ncpuids = cpuid.as_slice().len();
+            assert_eq!(cpuid.as_fam_struct_ref().len(), ncpuids);
             assert!(ncpuids <= KVM_MAX_CPUID_ENTRIES);
             let vcpu = vm.create_vcpu(0).unwrap();
 
@@ -1521,7 +1525,7 @@ mod tests {
             {
                 let entries = cpuid.as_mut_slice();
                 for entry in entries.iter_mut() {
-                    if entry.function == 0 {
+                    if entry.function == 0 && entry.index == 0 && entry.flags == 0 {
                         // " KVMKVMKVM "
                         entry.ebx = 0x4b4d564b;
                         entry.ecx = 0x564b4d56;
@@ -1531,8 +1535,12 @@ mod tests {
             }
             vcpu.set_cpuid2(&cpuid).unwrap();
             let cpuid_0 = vcpu.get_cpuid2(ncpuids).unwrap();
+            assert_eq!(
+                cpuid.as_fam_struct_ref().len(),
+                cpuid_0.as_fam_struct_ref().len()
+            );
             for entry in cpuid_0.as_slice() {
-                if entry.function == 0 {
+                if entry.function == 0 && entry.index == 0 && entry.flags == 0 {
                     assert_eq!(entry.ebx, 0x4b4d564b);
                     assert_eq!(entry.ecx, 0x564b4d56);
                     assert_eq!(entry.edx, 0x4d);
@@ -1561,6 +1569,41 @@ mod tests {
         }
     }
 
+    #[cfg(target_arch = "x86_64")]
+    #[test]
+    fn test_set_cpuid_one_entry() {
+        let kvm = Kvm::new().unwrap();
+        if kvm.check_extension(Cap::ExtCpuid) {
+            let vm = kvm.create_vm().unwrap();
+            let mut cpuid = CpuId::new(1).unwrap();
+            let ncpuids = cpuid.as_slice().len();
+            assert_eq!(ncpuids, 1);
+            assert_eq!(cpuid.as_fam_struct_ref().len(), ncpuids);
+            let vcpu = vm.create_vcpu(0).unwrap();
+
+            // Setting Manufacturer ID
+            {
+                let entries = cpuid.as_mut_slice();
+                // " KVMKVMKVM "
+                entries[0].ebx = 0x4b4d564b;
+                entries[0].ecx = 0x564b4d56;
+                entries[0].edx = 0x4d;
+            }
+            vcpu.set_cpuid2(&cpuid).unwrap();
+            let cpuid_0 = vcpu.get_cpuid2(ncpuids).unwrap();
+            assert_eq!(
+                cpuid.as_fam_struct_ref().len(),
+                cpuid_0.as_fam_struct_ref().len()
+            );
+            assert_eq!(cpuid, cpuid_0);
+            for entry in cpuid_0.as_slice() {
+                assert_eq!(entry.ebx, 0x4b4d564b);
+                assert_eq!(entry.ecx, 0x564b4d56);
+                assert_eq!(entry.edx, 0x4d);
+            }
+        }
+    }
+
     #[cfg(target_arch = "x86_64")]
     #[allow(non_snake_case)]
     #[test]