From 4ec67984f49461bb11465172a4cc19dd99f41768 Mon Sep 17 00:00:00 2001
From: Nuno Das Neves <nudasnev@microsoft.com>
Date: Thu, 11 Jan 2024 00:46:28 +0000
Subject: [PATCH] mshv-ioctls/vcpu: Use generic hypercall ioctl for set/get vp
 registers

Signed-off-by: Nuno Das Neves <nudasnev@microsoft.com>
---
 mshv-ioctls/src/ioctls/vcpu.rs | 81 ++++++++++++++++++++++------------
 1 file changed, 54 insertions(+), 27 deletions(-)

diff --git a/mshv-ioctls/src/ioctls/vcpu.rs b/mshv-ioctls/src/ioctls/vcpu.rs
index 6416bf0f..4edd0881 100644
--- a/mshv-ioctls/src/ioctls/vcpu.rs
+++ b/mshv-ioctls/src/ioctls/vcpu.rs
@@ -66,39 +66,66 @@ impl AsRawFd for VcpuFd {
 impl VcpuFd {
     /// Get the register values by providing an array of register names
     #[cfg(not(any(target_arch = "arm", target_arch = "aarch64")))]
-    pub fn get_reg(&self, reg_names: &mut [hv_register_assoc]) -> Result<()> {
+    pub fn get_reg(&self, reg_assocs: &mut [hv_register_assoc]) -> Result<()> {
+
+        if reg_assocs.len() == 0 {
+            return Err(errno::Error::new(libc::EINVAL));
+        }
+
         //TODO: Error if input register len is zero
-        let mut mshv_vp_register_args = mshv_vp_registers {
-            count: reg_names.len() as i32,
-            regs: reg_names.as_mut_ptr(),
-        };
-        // SAFETY: we know that our file is a vCPU fd, we know the kernel will only read the
-        // correct amount of memory from our pointer, and we verify the return result.
-        let ret = unsafe {
-            ioctl_with_mut_ref(self, MSHV_GET_VP_REGISTERS(), &mut mshv_vp_register_args)
-        };
-        if ret != 0 {
-            return Err(errno::Error::last());
+        let reg_names: Vec<hv_register_name> = reg_assocs
+            .iter()
+            .map(|assoc| assoc.name)
+            .collect();
+        let input = make_rep_input!(
+            hv_input_get_vp_registers {
+                vp_index: self.index,
+                ..Default::default()
+            },
+            names,
+            reg_names.as_slice()
+        );
+        let mut output: Vec<hv_register_value> = reg_names
+            .iter()
+            .map(|_| hv_register_value {
+                reg128: hv_u128 { ..Default::default() },
+            })
+            .collect();
+        let output_slice = output.as_mut_slice();
+
+        let mut args = make_rep_args!(HVCALL_GET_VP_REGISTERS, input, output_slice);
+        self.hvcall(&mut args)?;
+
+        if args.reps as usize != reg_assocs.len() {
+            // TODO better handling? partial success?
+            return Err(errno::Error::new(libc::EINTR));
+        }
+
+        for (assoc, value) in reg_assocs.iter_mut().zip(output.iter()) {
+            assoc.value = *value;
         }
+
         Ok(())
     }
-    /// Sets a vCPU register to input value.
-    ///
-    /// # Arguments
-    ///
-    /// * `reg_name` - general purpose register name.
-    /// * `reg_value` - register value.
+    /// Set the register values by providing an array of register assocs
     #[cfg(not(any(target_arch = "arm", target_arch = "aarch64")))]
-    pub fn set_reg(&self, regs: &[hv_register_assoc]) -> Result<()> {
-        let hv_vp_register_args = mshv_vp_registers {
-            count: regs.len() as i32,
-            regs: regs.as_ptr() as *mut hv_register_assoc,
-        };
-        // SAFETY: IOCTL call with correct types.
-        let ret = unsafe { ioctl_with_ref(self, MSHV_SET_VP_REGISTERS(), &hv_vp_register_args) };
-        if ret != 0 {
-            return Err(errno::Error::last());
+    pub fn set_reg(&self, reg_assocs: &[hv_register_assoc]) -> Result<()> {
+        let input = make_rep_input!(
+            hv_input_set_vp_registers {
+                vp_index: self.index,
+                ..Default::default()
+            },
+            elements,
+            reg_assocs
+        );
+        let mut args = make_rep_args!(HVCALL_SET_VP_REGISTERS, input);
+        self.hvcall(&mut args)?;
+
+        if args.reps as usize != reg_assocs.len() {
+            // TODO better handling? partial success?
+            return Err(errno::Error::new(libc::EINTR));
         }
+
         Ok(())
     }
     /// Sets the vCPU general purpose registers