Separate out splitstream writer into its own struct

Separating out the writer allows us to abstract the parallelism
of writing the splitstream.

Signed-off-by: Pragyan Poudyal <[email protected]>

Author: Johan-Liebert1

src/lib.rs

@@ -13,6 +13,7 @@ pub mod repository;
 pub mod selabel;
 pub mod splitstream;
 pub mod util;
+pub mod zstd_encoder;
 
 /// All files that contain 64 or fewer bytes (size <= INLINE_CONTENT_MAX) should be stored inline
 /// in the erofs image (and also in splitstreams).  All files with 65 or more bytes (size > MAX)

src/repository.rs

@@ -171,7 +171,7 @@ impl Repository {
         Ok(result)
     }
 
-    fn format_object_path(id: &Sha256HashValue) -> String {
+    pub fn format_object_path(id: &Sha256HashValue) -> String {
         format!("objects/{:02x}/{}", id[0], hex::encode(&id[1..]))
     }
 
@@ -236,7 +236,7 @@ impl Repository {
         writer: SplitStreamWriter,
         reference: Option<&str>,
     ) -> Result<Sha256HashValue> {
-        let Some((.., ref sha256)) = writer.sha256 else {
+        let Some((.., ref sha256)) = writer.get_sha_builder() else {
             bail!("Writer doesn't have sha256 enabled");
         };
         let stream_path = format!("streams/{}", hex::encode(sha256));

src/splitstream.rs

@@ -6,13 +6,14 @@
 use std::io::{BufReader, Read, Write};
 
 use anyhow::{bail, Result};
-use sha2::{Digest, Sha256};
-use zstd::stream::{read::Decoder, write::Encoder};
+use sha2::Sha256;
+use zstd::stream::read::Decoder;
 
 use crate::{
     fsverity::{FsVerityHashValue, Sha256HashValue},
     repository::Repository,
     util::read_exactish,
+    zstd_encoder::ZstdWriter,
 };
 
 #[derive(Debug)]
@@ -60,9 +61,8 @@ impl DigestMap {
 
 pub struct SplitStreamWriter<'a> {
     repo: &'a Repository,
-    inline_content: Vec<u8>,
-    writer: Encoder<'a, Vec<u8>>,
-    pub sha256: Option<(Sha256, Sha256HashValue)>,
+    pub(crate) inline_content: Vec<u8>,
+    writer: ZstdWriter,
 }
 
 impl std::fmt::Debug for SplitStreamWriter<'_> {
@@ -71,7 +71,6 @@ impl std::fmt::Debug for SplitStreamWriter<'_> {
         f.debug_struct("SplitStreamWriter")
             .field("repo", &self.repo)
             .field("inline_content", &self.inline_content)
-            .field("sha256", &self.sha256)
             .finish()
     }
 }
@@ -82,85 +81,46 @@ impl SplitStreamWriter<'_> {
         refs: Option<DigestMap>,
         sha256: Option<Sha256HashValue>,
     ) -> SplitStreamWriter {
-        // SAFETY: we surely can't get an error writing the header to a Vec<u8>
-        let mut writer = Encoder::new(vec![], 0).unwrap();
-
-        match refs {
-            Some(DigestMap { map }) => {
-                writer.write_all(&(map.len() as u64).to_le_bytes()).unwrap();
-                for ref entry in map {
-                    writer.write_all(&entry.body).unwrap();
-                    writer.write_all(&entry.verity).unwrap();
-                }
-            }
-            None => {
-                writer.write_all(&0u64.to_le_bytes()).unwrap();
-            }
-        }
-
         SplitStreamWriter {
             repo,
             inline_content: vec![],
-            writer,
-            sha256: sha256.map(|x| (Sha256::new(), x)),
+            writer: ZstdWriter::new(sha256, refs),
         }
     }
 
