From f5a2f8473bb00753983b85a6cc2b89c41d81db21 Mon Sep 17 00:00:00 2001
From: Jinank Jain <jinankjain@microsoft.com>
Date: Tue, 4 Feb 2025 15:26:00 +0000
Subject: [PATCH] hypervisor: Add a basic implementation of MshvVcpuState for
 aarch64

Currently we are just storing the StandardRegisters in the Vcpu state
which would be required for saving and restoring the ARM64 guest on
MSHV.

Signed-off-by: Jinank Jain <jinankjain@microsoft.com>
---
 hypervisor/src/lib.rs              |  2 +-
 hypervisor/src/mshv/aarch64/mod.rs | 22 ++++++++++++++++++++++
 hypervisor/src/mshv/mod.rs         |  7 +++++++
 3 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 hypervisor/src/mshv/aarch64/mod.rs

diff --git a/hypervisor/src/lib.rs b/hypervisor/src/lib.rs
index a644d5cd9..87dd9a9b6 100644
--- a/hypervisor/src/lib.rs
+++ b/hypervisor/src/lib.rs
@@ -163,7 +163,7 @@ pub enum IoEventAddress {
 pub enum CpuState {
     #[cfg(feature = "kvm")]
     Kvm(kvm::VcpuKvmState),
-    #[cfg(all(feature = "mshv", target_arch = "x86_64"))]
+    #[cfg(feature = "mshv")]
     Mshv(mshv::VcpuMshvState),
 }
 
diff --git a/hypervisor/src/mshv/aarch64/mod.rs b/hypervisor/src/mshv/aarch64/mod.rs
new file mode 100644
index 000000000..4634643c9
--- /dev/null
+++ b/hypervisor/src/mshv/aarch64/mod.rs
@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: Apache-2.0 OR BSD-3-Clause
+//
+// Copyright © 2025, Microsoft Corporation
+//
+use std::fmt;
+
+///
+/// Export generically-named wrappers of mshv_bindings for Unix-based platforms
+///
+pub use mshv_bindings::StandardRegisters as MshvStandardRegisters;
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct VcpuMshvState {
+    pub regs: MshvStandardRegisters,
+}
+
+impl fmt::Display for VcpuMshvState {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Standard registers: {:?}", self.regs)
+    }
+}
diff --git a/hypervisor/src/mshv/mod.rs b/hypervisor/src/mshv/mod.rs
index 15cfe8fbd..7ea07c687 100644
--- a/hypervisor/src/mshv/mod.rs
+++ b/hypervisor/src/mshv/mod.rs
@@ -29,10 +29,17 @@ mod snp_constants;
 // x86_64 dependencies
 #[cfg(target_arch = "x86_64")]
 pub mod x86_64;
+// aarch64 dependencies
+#[cfg(target_arch = "aarch64")]
+pub mod aarch64;
 #[cfg(target_arch = "x86_64")]
 use std::fs::File;
 use std::os::unix::io::AsRawFd;
+#[cfg(target_arch = "aarch64")]
+use std::sync::Mutex;
 
+#[cfg(target_arch = "aarch64")]
+pub use aarch64::VcpuMshvState;
 #[cfg(feature = "sev_snp")]
 use igvm_defs::IGVM_VHS_SNP_ID_BLOCK;
 #[cfg(feature = "sev_snp")]