Extend layer decompression based on mediatype label from containerd

src/tardev-snapshotter/src/snapshotter.rs

@@ -856,40 +856,64 @@ impl TarDevSnapshotter {
                     "missing target layer digest label",
                 ));
             };
+
+            // Determine image layer media type
             let Some(media_type) = labels.get(TARGET_LAYER_MEDIA_TYPE_LABEL) else {
                 error!("missing target layer media type label");
                 return Err(Status::invalid_argument(
                     "missing target layer media type label",
                 ));
             };
-
             trace!("layer digest {} media_type: {:?}", digest_str, media_type);
 
-            // TODO use the media_type:
-            // gzipped: "application/vnd.docker.image.rootfs.diff.tar.gzip", "application/vnd.oci.image.layer.v1.tar+gzip";
-            // uncompressed: "application/vnd.oci.image.layer.v1.tar", "application/vnd.docker.image.rootfs.diff.tar"
+            let layer_type = match media_type.as_str() {
+                "application/vnd.docker.image.rootfs.diff.tar.gzip" 
+                | "application/vnd.oci.image.layer.v1.tar+gzip" => "tar.gz",
+                "application/vnd.oci.image.layer.v1.tar" 
+                | "application/vnd.docker.image.rootfs.diff.tar" => "tar",
+                _ => {
+                    error!("unsupported media type: {}", media_type);
+                    return Err(Status::invalid_argument(format!(
+                        "unsupported media type: {}",
+                        media_type
+                    )));
+                }
+            };
 
+            // Fetch the layer image
             let name = dir.path().join(name_to_hash(&key));
-            let mut gzname = name.clone();
-            gzname.set_extension("gz");
-            trace!("Fetching layer image to {:?}", &gzname);
-            self.get_layer_image(&gzname, digest_str).await?;
-
-            // TODO: Decompress in stream instead of reopening.
-            // Decompress data.
-            trace!("Decompressing {:?} to {:?}", &gzname, &name);
+            trace!("Fetching {} layer image to {:?}", layer_type, name);
+            self.get_layer_image(&name, digest_str).await?;
+
+            // Rename the file with the correct extension
+            let target_name = {
+                let mut renamed = name.clone();
+                renamed.set_extension(layer_type);
+                trace!("Renaming {:?} to {:?}", &name, &renamed);
+                std::fs::rename(&name, &renamed)
+                    .map_err(|e| Status::internal(format!("Failed to rename file: {e}")))?;
+                renamed
+            };
+
+            // Decompress and process the layer
+            trace!("Decompressing {:?} to {:?}", &target_name, &name);
             let generated_root_hash = tokio::task::spawn_blocking(move || -> Result<_> {
-                let compressed = fs::File::open(&gzname)?;
                 let mut file = OpenOptions::new()
                     .read(true)
                     .write(true)
                     .create(true)
                     .truncate(true)
                     .open(&name)?;
-                let mut gz_decoder = flate2::read::GzDecoder::new(compressed);
-                std::io::copy(&mut gz_decoder, &mut file)
+                if layer_type == "tar.gz" {
+                    let compressed = fs::File::open(&target_name)?;
+                    let mut gz_decoder = flate2::read::GzDecoder::new(compressed);
+                    std::io::copy(&mut gz_decoder, &mut file)
                     .context("failed to copy payload from gz decoder")?;
-
+                } else {
+                    let mut tar_file = fs::File::open(&target_name)?;
+                    std::io::copy(&mut tar_file, &mut file).context("failed to copy payload from gz decoder")?;
+                }
+
                 trace!("Appending index to {:?}", &name);
                 file.rewind().context("failed to rewind the file handle")?;
                 tarindex::append_index(&mut file).context("failed to append tar index")?;