unencapsulate: use "inner" stream when finishing Decompressor

For things like GzDecoder, we don't want to read via the actual
decompression reader because we don't care about decompressing at this
point.  Plus, the inner reader may have encountered an error partway
through, and trying to decode via decompression will error with
UnexpectedEof.

Instead, wrap a reader for each content type which knows how to get a
mutable reference to the underlying data stream.  When we finish
decompressing, read directly from this "inner" stream to flush any
data.

Resolves: #1407
Signed-off-by: John Eckersberg <[email protected]>

Author: jeckersb

ostree-ext/src/container/unencapsulate.rs

@@ -191,8 +191,59 @@ pub async fn unencapsulate(repo: &ostree::Repo, imgref: &OstreeImageReference) -
     importer.unencapsulate().await
 }
 
+trait Decompressable: Read + Send + 'static {
+    fn get_inner_mut(&mut self) -> &mut (dyn Read);
+}
+
+// TransparentDecompressor
+
+struct TransparentDecompressor<R: Read + Send + 'static>(R);
+
+impl<R: Read + Send + 'static> Read for TransparentDecompressor<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        self.0.read(buf)
+    }
+}
+
+impl<R: Read + Send + 'static> Decompressable for TransparentDecompressor<R> {
+    fn get_inner_mut(&mut self) -> &mut (dyn Read) {
+        &mut self.0
+    }
+}
+
+// GzipDecompressor
+
+struct GzipDecompressor<R: std::io::BufRead>(flate2::bufread::GzDecoder<R>);
+
+impl<R: std::io::BufRead + Send + 'static> Read for GzipDecompressor<R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        self.0.read(buf)
+    }
+}
+
+impl<R: std::io::BufRead + Send + 'static> Decompressable for GzipDecompressor<R> {
+    fn get_inner_mut(&mut self) -> &mut (dyn Read) {
+        self.0.get_mut()
+    }
+}
+
+// ZstdDecompressor
+struct ZstdDecompressor<'a, R: std::io::BufRead>(zstd::stream::read::Decoder<'a, R>);
+
+impl<'a: 'static, R: std::io::BufRead + Send + 'static> Read for ZstdDecompressor<'a, R> {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        self.0.read(buf)
+    }
+}
+
+impl<'a: 'static, R: std::io::BufRead + Send + 'static> Decompressable for ZstdDecompressor<'a, R> {
+    fn get_inner_mut(&mut self) -> &mut (dyn Read) {
+        self.0.get_mut()
+    }
+}
+
 pub(crate) struct Decompressor {
-    inner: Box<dyn Read + Send + 'static>,
+    inner: Box<dyn Decompressable>,
     finished: bool,
 }
 
@@ -224,15 +275,17 @@ impl Decompressor {
         media_type: &oci_image::MediaType,
         src: impl Read + Send + 'static,
     ) -> Result<Self> {
-        let r: Box<dyn std::io::Read + Send + 'static> = match media_type {
+        let r: Box<dyn Decompressable> = match media_type {
             oci_image::MediaType::ImageLayerZstd => {
-                Box::new(zstd::stream::read::Decoder::new(src)?)
+                Box::new(ZstdDecompressor(zstd::stream::read::Decoder::new(src)?))
             }
-            oci_image::MediaType::ImageLayerGzip => Box::new(flate2::bufread::GzDecoder::new(
-                std::io::BufReader::new(src),
+            oci_image::MediaType::ImageLayerGzip => Box::new(GzipDecompressor(
+                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),
+            oci_image::MediaType::ImageLayer => Box::new(TransparentDecompressor(src)),
+            oci_image::MediaType::Other(t) if t.as_str() == DOCKER_TYPE_LAYER_TAR => {
+                Box::new(TransparentDecompressor(src))
+            }
             o => anyhow::bail!("Unhandled layer type: {}", o),
         };
         Ok(Self {
@@ -264,7 +317,7 @@ impl Decompressor {
         // https://github.com/bootc-dev/bootc/issues/1204
 
         let mut sink = std::io::sink();
-        let n = std::io::copy(&mut self.inner, &mut sink)?;
+        let n = std::io::copy(self.inner.get_inner_mut(), &mut sink)?;
 
         if n > 0 {
             tracing::debug!("Read extra {n} bytes at end of decompressor stream");
@@ -362,4 +415,18 @@ mod tests {
         let d = Decompressor::new(&oci_image::MediaType::ImageLayer, empty).unwrap();
         drop(d)
     }
+
+    #[test]
+    fn test_drop_decompressor_with_incomplete_gzip_data() {
+        let empty = std::io::empty();
+        let d = Decompressor::new(&oci_image::MediaType::ImageLayerGzip, empty).unwrap();
+        drop(d)
+    }
+
+    #[test]
+    fn test_drop_decompressor_with_incomplete_zstd_data() {
+        let empty = std::io::empty();
+        let d = Decompressor::new(&oci_image::MediaType::ImageLayerZstd, empty).unwrap();
+        drop(d)
+    }
 }