Fix alignment issues

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/src/hypervisor/regs/special_regs.rs

@@ -10,6 +10,8 @@ extern crate mshv_ioctls3 as mshv_ioctls;
 
 #[cfg(target_os = "windows")]
 use std::collections::HashSet;
+#[cfg(target_os = "windows")]
+use std::mem::offset_of;
 
 #[cfg(kvm)]
 use kvm_bindings::{kvm_dtable, kvm_segment, kvm_sregs};
@@ -20,6 +22,8 @@ use windows::Win32::System::Hypervisor::*;
 
 #[cfg(target_os = "windows")]
 use super::FromWhpRegisterError;
+#[cfg(target_os = "windows")]
+use crate::hypervisor::regs::{WHP_FPU_NAMES_LEN, WHP_REGS_NAMES_LEN};
 
 #[derive(Debug, Default, Copy, Clone, PartialEq)]
 pub(crate) struct CommonSpecialRegisters {
@@ -147,6 +151,23 @@ impl From<&CommonSpecialRegisters> for kvm_sregs {
     }
 }
 
+#[repr(C, align(16))]
+#[cfg(target_os = "windows")]
+pub(crate) struct AlignedRegisterValues<const N: usize>(pub(crate) [WHV_REGISTER_VALUE; N]);
+
+#[cfg(target_os = "windows")]
+const _: () = {
+    // WHP_SREGS_NAMES_LEN
+    assert!(std::mem::align_of::<AlignedRegisterValues<WHP_SREGS_NAMES_LEN>>() % 16 == 0);
+    assert!(offset_of!(AlignedRegisterValues<WHP_SREGS_NAMES_LEN>, 0) % 16 == 0);
+    // WHP_REGS_NAMES_LEN
+    assert!(std::mem::align_of::<AlignedRegisterValues<WHP_REGS_NAMES_LEN>>() % 16 == 0,);
+    assert!(offset_of!(AlignedRegisterValues<WHP_REGS_NAMES_LEN>, 0) % 16 == 0);
+    // WHP_FPU_NAMES_LEN
+    assert!(std::mem::align_of::<AlignedRegisterValues<WHP_FPU_NAMES_LEN>>() % 16 == 0,);
+    assert!(offset_of!(AlignedRegisterValues<WHP_FPU_NAMES_LEN>, 0) % 16 == 0);
+};
+
 #[cfg(target_os = "windows")]
 pub(crate) const WHP_SREGS_NAMES_LEN: usize = 17;
 #[cfg(target_os = "windows")]

src/hyperlight_host/src/hypervisor/whp.rs

@@ -26,8 +26,8 @@ use windows_result::HRESULT;
 #[cfg(gdb)]
 use super::handlers::DbgMemAccessHandlerWrapper;
 use super::regs::{
-    WHP_FPU_NAMES, WHP_FPU_NAMES_LEN, WHP_REGS_NAMES, WHP_REGS_NAMES_LEN, WHP_SREGS_NAMES,
-    WHP_SREGS_NAMES_LEN,
+    AlignedRegisterValues, WHP_FPU_NAMES, WHP_FPU_NAMES_LEN, WHP_REGS_NAMES, WHP_REGS_NAMES_LEN,
+    WHP_SREGS_NAMES, WHP_SREGS_NAMES_LEN,
 };
 use super::vm::HyperlightExit;
 use super::wrappers::HandleWrapper;
