AArch64: Implement AArch64 vCPU states save/restore

This commit adds methods to save/restore AArch64 vCPU registers,
including:

1. The AArch64 `VcpuKvmState` structure.

2. Some `Vcpu` trait methods of the `KvmVcpu` structure to
enable the save/restore of the AArch64 vCPU states.

Signed-off-by: Henry Wang <Henry.Wang@arm.com>

hypervisor/src/cpu.rs

@@ -10,6 +10,8 @@
 
 #[cfg(target_arch = "aarch64")]
 use crate::aarch64::VcpuInit;
+#[cfg(target_arch = "aarch64")]
+use crate::aarch64::{OneRegister, RegList, StandardRegisters};
 use crate::{CpuState, MpState};
 
 #[cfg(target_arch = "x86_64")]
@@ -154,10 +156,30 @@ pub enum HypervisorCpuError {
     #[error("Failed to enable HyperV SynIC")]
     EnableHyperVSynIC(#[source] anyhow::Error),
     ///
+    /// Getting AArch64 core register error
+    ///
+    #[error("Failed to get core register: {0}")]
+    GetCoreRegister(#[source] anyhow::Error),
+    ///
+    /// Setting AArch64 core register error
+    ///
+    #[error("Failed to set core register: {0}")]
+    SetCoreRegister(#[source] anyhow::Error),
+    ///
+    /// Getting AArch64 registers list error
+    ///
+    #[error("Failed to retrieve list of registers: {0}")]
+    GetRegList(#[source] anyhow::Error),
+    ///
     /// Getting AArch64 system register error
     ///
     #[error("Failed to get system register: {0}")]
     GetSysRegister(#[source] anyhow::Error),
+    ///
+    /// Setting AArch64 system register error
+    ///
+    #[error("Failed to set system register: {0}")]
+    SetSysRegister(#[source] anyhow::Error),
 }
 
 #[derive(Debug)]
@@ -311,6 +333,32 @@ pub trait Vcpu: Send + Sync {
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
     fn get_one_reg(&self, reg_id: u64) -> Result<u64>;
     ///
+    /// Gets a list of the guest registers that are supported for the
+    /// KVM_GET_ONE_REG/KVM_SET_ONE_REG calls.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn get_reg_list(&self, reg_list: &mut RegList) -> Result<()>;
+    ///
+    /// Save the state of the core registers.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn core_registers(&self, state: &mut StandardRegisters) -> Result<()>;
+    ///
+    /// Restore the state of the core registers.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn set_core_registers(&self, state: &StandardRegisters) -> Result<()>;
+    ///
+    /// Save the state of the system registers.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn system_registers(&self, state: &mut Vec<OneRegister>) -> Result<()>;
+    ///
+    /// Restore the state of the system registers.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn set_system_registers(&self, state: &[OneRegister]) -> Result<()>;
+    ///
     /// Read the MPIDR - Multiprocessor Affinity Register.
     ///
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]

hypervisor/src/kvm/aarch64/mod.rs

@@ -12,14 +12,16 @@
 /// Export generically-named wrappers of kvm-bindings for Unix-based platforms
 ///
 use crate::kvm::{KvmError, KvmResult};
-pub use kvm_bindings::kvm_vcpu_init as VcpuInit;
 use kvm_bindings::{
-    KVM_REG_ARM64, KVM_REG_ARM64_SYSREG, KVM_REG_ARM64_SYSREG_CRM_MASK,
+    kvm_mp_state, kvm_one_reg, kvm_regs, KVM_REG_ARM64, KVM_REG_ARM64_SYSREG,
-    KVM_REG_ARM64_SYSREG_CRM_SHIFT, KVM_REG_ARM64_SYSREG_CRN_MASK, KVM_REG_ARM64_SYSREG_CRN_SHIFT,
+    KVM_REG_ARM64_SYSREG_CRM_MASK, KVM_REG_ARM64_SYSREG_CRM_SHIFT, KVM_REG_ARM64_SYSREG_CRN_MASK,
-    KVM_REG_ARM64_SYSREG_OP0_MASK, KVM_REG_ARM64_SYSREG_OP0_SHIFT, KVM_REG_ARM64_SYSREG_OP1_MASK,
+    KVM_REG_ARM64_SYSREG_CRN_SHIFT, KVM_REG_ARM64_SYSREG_OP0_MASK, KVM_REG_ARM64_SYSREG_OP0_SHIFT,
-    KVM_REG_ARM64_SYSREG_OP1_SHIFT, KVM_REG_ARM64_SYSREG_OP2_MASK, KVM_REG_ARM64_SYSREG_OP2_SHIFT,
+    KVM_REG_ARM64_SYSREG_OP1_MASK, KVM_REG_ARM64_SYSREG_OP1_SHIFT, KVM_REG_ARM64_SYSREG_OP2_MASK,
-    KVM_REG_ARM_COPROC_MASK, KVM_REG_ARM_CORE, KVM_REG_SIZE_MASK, KVM_REG_SIZE_U32,
+    KVM_REG_ARM64_SYSREG_OP2_SHIFT, KVM_REG_ARM_COPROC_MASK, KVM_REG_ARM_CORE, KVM_REG_SIZE_MASK,
-    KVM_REG_SIZE_U64,
+    KVM_REG_SIZE_U32, KVM_REG_SIZE_U64,
+};
+pub use kvm_bindings::{
+    kvm_one_reg as OneRegister, kvm_regs as StandardRegisters, kvm_vcpu_init as VcpuInit, RegList,
 };
 use serde_derive::{Deserialize, Serialize};
 pub use {kvm_ioctls::Cap, kvm_ioctls::Kvm};
@@ -128,5 +130,13 @@ pub fn check_required_kvm_extensions(kvm: &Kvm) -> KvmResult<()> {
     Ok(())
 }
 
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone, Default, Serialize, Deserialize)]
-pub struct VcpuKvmState {}
+pub struct VcpuKvmState {
+    pub mp_state: kvm_mp_state,
+    pub core_regs: kvm_regs,
+    pub sys_regs: Vec<kvm_one_reg>,
+    // We will be using the mpidr for passing it to the VmState.
+    // The VmState will give this away for saving restoring the icc and redistributor
+    // registers.
+    pub mpidr: u64,
+}

hypervisor/src/kvm/mod.rs

@@ -10,12 +10,15 @@
 
 #[cfg(target_arch = "aarch64")]
 pub use crate::aarch64::{
-    check_required_kvm_extensions, VcpuInit, VcpuKvmState as CpuState, MPIDR_EL1,
+    check_required_kvm_extensions, is_system_register, VcpuInit, VcpuKvmState as CpuState,
+    MPIDR_EL1,
 };
 use crate::cpu;
 use crate::device;
 use crate::hypervisor;
 use crate::vm;
+#[cfg(target_arch = "aarch64")]
+use crate::{arm64_core_reg_id, offset__of};
 use kvm_ioctls::{NoDatamatch, VcpuFd, VmFd};
 use serde_derive::{Deserialize, Serialize};
 use std::os::unix::io::{AsRawFd, RawFd};
@@ -33,6 +36,9 @@ use x86_64::{
     check_required_kvm_extensions, FpuState, SpecialRegisters, StandardRegisters, KVM_TSS_ADDRESS,
 };
 
+#[cfg(target_arch = "aarch64")]
+use aarch64::{OneRegister, RegList, StandardRegisters};
+
 #[cfg(target_arch = "x86_64")]
 pub use x86_64::{
     CpuId, CpuIdEntry, ExtendedControlRegisters, LapicState, MsrEntries, VcpuKvmState as CpuState,
@@ -48,6 +54,13 @@ use crate::arch::x86::NUM_IOAPIC_PINS;
 // aarch64 dependencies
 #[cfg(target_arch = "aarch64")]
 pub mod aarch64;
+#[cfg(target_arch = "aarch64")]
+use kvm_bindings::{
+    kvm_regs, user_fpsimd_state, user_pt_regs, KVM_NR_SPSR, KVM_REG_ARM64, KVM_REG_ARM_CORE,
+    KVM_REG_SIZE_U128, KVM_REG_SIZE_U32, KVM_REG_SIZE_U64,
+};
+#[cfg(target_arch = "aarch64")]
+use std::mem;
 
 pub use kvm_bindings;
 pub use kvm_bindings::{
@@ -766,6 +779,241 @@ impl cpu::Vcpu for KvmVcpu {
             .map_err(|e| cpu::HypervisorCpuError::GetOneReg(e.into()))
     }
     ///
+    /// Gets a list of the guest registers that are supported for the
+    /// KVM_GET_ONE_REG/KVM_SET_ONE_REG calls.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn get_reg_list(&self, reg_list: &mut RegList) -> cpu::Result<()> {
+        self.fd
+            .get_reg_list(reg_list)
+            .map_err(|e| cpu::HypervisorCpuError::GetRegList(e.into()))
+    }
+    ///
+    /// Save the state of the core registers.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn core_registers(&self, state: &mut StandardRegisters) -> cpu::Result<()> {
+        let mut off = offset__of!(user_pt_regs, regs);
+        // There are 31 user_pt_regs:
+        // https://elixir.free-electrons.com/linux/v4.14.174/source/arch/arm64/include/uapi/asm/ptrace.h#L72
+        // These actually are the general-purpose registers of the Armv8-a
+        // architecture (i.e x0-x30 if used as a 64bit register or w0-30 when used as a 32bit register).
+        for i in 0..31 {
+            state.regs.regs[i] = self
+                .fd
+                .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off))
+                .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?;
+            off += std::mem::size_of::<u64>();
+        }
+
+        // We are now entering the "Other register" section of the ARMv8-a architecture.
+        // First one, stack pointer.
+        let off = offset__of!(user_pt_regs, sp);
+        state.regs.sp = self
+            .fd
+            .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off))
+            .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?;
+
+        // Second one, the program counter.
+        let off = offset__of!(user_pt_regs, pc);
+        state.regs.pc = self
+            .fd
+            .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off))
+            .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?;
+
+        // Next is the processor state.
+        let off = offset__of!(user_pt_regs, pstate);
+        state.regs.pstate = self
+            .fd
+            .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off))
+            .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?;
+
+        // The stack pointer associated with EL1
+        let off = offset__of!(kvm_regs, sp_el1);
+        state.sp_el1 = self
+            .fd
+            .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off))
+            .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?;
+
+        // Exception Link Register for EL1, when taking an exception to EL1, this register
+        // holds the address to which to return afterwards.
+        let off = offset__of!(kvm_regs, elr_el1);
+        state.elr_el1 = self
+            .fd
+            .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off))
+            .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?;
+
+        // Saved Program Status Registers, there are 5 of them used in the kernel.
+        let mut off = offset__of!(kvm_regs, spsr);
+        for i in 0..KVM_NR_SPSR as usize {
+            state.spsr[i] = self
+                .fd
+                .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off))
+                .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?;
+            off += std::mem::size_of::<u64>();
+        }
+
+        // Now moving on to floting point registers which are stored in the user_fpsimd_state in the kernel:
+        // https://elixir.free-electrons.com/linux/v4.9.62/source/arch/arm64/include/uapi/asm/kvm.h#L53
+        let mut off = offset__of!(kvm_regs, fp_regs) + offset__of!(user_fpsimd_state, vregs);
+        for i in 0..32 {
+            state.fp_regs.vregs[i][0] = self
+                .fd
+                .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U128, off))
+                .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?;
+            off += mem::size_of::<u128>();
+        }
+
+        // Floating-point Status Register
+        let off = offset__of!(kvm_regs, fp_regs) + offset__of!(user_fpsimd_state, fpsr);
+        state.fp_regs.fpsr = self
+            .fd
+            .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U32, off))
+            .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?
+            as u32;
+
+        // Floating-point Control Register
+        let off = offset__of!(kvm_regs, fp_regs) + offset__of!(user_fpsimd_state, fpcr);
+        state.fp_regs.fpcr = self
+            .fd
+            .get_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U32, off))
+            .map_err(|e| cpu::HypervisorCpuError::GetCoreRegister(e.into()))?
+            as u32;
+        Ok(())
+    }
+    ///
+    /// Restore the state of the core registers.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn set_core_registers(&self, state: &StandardRegisters) -> cpu::Result<()> {
+        // The function follows the exact identical order from `state`. Look there
+        // for some additional info on registers.
+        let mut off = offset__of!(user_pt_regs, regs);
+        for i in 0..31 {
+            self.fd
+                .set_one_reg(
+                    arm64_core_reg_id!(KVM_REG_SIZE_U64, off),
+                    state.regs.regs[i],
+                )
+                .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+            off += std::mem::size_of::<u64>();
+        }
+
+        let off = offset__of!(user_pt_regs, sp);
+        self.fd
+            .set_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off), state.regs.sp)
+            .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+
+        let off = offset__of!(user_pt_regs, pc);
+        self.fd
+            .set_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off), state.regs.pc)
+            .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+
+        let off = offset__of!(user_pt_regs, pstate);
+        self.fd
+            .set_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off), state.regs.pstate)
+            .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+
+        let off = offset__of!(kvm_regs, sp_el1);
+        self.fd
+            .set_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off), state.sp_el1)
+            .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+
+        let off = offset__of!(kvm_regs, elr_el1);
+        self.fd
+            .set_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off), state.elr_el1)
+            .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+
+        let mut off = offset__of!(kvm_regs, spsr);
+        for i in 0..KVM_NR_SPSR as usize {
+            self.fd
+                .set_one_reg(arm64_core_reg_id!(KVM_REG_SIZE_U64, off), state.spsr[i])
+                .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+            off += std::mem::size_of::<u64>();
+        }
+
+        let mut off = offset__of!(kvm_regs, fp_regs) + offset__of!(user_fpsimd_state, vregs);
+        for i in 0..32 {
+            self.fd
+                .set_one_reg(
+                    arm64_core_reg_id!(KVM_REG_SIZE_U128, off),
+                    state.fp_regs.vregs[i][0],
+                )
+                .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+            off += mem::size_of::<u128>();
+        }
+
+        let off = offset__of!(kvm_regs, fp_regs) + offset__of!(user_fpsimd_state, fpsr);
+        self.fd
+            .set_one_reg(
+                arm64_core_reg_id!(KVM_REG_SIZE_U32, off),
+                state.fp_regs.fpsr as u64,
+            )
+            .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+
+        let off = offset__of!(kvm_regs, fp_regs) + offset__of!(user_fpsimd_state, fpcr);
+        self.fd
+            .set_one_reg(
+                arm64_core_reg_id!(KVM_REG_SIZE_U32, off),
+                state.fp_regs.fpcr as u64,
+            )
+            .map_err(|e| cpu::HypervisorCpuError::SetCoreRegister(e.into()))?;
+        Ok(())
+    }
+    ///
+    /// Save the state of the system registers.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn system_registers(&self, state: &mut Vec<OneRegister>) -> cpu::Result<()> {
+        // Call KVM_GET_REG_LIST to get all registers available to the guest. For ArmV8 there are
+        // around 500 registers.
+        let mut reg_list = RegList::new(512);
+        self.fd
+            .get_reg_list(&mut reg_list)
+            .map_err(|e| cpu::HypervisorCpuError::GetRegList(e.into()))?;
+
+        // At this point reg_list should contain: core registers and system registers.
+        // The register list contains the number of registers and their ids. We will be needing to
+        // call KVM_GET_ONE_REG on each id in order to save all of them. We carve out from the list
+        // the core registers which are represented in the kernel by kvm_regs structure and for which
+        // we can calculate the id based on the offset in the structure.
+
+        reg_list.retain(|regid| *regid != 0);
+        reg_list.as_slice().to_vec().sort();
+
+        reg_list.retain(|regid| is_system_register(*regid));
+
+        // Now, for the rest of the registers left in the previously fetched register list, we are
+        // simply calling KVM_GET_ONE_REG.
+        let indices = reg_list.as_slice();
+        for (_pos, index) in indices.iter().enumerate() {
+            if _pos > 230 {
+                break;
+            }
+            state.push(kvm_bindings::kvm_one_reg {
+                id: *index,
+                addr: self
+                    .fd
+                    .get_one_reg(*index)
+                    .map_err(|e| cpu::HypervisorCpuError::GetSysRegister(e.into()))?,
+            });
+        }
+
+        Ok(())
+    }
+    ///
+    /// Restore the state of the system registers.
+    ///
+    #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
+    fn set_system_registers(&self, state: &[OneRegister]) -> cpu::Result<()> {
+        for reg in state {
+            self.fd
+                .set_one_reg(reg.id, reg.addr)
+                .map_err(|e| cpu::HypervisorCpuError::SetSysRegister(e.into()))?;
+        }
+        Ok(())
+    }
+    ///
     /// Read the MPIDR - Multiprocessor Affinity Register.
     ///
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
@@ -883,9 +1131,19 @@ impl cpu::Vcpu for KvmVcpu {
             mp_state,
         })
     }
