From 3b59e57001679f2f0f07cffa9f77c94116cb9047 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Tue, 7 Jun 2022 16:30:49 +0200
Subject: [PATCH] pci: vfio: Implement Migratable for VfioPciDevice

Based on the VfioCommon implementation, the VfioPciDevice now implements
the Migratable trait.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 pci/src/vfio.rs | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/pci/src/vfio.rs b/pci/src/vfio.rs
index ef8df9f49..88ad93b30 100644
--- a/pci/src/vfio.rs
+++ b/pci/src/vfio.rs
@@ -30,7 +30,9 @@ use vm_device::interrupt::{
 };
 use vm_device::{BusDevice, Resource};
 use vm_memory::{Address, GuestAddress, GuestUsize};
-use vm_migration::{MigratableError, Pausable, Snapshot, Snapshottable, VersionMapped};
+use vm_migration::{
+    Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable, VersionMapped,
+};
 use vmm_sys_util::eventfd::EventFd;
 
 #[derive(Debug, Error)]
@@ -1616,3 +1618,31 @@ impl PciDevice for VfioPciDevice {
         Some(self.id.clone())
     }
 }
+
+impl Pausable for VfioPciDevice {}
+
+impl Snapshottable for VfioPciDevice {
+    fn id(&self) -> String {
+        self.id.clone()
+    }
+
+    fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
+        let mut vfio_pci_dev_snapshot = Snapshot::new(&self.id);
+
+        // Snapshot VfioCommon
+        vfio_pci_dev_snapshot.add_snapshot(self.common.snapshot()?);
+
+        Ok(vfio_pci_dev_snapshot)
+    }
+
+    fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
+        // Restore VfioCommon
+        if let Some(vfio_common_snapshot) = snapshot.snapshots.get(&self.common.id()) {
+            self.common.restore(*vfio_common_snapshot.clone())?;
+        }
+
+        Ok(())
+    }
+}
+impl Transportable for VfioPciDevice {}
+impl Migratable for VfioPciDevice {}