@@ -86,6 +86,15 @@ pub(crate) struct WhpVm {
 unsafe impl Send for WhpVm {}
 unsafe impl Sync for WhpVm {}
 
+#[repr(C, align(16))]
+struct Align16<T>(T);
+const _: () = {
+    assert!(
+        std::mem::size_of::<Align16<WHV_REGISTER_VALUE>>()
+            == std::mem::size_of::<WHV_REGISTER_VALUE>()
+    );
+};
+
 impl WhpVm {
     pub(crate) fn new(mmap_file_handle: HandleWrapper) -> Result<Self> {
         const NUM_CPU: u32 = 1;
@@ -112,21 +121,29 @@ impl WhpVm {
     /// Helper for setting arbitrary registers.
     fn set_registers(&self, registers: &[(WHV_REGISTER_NAME, WHV_REGISTER_VALUE)]) -> Result<()> {
         let register_count = registers.len();
-        let mut register_names: Vec<WHV_REGISTER_NAME> = vec![];
-        let mut register_values: Vec<WHV_REGISTER_VALUE> = vec![];
+
+        // Prepare register names (no special alignment needed)
+        let mut register_names = Vec::with_capacity(register_count);
+        let mut register_values = Vec::with_capacity(register_count);
 
         for (key, value) in registers.iter() {
             register_names.push(*key);
-            register_values.push(*value);
+            register_values.push(Align16(*value));
         }
 
+        assert_eq!(
+            (register_values.as_ptr() as usize) % 16,
+            0,
+            "register_values is not 16-byte aligned"
+        );
+
         unsafe {
             WHvSetVirtualProcessorRegisters(
                 self.partition,
                 0,
                 register_names.as_ptr(),
                 register_count as u32,
-                register_values.as_ptr(),
+                register_values.as_ptr() as *const WHV_REGISTER_VALUE,
             )?;
         }
 
@@ -136,22 +153,24 @@ impl WhpVm {
 
 impl Vm for WhpVm {
     fn get_regs(&self) -> Result<CommonRegisters> {
-        let mut whv_regs_values: [WHV_REGISTER_VALUE; WHP_REGS_NAMES_LEN] =
-            unsafe { std::mem::zeroed() };
+        let mut whv_regs_values =
+            AlignedRegisterValues::<WHP_REGS_NAMES_LEN>(unsafe { std::mem::zeroed() });
+
+        assert_eq!((whv_regs_values.0.as_ptr() as usize) % 16, 0);
 
         unsafe {
             WHvGetVirtualProcessorRegisters(
                 self.partition,
                 0,
                 WHP_REGS_NAMES.as_ptr(),
-                WHP_REGS_NAMES_LEN as u32,
-                whv_regs_values.as_mut_ptr(),
+                whv_regs_values.0.len() as u32,
+                whv_regs_values.0.as_mut_ptr(),
             )?;
         }
 
         WHP_REGS_NAMES
             .into_iter()
-            .zip(whv_regs_values)
+            .zip(whv_regs_values.0)
             .collect::<Vec<(WHV_REGISTER_NAME, WHV_REGISTER_VALUE)>>()
             .as_slice()
             .try_into()
@@ -170,22 +189,24 @@ impl Vm for WhpVm {
     }
 
     fn get_sregs(&self) -> Result<CommonSpecialRegisters> {
-        let mut whp_sregs_values: [WHV_REGISTER_VALUE; WHP_SREGS_NAMES_LEN] =
-            unsafe { std::mem::zeroed() };
+        let mut whp_sregs_values =
+            AlignedRegisterValues::<WHP_SREGS_NAMES_LEN>(unsafe { std::mem::zeroed() });
+
+        assert_eq!((whp_sregs_values.0.as_ptr() as usize) % 16, 0);
 
         unsafe {
             WHvGetVirtualProcessorRegisters(
                 self.partition,
                 0,
                 WHP_SREGS_NAMES.as_ptr(),
-                whp_sregs_values.len() as u32,
-                whp_sregs_values.as_mut_ptr(),
+                whp_sregs_values.0.len() as u32,
+                whp_sregs_values.0.as_mut_ptr(),
             )?;
         }
 
         WHP_SREGS_NAMES
             .into_iter()
-            .zip(whp_sregs_values)
+            .zip(whp_sregs_values.0)
             .collect::<Vec<(WHV_REGISTER_NAME, WHV_REGISTER_VALUE)>>()
             .as_slice()
             .try_into()
@@ -204,22 +225,24 @@ impl Vm for WhpVm {
     }
 
     fn get_fpu(&self) -> Result<CommonFpu> {
-        let mut whp_fpu_values: [WHV_REGISTER_VALUE; WHP_FPU_NAMES_LEN] =
-            unsafe { std::mem::zeroed() };
+        let mut whp_fpu_values =
+            AlignedRegisterValues::<WHP_FPU_NAMES_LEN>(unsafe { std::mem::zeroed() });
+
+        assert_eq!((whp_fpu_values.0.as_ptr() as usize) % 16, 0);
 
         unsafe {
             WHvGetVirtualProcessorRegisters(
                 self.partition,
                 0,
                 WHP_FPU_NAMES.as_ptr(),
-                whp_fpu_values.len() as u32,
-                whp_fpu_values.as_mut_ptr(),
+                whp_fpu_values.0.len() as u32,
+                whp_fpu_values.0.as_mut_ptr(),
             )?;
         }
 
         WHP_FPU_NAMES
             .into_iter()
-            .zip(whp_fpu_values)
+            .zip(whp_fpu_values.0)
             .collect::<Vec<(WHV_REGISTER_NAME, WHV_REGISTER_VALUE)>>()
             .as_slice()
             .try_into()