Resolves Vector35#3714: Adds core initialization/shutdown wrapper obj…

…ect (to the Rust API) that you can initialize at the top of your script and it'll auto-cleanup at the end of your script. Also some misc doc fixes.
unknowntrojan · Dec 5, 2023 · ae14406 · ae14406
1 parent b02b658
commit ae14406
Show file tree

Hide file tree

Showing 7 changed files with 75 additions and 27 deletions.
diff --git a/rust/README.md b/rust/README.md
@@ -46,10 +46,13 @@ See the `./examples/`.  Plugin registration commands are in `binaryninja::comman
 binaryninja = { git = "https://github.com/Vector35/binaryninja-api.git", branch = "dev"}
 ```
 
-All standalone binaries should call both `binaryninja::headless::init()` and `binaryninja::headless::shutdown()`.  
 All standalone binaries need to provide a `build.rs`.  
 See [`examples/template`](examples/template) for details.  
 
+## Docs
+
+Docs can be found at https://dev-rust.binary.ninja/
+
 ---
 
 #### Attribution

diff --git a/rust/examples/basic_script/src/main.rs b/rust/examples/basic_script/src/main.rs
@@ -3,10 +3,12 @@ use binaryninja::binaryview::{BinaryViewBase, BinaryViewExt};
 
 fn main() {
     println!("Loading plugins..."); // This loads all the core architecture, platform, etc plugins
-    binaryninja::headless::init();
+    let headless_session = binaryninja::headless::Session::new();
 
     println!("Loading binary...");
-    let bv = binaryninja::load("/bin/cat").expect("Couldn't open `/bin/cat`");
+    let bv = headless_session
+        .load("/bin/cat")
+        .expect("Couldn't open `/bin/cat`");
 
     println!("Filename:  `{}`", bv.file().filename());
     println!("File size: `{:#x}`", bv.len());
@@ -32,7 +34,4 @@ fn main() {
             }
         }
     }
-
-    // Important!  You need to call shutdown or your script will hang forever
-    binaryninja::headless::shutdown();
 }
diff --git a/rust/examples/mlil_lifter/src/main.rs b/rust/examples/mlil_lifter/src/main.rs
@@ -14,7 +14,7 @@ fn main() {
     // This loads all the core architecture, platform, etc plugins
     // Standalone executables probably need to call this, but plugins do not
     println!("Loading plugins...");
-    binaryninja::headless::init();
+    let _headless_session = binaryninja::headless::Session::new();
 
     // Your code here...
     println!("Loading binary...");
@@ -46,7 +46,4 @@ fn main() {
         }
         println!();
     }
-
-    // Important!  Standalone executables need to call shutdown or they will hang forever
-    binaryninja::headless::shutdown();
 }
diff --git a/rust/examples/mlil_visitor/src/main.rs b/rust/examples/mlil_visitor/src/main.rs
@@ -235,7 +235,7 @@ fn main() {
     // This loads all the core architecture, platform, etc plugins
     // Standalone executables probably need to call this, but plugins do not
     println!("Loading plugins...");
-    binaryninja::headless::init();
+    let _headless_session = binaryninja::headless::Session::new();
 
     // Your code here...
     println!("Loading binary...");
@@ -261,7 +261,4 @@ fn main() {
         }
         println!();
     }
-
-    // Important!  Standalone executables need to call shutdown or they will hang forever
-    binaryninja::headless::shutdown();
 }
diff --git a/rust/examples/template/src/main.rs b/rust/examples/template/src/main.rs
@@ -7,11 +7,13 @@ fn main() {
     // This loads all the core architecture, platform, etc plugins
     // Standalone executables probably need to call this, but plugins do not
     println!("Loading plugins...");
-    binaryninja::headless::init();
+    let headless_session = binaryninja::headless::Session::new();
 
     // Your code here...
     println!("Loading binary...");
-    let bv = binaryninja::load("/bin/cat").expect("Couldn't open `/bin/cat`");
+    let bv = headless_session
+        .load("/bin/cat")
+        .expect("Couldn't open `/bin/cat`");
 
     println!("Filename:  `{}`", bv.file().filename());
     println!("File size: `{:#x}`", bv.len());
@@ -37,7 +39,4 @@ fn main() {
             }
         }
     }
-
-    // Important!  Standalone executables need to call shutdown or they will hang forever
-    binaryninja::headless::shutdown();
 }
diff --git a/rust/src/headless.rs b/rust/src/headless.rs
@@ -12,7 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::string::BnStrCompatible;
+use crate::{
+    binaryview,
+    metadata::Metadata,
+    rc::{self, Ref},
+    string::BnStrCompatible,
+};
 
 use std::env;
 use std::path::PathBuf;