+    ///
+    /// Get the current AArch64 CPU state
+    ///
     #[cfg(target_arch = "aarch64")]
     fn state(&self) -> cpu::Result<CpuState> {
-        unimplemented!();
+        let mut state = CpuState::default();
+        // Get this vCPUs multiprocessing state.
+        state.mp_state = self.get_mp_state()?;
+        self.core_registers(&mut state.core_regs)?;
+        self.system_registers(&mut state.sys_regs)?;
+        state.mpidr = self.read_mpidr()?;
+
+        Ok(state)
     }
     #[cfg(target_arch = "x86_64")]
     ///
@@ -969,10 +1227,15 @@ impl cpu::Vcpu for KvmVcpu {
 
         Ok(())
     }
-    #[allow(unused_variables)]
+    ///
+    /// Restore the previously saved AArch64 CPU state
+    ///
     #[cfg(target_arch = "aarch64")]
     fn set_state(&self, state: &CpuState) -> cpu::Result<()> {
-        warn!("CPU state was not restored");
+        self.set_core_registers(&state.core_regs)?;
+        self.set_system_registers(&state.sys_regs)?;
+        self.set_mp_state(state.mp_state)?;
+
         Ok(())
     }
 }

hypervisor/src/lib.rs

@@ -20,6 +20,7 @@
 
 #[macro_use]
 extern crate anyhow;
+#[cfg(target_arch = "x86_64")]
 #[macro_use]
 extern crate log;
 extern crate serde;