Improve headless initialization in Rust API

- Add option to disable loading of user and repo plugins
- Prevent shutting down the core while other sessions are still running

The reason we need to prevent shutting down the core is if one session is dropped and another session has worker actions queued than they would be cleared prematurely, and also the enterprise license would be revoked prematurely as well.

1 files changed, 56 insertions, 6 deletions

diff --git a/rust/src/headless.rs b/rust/src/headless.rs
index 78c04a96..2cc87056 100644
--- a/rust/src/headless.rs
+++ b/rust/src/headless.rs
@@ -18,6 +18,8 @@ use crate::{
 };
 use std::io;
 use std::path::{Path, PathBuf};
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering::SeqCst;
 use thiserror::Error;
 
 use crate::enterprise::release_license;
@@ -32,6 +34,9 @@ use std::time::Duration;
 
 static MAIN_THREAD_HANDLE: Mutex<Option<JoinHandle<()>>> = Mutex::new(None);
 
+/// Used to prevent shutting down Binary Ninja if there are other [`Session`]'s.
+static SESSION_COUNT: AtomicUsize = AtomicUsize::new(0);
+
 #[derive(Error, Debug)]
 pub enum InitializationError {
     #[error("main thread could not be started: {0}")]
@@ -91,6 +96,14 @@ pub struct InitializationOptions {
     pub floating_license_duration: Duration,
     /// The bundled plugin directory to use.
     pub bundled_plugin_directory: PathBuf,
+    /// Whether to initialize user plugins.
+    ///
+    /// Set this to false if your use might be impacted by a user installed plugin.
+    pub user_plugins: bool,
+    /// Whether to initialize repo plugins.
+    ///
+    /// Set this to false if your use might be impacted by a repo installed plugin.
+    pub repo_plugins: bool,
 }
 
 impl InitializationOptions {
@@ -98,6 +111,14 @@ impl InitializationOptions {
         Self::default()
     }
 
+    pub fn minimal() -> Self {
+        Self {
+            user_plugins: false,
+            repo_plugins: false,
+            ..Self::default()
+        }
+    }
+
     /// A license to override with, you can use this to make sure you initialize with a specific license.
     ///
     /// This takes the form of a JSON array. The string should be formed like:
@@ -112,7 +133,7 @@ impl InitializationOptions {
     /// If you need to make sure that you do not check out a license set this to false.
     ///
     /// This is really only useful if you have a headless license but are using an enterprise enabled core.
-    pub fn with_checkout_license(mut self, should_checkout: bool) -> Self {
+    pub fn with_license_checkout(mut self, should_checkout: bool) -> Self {
         self.checkout_license = should_checkout;
         self
     }
@@ -130,6 +151,18 @@ impl InitializationOptions {
         self.floating_license_duration = duration;
         self
     }
+
+    /// Set this to false if your use might be impacted by a user installed plugin.
+    pub fn with_user_plugins(mut self, should_initialize: bool) -> Self {
+        self.user_plugins = should_initialize;
+        self
+    }
+
+    /// Set this to false if your use might be impacted by a repo installed plugin.
+    pub fn with_repo_plugins(mut self, should_initialize: bool) -> Self {
+        self.repo_plugins = should_initialize;
+        self
+    }
 }
 
 impl Default for InitializationOptions {
@@ -141,6 +174,8 @@ impl Default for InitializationOptions {
             floating_license_duration: Duration::from_secs(900),
             bundled_plugin_directory: bundled_plugin_directory()
                 .expect("Failed to get bundled plugin directory"),
+            user_plugins: true,
+            repo_plugins: true,
         }
     }
 }
@@ -188,8 +223,11 @@ pub fn init_with_opts(options: InitializationOptions) -> Result<(), Initializati
     set_bundled_plugin_directory(options.bundled_plugin_directory);
 
     unsafe {
-        BNInitPlugins(true);
-        BNInitRepoPlugins();
+        BNInitPlugins(options.user_plugins);
+        if options.repo_plugins {
+            // We are allowed to initialize repo plugins, so do it!
+            BNInitRepoPlugins();
+        }
     }
 
     if !is_license_validated() {
@@ -249,6 +287,14 @@ pub fn license_location() -> Option<LicenseLocation> {
 pub struct Session {}
 
 impl Session {
+    /// Get a registered [`Session`] for use.
+    ///
+    /// This is required so that we can keep track of the [`SESSION_COUNT`].
+    fn registered_session() -> Self {
+        let _previous_count = SESSION_COUNT.fetch_add(1, SeqCst);
+        Self {}
+    }
+
     /// Before calling new you must make sure that the license is retrievable, otherwise the core won't be able to initialize.
     ///
     /// If you cannot otherwise provide a license via `BN_LICENSE_FILE` environment variable or the Binary Ninja user directory
@@ -257,7 +303,7 @@ impl Session {
         if license_location().is_some() {
             // We were able to locate a license, continue with initialization.
             init()?;
-            Ok(Self {})
+            Ok(Self::registered_session())
         } else {
             // There was no license that could be automatically retrieved, you must call [Self::new_with_license].
             Err(InitializationError::NoLicenseFound)
@@ -270,7 +316,7 @@ impl Session {
     /// can discover by itself, therefor it is expected that you know where your license is when calling this directly.
     pub fn new_with_opts(options: InitializationOptions) -> Result<Self, InitializationError> {
         init_with_opts(options)?;
-        Ok(Self {})
+        Ok(Self::registered_session())
     }
 
     /// ```no_run
@@ -364,6 +410,10 @@ impl Session {
 
 impl Drop for Session {
     fn drop(&mut self) {
-        shutdown()
+        let previous_count = SESSION_COUNT.fetch_sub(1, SeqCst);
+        if previous_count == 1 {
+            // We were the last session, therefor we can safely shut down.
+            shutdown();
+        }
     }
 }