Merge pull request #1248 from jeckersb/more-decompressor

unencapsulate: more Decompressor improvements

Author: cgwalters

ostree-ext/src/container/store.rs

@@ -7,6 +7,7 @@
 
 use super::*;
 use crate::chunking::{self, Chunk};
+use crate::container::Decompressor;
 use crate::logging::system_repo_journal_print;
 use crate::refescape;
 use crate::sysroot::SysrootLock;
@@ -779,8 +780,8 @@ impl ImageImporter {
                     let txn = repo.auto_transaction(Some(cancellable))?;
                     let mut importer = crate::tar::Importer::new_for_object_set(&repo);
                     let blob = tokio_util::io::SyncIoBridge::new(blob);
-                    let blob = super::unencapsulate::decompressor(&media_type, blob)?;
-                    let mut archive = tar::Archive::new(blob);
+                    let mut blob = Decompressor::new(&media_type, blob)?;
+                    let mut archive = tar::Archive::new(&mut blob);
                     importer.import_objects(&mut archive, Some(cancellable))?;
                     let commit = if write_refs {
                         let commit = importer.finish_import_object_set()?;
@@ -791,6 +792,7 @@ impl ImageImporter {
                         None
                     };
                     txn.commit(Some(cancellable))?;
+                    blob.finish()?;
                     Ok::<_, anyhow::Error>(commit)
                 })
                 .map_err(|e| e.context(format!("Layer {}", layer.layer.digest())));
@@ -825,8 +827,8 @@ impl ImageImporter {
                     let txn = repo.auto_transaction(Some(cancellable))?;
                     let mut importer = crate::tar::Importer::new_for_commit(&repo, remote);
                     let blob = tokio_util::io::SyncIoBridge::new(blob);
-                    let blob = super::unencapsulate::decompressor(&media_type, blob)?;
-                    let mut archive = tar::Archive::new(blob);
+                    let mut blob = Decompressor::new(&media_type, blob)?;
+                    let mut archive = tar::Archive::new(&mut blob);
                     importer.import_commit(&mut archive, Some(cancellable))?;
                     let (commit, verify_text) = importer.finish_import_commit();
                     if write_refs {
@@ -835,6 +837,7 @@ impl ImageImporter {
                     }
                     repo.mark_commit_partial(&commit, false)?;
                     txn.commit(Some(cancellable))?;
+                    blob.finish()?;
                     Ok::<_, anyhow::Error>((commit, verify_text))
                 });
             let (commit, verify_text) =

ostree-ext/src/container/unencapsulate.rs

