From e8c324687a472243d4dedadafd7fda63044aa8ac Mon Sep 17 00:00:00 2001
From: zyfjeff <zyfjeff@linux.alibaba.com>
Date: Mon, 28 Aug 2023 16:17:36 +0800
Subject: [PATCH] add --original-blob-ids args for merge

the default merge command is to get the name of the original
blob from the bootstrap name, and add a cli args for it

Signed-off-by: zyfjeff <zyfjeff@linux.alibaba.com>
---
 src/bin/nydus-image/main.rs  | 13 ++++++++++++
 src/bin/nydus-image/merge.rs | 40 +++++++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/src/bin/nydus-image/main.rs b/src/bin/nydus-image/main.rs
index b52b11a6a39..70aa84e5097 100644
--- a/src/bin/nydus-image/main.rs
+++ b/src/bin/nydus-image/main.rs
@@ -360,6 +360,12 @@ fn prepare_cmd_args(bti_string: &'static str) -> App {
                         .required(false)
                         .help("RAFS blob digest list separated by comma"),
                 )
+                .arg(
+                    Arg::new("original-blob-ids")
+                        .long("original-blob-ids")
+                        .required(false)
+                        .help("original blob id list separated by comma, it may usually be a sha256 hex string"),
+                )
                 .arg(
                     Arg::new("blob-sizes")
                     .long("blob-sizes")
@@ -903,6 +909,12 @@ impl Command {
                     .map(|item| item.trim().to_string())
                     .collect()
             });
+        let original_blob_ids: Option<Vec<String>> =
+            matches.get_one::<String>("original-blob-ids").map(|list| {
+                list.split(',')
+                    .map(|item| item.trim().to_string())
+                    .collect()
+            });
         let blob_toc_sizes: Option<Vec<u64>> =
             matches.get_one::<String>("blob-toc-sizes").map(|list| {
                 list.split(',')
@@ -943,6 +955,7 @@ impl Command {
             parent_bootstrap_path,
             source_bootstrap_paths,
             blob_digests,
+            original_blob_ids,
             blob_sizes,
             blob_toc_digests,
             blob_toc_sizes,
diff --git a/src/bin/nydus-image/merge.rs b/src/bin/nydus-image/merge.rs
index 28223b849fe..6e4be8f1ab7 100644
--- a/src/bin/nydus-image/merge.rs
+++ b/src/bin/nydus-image/merge.rs
@@ -32,6 +32,20 @@ use crate::core::tree::{MetadataTreeBuilder, Tree};
 pub struct Merger {}
 
 impl Merger {
+    fn get_string_from_list(
+        original_ids: &Option<Vec<String>>,
+        idx: usize,
+    ) -> Result<Option<String>> {
+        Ok(if let Some(id) = &original_ids {
+            let id_string = id
+                .get(idx)
+                .ok_or_else(|| anyhow!("unmatched digest index {}", idx))?;
+            Some(id_string.clone())
+        } else {
+            None
+        })
+    }
+
     fn get_digest_from_list(digests: &Option<Vec<String>>, idx: usize) -> Result<Option<[u8; 32]>> {
         Ok(if let Some(digests) = &digests {
             let digest = digests
@@ -65,6 +79,7 @@ impl Merger {
         parent_bootstrap_path: Option<String>,
         sources: Vec<PathBuf>,
         blob_digests: Option<Vec<String>>,
+        original_blob_ids: Option<Vec<String>>,
         blob_sizes: Option<Vec<u64>>,
         blob_toc_digests: Option<Vec<String>>,
         blob_toc_sizes: Option<Vec<u64>>,
@@ -83,6 +98,22 @@ impl Merger {
                 sources.len(),
             );
         }
+        if let Some(original_ids) = original_blob_ids.as_ref() {
+            ensure!(
+                original_ids.len() == sources.len(),
+                "number of original blob id entries {} doesn't match number of sources {}",
+                original_ids.len(),
+                sources.len(),
+            );
+        }
+        if let Some(sizes) = blob_sizes.as_ref() {
+            ensure!(
+                sizes.len() == sources.len(),
+                "number of blob size entries {} doesn't match number of sources {}",
+                sizes.len(),
+                sources.len(),
+            );
+        }
         if let Some(toc_digests) = blob_toc_digests.as_ref() {
             ensure!(
                 toc_digests.len() == sources.len(),
@@ -186,7 +217,14 @@ impl Merger {
                     } else {
                         // The blob id (blob sha256 hash) in parent bootstrap is invalid for nydusd
                         // runtime, should change it to the hash of whole tar blob.
-                        blob_ctx.blob_id = BlobInfo::get_blob_id_from_meta_path(bootstrap_path)?;
+                        if let Some(original_id) =
+                            Self::get_string_from_list(&original_blob_ids, layer_idx)?
+                        {
+                            blob_ctx.blob_id = original_id;
+                        } else {
+                            blob_ctx.blob_id =
+                                BlobInfo::get_blob_id_from_meta_path(bootstrap_path)?;
+                        }
                     }
                     if let Some(digest) = Self::get_digest_from_list(&blob_digests, layer_idx)? {
                         if blob.has_feature(BlobFeatures::SEPARATE) {