Author: beamandala

src/snapshot-editor/src/utils.rs

@@ -26,13 +26,22 @@ pub enum UtilsError {
 }
 
 #[allow(unused)]
-pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<Snapshot<MicrovmState>, UtilsError> {
+pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<(MicrovmState, Version), UtilsError> {
     let mut snapshot_reader = File::open(snapshot_path).map_err(UtilsError::VmStateFileOpen)?;
     let metadata = std::fs::metadata(snapshot_path).map_err(UtilsError::VmStateFileMeta)?;
     let snapshot_len = u64_to_usize(metadata.len());
 
     let snapshot: Result<Snapshot<MicrovmState>, UtilsError> = Snapshot::load(&mut snapshot_reader, snapshot_len).map_err(UtilsError::VmStateLoad);
-    snapshot
+    match snapshot {
+        Ok(snapshot) => {
+            let version = snapshot.version();
+            Ok((snapshot.data, version.to_owned()))
+        }
+        Err(e) => {
+            return Err(e);
+        }
+    }
+    
 }
 
 // This method is used only in aarch64 code so far

src/vmm/src/persist.rs

@@ -731,6 +731,8 @@ mod tests {
             #[cfg(target_arch = "x86_64")]
             acpi_dev_state: vmm.acpi_device_manager.save(),
         };
+        let vm_info = microvm_state.vm_info.clone();
+        let device_states = microvm_state.device_states.clone();
 
         let mut buf = vec![0; 10000];
 
@@ -742,10 +744,10 @@ mod tests {
             Snapshot::load(&mut buf.as_slice(), buf.len()).unwrap();
         let restored_microvm_state = restored_snapshot.data;
 
-        assert_eq!(restored_microvm_state.vm_info, microvm_state.vm_info);
+        assert_eq!(restored_microvm_state.vm_info, vm_info);
         assert_eq!(
             restored_microvm_state.device_states,
-            microvm_state.device_states
+            device_states
         )
     }
 

src/vmm/src/snapshot/mod.rs

@@ -82,19 +82,16 @@ impl SnapshotHdr {
     }
 
     pub fn load<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
-        let hdr: SnapshotHdr = bincode::DefaultOptions::new()
-            .with_limit(VM_STATE_DESERIALIZE_LIMIT)
-            .with_fixint_encoding()
-            .allow_trailing_bytes() // need this because we deserialize header and snapshot from the same file, so after
-            // reading the header, there will be trailing bytes.
-            .deserialize_from(reader)
-            .map_err(|err| SnapshotError::Serde(err.to_string()))?;
+        let hdr: SnapshotHdr = deserialize(reader)?;
 
         Ok(hdr)
     }
 
     pub fn store<W: Write>(&self, writer: &mut W) -> Result<(), SnapshotError> {
-        bincode::serialize_into(writer, self).map_err(|err| SnapshotError::Serde(err.to_string()))
+        match serialize(writer, self) {
+            Ok(_) => Ok(()),
+            Err(e) => Err(e)
+        }
     }
 }
 
@@ -105,33 +102,17 @@ pub struct Snapshot<Data> {
     pub data: Data,
 }
 