@@ -193,6 +193,7 @@ pub async fn unencapsulate(repo: &ostree::Repo, imgref: &OstreeImageReference) -
 
 pub(crate) struct Decompressor {
     inner: Box<dyn Read + Send + 'static>,
+    finished: bool,
 }
 
 impl Read for Decompressor {
@@ -203,6 +204,50 @@ impl Read for Decompressor {
 
 impl Drop for Decompressor {
     fn drop(&mut self) {
+        if self.finished {
+            return;
+        }
+
+        // We really should not get here; users are required to call
+        // `finish()` to clean up the stream.  But we'll give
+        // best-effort to clean things up nonetheless.  If things go
+        // wrong, then panic, because we're in a bad state and it's
+        // likely that we end up with a broken pipe error or a
+        // deadlock.
+        self._finish().expect("Decompressor::finish MUST be called")
+    }
+}
+
+impl Decompressor {
+    /// Create a decompressor for this MIME type, given a stream of input.
+    pub(crate) fn new(
+        media_type: &oci_image::MediaType,
+        src: impl Read + Send + 'static,
+    ) -> Result<Self> {
+        let r: Box<dyn std::io::Read + Send + 'static> = match media_type {
+            oci_image::MediaType::ImageLayerZstd => {
+                Box::new(zstd::stream::read::Decoder::new(src)?)
+            }
+            oci_image::MediaType::ImageLayerGzip => Box::new(flate2::bufread::GzDecoder::new(
+                std::io::BufReader::new(src),
+            )),
+            oci_image::MediaType::ImageLayer => Box::new(src),
+            oci_image::MediaType::Other(t) if t.as_str() == DOCKER_TYPE_LAYER_TAR => Box::new(src),
+            o => anyhow::bail!("Unhandled layer type: {}", o),
+        };
+        Ok(Self {
+            inner: r,
+            finished: false,
+        })
+    }
+
+    pub(crate) fn finish(mut self) -> Result<()> {
+        self._finish()
+    }
+
+    fn _finish(&mut self) -> Result<()> {
+        self.finished = true;
+
         // We need to make sure to flush out the decompressor and/or
         // tar stream here.  For tar, we might not read through the
         // entire stream, because the archive has zero-block-markers
@@ -219,29 +264,14 @@ impl Drop for Decompressor {
         // https://github.com/bootc-dev/bootc/issues/1204
 
         let mut sink = std::io::sink();
-        match std::io::copy(&mut self.inner, &mut sink) {
-            Err(e) => tracing::debug!("Ignoring error while dropping decompressor: {e}"),
-            Ok(0) => { /* We already read everything and are happy */ }
-            Ok(n) => tracing::debug!("Read extra {n} bytes at end of decompressor stream"),
+        let n = std::io::copy(&mut self.inner, &mut sink)?;
+
+        if n > 0 {
+            tracing::debug!("Read extra {n} bytes at end of decompressor stream");
         }
-    }
-}
 
-/// Create a decompressor for this MIME type, given a stream of input.
-pub(crate) fn decompressor(
-    media_type: &oci_image::MediaType,
-    src: impl Read + Send + 'static,
-) -> Result<Decompressor> {
-    let r: Box<dyn std::io::Read + Send + 'static> = match media_type {
-        oci_image::MediaType::ImageLayerZstd => Box::new(zstd::stream::read::Decoder::new(src)?),
-        oci_image::MediaType::ImageLayerGzip => Box::new(flate2::bufread::GzDecoder::new(
-            std::io::BufReader::new(src),
-        )),
-        oci_image::MediaType::ImageLayer => Box::new(src),
-        oci_image::MediaType::Other(t) if t.as_str() == DOCKER_TYPE_LAYER_TAR => Box::new(src),
-        o => anyhow::bail!("Unhandled layer type: {}", o),
-    };
-    Ok(Decompressor { inner: r })
+        Ok(())
+    }
 }
 
 /// A wrapper for [`get_blob`] which fetches a layer and decompresses it.
@@ -305,3 +335,31 @@ pub(crate) async fn fetch_layer<'a>(
         Ok((Box::new(blob), Either::Right(driver), media_type))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    struct BrokenPipe;
+
+    impl Read for BrokenPipe {
+        fn read(&mut self, _buf: &mut [u8]) -> std::io::Result<usize> {
+            std::io::Result::Err(std::io::ErrorKind::BrokenPipe.into())
+        }
+    }
+
+    #[test]
+    #[should_panic(expected = "Decompressor::finish MUST be called")]
+    fn test_drop_decompressor_with_finish_error_should_panic() {
+        let broken = BrokenPipe;
+        let d = Decompressor::new(&oci_image::MediaType::ImageLayer, broken).unwrap();
+        drop(d)
+    }
+
+    #[test]
+    fn test_drop_decompressor_with_successful_finish() {
+        let empty = std::io::empty();
+        let d = Decompressor::new(&oci_image::MediaType::ImageLayer, empty).unwrap();
+        drop(d)
+    }
+}

ostree-ext/src/tar/write.rs

@@ -7,6 +7,7 @@
 //! In the future, this may also evolve into parsing the tar
 //! stream in Rust, not in C.
 
+use crate::container::Decompressor;
 use crate::Result;
 use anyhow::{anyhow, Context};
 use camino::{Utf8Component, Utf8Path, Utf8PathBuf};
@@ -362,11 +363,13 @@ async fn filter_tar_async(
     let config = config.clone();
     let tar_transformer = crate::tokio_util::spawn_blocking_flatten(move || {
         let src = tokio_util::io::SyncIoBridge::new(src);
-        let mut src = crate::container::decompressor(&media_type, src)?;
+        let mut src = Decompressor::new(&media_type, src)?;
         let dest = tokio_util::io::SyncIoBridge::new(tx_buf);
 
         let r = filter_tar(&mut src, dest, &config, &repo_tmpdir);
 
+        src.finish()?;
+
         Ok(r)
     });
     let copier = tokio::io::copy(&mut rx_buf, &mut dest);