@@ -70,6 +75,8 @@ use binaryninjacore_sys::{BNInitPlugins, BNInitRepoPlugins, BNSetBundledPluginDi
 /// Loads plugins, core architecture, platform, etc.
 ///
 /// ⚠️ Important! Must be called at the beginning of scripts.  Plugins do not need to call this. ⚠️
+///
+/// You can instead call this through [`Session`] or [`script_helper`]
 pub fn init() {
     unsafe {
         let path = binja_path().join("plugins").into_os_string();
@@ -102,3 +109,50 @@ pub fn script_helper(func: fn()) {
     func();
     shutdown();
 }
+
+/// Wrapper for [`init`] and [`shutdown`]. Instantiating this at the top of your script will initialize everything correctly and then clean itself up at exit as well.
+pub struct Session {}
+
+impl Session {
+    pub fn new() -> Self {
+        init();
+        Self {}
+    }
+
+    /// ```rust
+    /// let headless_session = binaryninja::headless::Session::new();
+    ///
+    /// let bv = headless_session.load("/bin/cat").expect("Couldn't open `/bin/cat`");
+    /// ```
+    pub fn load(&self, filename: &str) -> Option<rc::Ref<binaryview::BinaryView>> {
+        crate::load(filename)
+    }
+
+    /// ```rust
+    /// let settings = [("analysis.linearSweep.autorun", false)].into();
+    /// let headless_session = binaryninja::headless::Session::new();
+    ///
+    /// let bv = headless_session.load_with_options("/bin/cat", true, Some(settings))
+    ///     .expect("Couldn't open `/bin/cat`");
+    /// ```
+    pub fn load_with_options(
+        &self,
+        filename: &str,
+        update_analysis_and_wait: bool,
+        options: Option<Ref<Metadata>>,
+    ) -> Option<rc::Ref<binaryview::BinaryView>> {
+        crate::load_with_options(filename, update_analysis_and_wait, options)
+    }
+}
+
+impl Default for Session {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Drop for Session {
+    fn drop(&mut self) {
+        shutdown()
+    }
+}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
@@ -72,20 +72,17 @@
 //! The most up-to-date version of the suggested [`build.rs` is here].
 //!
 //! ### `main.rs`
-//! All standalone binaries need to call [`headless::init()`] at start and [`headless::shutdown()`] at shutdown.
+//! Standalone binaries need to initialize Binary Ninja before they can work. You can do this through [`headless::Session`], [`headless::script_helper`], or [`headless::init()`] at start and [`headless::shutdown()`] at shutdown.
 //! ```rust
 //! fn main() {
 //!     // This loads all the core architecture, platform, etc plugins
 //!     // Standalone executables need to call this, but plugins do not
-//!     binaryninja::headless::init();
+//!     let headless_session = binaryninja::headless::Session::new();
 //!
 //!     println!("Loading binary...");
-//!     let bv = binaryninja::load("/bin/cat").expect("Couldn't open `/bin/cat`");
+//!     let bv = headless_session.load("/bin/cat").expect("Couldn't open `/bin/cat`");
 //!
 //!     // Your code here...
-//!
-//!     // Important!  Standalone executables need to call shutdown or they will hang forever
-//!     binaryninja::headless::shutdown();
 //! }
 //! ```
 //!
@@ -123,7 +120,6 @@ extern crate rayon;
 // cc possible values
 // bv reorg
 // core fileaccessor (for bv saving)
-// headless wrapper for shutdown
 // platform cc
 
 #[macro_use]
@@ -196,6 +192,7 @@ use string::BnStrCompatible;
 const BN_FULL_CONFIDENCE: u8 = 255;
 const BN_INVALID_EXPR: usize = usize::MAX;
 
+/// The main way to open and load files into Binary Ninja. Make sure you've properly initialized the core before calling this function. See [`crate::headless::init()`]
 pub fn load<S: BnStrCompatible>(filename: S) -> Option<rc::Ref<binaryview::BinaryView>> {
     let filename = filename.into_bytes_with_nul();
     let metadata = Metadata::new_of_type(MetadataType::KeyValueDataType);
@@ -216,6 +213,8 @@ pub fn load<S: BnStrCompatible>(filename: S) -> Option<rc::Ref<binaryview::Binar
     }
 }
 
+/// The main way to open and load files (with options) into Binary Ninja. Make sure you've properly initialized the core before calling this function. See [`crate::headless::init()`]
+///
 /// ```rust
 /// let settings = [("analysis.linearSweep.autorun", false)].into();
 ///