-impl<Data: for<'a> Deserialize<'a>> Snapshot<Data> {
-    /// Helper function to deserialize an object from a reader
-    pub fn deserialize<T, O>(reader: &mut T) -> Result<O, SnapshotError>
-    where
-        T: Read,
-        O: DeserializeOwned + Debug,
-    {
-        // flags below are those used by default by bincode::deserialize_from, plus `with_limit`.
-        bincode::DefaultOptions::new()
-            .with_limit(VM_STATE_DESERIALIZE_LIMIT)
-            .with_fixint_encoding()
-            .allow_trailing_bytes() // need this because we deserialize header and snapshot from the same file, so after
-            // reading the header, there will be trailing bytes.
-            .deserialize_from(reader)
-            .map_err(|err| SnapshotError::Serde(err.to_string()))
-    }
-
+impl<Data: DeserializeOwned> Snapshot<Data> {
     pub fn load_unchecked<R: Read>(reader: &mut R) -> Result<Self, SnapshotError>
     where
         Data: DeserializeOwned + Debug,
     {
-        let hdr: SnapshotHdr = Self::deserialize(reader)?;
+        let hdr: SnapshotHdr = deserialize(reader)?;
         if hdr.magic != SNAPSHOT_MAGIC_ID {
             return Err(SnapshotError::InvalidMagic(hdr.magic));
         }
 
-        let data: Data = Self::deserialize(reader)?;
+        let data: Data = deserialize(reader)?;
         Ok(Self { header: hdr, data })
     }
 
@@ -156,7 +137,7 @@ impl<Data: for<'a> Deserialize<'a>> Snapshot<Data> {
         // 2 statements is important, we first get the checksum computed on the read bytes
         // then read the stored checksum.
         let computed_checksum = crc_reader.checksum();
-        let stored_checksum: u64 = Self::deserialize(&mut crc_reader)?;
+        let stored_checksum: u64 = deserialize(&mut crc_reader)?;
         if computed_checksum != stored_checksum {
             return Err(SnapshotError::Crc64(computed_checksum));
         }
@@ -167,29 +148,11 @@ impl<Data: for<'a> Deserialize<'a>> Snapshot<Data> {
 }
 
 impl<Data: Serialize + Debug> Snapshot<Data> {
-    /// Helper function to serialize an object to a writer
-    pub fn serialize<T, O>(writer: &mut T, data: &O) -> Result<usize, SnapshotError>
-    where
-        T: Write,
-        O: Serialize + Debug,
-    {
-        let mut buffer = Vec::new();
-        bincode::serialize_into(&mut buffer, data)
-            .map_err(|err| SnapshotError::Serde(err.to_string()))?;
-
-        writer
-            .write_all(&buffer)
-            .map_err(|err| SnapshotError::Serde(err.to_string()))?;
-
-        Ok(buffer.len())
-        // bincode::serialize_into(writer, data).map_err(|err| SnapshotError::Serde(err.to_string()))
-    }
-
     pub fn save<W: Write>(&self, mut writer: &mut W) -> Result<usize, SnapshotError> {
         // Write magic value and snapshot version
-        Self::serialize(&mut writer, &SnapshotHdr::new(self.header.version.clone()))?;
+        serialize(&mut writer, &SnapshotHdr::new(self.header.version.clone()))?;
         // Write data
-        Self::serialize(&mut writer, &self.data)
+        serialize(&mut writer, &self.data)
     }
 
     pub fn save_with_crc<W: Write>(&self, writer: &mut W) -> Result<usize, SnapshotError> {
@@ -198,14 +161,52 @@ impl<Data: Serialize + Debug> Snapshot<Data> {
 
         // Now write CRC value
         let checksum = crc_writer.checksum();
-        Self::serialize(&mut crc_writer, &checksum)
+        serialize(&mut crc_writer, &checksum)
     }
 }
 
 impl<Data> Snapshot<Data> {
     pub fn new(header: SnapshotHdr, data: Data) -> Self {
         Snapshot { header, data }
     }
+
+    pub fn version(&self) -> Version {
+        self.header.version.clone()
+    }
+}
+
+/// Helper function to deserialize an object from a reader
+fn deserialize<T, O>(reader: &mut T) -> Result<O, SnapshotError>
+where
+    T: Read,
+    O: DeserializeOwned + Debug,
+{
+    // flags below are those used by default by bincode::deserialize_from, plus `with_limit`.
+    bincode::DefaultOptions::new()
+        .with_limit(VM_STATE_DESERIALIZE_LIMIT)
+        .with_fixint_encoding()
+        .allow_trailing_bytes() // need this because we deserialize header and snapshot from the same file, so after
+        // reading the header, there will be trailing bytes.
+        .deserialize_from(reader)
+        .map_err(|err| SnapshotError::Serde(err.to_string()))
+}
+
+/// Helper function to serialize an object to a writer
+fn serialize<T, O>(writer: &mut T, data: &O) -> Result<usize, SnapshotError>
+where
+    T: Write,
+    O: Serialize + Debug,
+{
+    let mut buffer = Vec::new();
+    bincode::serialize_into(&mut buffer, data)
+        .map_err(|err| SnapshotError::Serde(err.to_string()))?;
+
+    writer
+        .write_all(&buffer)
+        .map_err(|err| SnapshotError::Serde(err.to_string()))?;
+
+    Ok(buffer.len())
+    // bincode::serialize_into(writer, data).map_err(|err| SnapshotError::Serde(err.to_string()))
 }
 
 #[cfg(test)]

src/vmm/src/vmm_config/boot_source.rs

@@ -132,7 +132,7 @@ pub(crate) mod tests {
         };
 
         let mut snapshot_data = vec![0u8; 1000];
-        Snapshot::serialize(&mut snapshot_data.as_mut_slice(), &boot_src_cfg).unwrap();
+        Snapshot::serialize::<&BootSourceConfig>(&mut snapshot_data.as_mut_slice(), &boot_src_cfg).unwrap();
         let restored_boot_cfg = Snapshot::deserialize(&mut snapshot_data.as_slice()).unwrap();
         assert_eq!(boot_src_cfg, restored_boot_cfg);
     }

src/vmm/src/vstate/memory.rs

@@ -603,7 +603,7 @@ mod tests {
     fn check_serde(guest_memory: &GuestMemoryMmap) {
         let mut snapshot_data = vec![0u8; 10000];
         let original_state = guest_memory.describe();
-        Snapshot::serialize(&mut snapshot_data.as_mut_slice(), &original_state).unwrap();
+        Snapshot::serialize::<GuestMemoryState>(&mut snapshot_data.as_mut_slice(), &original_state).unwrap();
         let restored_state = Snapshot::deserialize(&mut snapshot_data.as_slice()).unwrap();
         assert_eq!(original_state, restored_state);
     }

src/vmm/src/vstate/vm.rs

@@ -579,7 +579,7 @@ pub(crate) mod tests {
         let (mut vm, _) = setup_vm(0x1000);
         vm.setup_irqchip().unwrap();
         let state = vm.save_state().unwrap();
-        Snapshot::serialize(&mut snapshot_data.as_mut_slice(), &state).unwrap();
+        Snapshot::serialize::<&mut [u8], VmState>(&mut snapshot_data.as_mut_slice(), &state).unwrap();
         let restored_state: VmState = Snapshot::deserialize(&mut snapshot_data.as_slice()).unwrap();
 
         vm.restore_state(&restored_state).unwrap();