refactor: cleanup vmm::snapshot module

src/firecracker/src/main.rs

@@ -28,7 +28,7 @@ use vmm::persist::SNAPSHOT_VERSION;
 use vmm::resources::VmResources;
 use vmm::seccomp::BpfThreadMap;
 use vmm::signal_handler::register_signal_handlers;
-use vmm::snapshot::{Snapshot, SnapshotError};
+use vmm::snapshot::{SnapshotError, get_format_version};
 use vmm::vmm_config::instance_info::{InstanceInfo, VmState};
 use vmm::vmm_config::metrics::{MetricsConfig, MetricsConfigError, init_metrics};
 use vmm::{EventManager, FcExitCode, HTTP_MAX_PAYLOAD_SIZE};
@@ -546,8 +546,8 @@ 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)?;
+    let data_format_version =
+        get_format_version(&mut snapshot_reader).map_err(SnapshotVersionError::SnapshotVersion)?;
 
     println!("v{}", data_format_version);
     Ok(())

src/snapshot-editor/src/edit_vmstate.rs

@@ -51,9 +51,9 @@ fn edit(
     output_path: &PathBuf,
     f: impl Fn(MicrovmState) -> Result<MicrovmState, EditVmStateError>,
 ) -> Result<(), EditVmStateError> {
-    let (microvm_state, version) = open_vmstate(vmstate_path)?;
-    let microvm_state = f(microvm_state)?;
-    save_vmstate(microvm_state, output_path, version)?;
+    let snapshot = open_vmstate(vmstate_path)?;
+    let microvm_state = f(snapshot.data)?;
+    save_vmstate(microvm_state, output_path)?;
     Ok(())
 }
 
