[PATCH proxmox 2/3] oci: reduce redundant layer decompression during extraction

[Filip Schauer <f.schauer@proxmox.com>]

According to the OCI Image Specification, the order of layers in the
image manifest matches the order of the rootfs.diff_ids in the image
configuration.

Instead of redundantly decompressing and hashing every layer upfront to
build a lookup map, pair layer blobs directly with their corresponding
diff_ids and verify the digests on the fly.

Signed-off-by: Filip Schauer <f.schauer@proxmox.com>
---
 proxmox-oci/src/lib.rs | 112 ++++++++---------------------------------
 1 file changed, 21 insertions(+), 91 deletions(-)

diff --git a/proxmox-oci/src/lib.rs b/proxmox-oci/src/lib.rs
index 71970cd8..7182b008 100644
--- a/proxmox-oci/src/lib.rs
+++ b/proxmox-oci/src/lib.rs
@@ -1,4 +1,3 @@
-use std::collections::HashMap;
 use std::ffi::OsStr;
 use std::fs::{File, read_dir, remove_dir_all, remove_file};
 use std::io::{self, ErrorKind, Read, Seek, SeekFrom};
@@ -10,12 +9,11 @@ use flate2::read::GzDecoder;
 use oci_spec::OciSpecError;
 pub use oci_spec::image::{Arch, Config};
 use oci_spec::image::{ImageConfiguration, ImageManifest, MediaType};
-use sha2::digest::generic_array::GenericArray;
-use sha2::{Digest, Sha256};
 use tar::{Archive, EntryType};
 use thiserror::Error;
 
 mod digest;
+use digest::{DigestError, VerifyingReader};
 
 mod oci_tar_image;
 use oci_tar_image::{OciTarImage, OciTarImageBlob};
@@ -23,80 +21,6 @@ use oci_tar_image::{OciTarImage, OciTarImageBlob};
 const WHITEOUT_PREFIX: &[u8] = b".wh.";
 const OPAQUE_WHITEOUT_NAME: &str = ".wh..wh..opq";
 