-    fn write_fragment(writer: &mut impl Write, size: usize, data: &[u8]) -> Result<()> {
-        writer.write_all(&(size as u64).to_le_bytes())?;
-        Ok(writer.write_all(data)?)
+    pub fn get_sha_builder(&self) -> &Option<(Sha256, Sha256HashValue)> {
+        &self.writer.sha256_builder
     }
 
     /// flush any buffered inline data, taking new_value as the new value of the buffer
     fn flush_inline(&mut self, new_value: Vec<u8>) -> Result<()> {
-        if !self.inline_content.is_empty() {
-            SplitStreamWriter::write_fragment(
-                &mut self.writer,
-                self.inline_content.len(),
-                &self.inline_content,
-            )?;
-            self.inline_content = new_value;
-        }
+        self.writer.flush_inline(&self.inline_content)?;
+        self.inline_content = new_value;
         Ok(())
     }
 
     /// really, "add inline content to the buffer"
     /// you need to call .flush_inline() later
     pub fn write_inline(&mut self, data: &[u8]) {
-        if let Some((ref mut sha256, ..)) = self.sha256 {
-            sha256.update(data);
-        }
+        self.writer.update_sha(data);
         self.inline_content.extend(data);
     }
 
-    /// write a reference to external data to the stream.  If the external data had padding in the
-    /// stream which is not stored in the object then pass it here as well and it will be stored
-    /// inline after the reference.
-    fn write_reference(&mut self, reference: Sha256HashValue, padding: Vec<u8>) -> Result<()> {
-        // Flush the inline data before we store the external reference.  Any padding from the
-        // external data becomes the start of a new inline block.
-        self.flush_inline(padding)?;
+    pub fn write_external(&mut self, data: &[u8], padding: Vec<u8>) -> Result<()> {
+        let id = self.repo.ensure_object(&data)?;
 
-        SplitStreamWriter::write_fragment(&mut self.writer, 0, &reference)
-    }
+        self.writer.update_sha(data);
+        self.writer.update_sha(&padding);
+        self.writer.flush_inline(&padding)?;
 
-    pub fn write_external(&mut self, data: &[u8], padding: Vec<u8>) -> Result<()> {
-        if let Some((ref mut sha256, ..)) = self.sha256 {
-            sha256.update(data);
-            sha256.update(&padding);
-        }
-        let id = self.repo.ensure_object(data)?;
-        self.write_reference(id, padding)
+        self.writer.write_fragment(0, &id)?;
+        Ok(())
     }
 
     pub fn done(mut self) -> Result<Sha256HashValue> {
         self.flush_inline(vec![])?;
 
-        if let Some((context, expected)) = self.sha256 {
-            if Into::<Sha256HashValue>::into(context.finalize()) != expected {
-                bail!("Content doesn't have expected SHA256 hash value!");
-            }
-        }
+        self.writer.finalize_sha256_builder()?;
 
         self.repo.ensure_object(&self.writer.finish()?)
     }

src/zstd_encoder.rs

@@ -0,0 +1,92 @@
+use std::io::{self, Write};
+
+use sha2::{Digest, Sha256};
+
+use anyhow::{bail, Result};
+
+use crate::{
+    fsverity::{FsVerityHashValue, Sha256HashValue},
+    splitstream::DigestMap,
+};
+
+pub(crate) struct ZstdWriter {
+    writer: zstd::Encoder<'static, Vec<u8>>,
+    pub(crate) sha256_builder: Option<(Sha256, Sha256HashValue)>,
+}
+
+impl ZstdWriter {
+    pub fn new(sha256: Option<Sha256HashValue>, refs: Option<DigestMap>) -> Self {
+        Self {
+            writer: ZstdWriter::instantiate_writer(refs),
+            sha256_builder: sha256.map(|x| (Sha256::new(), x)),
+        }
+    }
+
+    fn instantiate_writer(refs: Option<DigestMap>) -> zstd::Encoder<'static, Vec<u8>> {
+        let mut writer = zstd::Encoder::new(vec![], 0).unwrap();
+
+        match refs {
+            Some(DigestMap { map }) => {
+                writer.write_all(&(map.len() as u64).to_le_bytes()).unwrap();
+
+                for ref entry in map {
+                    writer.write_all(&entry.body).unwrap();
+                    writer.write_all(&entry.verity).unwrap();
+                }
+            }
+
+            None => {
+                writer.write_all(&0u64.to_le_bytes()).unwrap();
+            }
+        }
+
+        return writer;
+    }
+
+    pub(crate) fn write_fragment(&mut self, size: usize, data: &[u8]) -> Result<()> {
+        self.writer.write_all(&(size as u64).to_le_bytes())?;
+        Ok(self.writer.write_all(data)?)
+    }
+
+    pub(crate) fn update_sha(&mut self, data: &[u8]) {
+        if let Some((sha256, ..)) = &mut self.sha256_builder {
+            sha256.update(&data);
+        }
+    }
+
+    pub(crate) fn flush_inline(&mut self, inline_content: &Vec<u8>) -> Result<()> {
+        if inline_content.is_empty() {
+            return Ok(());
+        }
+
+        self.write_fragment(inline_content.len(), &inline_content)?;
+
+        Ok(())
+    }
+
+    pub(crate) fn finalize_sha256_builder(&mut self) -> Result<Sha256HashValue> {
+        let sha256_builder = std::mem::replace(&mut self.sha256_builder, None);
+
+        let mut sha = Sha256HashValue::EMPTY;
+
+        if let Some((context, expected)) = sha256_builder {
+            let final_sha = Into::<Sha256HashValue>::into(context.finalize());
+
+            if final_sha != expected {
+                bail!(
+                    "Content doesn't have expected SHA256 hash value!\nExpected: {}, final: {}",
+                    hex::encode(expected),
+                    hex::encode(final_sha)
+                );
+            }
+
+            sha = final_sha;
+        }
+
+        return Ok(sha);
+    }
+
+    pub(crate) fn finish(self) -> io::Result<Vec<u8>> {
+        self.writer.finish()
+    }
+}