From de248169ce7b734ea9dcffd18d0cad6c2266c0f5 Mon Sep 17 00:00:00 2001
From: xuejun-xj <jiyunxue@linux.alibaba.com>
Date: Thu, 12 Oct 2023 15:08:35 +0800
Subject: [PATCH] Add Cap::KvmCapStealTime capability

This capability is used for kvm steal time feature.

Signed-off-by: xuejun-xj <jiyunxue@linux.alibaba.com>
---
 CHANGELOG.md     |  2 ++
 src/cap.rs       |  2 ++
 src/ioctls/vm.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5bc8445..94c03e6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@
 ## Changed
 - [[#234](https://github.com/rust-vmm/kvm-ioctls/issues/234)] vcpu: export
 reg_size as a public method.
+- [[#239](https://github.com/rust-vmm/kvm-ioctls/pull/239)] Add Cap::KvmCapStealTime
+capability.
 
 # v0.15.0
 
diff --git a/src/cap.rs b/src/cap.rs
index 245ae74..cd427b5 100644
--- a/src/cap.rs
+++ b/src/cap.rs
@@ -162,4 +162,6 @@ pub enum Cap {
     ArmPtrAuthAddress = KVM_CAP_ARM_PTRAUTH_ADDRESS,
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
     ArmPtrAuthGeneric = KVM_CAP_ARM_PTRAUTH_GENERIC,
+    #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
+    KvmCapStealTime = KVM_CAP_STEAL_TIME,
 }
diff --git a/src/ioctls/vm.rs b/src/ioctls/vm.rs
index 6b9dc9c..232ac8c 100644
--- a/src/ioctls/vm.rs
+++ b/src/ioctls/vm.rs
@@ -2251,4 +2251,48 @@ mod tests {
         assert!(vm.register_enc_memory_region(&memory_region).is_ok());
         assert!(vm.unregister_enc_memory_region(&memory_region).is_ok());
     }
+
+    #[test]
+    #[cfg(target_arch = "aarch64")]
+    fn test_kvm_steal_time_capability() {
+        let kvm = Kvm::new().unwrap();
+        let vm = kvm.create_vm().unwrap();
+        let vcpu = vm.create_vcpu(0).unwrap();
+        create_gic_device(&vm, 0);
+
+        if kvm.check_extension(Cap::KvmCapStealTime) {
+            const PVTIME_REGION_SIZE: u64 = 0x10000;
+            const TEST_KVM_MEM_SLOT: u32 = 0;
+            const TEST_IPA: u64 = 0x20000000;
+            let hva = unsafe {
+                libc::mmap(
+                    std::ptr::null_mut(),
+                    PVTIME_REGION_SIZE as usize,
+                    libc::PROT_READ | libc::PROT_WRITE,
+                    libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
+                    0,
+                    0,
+                )
+            } as *const u8 as u64;
+
+            let mem_region = kvm_bindings::kvm_userspace_memory_region {
+                slot: TEST_KVM_MEM_SLOT,
+                guest_phys_addr: TEST_IPA,
+                memory_size: PVTIME_REGION_SIZE,
+                userspace_addr: hva,
+                flags: 0,
+            };
+            unsafe {
+                vm.set_user_memory_region(mem_region).unwrap();
+            }
+
+            let device_attribute = kvm_device_attr {
+                group: KVM_ARM_VCPU_PVTIME_CTRL,
+                attr: KVM_ARM_VCPU_PVTIME_IPA as u64,
+                addr: &TEST_IPA as *const _ as u64,
+                ..Default::default()
+            };
+            assert!(vcpu.set_device_attr(&device_attribute).is_ok())
+        }
+    }
 }