From ff99a6766621d5a5f5e52b04750ff363e9a7275c 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>
---
 builder/src/merge.rs        | 32 +++++++++++++++++++++++++++++++-
 src/bin/nydus-image/main.rs | 14 ++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/builder/src/merge.rs b/builder/src/merge.rs
index 8301b32cf8d..56b0c8d3a94 100644
--- a/builder/src/merge.rs
+++ b/builder/src/merge.rs
@@ -29,6 +29,20 @@ use super::{
 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
@@ -62,6 +76,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>>,
@@ -80,6 +95,14 @@ 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(),
@@ -194,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) {
diff --git a/src/bin/nydus-image/main.rs b/src/bin/nydus-image/main.rs
index f551930b78b..090b7c77c13 100644
--- a/src/bin/nydus-image/main.rs
+++ b/src/bin/nydus-image/main.rs
@@ -431,6 +431,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")
@@ -1194,6 +1200,13 @@ impl Command {
                     .map(|item| item.trim().to_string())
                     .collect()
             });
+        let original_blob_digest: Option<Vec<String>> = matches
+            .get_one::<String>("original-blob-digest")
+            .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(',')
@@ -1234,6 +1247,7 @@ impl Command {
             parent_bootstrap_path,
             source_bootstrap_paths,
             blob_digests,
+            original_blob_digest,
             blob_sizes,
             blob_toc_digests,
             blob_toc_sizes,