splitstream: Clean up sha256 argument to SplitStreaWriter::new()

This separates out the two cases where we have an known expected
digest, and the case were we maybe don't know it but want to compute
it.

Author: alexlarsson

crates/composefs-oci/src/tar.rs

@@ -344,7 +344,7 @@ mod tests {
         let mut writer = repo.create_stream(TAR_LAYER_CONTENT_TYPE, None);
 
         split(&mut tar_cursor, &mut writer)?;
-        let object_id = writer.done()?;
+        let (object_id, _) = writer.done()?;
 
         let mut reader: SplitStreamReader<std::fs::File, Sha256HashValue> = SplitStreamReader::new(
             repo.open_object(&object_id)?.into(),
@@ -371,7 +371,7 @@ mod tests {
         let mut writer = repo.create_stream(TAR_LAYER_CONTENT_TYPE, None);
 
         split(&mut tar_cursor, &mut writer).unwrap();
-        let object_id = writer.done().unwrap();
+        let (object_id, _) = writer.done().unwrap();
 
         let mut reader: SplitStreamReader<std::fs::File, Sha256HashValue> = SplitStreamReader::new(
             repo.open_object(&object_id).unwrap().into(),
@@ -400,7 +400,7 @@ mod tests {
         let mut writer = repo.create_stream(TAR_LAYER_CONTENT_TYPE, None);
 
         split(&mut tar_cursor, &mut writer).unwrap();
-        let object_id = writer.done().unwrap();
+        let (object_id, _) = writer.done().unwrap();
 
         let mut reader: SplitStreamReader<std::fs::File, Sha256HashValue> = SplitStreamReader::new(
             repo.open_object(&object_id).unwrap().into(),
@@ -458,7 +458,7 @@ mod tests {
         let mut writer = repo.create_stream(TAR_LAYER_CONTENT_TYPE, None);
 
         split(&mut tar_cursor, &mut writer).unwrap();
-        let object_id = writer.done().unwrap();
+        let (object_id, _) = writer.done().unwrap();
 
         let mut reader: SplitStreamReader<std::fs::File, Sha256HashValue> = SplitStreamReader::new(
             repo.open_object(&object_id).unwrap().into(),
@@ -524,7 +524,7 @@ mod tests {
         let repo = create_test_repository().unwrap();
         let mut writer = repo.create_stream(TAR_LAYER_CONTENT_TYPE, None);
         split(&mut tar_cursor, &mut writer).unwrap();
-        let object_id = writer.done().unwrap();
+        let (object_id, _) = writer.done().unwrap();
 
         // Read back entries and compare with original headers
         let mut reader: SplitStreamReader<std::fs::File, Sha256HashValue> = SplitStreamReader::new(

crates/composefs/src/repository.rs

@@ -238,7 +238,7 @@ impl<ObjectID: FsVerityHashValue> Repository<ObjectID> {
         content_type: u64,
         sha256: Option<Sha256Digest>,
     ) -> SplitStreamWriter<ObjectID> {
-        SplitStreamWriter::new(self, content_type, sha256)
+        SplitStreamWriter::new(self, content_type, false, sha256)
     }
 
     fn format_object_path(id: &ObjectID) -> String {
@@ -319,11 +319,11 @@ impl<ObjectID: FsVerityHashValue> Repository<ObjectID> {
         writer: SplitStreamWriter<ObjectID>,
         reference: Option<&str>,
     ) -> Result<ObjectID> {
-        let Some((.., ref sha256)) = writer.sha256 else {
+        let Some(sha256) = writer.expected_sha256 else {
             bail!("Writer doesn't have sha256 enabled");
         };
         let stream_path = format!("streams/{}", hex::encode(sha256));
-        let object_id = writer.done()?;
+        let (object_id, _) = writer.done()?;
         let object_path = Self::format_object_path(&object_id);
         self.symlink(&stream_path, &object_path)?;
 
@@ -382,7 +382,7 @@ impl<ObjectID: FsVerityHashValue> Repository<ObjectID> {
             None => {
                 let mut writer = self.create_stream(content_type, Some(*sha256));
                 callback(&mut writer)?;
-                let object_id = writer.done()?;
+                let (object_id, _) = writer.done()?;
 
                 let object_path = Self::format_object_path(&object_id);
                 self.symlink(&stream_path, &object_path)?;

crates/composefs/src/splitstream.rs

@@ -97,7 +97,8 @@ pub struct SplitStreamWriter<ObjectID: FsVerityHashValue> {
     total_size: u64,
     writer: Encoder<'static, Vec<u8>>,
     pub content_type: u64,
-    pub sha256: Option<(Sha256, Sha256Digest)>,
+    pub sha256: Option<Sha256>,
+    pub expected_sha256: Option<Sha256Digest>,
 }
 
 impl<ObjectID: FsVerityHashValue> std::fmt::Debug for SplitStreamWriter<ObjectID> {
@@ -106,6 +107,7 @@ impl<ObjectID: FsVerityHashValue> std::fmt::Debug for SplitStreamWriter<ObjectID
         f.debug_struct("SplitStreamWriter")
             .field("repo", &self.repo)
             .field("inline_content", &self.inline_content)
+            .field("expected_sha256", &self.expected_sha256)
             .field("sha256", &self.sha256)
             .finish()
     }
@@ -115,7 +117,8 @@ impl<ObjectID: FsVerityHashValue> SplitStreamWriter<ObjectID> {
     pub fn new(
         repo: &Arc<Repository<ObjectID>>,
         content_type: u64,
-        sha256: Option<Sha256Digest>,
+        compute_sha256: bool,
+        expected_sha256: Option<Sha256Digest>,
     ) -> Self {
         // SAFETY: we surely can't get an error writing the header to a Vec<u8>
         let writer = Encoder::new(vec![], 0).unwrap();
@@ -128,7 +131,12 @@ impl<ObjectID: FsVerityHashValue> SplitStreamWriter<ObjectID> {
             total_size: 0,
             mappings: DigestMap::new(),
             writer,
-            sha256: sha256.map(|x| (Sha256::new(), x)),
+            sha256: if compute_sha256 || expected_sha256.is_some() {
+                Some(Sha256::new())
+            } else {
+                None
+            },
+            expected_sha256,
         }
     }
 
@@ -177,7 +185,7 @@ impl<ObjectID: FsVerityHashValue> SplitStreamWriter<ObjectID> {
     /// 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 {
+        if let Some(ref mut sha256) = self.sha256 {
             sha256.update(data);
         }
         self.inline_content.extend(data);
@@ -195,7 +203,7 @@ impl<ObjectID: FsVerityHashValue> SplitStreamWriter<ObjectID> {
     }
 
     pub fn write_external(&mut self, data: &[u8], padding: Vec<u8>) -> Result<()> {
-        if let Some((ref mut sha256, ..)) = self.sha256 {
+        if let Some(ref mut sha256, ..) = self.sha256 {
             sha256.update(data);
             sha256.update(&padding);
         }
@@ -207,7 +215,7 @@ impl<ObjectID: FsVerityHashValue> SplitStreamWriter<ObjectID> {
     }
 
     pub async fn write_external_async(&mut self, data: Vec<u8>, padding: Vec<u8>) -> Result<()> {
-        if let Some((ref mut sha256, ..)) = self.sha256 {
+        if let Some(ref mut sha256, ..) = self.sha256 {
             sha256.update(&data);
             sha256.update(&padding);
         }
@@ -217,14 +225,20 @@ impl<ObjectID: FsVerityHashValue> SplitStreamWriter<ObjectID> {
         self.write_reference(&id, padding)
     }
 
-    pub fn done(mut self) -> Result<ObjectID> {
+    pub fn done(mut self) -> Result<(ObjectID, Option<Sha256Digest>)> {
         self.flush_inline(vec![])?;
 
-        if let Some((context, expected)) = self.sha256 {
-            if Into::<Sha256Digest>::into(context.finalize()) != expected {
-                bail!("Content doesn't have expected SHA256 hash value!");
+        let sha256_digest = if let Some(sha256) = self.sha256 {
+            let actual = Into::<Sha256Digest>::into(sha256.finalize());
+            if let Some(expected) = self.expected_sha256 {
+                if actual != expected {
+                    bail!("Content doesn't have expected SHA256 hash value!");
+                }
             }
-        }
+            Some(actual)
+        } else {
+            None
+        };
 
         let mut buf = vec![];
         let header = SplitstreamHeader {
@@ -252,7 +266,7 @@ impl<ObjectID: FsVerityHashValue> SplitStreamWriter<ObjectID> {
 
         buf.extend_from_slice(&self.writer.finish()?);
 
-        self.repo.ensure_object(&buf)
+        Ok((self.repo.ensure_object(&buf)?, sha256_digest))
     }
 }