Rustify snapshot module

src/firecracker/src/main.rs

@@ -533,10 +533,11 @@ fn print_snapshot_data_format(snapshot_path: &str) -> Result<(), SnapshotVersion
     let mut snapshot_reader =
         File::open(snapshot_path).map_err(SnapshotVersionError::OpenSnapshot)?;
 
-    let data_format_version = Snapshot::get_format_version(&mut snapshot_reader)
-        .map_err(SnapshotVersionError::SnapshotVersion)?;
+    // TODO: Need to fix this
+    // let data_format_version = Snapshot::get_format_version(&mut snapshot_reader)
+    //     .map_err(SnapshotVersionError::SnapshotVersion)?;
 
-    println!("v{}", data_format_version);
+    // println!("v{}", data_format_version);
     Ok(())
 }
 

src/snapshot-editor/src/utils.rs

@@ -6,7 +6,7 @@ use std::path::PathBuf;
 
 use semver::Version;
 use vmm::persist::MicrovmState;
-use vmm::snapshot::Snapshot;
+use vmm::snapshot::{Snapshot, SnapshotError, SnapshotHdr};
 use vmm::utils::u64_to_usize;
 
 // Some errors are only used in aarch64 code
@@ -30,7 +30,18 @@ pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<(MicrovmState, Version),
     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());
-    Snapshot::load(&mut snapshot_reader, snapshot_len).map_err(UtilsError::VmStateLoad)
+
+    let snapshot: Result<Snapshot<MicrovmState>, UtilsError> = Snapshot::load(&mut snapshot_reader, snapshot_len).map_err(UtilsError::VmStateLoad);
+    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
@@ -46,9 +57,10 @@ pub fn save_vmstate(
         .truncate(true)
         .open(output_path)
         .map_err(UtilsError::OutputFileOpen)?;
-    let mut snapshot = Snapshot::new(version);
+    let snapshot_hdr = SnapshotHdr::new(version);
+    let snapshot = Snapshot::new(snapshot_hdr, microvm_state);
     snapshot
-        .save(&mut output_file, &microvm_state)
+        .save(&mut output_file)
         .map_err(UtilsError::VmStateSave)?;
     Ok(())
 }

src/vmm/src/persist.rs

@@ -28,7 +28,7 @@ use crate::cpu_config::x86_64::cpuid::CpuidTrait;
 use crate::device_manager::persist::{ACPIDeviceManagerState, DevicePersistError, DeviceStates};
 use crate::logger::{info, warn};
 use crate::resources::VmResources;
-use crate::snapshot::Snapshot;
+use crate::snapshot::{Snapshot, SnapshotHdr};
 use crate::utils::u64_to_usize;
 use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
@@ -188,9 +188,10 @@ fn snapshot_state_to_file(
         .open(snapshot_path)
         .map_err(|err| SnapshotBackingFile("open", err))?;
 
-    let snapshot = Snapshot::new(SNAPSHOT_VERSION);
+    let snapshot_hdr = SnapshotHdr::new(SNAPSHOT_VERSION);
+    let snapshot = Snapshot::new(snapshot_hdr, microvm_state);
     snapshot
-        .save(&mut snapshot_file, microvm_state)
+        .save(&mut snapshot_file)
         .map_err(SerializeMicrovmState)?;
     snapshot_file
         .flush()
@@ -490,15 +491,14 @@ pub enum SnapshotStateFromFileError {
 fn snapshot_state_from_file(
     snapshot_path: &Path,
 ) -> Result<MicrovmState, SnapshotStateFromFileError> {
-    let snapshot = Snapshot::new(SNAPSHOT_VERSION);
     let mut snapshot_reader =
         File::open(snapshot_path).map_err(SnapshotStateFromFileError::Open)?;
     let metadata = std::fs::metadata(snapshot_path).map_err(SnapshotStateFromFileError::Meta)?;
     let snapshot_len = u64_to_usize(metadata.len());
-    let state: MicrovmState = snapshot
-        .load_with_version_check(&mut snapshot_reader, snapshot_len)
+    let state: Snapshot<MicrovmState> = Snapshot::load(&mut snapshot_reader, snapshot_len)
         .map_err(SnapshotStateFromFileError::Load)?;
-    Ok(state)
+
+    Ok(state.data)
 }
 
 /// Error type for [`guest_memory_from_file`].
@@ -744,17 +744,23 @@ mod tests {
             vm_state: vmm.vm.save_state().unwrap(),
             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];
-        Snapshot::serialize(&mut buf.as_mut_slice(), &microvm_state).unwrap();
 
-        let restored_microvm_state: MicrovmState =
-            Snapshot::deserialize(&mut buf.as_slice()).unwrap();
+        let snapshot_hdr = SnapshotHdr::new(Version::new(1, 0, 42));
+        let snapshot = Snapshot::new(snapshot_hdr, microvm_state);
+        snapshot.save(&mut buf.as_mut_slice()).unwrap();
+
+        let restored_snapshot: Snapshot<MicrovmState> =
+            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
         )
     }