@@ -78,7 +78,7 @@ fn remove_regs(
         }
         vcpu_state.regs = new_regs;
         for (reg, removed) in remove_regs.iter().zip(removed.iter()) {
-            print!("Regsiter {reg:#x}: ");
+            print!("Register {reg:#x}: ");
             match removed {
                 true => println!("removed"),
                 false => println!("not present"),

src/snapshot-editor/src/info.rs

@@ -4,8 +4,8 @@
 use std::path::PathBuf;
 
 use clap::Subcommand;
-use semver::Version;
 use vmm::persist::MicrovmState;
+use vmm::snapshot::Snapshot;
 
 use crate::utils::*;
 
@@ -50,27 +50,27 @@ pub fn info_vmstate_command(command: InfoVmStateSubCommand) -> Result<(), InfoVm
 
 fn info(
     vmstate_path: &PathBuf,
-    f: impl Fn(&MicrovmState, Version) -> Result<(), InfoVmStateError>,
+    f: impl Fn(&Snapshot<MicrovmState>) -> Result<(), InfoVmStateError>,
 ) -> Result<(), InfoVmStateError> {
-    let (vmstate, version) = open_vmstate(vmstate_path)?;
-    f(&vmstate, version)?;
+    let snapshot = open_vmstate(vmstate_path)?;
+    f(&snapshot)?;
     Ok(())
 }
 
-fn info_version(_: &MicrovmState, version: Version) -> Result<(), InfoVmStateError> {
-    println!("v{version}");
+fn info_version(snapshot: &Snapshot<MicrovmState>) -> Result<(), InfoVmStateError> {
+    println!("v{}", snapshot.version());
     Ok(())
 }
 
-fn info_vcpu_states(state: &MicrovmState, _: Version) -> Result<(), InfoVmStateError> {
-    for (i, state) in state.vcpu_states.iter().enumerate() {
+fn info_vcpu_states(snapshot: &Snapshot<MicrovmState>) -> Result<(), InfoVmStateError> {
+    for (i, state) in snapshot.data.vcpu_states.iter().enumerate() {
         println!("vcpu {i}:");
         println!("{state:#?}");
     }
     Ok(())
 }
 
-fn info_vmstate(vmstate: &MicrovmState, _version: Version) -> Result<(), InfoVmStateError> {
-    println!("{vmstate:#?}");
+fn info_vmstate(snapshot: &Snapshot<MicrovmState>) -> Result<(), InfoVmStateError> {
+    println!("{:#?}", snapshot.data);
     Ok(())
 }

src/snapshot-editor/src/utils.rs

@@ -4,19 +4,15 @@
 use std::fs::{File, OpenOptions};
 use std::path::PathBuf;
 
-use semver::Version;
 use vmm::persist::MicrovmState;
 use vmm::snapshot::Snapshot;
-use vmm::utils::u64_to_usize;
 
 // Some errors are only used in aarch64 code
 #[allow(unused)]
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum UtilsError {
     /// Can not open snapshot file: {0}
     VmStateFileOpen(std::io::Error),
-    /// Can not retrieve metadata for snapshot file: {0}
-    VmStateFileMeta(std::io::Error),
     /// Can not load snapshot: {0}
     VmStateLoad(vmm::snapshot::SnapshotError),
     /// Can not open output file: {0}
@@ -26,29 +22,23 @@ pub enum UtilsError {
 }
 
 #[allow(unused)]
-pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<(MicrovmState, Version), UtilsError> {
+pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<Snapshot<MicrovmState>, 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());
-    Snapshot::load(&mut snapshot_reader, snapshot_len).map_err(UtilsError::VmStateLoad)
+    Snapshot::load(&mut snapshot_reader).map_err(UtilsError::VmStateLoad)
 }
 
 // This method is used only in aarch64 code so far
 #[allow(unused)]
-pub fn save_vmstate(
-    microvm_state: MicrovmState,
-    output_path: &PathBuf,
-    version: Version,
-) -> Result<(), UtilsError> {
+pub fn save_vmstate(microvm_state: MicrovmState, output_path: &PathBuf) -> Result<(), UtilsError> {
     let mut output_file = OpenOptions::new()
         .create(true)
         .write(true)
         .truncate(true)
         .open(output_path)
         .map_err(UtilsError::OutputFileOpen)?;
-    let mut snapshot = Snapshot::new(version);
+    let mut snapshot = Snapshot::new(microvm_state);
     snapshot
-        .save(&mut output_file, &microvm_state)
+        .save(&mut output_file)
         .map_err(UtilsError::VmStateSave)?;
     Ok(())
 }

src/vmm/src/arch/aarch64/regs.rs

@@ -546,8 +546,8 @@ mod tests {
 
         let mut buf = vec![0; 10000];
 
-        Snapshot::serialize(&mut buf.as_mut_slice(), &v).unwrap();
-        let restored: Aarch64RegisterVec = Snapshot::deserialize(&mut buf.as_slice()).unwrap();
+        Snapshot::new(&v).save(&mut buf.as_mut_slice()).unwrap();
+        let restored: Aarch64RegisterVec = Snapshot::load(&mut buf.as_slice()).unwrap().data;
 
         for (old, new) in v.iter().zip(restored.iter()) {
             assert_eq!(old, new);
@@ -572,11 +572,11 @@ mod tests {
 
         let mut buf = vec![0; 10000];
 
-        Snapshot::serialize(&mut buf.as_mut_slice(), &v).unwrap();
+        Snapshot::new(&v).save(&mut buf.as_mut_slice()).unwrap();
 
         // Total size of registers according IDs are 16 + 16 = 32,
         // but actual data size is 8 + 16 = 24.
-        Snapshot::deserialize::<_, Aarch64RegisterVec>(&mut buf.as_slice()).unwrap_err();
+        Snapshot::<Aarch64RegisterVec>::load(&mut buf.as_slice()).unwrap_err();
     }
 
     #[test]
@@ -595,10 +595,10 @@ mod tests {
 
         let mut buf = vec![0; 10000];
 
-        Snapshot::serialize(&mut buf.as_mut_slice(), &v).unwrap();
+        Snapshot::new(v).save(&mut buf.as_mut_slice()).unwrap();
 
         // 4096 bit wide registers are not supported.
-        Snapshot::deserialize::<_, Aarch64RegisterVec>(&mut buf.as_slice()).unwrap_err();
+        Snapshot::<Aarch64RegisterVec>::load(&mut buf.as_slice()).unwrap_err();
     }
 
     #[test]

src/vmm/src/arch/x86_64/vm.rs

@@ -318,8 +318,10 @@ mod tests {
         let (_, mut vm) = setup_vm_with_memory(0x1000);
         vm.setup_irqchip().unwrap();
         let state = vm.save_state().unwrap();
-        Snapshot::serialize(&mut snapshot_data.as_mut_slice(), &state).unwrap();
-        let restored_state: VmState = Snapshot::deserialize(&mut snapshot_data.as_slice()).unwrap();
+        Snapshot::new(state)
+            .save(&mut snapshot_data.as_mut_slice())
+            .unwrap();
+        let restored_state: VmState = Snapshot::load(&mut snapshot_data.as_slice()).unwrap().data;
 
         vm.restore_state(&restored_state).unwrap();
     }

src/vmm/src/device_manager/pci_mngr.rs

@@ -645,7 +645,9 @@ mod tests {
             let entropy_config = EntropyDeviceConfig::default();
             insert_entropy_device(&mut vmm, &mut cmdline, &mut event_manager, entropy_config);
 
-            Snapshot::serialize(&mut buf.as_mut_slice(), &vmm.device_manager.save()).unwrap();
+            Snapshot::new(vmm.device_manager.save())
+                .save(&mut buf.as_mut_slice())
+                .unwrap();
         }
 
         tmp_sock_file.remove().unwrap();
@@ -657,7 +659,7 @@ mod tests {
         // object and calling default_vmm() is the easiest way to create one.
         let vmm = default_vmm();
         let device_manager_state: device_manager::DevicesState =
-            Snapshot::deserialize(&mut buf.as_slice()).unwrap();
+            Snapshot::load(&mut buf.as_slice()).unwrap().data;
         let vm_resources = &mut VmResources::default();
         let restore_args = PciDevicesConstructorArgs {
             vm: vmm.vm.clone(),

src/vmm/src/device_manager/persist.rs

@@ -676,15 +676,17 @@ mod tests {
             let entropy_config = EntropyDeviceConfig::default();
             insert_entropy_device(&mut vmm, &mut cmdline, &mut event_manager, entropy_config);
 
-            Snapshot::serialize(&mut buf.as_mut_slice(), &vmm.device_manager.save()).unwrap();
+            Snapshot::new(vmm.device_manager.save())
+                .save(&mut buf.as_mut_slice())
+                .unwrap();
         }
 
         tmp_sock_file.remove().unwrap();
 
         let mut event_manager = EventManager::new().expect("Unable to create EventManager");
         let vmm = default_vmm();
         let device_manager_state: device_manager::DevicesState =
-            Snapshot::deserialize(&mut buf.as_slice()).unwrap();
+            Snapshot::load(&mut buf.as_slice()).unwrap().data;
         let vm_resources = &mut VmResources::default();
         let restore_args = MMIODevManagerConstructorArgs {
             mem: vmm.vm.guest_memory(),

src/vmm/src/devices/virtio/balloon/persist.rs

@@ -187,15 +187,17 @@ mod tests {
         // Create and save the balloon device.
         let balloon = Balloon::new(0x42, false, 2, false).unwrap();
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &balloon.save()).unwrap();
+        Snapshot::new(balloon.save())
+            .save(&mut mem.as_mut_slice())
+            .unwrap();
 
         // Deserialize and restore the balloon device.
         let restored_balloon = Balloon::restore(
             BalloonConstructorArgs {
                 mem: guest_mem,
                 restored_from_file: true,
             },
-            &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
+            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
         )
         .unwrap();
 

src/vmm/src/devices/virtio/block/virtio/persist.rs

@@ -171,7 +171,9 @@ mod tests {
         // Save the block device.
         let mut mem = vec![0; 4096];
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &block.save()).unwrap();
+        Snapshot::new(block.save())
+            .save(&mut mem.as_mut_slice())
+            .unwrap();
     }
 
     #[test]
@@ -214,12 +216,14 @@ mod tests {
         // Save the block device.
         let mut mem = vec![0; 4096];
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &block.save()).unwrap();
+        Snapshot::new(block.save())
+            .save(&mut mem.as_mut_slice())
+            .unwrap();
 
         // Restore the block device.
         let restored_block = VirtioBlock::restore(
             BlockConstructorArgs { mem: guest_mem },
-            &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
+            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
         )
         .unwrap();
 

src/vmm/src/devices/virtio/net/persist.rs

@@ -153,7 +153,9 @@ mod tests {
 
         // Create and save the net device.
         {
-            Snapshot::serialize(&mut mem.as_mut_slice(), &net.save()).unwrap();
+            Snapshot::new(net.save())
+                .save(&mut mem.as_mut_slice())
+                .unwrap();
 
             // Save some fields that we want to check later.
             id = net.id.clone();
@@ -173,7 +175,7 @@ mod tests {
                     mem: guest_mem,
                     mmds: mmds_ds,
                 },
-                &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
+                &Snapshot::load(&mut mem.as_slice()).unwrap().data,
             ) {
                 Ok(restored_net) => {
                     // Test that virtio specific fields are the same.

src/vmm/src/devices/virtio/persist.rs

@@ -358,14 +358,16 @@ mod tests {
 
         let mut bytes = vec![0; 4096];
 
-        Snapshot::serialize(&mut bytes.as_mut_slice(), &queue.save()).unwrap();
+        Snapshot::new(queue.save())
+            .save(&mut bytes.as_mut_slice())
+            .unwrap();
 
         let ca = QueueConstructorArgs {
             mem,
             is_activated: true,
         };
         let restored_queue =
-            Queue::restore(ca, &Snapshot::deserialize(&mut bytes.as_slice()).unwrap()).unwrap();
+            Queue::restore(ca, &Snapshot::load(&mut bytes.as_slice()).unwrap().data).unwrap();
 
         assert_eq!(restored_queue, queue);
     }
@@ -376,9 +378,9 @@ mod tests {
         let mut mem = vec![0; 4096];
 
         let state = VirtioDeviceState::from_device(&dummy);
-        Snapshot::serialize(&mut mem.as_mut_slice(), &state).unwrap();
+        Snapshot::new(&state).save(&mut mem.as_mut_slice()).unwrap();
 
-        let restored_state: VirtioDeviceState = Snapshot::deserialize(&mut mem.as_slice()).unwrap();
+        let restored_state: VirtioDeviceState = Snapshot::load(&mut mem.as_slice()).unwrap().data;
         assert_eq!(restored_state, state);
     }
 
@@ -405,7 +407,9 @@ mod tests {
     ) {
         let mut buf = vec![0; 4096];
 
-        Snapshot::serialize(&mut buf.as_mut_slice(), &mmio_transport.save()).unwrap();
+        Snapshot::new(mmio_transport.save())
+            .save(&mut buf.as_mut_slice())
+            .unwrap();
 
         let restore_args = MmioTransportConstructorArgs {
             mem,
@@ -415,7 +419,7 @@ mod tests {
         };
         let restored_mmio_transport = MmioTransport::restore(
             restore_args,
-            &Snapshot::deserialize(&mut buf.as_slice()).unwrap(),
+            &Snapshot::load(&mut buf.as_slice()).unwrap().data,
         )
         .unwrap();
 

src/vmm/src/devices/virtio/rng/persist.rs

@@ -85,12 +85,14 @@ mod tests {
         let mut mem = vec![0u8; 4096];
         let entropy = Entropy::new(RateLimiter::default()).unwrap();
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &entropy.save()).unwrap();
+        Snapshot::new(entropy.save())
+            .save(&mut mem.as_mut_slice())
+            .unwrap();
 
         let guest_mem = create_virtio_mem();
         let restored = Entropy::restore(
             EntropyConstructorArgs { mem: guest_mem },
-            &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
+            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
         )
         .unwrap();
 

src/vmm/src/devices/virtio/vsock/persist.rs

@@ -178,9 +178,9 @@ pub(crate) mod tests {
             frontend: ctx.device.save(),
         };
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &state).unwrap();
+        Snapshot::new(&state).save(&mut mem.as_mut_slice()).unwrap();
 
-        let restored_state: VsockState = Snapshot::deserialize(&mut mem.as_slice()).unwrap();
+        let restored_state: VsockState = Snapshot::load(&mut mem.as_slice()).unwrap().data;
         let mut restored_device = Vsock::restore(
             VsockConstructorArgs {
                 mem: ctx.mem.clone(),

src/vmm/src/mmds/persist.rs

@@ -62,11 +62,13 @@ mod tests {
 
         let mut mem = vec![0; 4096];
 
-        Snapshot::serialize(&mut mem.as_mut_slice(), &ns.save()).unwrap();
+        Snapshot::new(ns.save())
+            .save(&mut mem.as_mut_slice())
+            .unwrap();
 
         let restored_ns = MmdsNetworkStack::restore(
             Arc::new(Mutex::new(Mmds::default())),
-            &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
+            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
         )
         .unwrap();