Support custom image compression

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "ocidir"
 description = "A Rust library for reading and writing OCI (opencontainers) layout directories"
-version = "0.4.0"
+version = "0.5.0"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 repository = "https://github.com/containers/ocidir-rs"

examples/custom_compressor.rs

@@ -0,0 +1,61 @@
+/// Example that shows how to use a custom compression and media type for image layers.
+/// The example below does no compression.
+use std::{env, io, path::PathBuf};
+
+use oci_spec::image::Platform;
+use ocidir::{cap_std::fs::Dir, BlobWriter, OciDir, WriteComplete};
+
+struct NoCompression<'a>(BlobWriter<'a>);
+
+impl io::Write for NoCompression<'_> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.0.write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.0.flush()
+    }
+}
+
+impl<'a> WriteComplete<BlobWriter<'a>> for NoCompression<'a> {
+    fn complete(self) -> io::Result<BlobWriter<'a>> {
+        Ok(self.0)
+    }
+}
+
+fn main() {
+    let dir = Dir::open_ambient_dir(env::temp_dir(), ocidir::cap_std::ambient_authority()).unwrap();
+    let oci_dir = OciDir::ensure(dir).unwrap();
+
+    let mut manifest = oci_dir.new_empty_manifest().unwrap().build().unwrap();
+    let mut config = ocidir::oci_spec::image::ImageConfigurationBuilder::default()
+        .build()
+        .unwrap();
+
+    // Add the src as a layer
+    let writer = oci_dir
+        .create_custom_layer(
+            |bw| Ok(NoCompression(bw)),
+            oci_spec::image::MediaType::ImageLayer,
+        )
+        .unwrap();
+    let mut builder = tar::Builder::new(writer);
+    builder.follow_symlinks(false);
+
+    builder
+        .append_dir_all(".", PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("src"))
+        .unwrap();
+
+    let layer = builder.into_inner().unwrap().complete().unwrap();
+    oci_dir.push_layer(&mut manifest, &mut config, layer, "src", None);
+
+    println!(
+        "Created image with manifest: {}",
+        manifest.to_string_pretty().unwrap()
+    );
+
+    // Add the image manifest
+    let _descriptor = oci_dir
+        .insert_manifest_and_config(manifest.clone(), config, None, Platform::default())
+        .unwrap();
+}

src/lib.rs

@@ -15,7 +15,8 @@ use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
 use std::fs::File;
-use std::io::{prelude::*, BufReader};
+use std::io::{prelude::*, BufReader, BufWriter};
+use std::marker::PhantomData;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 use thiserror::Error;
@@ -141,7 +142,7 @@ pub struct BlobWriter<'a> {
     /// Compute checksum
     hash: Hasher,
     /// Target file
-    target: Option<cap_tempfile::TempFile<'a>>,
+    target: Option<BufWriter<cap_tempfile::TempFile<'a>>>,
     size: u64,
 }
 
@@ -154,13 +155,6 @@ impl Debug for BlobWriter<'_> {
     }
 }
 
