fix: use temporary buffer for loading vmstate file

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

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

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

@@ -321,7 +321,9 @@ mod tests {
         Snapshot::new(state)
             .save(&mut snapshot_data.as_mut_slice())
             .unwrap();
-        let restored_state: VmState = Snapshot::load(&mut snapshot_data.as_slice()).unwrap().data;
+        let restored_state: VmState = Snapshot::load_without_crc_check(snapshot_data.as_slice())
+            .unwrap()
+            .data;
 
         vm.restore_state(&restored_state).unwrap();
     }

src/vmm/src/device_manager/pci_mngr.rs

@@ -659,7 +659,9 @@ 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::load(&mut buf.as_slice()).unwrap().data;
+            Snapshot::load_without_crc_check(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

@@ -680,7 +680,9 @@ mod tests {
         let mut event_manager = EventManager::new().expect("Unable to create EventManager");
         let vmm = default_vmm();
         let device_manager_state: device_manager::DevicesState =
-            Snapshot::load(&mut buf.as_slice()).unwrap().data;
+            Snapshot::load_without_crc_check(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

@@ -196,7 +196,9 @@ mod tests {
                 mem: guest_mem,
                 restored_from_file: true,
             },
-            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 

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

@@ -223,7 +223,9 @@ mod tests {
         // Restore the block device.
         let restored_block = VirtioBlock::restore(
             BlockConstructorArgs { mem: guest_mem },
-            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 

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

@@ -175,7 +175,9 @@ mod tests {
                     mem: guest_mem,
                     mmds: mmds_ds,
                 },
-                &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+                &Snapshot::load_without_crc_check(mem.as_slice())
+                    .unwrap()
+                    .data,
             ) {
                 Ok(restored_net) => {
                     // Test that virtio specific fields are the same.

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

@@ -366,8 +366,13 @@ mod tests {
             mem,
             is_activated: true,
         };
-        let restored_queue =
-            Queue::restore(ca, &Snapshot::load(&mut bytes.as_slice()).unwrap().data).unwrap();
+        let restored_queue = Queue::restore(
+            ca,
+            &Snapshot::load_without_crc_check(bytes.as_slice())
+                .unwrap()
+                .data,
+        )
+        .unwrap();
 
         assert_eq!(restored_queue, queue);
     }
@@ -380,7 +385,9 @@ mod tests {
         let state = VirtioDeviceState::from_device(&dummy);
         Snapshot::new(&state).save(&mut mem.as_mut_slice()).unwrap();
 
-        let restored_state: VirtioDeviceState = Snapshot::load(&mut mem.as_slice()).unwrap().data;
+        let restored_state: VirtioDeviceState = Snapshot::load_without_crc_check(mem.as_slice())
+            .unwrap()
+            .data;
         assert_eq!(restored_state, state);
     }
 
@@ -419,7 +426,9 @@ mod tests {
         };
         let restored_mmio_transport = MmioTransport::restore(
             restore_args,
-            &Snapshot::load(&mut buf.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(buf.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 

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

@@ -92,7 +92,9 @@ mod tests {
         let guest_mem = create_virtio_mem();
         let restored = Entropy::restore(
             EntropyConstructorArgs { mem: guest_mem },
-            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 

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

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

src/vmm/src/mmds/persist.rs

@@ -68,7 +68,9 @@ mod tests {
 
         let restored_ns = MmdsNetworkStack::restore(
             Arc::new(Mutex::new(Mmds::default())),
-            &Snapshot::load(&mut mem.as_slice()).unwrap().data,
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
         )
         .unwrap();
 

src/vmm/src/persist.rs

@@ -676,8 +676,9 @@ mod tests {
             .save(&mut buf.as_mut_slice())
             .unwrap();
 
-        let restored_microvm_state: MicrovmState =
-            Snapshot::load(&mut buf.as_slice()).unwrap().data;
+        let restored_microvm_state: MicrovmState = Snapshot::load_without_crc_check(buf.as_slice())
+            .unwrap()
+            .data;
 
         assert_eq!(restored_microvm_state.vm_info, microvm_state.vm_info);
         assert_eq!(

src/vmm/src/rate_limiter/persist.rs

@@ -120,8 +120,13 @@ mod tests {
             .save(&mut mem.as_mut_slice())
             .unwrap();
 
-        let restored_tb =
-            TokenBucket::restore((), &Snapshot::load(&mut mem.as_slice()).unwrap().data).unwrap();
+        let restored_tb = TokenBucket::restore(
+            (),
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
+        )
+        .unwrap();
         assert!(tb.partial_eq(&restored_tb));
     }
 
@@ -197,8 +202,13 @@ mod tests {
         Snapshot::new(rate_limiter.save())
             .save(&mut mem.as_mut_slice())
             .unwrap();
-        let restored_rate_limiter =
-            RateLimiter::restore((), &Snapshot::load(&mut mem.as_slice()).unwrap().data).unwrap();
+        let restored_rate_limiter = RateLimiter::restore(
+            (),
+            &Snapshot::load_without_crc_check(mem.as_slice())
+                .unwrap()
+                .data,
+        )
+        .unwrap();
 
         assert!(
             rate_limiter

src/vmm/src/snapshot/mod.rs

@@ -26,7 +26,7 @@
 pub mod crc;
 mod persist;
 use std::fmt::Debug;
-use std::io::{Read, Write};
+use std::io::{Read, Seek, SeekFrom, Write};
 
 use bincode::config;
 use bincode::config::{Configuration, Fixint, Limit, LittleEndian};
@@ -38,7 +38,7 @@ use serde::{Deserialize, Serialize};
 use crate::persist::SNAPSHOT_VERSION;
 use crate::snapshot::crc::{CRC64Reader, CRC64Writer};
 pub use crate::snapshot::persist::Persist;
-use crate::utils::mib_to_bytes;
+use crate::utils::{mib_to_bytes, u64_to_usize};
 
 #[cfg(target_arch = "x86_64")]
 const SNAPSHOT_MAGIC_ID: u64 = 0x0710_1984_8664_0000u64;
@@ -64,10 +64,14 @@ pub enum SnapshotError {
     InvalidFormatVersion(Version),
     /// Magic value does not match arch: {0}
     InvalidMagic(u64),
+    /// Snapshot file is not long enough to even contain the CRC
+    TooShort,
     /// An error occured during bincode encoding: {0}
     Encode(#[from] EncodeError),
     /// An error occured during bincode decoding: {0}
     Decode(#[from] DecodeError),
+    /// IO Error: {0}
+    Io(#[from] std::io::Error),
 }
 
 fn serialize<S: Serialize, W: Write>(data: &S, write: &mut W) -> Result<(), SnapshotError> {
@@ -86,8 +90,8 @@ struct SnapshotHdr {
 }
 
 impl SnapshotHdr {
-    fn load<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
-        let hdr: SnapshotHdr = bincode::serde::decode_from_std_read(reader, BINCODE_CONFIG)?;
+    fn load(buf: &mut &[u8]) -> Result<Self, SnapshotError> {
+        let (hdr, bytes_read) = bincode::serde::decode_from_slice::<Self, _>(buf, BINCODE_CONFIG)?;
 
         if hdr.magic != SNAPSHOT_MAGIC_ID {
             return Err(SnapshotError::InvalidMagic(hdr.magic));
@@ -98,6 +102,8 @@ impl SnapshotHdr {
             return Err(SnapshotError::InvalidFormatVersion(hdr.version));
         }
 
+        *buf = &buf[bytes_read..];
+
         Ok(hdr)
     }
 }
@@ -138,17 +144,27 @@ impl<Data> Snapshot<Data> {
 }
 
 impl<Data: DeserializeOwned> Snapshot<Data> {
-    fn load_without_crc_check<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
-        let header = SnapshotHdr::load(reader)?;
-        let data = bincode::serde::decode_from_std_read(reader, BINCODE_CONFIG)?;
+    pub(crate) fn load_without_crc_check(mut buf: &[u8]) -> Result<Self, SnapshotError> {
+        let header = SnapshotHdr::load(&mut buf)?;
+        let data = bincode::serde::decode_from_slice(buf, BINCODE_CONFIG)?.0;
         Ok(Self { header, data })
     }
 
     /// Loads a snapshot from the given [`Read`] instance, performing all validations
     /// (CRC, snapshot magic value, snapshot version).
-    pub fn load<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
+    pub fn load<R: Read + Seek>(reader: &mut R) -> Result<Self, SnapshotError> {
+        let snapshot_size = reader.seek(SeekFrom::End(0))?;
+        reader.seek(SeekFrom::Start(0))?;
+        // dont read the CRC yet.
+        let mut buf = vec![
+            0;
+            u64_to_usize(snapshot_size)
+                .checked_sub(size_of::<u64>())
+                .ok_or(SnapshotError::TooShort)?
+        ];
         let mut crc_reader = CRC64Reader::new(reader);
-        let snapshot = Self::load_without_crc_check(&mut crc_reader)?;
+        crc_reader.read_exact(buf.as_mut_slice())?;
+        let snapshot = Self::load_without_crc_check(buf.as_slice())?;
         let computed_checksum = crc_reader.checksum();
         let stored_checksum: u64 =
             bincode::serde::decode_from_std_read(&mut crc_reader.reader, BINCODE_CONFIG)?;
@@ -171,7 +187,20 @@ impl<Data: Serialize> Snapshot<Data> {
 
 #[cfg(test)]
 mod tests {
+    use vmm_sys_util::tempfile::TempFile;
+
     use super::*;
+    use crate::persist::MicrovmState;
+
+    #[test]
+    fn test_snapshot_restore() {
+        let state = MicrovmState::default();
+        let file = TempFile::new().unwrap();
+
+        Snapshot::new(state).save(&mut file.as_file()).unwrap();
+        file.as_file().seek(SeekFrom::Start(0)).unwrap();
+        Snapshot::<MicrovmState>::load(&mut file.as_file()).unwrap();
+    }
 
     #[test]
     fn test_parse_version_from_file() {
@@ -199,10 +228,16 @@ mod tests {
             }
         }
 
+        impl Seek for BadReader {
+            fn seek(&mut self, _pos: SeekFrom) -> std::io::Result<u64> {
+                Ok(9) // needs to be long enough to prevent to have a CRC
+            }
+        }
+
         let mut reader = BadReader {};
 
         assert!(
-            matches!(Snapshot::<()>::load(&mut reader), Err(SnapshotError::Decode(DecodeError::Io {inner, ..})) if inner.kind() == std::io::ErrorKind::InvalidInput)
+            matches!(Snapshot::<()>::load(&mut reader), Err(SnapshotError::Io(inner)) if inner.kind() == std::io::ErrorKind::InvalidInput)
         );
     }
 
@@ -239,7 +274,7 @@ mod tests {
         serialize(&snapshot, &mut data.as_mut_slice()).unwrap();
 
         assert!(matches!(
-            Snapshot::<()>::load(&mut data.as_slice()),
+            Snapshot::<()>::load(&mut std::io::Cursor::new(data.as_slice())),
             Err(SnapshotError::Crc64(_))
         ));
     }
@@ -254,7 +289,7 @@ mod tests {
         snapshot.save(&mut data.as_mut_slice()).unwrap();
 
         assert!(matches!(
-            Snapshot::<()>::load(&mut data.as_slice()),
+            Snapshot::<()>::load_without_crc_check(data.as_slice()),
             Err(SnapshotError::InvalidFormatVersion(v)) if v.major == SNAPSHOT_VERSION.major + 1
         ));
 
@@ -263,36 +298,36 @@ mod tests {
         snapshot.header.version.minor = SNAPSHOT_VERSION.minor + 1;
         snapshot.save(&mut data.as_mut_slice()).unwrap();
         assert!(matches!(
-            Snapshot::<()>::load(&mut data.as_slice()),
+            Snapshot::<()>::load_without_crc_check(data.as_slice()),
             Err(SnapshotError::InvalidFormatVersion(v)) if v.minor == SNAPSHOT_VERSION.minor + 1
         ));
 
         // But we can support minor versions smaller or equal to ours. We also support
         // all patch versions within our supported major.minor version.
         let snapshot = Snapshot::new(());
         snapshot.save(&mut data.as_mut_slice()).unwrap();
-        Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+        Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
 
         if SNAPSHOT_VERSION.minor != 0 {
             let mut snapshot = Snapshot::new(());
             snapshot.header.version.minor = SNAPSHOT_VERSION.minor - 1;
             snapshot.save(&mut data.as_mut_slice()).unwrap();
-            Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+            Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
         }
 
         let mut snapshot = Snapshot::new(());
         snapshot.header.version.patch = 0;
         snapshot.save(&mut data.as_mut_slice()).unwrap();
-        Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+        Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
 
         let mut snapshot = Snapshot::new(());
         snapshot.header.version.patch = SNAPSHOT_VERSION.patch + 1;
         snapshot.save(&mut data.as_mut_slice()).unwrap();
-        Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+        Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
 
         let mut snapshot = Snapshot::new(());
         snapshot.header.version.patch = 1024;
         snapshot.save(&mut data.as_mut_slice()).unwrap();
-        Snapshot::<()>::load(&mut data.as_slice()).unwrap();
+        Snapshot::<()>::load_without_crc_check(data.as_slice()).unwrap();
     }
 }