Snapshot refactoring

Author: gjkeller

src/vmm/src/snapshot/mod.rs

@@ -31,6 +31,7 @@ use std::io::{Read, Write};
 
 use bincode::config;
 use bincode::config::{Configuration, Fixint, Limit, LittleEndian};
+use bincode::error::{DecodeError, EncodeError};
 use semver::Version;
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
@@ -73,7 +74,7 @@ pub enum SnapshotError {
 
 /// Firecracker snapshot header
 #[derive(Debug, Serialize, Deserialize)]
-pub struct SnapshotHdr {
+struct SnapshotHdr {
     /// Magic value
     magic: u64,
     /// Snapshot data version
@@ -82,21 +83,26 @@ pub struct SnapshotHdr {
 
 impl SnapshotHdr {
     /// Create a new header for writing snapshots
-    pub fn new(version: Version) -> Self {
+    fn new(version: Version) -> Self {
         Self {
             magic: SNAPSHOT_MAGIC_ID,
             version,
         }
     }
 
     /// Load and deserialize just the header (magic + version)
-    pub fn load<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
+    fn load<R: Read>(reader: &mut R) -> Result<Self, SnapshotError> {
         let hdr: SnapshotHdr = deserialize(reader)?;
-        Ok(hdr)
+        if hdr.magic != SNAPSHOT_MAGIC_ID {
+            Err(SnapshotError::InvalidMagic(hdr.magic))
+        }
+        else {
+            Ok(hdr)
+        }
     }
 
     /// Serialize and write just the header
-    pub fn store<W: Write>(&self, writer: &mut W) -> Result<(), SnapshotError> {
+    fn store<W: Write>(&self, writer: &mut W) -> Result<(), SnapshotError> {
         serialize(writer, self)?;
         Ok(())
     }
@@ -109,44 +115,44 @@ pub struct Snapshot<Data> {
     pub data: Data,
 }
 
+// Implementations for deserializing snapshots
+// Publicly exposed functions:
+// - load_unchecked()
+//- load()
 impl<Data: DeserializeOwned + Debug> Snapshot<Data> {
-    /// Load without CRC or version‐check, but verify the magic
+    /// Load without CRC or version‐check, but verify magic via `SnapshotHdr::load`.
     pub fn load_unchecked<R: Read + Debug>(reader: &mut R) -> Result<Self, SnapshotError> {
-        let hdr: SnapshotHdr = deserialize(reader)?;
-        if hdr.magic != SNAPSHOT_MAGIC_ID {
-            return Err(SnapshotError::InvalidMagic(hdr.magic));
-        }
-        let data: Data = deserialize(reader)?;
+        // this calls `deserialize` + checks magic internally
+        let hdr: SnapshotHdr = SnapshotHdr::load(reader)?;
+        let data: Data      = deserialize(reader)?;
         Ok(Self { header: hdr, data })
     }
 
-    /// Load with CRC64 validation
-    pub fn load<R: Read + Debug>(reader: &mut R, snapshot_len: usize) -> Result<Self, SnapshotError> {
+    /// Load with CRC64 validation in one pass, using `load_unchecked` for header+data.
+    pub fn load<R: Read + Debug>(reader: &mut R) -> Result<Self, SnapshotError> {
+        // 1) Wrap in CRC reader
         let mut crc_reader = CRC64Reader::new(reader);
 
-        // Snapshot must be at least (len of CRC64)
-        let raw_snapshot_len = snapshot_len
-            .checked_sub(std::mem::size_of::<u64>())
-            .ok_or(SnapshotError::InvalidSnapshotSize)?;
-
-        let mut snapshot_buf = vec![0u8; raw_snapshot_len];
-        crc_reader
-            .read_exact(&mut snapshot_buf)
-            .map_err(|err| SnapshotError::Io(err.raw_os_error().unwrap_or(libc::EINVAL)))?;
+        // 2) Parse header + payload & magic‐check
+        let snapshot = Snapshot::load_unchecked(&mut crc_reader)?;
 
-        // Compute then read stored checksum
+        // 3) Grab the computed CRC over everything read so far
         let computed = crc_reader.checksum();
+
+        // 4) Deserialize the trailing u64 and compare
         let stored: u64 = deserialize(&mut crc_reader)?;
-        if computed != stored {
+        if stored != computed {
             return Err(SnapshotError::Crc64(computed));
         }
 
-        // Now parse header+data from the buffered bytes
-        let mut slice: &[u8] = snapshot_buf.as_slice();
-        Snapshot::load_unchecked(&mut slice)
+        Ok(snapshot)
     }
 }
 
+// Implementations for serializing snapshots
+// Publicly-exposed *methods*:
+// - save(self,...)
+// - save_with_crc(self,...)
 impl<Data: Serialize + Debug> Snapshot<Data> {
     /// Save without CRC64
     pub fn save_without_crc<W: Write + Debug>(&self, writer: &mut W) -> Result<(), SnapshotError> {
@@ -166,9 +172,12 @@ impl<Data: Serialize + Debug> Snapshot<Data> {
     }
 }
 
+// General methods for snapshots (related to serialization, see above, since an
+// instance is needed to serialize)
 impl<Data> Snapshot<Data> {
     /// Construct from a pre‐built header + payload
-    pub fn new(header: SnapshotHdr, data: Data) -> Self {
+    pub fn new(version: Version, data: Data) -> Self {
+        header = SnapshotHdr::new(version);
         Snapshot { header, data }
     }
 
@@ -178,31 +187,56 @@ impl<Data> Snapshot<Data> {
     }
 }
 
-/// Deserialize any `O: DeserializeOwned + Debug` via bincode + our config
+/// Deserialize any `O: DeserializeOwned + Debug` via bincode + our config,
 fn deserialize<T, O>(reader: &mut T) -> Result<O, SnapshotError>
 where
     T: Read,
     O: DeserializeOwned + Debug,
 {
     bincode::serde::decode_from_std_read(reader, BINCODE_CONFIG)
-        .map_err(|err| SnapshotError::Serde(err.to_string()))
+        .map_err(|err| match err {
+            // The reader hit an actual IO error.
+            DecodeError::Io { inner, .. } =>
+                SnapshotError::Io(inner.raw_os_error().unwrap_or(EIO)),
+
+            // Not enough bytes in the input for what we expected.
+            DecodeError::UnexpectedEnd { .. } |
+            DecodeError::LimitExceeded =>
+                SnapshotError::InvalidSnapshotSize,
+
+            // Anything else is a ser/de format issue.
+            other =>
+                SnapshotError::Serde(other.to_string()),
+        })
 }
 
-/// Serialize any `O: Serialize + Debug` into a Vec, write it, and return the byte‐count
+/// Serialize any `O: Serialize + Debug` into a Vec, write it, and return the byte‐count,
 fn serialize<T, O>(writer: &mut T, data: &O) -> Result<usize, SnapshotError>
 where
     T: Write,
     O: Serialize + Debug,
 {
-    // First serialize into an in‐memory buffer using our config
+    // 1) Encode into an in-memory buffer
     let mut buf = Vec::new();
     bincode::serde::encode_into_std_write(data, &mut buf, BINCODE_CONFIG)
-        .map_err(|err| SnapshotError::Serde(err.to_string()))?;
+        .map_err(|err| match err {
+            // Ran out of room while encoding
+            EncodeError::UnexpectedEnd =>
+                SnapshotError::Io(libc::EIO),
+
+            // Underlying IO failure during encode (index tells how many bytes got written)
+            EncodeError::Io { inner, .. } =>
+                SnapshotError::Io(inner.raw_os_error().unwrap_or(libc::EIO)),
+
+            // Any other encode error we surface as Serde
+            other =>
+                SnapshotError::Serde(other.to_string()),
+        })?;
 
-    // Then write it out
+    // 2) Flush that buffer to the target writer
     writer
         .write_all(&buf)
-        .map_err(|err| SnapshotError::Io(err.raw_os_error().unwrap_or(libc::EIO)))?;
+        .map_err(|io_err| SnapshotError::Io(io_err.raw_os_error().unwrap_or(libc::EIO)))?;
 
     Ok(buf.len())
 }