-/// Create an OCI tar+gzip layer.
-pub struct GzipLayerWriter<'a>(Sha256Writer<GzEncoder<BlobWriter<'a>>>);
-
-#[cfg(feature = "zstd")]
-/// Writer for a OCI tar+zstd layer.
-pub struct ZstdLayerWriter<'a>(Sha256Writer<zstd::Encoder<'static, BlobWriter<'a>>>);
-
 #[derive(Debug)]
 /// An opened OCI directory.
 pub struct OciDir {
@@ -280,17 +274,32 @@ impl OciDir {
         BlobWriter::new(&self.dir)
     }
 
+    /// Create a layer writer with a custom encoder and
+    /// media type
+    pub fn create_custom_layer<'a, W: WriteComplete<BlobWriter<'a>>>(
+        &'a self,
+        create: impl FnOnce(BlobWriter<'a>) -> std::io::Result<W>,
+        media_type: MediaType,
+    ) -> Result<LayerWriter<'a, W>> {
+        let bw = BlobWriter::new(&self.dir)?;
+        Ok(LayerWriter::new(create(bw)?, media_type))
+    }
+
     /// Create a writer for a new gzip+tar blob; the contents
     /// are not parsed, but are expected to be a tarball.
-    pub fn create_gzip_layer(&self, c: Option<flate2::Compression>) -> Result<GzipLayerWriter> {
-        GzipLayerWriter::new(&self.dir, c)
+    pub fn create_gzip_layer<'a>(
+        &'a self,
+        c: Option<flate2::Compression>,
+    ) -> Result<LayerWriter<'a, GzEncoder<BlobWriter<'a>>>> {
+        let creator = |bw: BlobWriter<'a>| Ok(GzEncoder::new(bw, c.unwrap_or_default()));
+        self.create_custom_layer(creator, MediaType::ImageLayerGzip)
     }
 
     /// Create a tar output stream, backed by a blob
     pub fn create_layer(
         &self,
         c: Option<flate2::Compression>,
-    ) -> Result<tar::Builder<GzipLayerWriter>> {
+    ) -> Result<tar::Builder<LayerWriter<GzEncoder<BlobWriter>>>> {
         Ok(tar::Builder::new(self.create_gzip_layer(c)?))
     }
 
@@ -299,8 +308,12 @@ impl OciDir {
     /// are not parsed, but are expected to be a tarball.
     ///
     /// This method is only available when the `zstd` feature is enabled.
-    pub fn create_layer_zstd(&self, compression_level: Option<i32>) -> Result<ZstdLayerWriter> {
-        ZstdLayerWriter::new(&self.dir, compression_level)
+    pub fn create_layer_zstd<'a>(
+        &'a self,
+        compression_level: Option<i32>,
+    ) -> Result<LayerWriter<'a, zstd::Encoder<'static, BlobWriter<'a>>>> {
+        let creator = |bw: BlobWriter<'a>| zstd::Encoder::new(bw, compression_level.unwrap_or(0));
+        self.create_custom_layer(creator, MediaType::ImageLayerZstd)
     }
 
     #[cfg(feature = "zstdmt")]
@@ -312,12 +325,17 @@ impl OciDir {
     /// [zstd::Encoder::multithread]]
     ///
     /// This method is only available when the `zstdmt` feature is enabled.
-    pub fn create_layer_zstd_multithread(
-        &self,
+    pub fn create_layer_zstd_multithread<'a>(
+        &'a self,
         compression_level: Option<i32>,
         n_workers: u32,
-    ) -> Result<ZstdLayerWriter> {
-        ZstdLayerWriter::multithread(&self.dir, compression_level, n_workers)
+    ) -> Result<LayerWriter<'a, zstd::Encoder<'static, BlobWriter<'a>>>> {
+        let creator = |bw: BlobWriter<'a>| {
+            let mut encoder = zstd::Encoder::new(bw, compression_level.unwrap_or(0))?;
+            encoder.multithread(n_workers)?;
+            Ok(encoder)
+        };
+        self.create_custom_layer(creator, MediaType::ImageLayerZstd)
     }
 
     /// Add a layer to the top of the image stack.  The firsh pushed layer becomes the root.
@@ -655,7 +673,7 @@ impl<'a> BlobWriter<'a> {
         Ok(Self {
             hash: Hasher::new(MessageDigest::sha256())?,
             // FIXME add ability to choose filename after completion
-            target: Some(cap_tempfile::TempFile::new(ocidir)?),
+            target: Some(BufWriter::new(cap_tempfile::TempFile::new(ocidir)?)),
             size: 0,
         })
     }
