make idb_import a background task

Author: rbran

plugins/idb_import/src/lib.rs

@@ -2,6 +2,7 @@ mod types;
 use std::fs::File;
 use std::io::BufReader;
 
+use binaryninja::background_task::BackgroundTask;
 use binaryninja::command::Command;
 use binaryninja::string::BnStrCompatible;
 use binaryninja::type_library::TypeLibrary;
@@ -84,7 +85,7 @@ struct LoadTilFile;
 
 impl Command for LoadTilFile {
     fn action(&self, view: &BinaryView) {
-        if let Err(error) = interactive_import_til(view) {
+        if let Err(error) = background_import_til(view) {
             error!("Unable to convert TIL file: {error}");
         }
     }
@@ -126,10 +127,6 @@ impl std::io::Seek for BinaryViewReader<'_> {
     }
 }
 
-fn dummy_progress(_: usize, _: usize) -> Result<(), ()> {
-    Ok(())
-}
-
 fn import_idb_info<P: Fn(usize, usize) -> Result<(), ()>>(
     debug_info: &mut DebugInfo,
     bv: &BinaryView,
@@ -181,7 +178,18 @@ fn import_til_info_from_debug_file<P: Fn(usize, usize) -> Result<(), ()>>(
     import_til_to_type_library(til, filename, bv, progress)
 }
 
+fn background_import_til(view: &BinaryView) -> Result<()> {
+    let moved_view = view.to_owned();
+    binaryninja::worker_thread::execute_on_worker_thread_interactive(c"Til Import", move || {
+        if let Err(err) = interactive_import_til(&moved_view) {
+            error!("Unable to import TIL: {err}");
+        }
+    });
+    Ok(())
+}
+
 fn interactive_import_til(view: &BinaryView) -> Result<()> {
+    let bt = BackgroundTask::new("Import TIL", true);
     let Some(file) =
         binaryninja::interaction::get_open_filename_input("Select a .til file", "*.til")
     else {
@@ -192,8 +200,19 @@ fn interactive_import_til(view: &BinaryView) -> Result<()> {
     let mut file = BufReader::new(File::open(&file)?);
     let til = TILSection::read(&mut file, idb_rs::IDBSectionCompression::None)?;
 
-    // TODO remove the extension?
-    import_til_to_type_library(til, filename, view, dummy_progress)
+    let progress = |current, total| {
+        if bt.is_cancelled() {
+            return Err(());
+        }
+        bt.set_progress_text(format!(
+            "Import TIL progress: {}%",
+            ((current as f32 / total as f32) * 100f32) as u32
+        ));
+        Ok(())
+    };
+    import_til_to_type_library(til, filename, view, progress)?;
+    bt.finish();
+    Ok(())
 }
 
 fn import_til_to_type_library<S, P>(
@@ -271,7 +290,7 @@ fn import_til_section<P: Fn(usize, usize) -> Result<(), ()>>(
     let default_arch = view
         .default_arch()
         .ok_or_else(|| anyhow!("Unable to get the default arch"))?;
-    let types = types::translate_tils_types(default_arch, tils, progress)?;
+    let types = types::translate_til_types(default_arch, tils, progress)?;
 
     // print any errors
     print_til_convertsion_errors(&types)?;
@@ -437,11 +456,11 @@ pub extern "C" fn CorePluginInit() -> bool {
         .with_level(LevelFilter::Error)
         .init();
     binaryninja::command::register_command(
-        "Import TIL types",
-        "Convert and import a TIL file into a TypeLibrary",
+        c"Import TIL types",
+        c"Convert and import a TIL file into a TypeLibrary",
         LoadTilFile,
     );
-    DebugInfoParser::register("idb_parser", IDBDebugInfoParser);
-    DebugInfoParser::register("til_parser", TILDebugInfoParser);
+    DebugInfoParser::register(c"idb_parser", IDBDebugInfoParser);
+    DebugInfoParser::register(c"til_parser", TILDebugInfoParser);
     true
 }

plugins/idb_import/src/types.rs

@@ -114,7 +114,6 @@ pub struct TranslatesIDBType<'a> {
 
 pub struct TranslateIDBTypes<'a, F> {
     pub arch: CoreArchitecture,
-    pub total: usize,
     pub progress: F,
     // key til addr and type idx
     pub types_by_idx: HashMap<(usize, usize), usize>,
@@ -129,19 +128,16 @@ where
     F: Fn(usize, usize) -> Result<(), ()>,
 {
     pub fn from_types(arch: CoreArchitecture, tils: &'a [TILSection], progress: F) -> Result<Self> {
-        let total = tils
-            .iter()
-            .map(|til| til.symbols.len() + til.types.len())
-            .sum();
-        let mut types = Vec::with_capacity(total);
-        let mut types_by_idx = HashMap::with_capacity(total);
-        let mut types_by_ord = HashMap::with_capacity(total);
-        let mut types_by_name = HashMap::with_capacity(total);
         let all_types = tils.iter().flat_map(|til| {
             core::iter::repeat(til).zip(til.types.iter().enumerate().zip(core::iter::repeat(false)))
             // TODO: it's unclear how the demangle symbols and types names/ord, for now only parse types
             //.chain(til.symbols.iter().zip(core::iter::repeat(true)))
         });
+        let total = all_types.clone().count();
+        let mut types = Vec::with_capacity(total);
+        let mut types_by_idx = HashMap::with_capacity(total);
+        let mut types_by_ord = HashMap::with_capacity(total);
+        let mut types_by_name = HashMap::with_capacity(total);
         for (i, (til, ((ty_idx, ty), is_symbol))) in all_types.enumerate() {
             // TODO sanitized the input
             // TODO find out how the namespaces used by TIL works
@@ -189,7 +185,6 @@ where
         }
 
         Ok(TranslateIDBTypes {
-            total,
             arch,
             progress,
             tils,
@@ -201,7 +196,7 @@ where
     }
 
     fn resolve(&mut self) -> Result<()> {
-        if (self.progress)(0, self.total).is_err() {
+        if (self.progress)(0, self.types.len()).is_err() {
             return Err(anyhow!("IDB import aborted"));
         }
 
@@ -238,14 +233,15 @@ where
                 }
 
                 // count the number of finished types
-                if let TranslateTypeResult::Translated(_) = &self.types[i].ty {
-                    num_translated += 1
+                let was_translated =
+                    matches!(&self.types[i].ty, TranslateTypeResult::Translated(_));
+                if was_translated {
+                    num_translated += 1;
                 }
-
-                all_done &= matches!(&self.types[i].ty, TranslateTypeResult::Translated(_));
+                all_done &= was_translated;
             }
 
-            if (self.progress)(num_translated, self.total).is_err() {
+            if (self.progress)(num_translated, self.types.len()).is_err() {
                 // error means the user aborted the progress
                 return Err(anyhow!("User aborted during processing"));
             }
@@ -770,7 +766,6 @@ pub fn translate_ephemeral_type(debug_file: &BinaryView, ty: &TILType) -> Transl
             macros: None,
         }],
         types: vec![],
-        total: 0,
         types_by_idx: HashMap::new(),
         types_by_ord: HashMap::new(),
         types_by_name: HashMap::new(),
@@ -779,14 +774,6 @@ pub fn translate_ephemeral_type(debug_file: &BinaryView, ty: &TILType) -> Transl
     translator.translate_type(&translator.tils[0], ty)
 }
 
-pub fn translate_tils_types(
-    arch: CoreArchitecture,
-    tils: &[TILSection],
-    progress: impl Fn(usize, usize) -> Result<(), ()>,
-) -> Result<Vec<TranslatesIDBType<'_>>> {
-    translate_til_types(arch, tils, progress)
-}
-
 pub fn translate_til_types(
     arch: CoreArchitecture,
     tils: &[TILSection],