From f84ddedb1ac78b14df8fc8e1ecb23f30d75e8473 Mon Sep 17 00:00:00 2001
From: Wei Liu <liuwe@microsoft.com>
Date: Wed, 20 Jul 2022 11:42:24 +0000
Subject: [PATCH] hypervisor, vmm: introduce trait functions for aarch64 PMU

The original code uses kvm_device_attr directly outside of the
hyeprvisor crate. That leaks hypervisor details.

No functional change intended.

Signed-off-by: Wei Liu <liuwe@microsoft.com>
---
 hypervisor/src/cpu.rs     | 16 ++++++++++++++++
 hypervisor/src/kvm/mod.rs | 29 +++++++++++++++++++++++++++++
 vmm/src/cpu.rs            | 34 ++++------------------------------
 3 files changed, 49 insertions(+), 30 deletions(-)

diff --git a/hypervisor/src/cpu.rs b/hypervisor/src/cpu.rs
index eb2c24460..ad540eda8 100644
--- a/hypervisor/src/cpu.rs
+++ b/hypervisor/src/cpu.rs
@@ -236,6 +236,12 @@ pub enum HypervisorCpuError {
     #[cfg(feature = "tdx")]
     #[error("Unknown TDX VM call")]
     UnknownTdxVmCall,
+    #[cfg(target_arch = "aarch64")]
+    ///
+    /// Failed to intialize PMU
+    ///
+    #[error("Failed to initialize PMU")]
+    InitializePmu,
 }
 
 #[derive(Debug)]
@@ -400,6 +406,16 @@ pub trait Vcpu: Send + Sync {
     #[cfg(target_arch = "aarch64")]
     fn setup_regs(&self, cpu_id: u8, boot_ip: u64, fdt_start: u64) -> Result<()>;
     ///
+    /// Check if the CPU supports PMU
+    ///
+    #[cfg(target_arch = "aarch64")]
+    fn has_pmu_support(&self) -> bool;
+    ///
+    /// Initialize PMU
+    ///
+    #[cfg(target_arch = "aarch64")]
+    fn init_pmu(&self, irq: u32) -> Result<()>;
+    ///
     /// Retrieve the vCPU state.
     /// This function is necessary to snapshot the VM
     ///
diff --git a/hypervisor/src/kvm/mod.rs b/hypervisor/src/kvm/mod.rs
index 6e311a1c2..aa04251d5 100644
--- a/hypervisor/src/kvm/mod.rs
+++ b/hypervisor/src/kvm/mod.rs
@@ -2063,6 +2063,35 @@ impl cpu::Vcpu for KvmVcpu {
         ]
         .to_vec()
     }
+    #[cfg(target_arch = "aarch64")]
+    fn has_pmu_support(&self) -> bool {
+        let cpu_attr = kvm_bindings::kvm_device_attr {
+            group: kvm_bindings::KVM_ARM_VCPU_PMU_V3_CTRL,
+            attr: u64::from(kvm_bindings::KVM_ARM_VCPU_PMU_V3_INIT),
+            addr: 0x0,
+            flags: 0,
+        };
+        self.has_vcpu_attr(&cpu_attr).is_ok()
+    }
+    #[cfg(target_arch = "aarch64")]
+    fn init_pmu(&self, irq: u32) -> cpu::Result<()> {
+        let cpu_attr = kvm_bindings::kvm_device_attr {
+            group: kvm_bindings::KVM_ARM_VCPU_PMU_V3_CTRL,
+            attr: u64::from(kvm_bindings::KVM_ARM_VCPU_PMU_V3_INIT),
+            addr: 0x0,
+            flags: 0,
+        };
+        let cpu_attr_irq = kvm_bindings::kvm_device_attr {
+            group: kvm_bindings::KVM_ARM_VCPU_PMU_V3_CTRL,
+            attr: u64::from(kvm_bindings::KVM_ARM_VCPU_PMU_V3_IRQ),
+            addr: &irq as *const u32 as u64,
+            flags: 0,
+        };
+        self.set_vcpu_attr(&cpu_attr_irq)
+            .map_err(|_| cpu::HypervisorCpuError::InitializePmu)?;
+        self.set_vcpu_attr(&cpu_attr)
+            .map_err(|_| cpu::HypervisorCpuError::InitializePmu)
+    }
 }
 
 impl KvmVcpu {
diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index bd3f87d12..2865d9c46 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -787,41 +787,15 @@ impl CpuManager {
 
     #[cfg(target_arch = "aarch64")]
     pub fn init_pmu(&self, irq: u32) -> Result<bool> {
-        let cpu_attr = kvm_bindings::kvm_device_attr {
-            group: kvm_bindings::KVM_ARM_VCPU_PMU_V3_CTRL,
-            attr: u64::from(kvm_bindings::KVM_ARM_VCPU_PMU_V3_INIT),
-            addr: 0x0,
-            flags: 0,
-        };
-
         for cpu in self.vcpus.iter() {
-            let tmp = irq;
-            let cpu_attr_irq = kvm_bindings::kvm_device_attr {
-                group: kvm_bindings::KVM_ARM_VCPU_PMU_V3_CTRL,
-                attr: u64::from(kvm_bindings::KVM_ARM_VCPU_PMU_V3_IRQ),
-                addr: &tmp as *const u32 as u64,
-                flags: 0,
-            };
-
+            let cpu = cpu.lock().unwrap();
             // Check if PMU attr is available, if not, log the information.
-            if cpu.lock().unwrap().vcpu.has_vcpu_attr(&cpu_attr).is_ok() {
-                // Set irq for PMU
-                cpu.lock()
-                    .unwrap()
-                    .vcpu
-                    .set_vcpu_attr(&cpu_attr_irq)
-                    .map_err(Error::InitPmu)?;
-
-                // Init PMU
-                cpu.lock()
-                    .unwrap()
-                    .vcpu
-                    .set_vcpu_attr(&cpu_attr)
-                    .map_err(Error::InitPmu)?;
+            if cpu.vcpu.has_pmu_support() {
+                cpu.vcpu.init_pmu(irq).map_err(Error::InitPmu)?;
             } else {
                 debug!(
                     "PMU attribute is not supported in vCPU{}, skip PMU init!",
-                    cpu.lock().unwrap().id
+                    cpu.id
                 );
                 return Ok(false);
             }