vmm: AArch64: enable VM states save/restore for AArch64

The states of GIC should be part of the VM states. This commit
enables the AArch64 VM states save/restore by adding save/restore
of GIC states.

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

arch/src/aarch64/gic/gicv2.rs

@@ -4,6 +4,7 @@
 pub mod kvm {
     use crate::aarch64::gic::kvm::KvmGICDevice;
     use crate::aarch64::gic::{Error, GICDevice};
+    use std::any::Any;
     use std::{boxed::Box, result};
     type Result<T> = result::Result<T, Error>;
     use crate::layout;
@@ -82,6 +83,10 @@ pub mod kvm {
         fn set_gicr_typers(&mut self, gicr_typers: Vec<u64>) {
             self.gicr_typers = gicr_typers;
         }
+
+        fn as_any_concrete_mut(&mut self) -> &mut dyn Any {
+            self
+        }
     }
 
     impl KvmGICDevice for KvmGICv2 {

arch/src/aarch64/gic/gicv3.rs

@@ -10,6 +10,7 @@ pub mod kvm {
     use crate::layout;
     use anyhow::anyhow;
     use hypervisor::kvm::kvm_bindings;
+    use std::any::Any;
     use std::convert::TryInto;
     use std::sync::Arc;
     use std::{boxed::Box, result};
@@ -163,6 +164,10 @@ pub mod kvm {
         fn set_gicr_typers(&mut self, gicr_typers: Vec<u64>) {
             self.gicr_typers = gicr_typers;
         }
+
+        fn as_any_concrete_mut(&mut self) -> &mut dyn Any {
+            self
+        }
     }
 
     impl KvmGICDevice for KvmGICv3 {
@@ -217,7 +222,7 @@ pub mod kvm {
         }
     }
 
-    const GIC_V3_SNAPSHOT_ID: &str = "gic-v3";
+    pub const GIC_V3_SNAPSHOT_ID: &str = "gic-v3";
     impl Snapshottable for KvmGICv3 {
         fn id(&self) -> String {
             GIC_V3_SNAPSHOT_ID.to_string()

arch/src/aarch64/gic/gicv3_its.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 pub mod kvm {
+    use std::any::Any;
     use std::convert::TryInto;
     use std::sync::Arc;
     use std::{boxed::Box, result};
@@ -76,6 +77,10 @@ pub mod kvm {
         fn set_gicr_typers(&mut self, gicr_typers: Vec<u64>) {
             self.gicr_typers = gicr_typers;
         }
+
+        fn as_any_concrete_mut(&mut self) -> &mut dyn Any {
+            self
+        }
     }
 
     impl KvmGICDevice for KvmGICv3ITS {

arch/src/aarch64/gic/mod.rs

@@ -2,15 +2,16 @@
 // SPDX-License-Identifier: Apache-2.0
 
 pub mod dist_regs;
-mod gicv2;
-mod gicv3;
-mod gicv3_its;
+pub mod gicv2;
+pub mod gicv3;
+pub mod gicv3_its;
 pub mod icc_regs;
 pub mod redist_regs;
 
 pub use self::dist_regs::{get_dist_regs, read_ctlr, set_dist_regs, write_ctlr};
 pub use self::icc_regs::{get_icc_regs, set_icc_regs};
 pub use self::redist_regs::{get_redist_regs, set_redist_regs};
+use std::any::Any;
 use std::result;
 use std::sync::Arc;
 
@@ -59,6 +60,9 @@ pub trait GICDevice: Send {
 
     /// Get the values of GICR_TYPER for each vCPU.
     fn set_gicr_typers(&mut self, gicr_typers: Vec<u64>);
+
+    /// Downcast the trait object to its concrete type.
+    fn as_any_concrete_mut(&mut self) -> &mut dyn Any;
 }
 
 pub mod kvm {

vmm/src/vm.rs

@@ -72,6 +72,11 @@ use vm_migration::{
 use vmm_sys_util::eventfd::EventFd;
 use vmm_sys_util::terminal::Terminal;
 
+#[cfg(target_arch = "aarch64")]
+use arch::aarch64::gic::gicv3::kvm::{KvmGICv3, GIC_V3_SNAPSHOT_ID};
+#[cfg(target_arch = "aarch64")]
+use arch::aarch64::gic::kvm::create_gic;
+
 // 64 bit direct boot entry offset for bzImage
 #[cfg(target_arch = "x86_64")]
 const KERNEL_64BIT_ENTRY_OFFSET: u64 = 0x200;
@@ -1386,6 +1391,94 @@ impl Vm {
             .map_err(|_| Error::PoisonedState)
             .map(|state| *state)
     }
+
+    #[cfg(target_arch = "aarch64")]
+    /// Add the vGIC section to the VM snapshot.
+    fn add_vgic_snapshot_section(
+        &self,
+        vm_snapshot: &mut Snapshot,
+    ) -> std::result::Result<(), MigratableError> {
+        let saved_vcpu_states = self.cpu_manager.lock().unwrap().get_saved_states();
+        self.device_manager
+            .lock()
+            .unwrap()
+            .construct_gicr_typers(&saved_vcpu_states);
+
+        vm_snapshot.add_snapshot(
+            self.device_manager
+                .lock()
+                .unwrap()
+                .get_gic_device_entity()
+                .unwrap()
+                .lock()
+                .unwrap()
+                .as_any_concrete_mut()
+                .downcast_mut::<KvmGICv3>()
+                .unwrap()
+                .snapshot()?,
+        );
+
+        Ok(())
+    }
+
+    #[cfg(target_arch = "aarch64")]
+    /// Restore the vGIC from the VM snapshot and enable the interrupt controller routing.
+    fn restore_vgic_and_enable_interrupt(
+        &self,
+        vm_snapshot: &Snapshot,
+    ) -> std::result::Result<(), MigratableError> {
+        let saved_vcpu_states = self.cpu_manager.lock().unwrap().get_saved_states();
+        // The number of vCPUs is the same as the number of saved vCPU states.
+        let vcpu_numbers = saved_vcpu_states.len();
+
+        // Creating a GIC device here, as the GIC will not be created when
+        // restoring the device manager. Note that currently only the bare GICv3
+        // without ITS is supported.
+        let gic_device = create_gic(&self.vm, vcpu_numbers.try_into().unwrap(), false)
+            .map_err(|e| MigratableError::Restore(anyhow!("Could not create GIC: {:#?}", e)))?;
+
+        // Update the GIC entity in device manager
+        self.device_manager
+            .lock()
+            .unwrap()
+            .set_gic_device_entity(Arc::new(Mutex::new(gic_device)));
+
+        // Here we prepare the GICR_TYPER registers from the restored vCPU states.
+        self.device_manager
+            .lock()
+            .unwrap()
+            .construct_gicr_typers(&saved_vcpu_states);
+
+        // Restore GIC states.
+        if let Some(gic_v3_snapshot) = vm_snapshot.snapshots.get(GIC_V3_SNAPSHOT_ID) {
+            self.device_manager
+                .lock()
+                .unwrap()
+                .get_gic_device_entity()
+                .unwrap()
+                .lock()
+                .unwrap()
+                .as_any_concrete_mut()
+                .downcast_mut::<KvmGICv3>()
+                .unwrap()
+                .restore(*gic_v3_snapshot.clone())?;
+        } else {
+            return Err(MigratableError::Restore(anyhow!("Missing GICv3 snapshot")));
+        }
+
+        self.device_manager
+            .lock()
+            .unwrap()
+            .enable_interrupt_controller()
+            .map_err(|e| {
+                MigratableError::Restore(anyhow!(
+                    "Could not enable interrupt controller routing: {:#?}",
+                    e
+                ))
+            })?;
+
+        Ok(())
+    }
 }
 
 impl Pausable for Vm {
@@ -1484,6 +1577,11 @@ impl Snapshottable for Vm {
 
         vm_snapshot.add_snapshot(self.cpu_manager.lock().unwrap().snapshot()?);
         vm_snapshot.add_snapshot(self.memory_manager.lock().unwrap().snapshot()?);
+
+        #[cfg(target_arch = "aarch64")]
+        self.add_vgic_snapshot_section(&mut vm_snapshot)
+            .map_err(|e| MigratableError::Snapshot(e.into()))?;
+
         vm_snapshot.add_snapshot(self.device_manager.lock().unwrap().snapshot()?);
         vm_snapshot.add_data_section(SnapshotDataSection {
             id: format!("{}-section", VM_SNAPSHOT_ID),
@@ -1535,6 +1633,9 @@ impl Snapshottable for Vm {
             )));
         }
 
+        #[cfg(target_arch = "aarch64")]
+        self.restore_vgic_and_enable_interrupt(&snapshot)?;
+
         // Now we can start all vCPUs from here.
         self.cpu_manager
             .lock()