@@ -684,7 +702,7 @@ impl<'a> BlobWriter<'a> {
     fn complete_as(mut self, sha256_digest: &str) -> Result<Blob> {
         let destname = &format!("{}/{}", BLOBDIR, sha256_digest);
         let target = self.target.take().unwrap();
-        target.replace(destname)?;
+        target.into_inner().unwrap().replace(destname)?;
         Ok(Blob {
             sha256: Sha256Digest::from_str(sha256_digest).unwrap(),
             size: self.size,
@@ -700,94 +718,84 @@ impl<'a> BlobWriter<'a> {
 
 impl std::io::Write for BlobWriter<'_> {
     fn write(&mut self, srcbuf: &[u8]) -> std::io::Result<usize> {
-        self.hash.update(srcbuf)?;
-        self.target
-            .as_mut()
-            .unwrap()
-            .as_file_mut()
-            .write_all(srcbuf)?;
-        self.size += srcbuf.len() as u64;
-        Ok(srcbuf.len())
+        let written = self.target.as_mut().unwrap().write(srcbuf)?;
+        self.hash.update(&srcbuf[..written])?;
+        self.size += written as u64;
+        Ok(written)
     }
 
     fn flush(&mut self) -> std::io::Result<()> {
         Ok(())
     }
 }
 
-impl<'a> GzipLayerWriter<'a> {
-    /// Create a writer for a gzip compressed layer blob.
-    fn new(ocidir: &'a Dir, c: Option<flate2::Compression>) -> Result<Self> {
-        let bw = BlobWriter::new(ocidir)?;
-        let enc = flate2::write::GzEncoder::new(bw, c.unwrap_or_default());
-        Ok(Self(Sha256Writer::new(enc)))
-    }
+/// A writer that can be finalized to return an inner writer.
+pub trait WriteComplete<W>: Write {
+    fn complete(self) -> std::io::Result<W>;
+}
 
-    /// Consume this writer, flushing buffered data and put the blob in place.
-    pub fn complete(self) -> Result<Layer> {
-        let (uncompressed_sha256, enc) = self.0.finish();
-        let blob = enc.finish()?.complete()?;
-        Ok(Layer {
-            blob,
-            uncompressed_sha256,
-            media_type: MediaType::ImageLayerGzip,
-        })
+impl<W> WriteComplete<W> for GzEncoder<W>
+where
+    W: Write,
+{
+    fn complete(self) -> std::io::Result<W> {
+        self.finish()
     }
 }
 
-impl std::io::Write for GzipLayerWriter<'_> {
-    fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
-        self.0.write(data)
+#[cfg(feature = "zstd")]
+impl<W> WriteComplete<W> for zstd::Encoder<'_, W>
+where
+    W: Write,
+{
+    fn complete(self) -> std::io::Result<W> {
+        self.finish()
     }
+}
 
-    fn flush(&mut self) -> std::io::Result<()> {
-        self.0.flush()
-    }
+/// A writer for a layer.
+pub struct LayerWriter<'a, W>
+where
+    W: WriteComplete<BlobWriter<'a>>,
+{
+    inner: Sha256Writer<W>,
+    media_type: MediaType,
+    marker: PhantomData<&'a ()>,
 }
 
-#[cfg(feature = "zstd")]
-impl<'a> ZstdLayerWriter<'a> {
-    /// Create a writer for a gzip compressed layer blob.
-    fn new(ocidir: &'a Dir, c: Option<i32>) -> Result<Self> {
-        let bw = BlobWriter::new(ocidir)?;
-        let encoder = zstd::Encoder::new(bw, c.unwrap_or(0))?;
-        Ok(Self(Sha256Writer::new(encoder)))
+impl<'a, W> LayerWriter<'a, W>
+where
+    W: WriteComplete<BlobWriter<'a>>,
+{
+    pub fn new(inner: W, media_type: oci_image::MediaType) -> Self {
+        Self {
+            inner: Sha256Writer::new(inner),
+            media_type,
+            marker: PhantomData,
+        }
     }
 
-    /// Consume this writer, flushing buffered data and put the blob in place.
     pub fn complete(self) -> Result<Layer> {
-        let (uncompressed_sha256, enc) = self.0.finish();
-        let blob = enc.finish()?.complete()?;
+        let (uncompressed_sha256, enc) = self.inner.finish();
+        let blob = enc.complete()?.complete()?;
         Ok(Layer {
             blob,
             uncompressed_sha256,
-            media_type: MediaType::ImageLayerZstd,
+            media_type: self.media_type,
         })
     }
 }
 
-#[cfg(feature = "zstdmt")]
-impl<'a> ZstdLayerWriter<'a> {
-    /// Create a writer for a zstd compressed layer blob, with multithreaded compression enabled.
-    ///
-    /// The `n_workers` parameter specifies the number of threads to use for compression, per
-    /// [Encoder::multithread]]
-    fn multithread(ocidir: &'a Dir, c: Option<i32>, n_workers: u32) -> Result<Self> {
-        let bw = BlobWriter::new(ocidir)?;
-        let mut encoder = zstd::Encoder::new(bw, c.unwrap_or(0))?;
-        encoder.multithread(n_workers)?;
-        Ok(Self(Sha256Writer::new(encoder)))
-    }
-}
-
-#[cfg(feature = "zstd")]
-impl std::io::Write for ZstdLayerWriter<'_> {
+impl<'a, W> std::io::Write for LayerWriter<'a, W>
+where
+    W: WriteComplete<BlobWriter<'a>>,
+{
     fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
-        self.0.write(data)
+        self.inner.write(data)
     }
 
     fn flush(&mut self) -> std::io::Result<()> {
-        self.0.flush()
+        self.inner.flush()
     }
 }