From 9b19c896560f3422d6c5f6c5da63211f48e5ab66 Mon Sep 17 00:00:00 2001
From: Sida Chen <sidachen@google.com>
Date: Wed, 9 Oct 2024 22:14:04 +0000
Subject: [PATCH] Introduce check_extension_int and check_extension_raw

The two functions are used to return integer value from
KVM_CHECK_EXTENSION ioctl commands. This is useful for capabilities
returning an integer with each bits representing different configs.

Signed-off-by: Sida Chen <sidachen@google.com>
---
 kvm-ioctls/CHANGELOG.md     |  1 +
 kvm-ioctls/src/ioctls/vm.rs | 48 +++++++++++++++++++++++++++++++++----
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/kvm-ioctls/CHANGELOG.md b/kvm-ioctls/CHANGELOG.md
index a8affaf..2657c80 100644
--- a/kvm-ioctls/CHANGELOG.md
+++ b/kvm-ioctls/CHANGELOG.md
@@ -6,6 +6,7 @@
 
 - [[#288](https://github.com/rust-vmm/kvm-ioctls/pull/288)]: Introduce `Cap::GuestMemfd`, `Cap::MemoryAttributes` and 
    `Cap::UserMemory2` capabilities enum variants for use with `VmFd::check_extension`.
+- [[#288](https://github.com/rust-vmm/kvm-ioctls/pull/288)]: Introduce `VmFd::check_extension_raw` and `VmFd::check_extension_int` to allow `KVM_CHECK_EXTENSION` to return integer.
 
 ### Fixed
 
diff --git a/kvm-ioctls/src/ioctls/vm.rs b/kvm-ioctls/src/ioctls/vm.rs
index d293c6a..c5cec2f 100644
--- a/kvm-ioctls/src/ioctls/vm.rs
+++ b/kvm-ioctls/src/ioctls/vm.rs
@@ -1402,10 +1402,50 @@ impl VmFd {
     /// Wrapper over `KVM_CHECK_EXTENSION`.
     ///
     /// Returns 0 if the capability is not available and a positive integer otherwise.
-    fn check_extension_int(&self, c: Cap) -> i32 {
-        // SAFETY: Safe because we know that our file is a VM fd and that the extension is one of
-        // the ones defined by kernel.
-        unsafe { ioctl_with_val(self, KVM_CHECK_EXTENSION(), c as c_ulong) }
+    /// See the documentation for `KVM_CHECK_EXTENSION`.
+    ///
+    /// # Arguments
+    ///
+    /// * `c` - KVM capability to check.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use kvm_ioctls::Kvm;
+    /// use kvm_ioctls::Cap;
+    ///
+    /// let kvm = Kvm::new().unwrap();
+    /// let vm = kvm.create_vm().unwrap();
+    /// assert!(vm.check_extension_int(Cap::MaxVcpus) > 0);
+    /// ```
+    pub fn check_extension_int(&self, c: Cap) -> i32 {
+        self.check_extension_raw(c as c_ulong)
+    }
+
+    /// Wrapper over `KVM_CHECK_EXTENSION`.
+    ///
+    /// Returns 0 if the capability is not available and a positive integer otherwise.
+    /// See the documentation for `KVM_CHECK_EXTENSION`.
+    ///
+    /// # Arguments
+    ///
+    /// * `c` - KVM capability to check in a form of a raw integer.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use kvm_ioctls::Kvm;
+    /// # use std::os::raw::c_ulong;
+    /// use kvm_ioctls::Cap;
+    ///
+    /// let kvm = Kvm::new().unwrap();
+    /// let vm = kvm.create_vm().unwrap();
+    /// assert!(vm.check_extension_raw(Cap::MaxVcpus as c_ulong) > 0);
+    /// ```
+    pub fn check_extension_raw(&self, c: c_ulong) -> i32 {
+        // SAFETY: Safe because we know that our file is a KVM fd.
+        // If `c` is not a known kernel extension, kernel will return 0.
+        unsafe { ioctl_with_val(self, KVM_CHECK_EXTENSION(), c) }
     }
 
     /// Checks if a particular `Cap` is available.