[WARP] Only merge chunks in create from view command if an existing file was given

Fixes the case where we are generating many smaller chunks for a given view and then unintentionally merging
them back together, throwing away all data.

1 files changed, 20 insertions, 6 deletions

diff --git a/plugins/warp/src/plugin/create.rs b/plugins/warp/src/plugin/create.rs
index 0410e74a..a74310a5 100644
--- a/plugins/warp/src/plugin/create.rs
+++ b/plugins/warp/src/plugin/create.rs
@@ -164,8 +164,8 @@ impl CreateFromCurrentView {
 
         if let Err(err) = file {
             binaryninja::interaction::show_message_box(
-                "Error",
-                &format!("Failed to create signature file: {}", err),
+                "Failed to create signature file",
+                &err.to_string(),
                 MessageBoxButtonSet::OKButtonSet,
                 MessageBoxIcon::ErrorIcon,
             );
@@ -173,16 +173,30 @@ impl CreateFromCurrentView {
             return None;
         }
 
+        let background_task = BackgroundTask::new("Creating WARP File...", false);
         let mut file = file.unwrap();
         // Add back the existing chunks if the user selected to keep them.
-        file.chunks.extend(existing_chunks);
-        // TODO: Make merging optional?
-        file.chunks = Chunk::merge(&file.chunks, compression_type.into());
+        if !existing_chunks.is_empty() {
+            file.chunks.extend(existing_chunks);
+            // TODO: Make merging optional?
+            // TODO: Merging can lose chunk data if it goes above the maximum table count.
+            // TODO: We should probably solve that in the warp crate itself?
+            file.chunks = Chunk::merge(&file.chunks, compression_type.into());
 
-        if std::fs::write(&file_path, file.to_bytes()).is_err() {
+            // After merging, we should have at least one chunk. If not, merging actually removed data.
+            if file.chunks.len() < 1 {
+                log::error!("Failed to merge chunks! Please report this, it should not happen.");
+                return None;
+            }
+        }
+
+        let file_bytes = file.to_bytes();
+        let file_size = file_bytes.len();
+        if std::fs::write(&file_path, file_bytes).is_err() {
             log::error!("Failed to write data to signature file!");
         }
         log::info!("Saved signature file to: '{}'", file_path.display());
+        background_task.finish();
 
         // Show a report of the generate signatures, if desired.
         let report_generator = ReportGenerator::new();