-fn compute_digest<R: Read, H: Digest>(
-    mut reader: R,
-    mut hasher: H,
-) -> io::Result<GenericArray<u8, H::OutputSize>> {
-    let mut buf = proxmox_io::boxed::zeroed(32768);
-
-    loop {
-        let bytes_read = reader.read(&mut buf)?;
-        if bytes_read == 0 {
-            break Ok(hasher.finalize());
-        }
-
-        hasher.update(&buf[..bytes_read]);
-    }
-}
-
-fn compute_sha256<R: Read>(reader: R) -> io::Result<oci_spec::image::Sha256Digest> {
-    let digest = compute_digest(reader, Sha256::new())?;
-    Ok(oci_spec::image::Sha256Digest::from_str(&format!("{digest:x}")).expect("valid digest"))
-}
-
-/// Build a mapping from uncompressed layer digests (as found in the image config's `rootfs.diff_ids`)
-/// to their corresponding compressed-layer digests (i.e. the filenames under `blobs/<algorithm>/<digest>`)
-fn build_layer_map<R: Read + Seek>(
-    mut oci_tar_image: OciTarImage<R>,
-    image_manifest: &ImageManifest,
-) -> Result<
-    (
-        OciTarImage<R>,
-        HashMap<oci_spec::image::Digest, oci_spec::image::Descriptor>,
-    ),
-    ExtractError,
-> {
-    let mut layer_mapping = HashMap::new();
-
-    for layer in image_manifest.layers() {
-        let digest = match layer.media_type() {
-            MediaType::ImageLayer | MediaType::ImageLayerNonDistributable => layer.digest().clone(),
-            MediaType::ImageLayerGzip | MediaType::ImageLayerNonDistributableGzip => {
-                let mut compressed_blob = oci_tar_image
-                    .open_blob(layer.digest())
-                    .ok_or(ExtractError::MissingLayerFile(layer.digest().clone()))?;
-                let decoder = GzDecoder::new(&mut compressed_blob);
-                let hash = compute_sha256(decoder)?.into();
-                oci_tar_image = compressed_blob.into_oci_tar_image();
-                hash
-            }
-            MediaType::ImageLayerZstd | MediaType::ImageLayerNonDistributableZstd => {
-                let mut compressed_blob = oci_tar_image
-                    .open_blob(layer.digest())
-                    .ok_or(ExtractError::MissingLayerFile(layer.digest().clone()))?;
-                let decoder = zstd::Decoder::new(&mut compressed_blob)?;
-                let hash = compute_sha256(decoder)?.into();
-                oci_tar_image = compressed_blob.into_oci_tar_image();
-                hash
-            }
-            // Skip any other non-ImageLayer related media types.
-            // Match explicitly to avoid missing new image layer types when oci-spec updates.
-            MediaType::Descriptor
-            | MediaType::LayoutHeader
-            | MediaType::ImageManifest
-            | MediaType::ImageIndex
-            | MediaType::ImageConfig
-            | MediaType::ArtifactManifest
-            | MediaType::EmptyJSON
-            | MediaType::Other(_) => continue,
-        };
-
-        layer_mapping.insert(digest, layer.clone());
-    }
-
-    Ok((oci_tar_image, layer_mapping))
-}
-
 #[derive(Debug, Error)]
 pub enum ProxmoxOciError {
     #[error("Error while parsing OCI image: {0}")]
@@ -177,8 +101,8 @@ fn parse_image<P: AsRef<Path>>(
 pub enum ExtractError {
     #[error("Incorrectly formatted digest: \"{0}\"")]
     InvalidDigest(String),
-    #[error("Unknown layer digest {0} found in rootfs.diff_ids")]
-    UnknownLayerDigest(oci_spec::image::Digest),
+    #[error(transparent)]
+    DigestError(#[from] DigestError),
     #[error("Layer file {0} mentioned in image manifest is missing")]
     MissingLayerFile(oci_spec::image::Digest),
     #[error("IO error: {0}")]
@@ -188,23 +112,23 @@ pub enum ExtractError {
 }
 
 fn extract_image_rootfs<R: Read + Seek, P: AsRef<Path>>(
-    oci_tar_image: OciTarImage<R>,
+    mut oci_tar_image: OciTarImage<R>,
     image_manifest: &ImageManifest,
     image_config: &ImageConfiguration,
     target_path: P,
 ) -> Result<(), ExtractError> {
-    let (mut oci_tar_image, layer_map) = build_layer_map(oci_tar_image, image_manifest)?;
     let target_path = target_path.as_ref().canonicalize()?;
 
-    for layer in image_config.rootfs().diff_ids() {
-        let layer_digest = oci_spec::image::Digest::from_str(layer)
-            .map_err(|_| ExtractError::InvalidDigest(layer.to_string()))?;
-        let layer_descriptor = layer_map
-            .get(&layer_digest)
-            .ok_or(ExtractError::UnknownLayerDigest(layer_digest.clone()))?;
-        let mut layer_file = oci_tar_image
-            .open_blob(layer_descriptor.digest())
-            .ok_or(ExtractError::MissingLayerFile(layer_digest))?;
+    for (layer_descriptor, diffid_str) in image_manifest
+        .layers()
+        .iter()
+        .zip(image_config.rootfs().diff_ids())
+    {
+        let layer_diffid = oci_spec::image::Digest::from_str(diffid_str)
+            .map_err(|_| ExtractError::InvalidDigest(diffid_str.to_string()))?;
+        let mut layer_file = oci_tar_image.open_blob(layer_descriptor.digest()).ok_or(
+            ExtractError::MissingLayerFile(layer_descriptor.digest().clone()),
+        )?;
 
         type DecodeFn<T> = Box<dyn for<'a> Fn(&'a mut T) -> io::Result<Box<dyn Read + 'a>>>;
         let decode_fn: DecodeFn<OciTarImageBlob<R>> = match layer_descriptor.media_type() {
@@ -231,7 +155,13 @@ fn extract_image_rootfs<R: Read + Seek, P: AsRef<Path>>(
             }
         };
 
-        remove_replaced_paths(&mut decode_fn(&mut layer_file)?, &target_path)?;
+        let mut verified = VerifyingReader::new(decode_fn(&mut layer_file)?, layer_diffid)?;
+        remove_replaced_paths(&mut verified, &target_path)?;
+
+        // Force an EOF so VerifyingReader actually verifies.
+        io::copy(&mut verified, &mut io::sink())?;
+        drop(verified);
+
         layer_file.seek(SeekFrom::Start(0))?;
         extract_archive(&mut decode_fn(&mut layer_file)?, &target_path)?;
 
-- 
2.47.3