[Rust] Simplify enterprise initialization

We don't need to introduce extra global state when the goal is to not release floating licenses in shutdown.

2 files changed, 88 insertions, 96 deletions

diff --git a/rust/src/enterprise.rs b/rust/src/enterprise.rs
index 9b97a387..d09ac333 100644
--- a/rust/src/enterprise.rs
+++ b/rust/src/enterprise.rs
@@ -19,13 +19,23 @@ pub enum EnterpriseCheckoutError {
     RefreshExpiredLicenseFailed(String),
 }
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub enum EnterpriseCheckoutStatus {
+    /// The UI is managing the enterprise checkout.
+    AlreadyManaged,
+    /// Checkout was successful, attached duration is the duration of the license, if any.
+    Success(Option<Duration>),
+}
+
 /// Initialize the enterprise server connection to check out a floating license.
 /// Result value is if we actually checked out a license (i.e. Ok(false) means we already have a
 /// license checked out and will not need to release it later)
-pub fn checkout_license(duration: Duration) -> Result<bool, EnterpriseCheckoutError> {
+pub fn checkout_license(
+    duration: Duration,
+) -> Result<EnterpriseCheckoutStatus, EnterpriseCheckoutError> {
     if crate::is_ui_enabled() {
         // We only need to check out a license if running headlessly.
-        return Ok(false);
+        return Ok(EnterpriseCheckoutStatus::AlreadyManaged);
     }
 
     // The disparate core functions we call here might already have mutexes to guard.
@@ -67,20 +77,19 @@ pub fn checkout_license(duration: Duration) -> Result<bool, EnterpriseCheckoutEr
     if !is_server_license_still_activated()
         || (!is_server_floating_license() && crate::license_expiration_time() < SystemTime::now())
     {
-        // If the license is expired we should refresh the license.
+        // If the license is expired, we should refresh the license.
         if !update_server_license(duration) {
             let last_error = server_last_error().to_string();
             return Err(EnterpriseCheckoutError::RefreshExpiredLicenseFailed(
                 last_error,
             ));
         }
-        Ok(true)
-    } else {
-        Ok(false)
     }
+
+    Ok(EnterpriseCheckoutStatus::Success(license_duration()))
 }
 
-pub fn release_license() {
+pub fn release_license(release_floating: bool) {
     if !crate::is_ui_enabled() {
         // This might look dumb, why would we want to connect to the server, would that not just mean
         // we don't need to release the license? Well no, you could have run a script, acquired a license for 10 hours
@@ -92,6 +101,10 @@ pub fn release_license() {
         if !is_server_connected() {
             connect_server();
         }
+        // We optionally release floating licenses as users typically want to keep them around.
+        if is_server_floating_license() && !release_floating {
+            return;
+        }
         // We should only release the license if we are running headlessly.
         release_server_license();
     }
@@ -141,8 +154,14 @@ pub fn server_token() -> String {
     unsafe { BnString::into_string(binaryninjacore_sys::BNGetEnterpriseServerToken()) }
 }
 
-pub fn license_duration() -> Duration {
-    Duration::from_secs(unsafe { binaryninjacore_sys::BNGetEnterpriseServerLicenseDuration() })
+pub fn license_duration() -> Option<Duration> {
+    let duration =
+        Duration::from_secs(unsafe { binaryninjacore_sys::BNGetEnterpriseServerLicenseDuration() });
+    match duration {
+        // If the core returns 0 there is no license duration.
+        Duration::ZERO => None,
+        _ => Some(duration),
+    }
 }
 
 pub fn license_expiration_time() -> SystemTime {
diff --git a/rust/src/headless.rs b/rust/src/headless.rs
index 32286331..7279a459 100644
--- a/rust/src/headless.rs
+++ b/rust/src/headless.rs
@@ -14,14 +14,15 @@
 
 use crate::{
     binary_view, bundled_plugin_directory, enterprise, is_license_validated, is_main_thread,
-    license_path, set_bundled_plugin_directory, set_license, string::IntoJson,
+    is_ui_enabled, license_path, set_bundled_plugin_directory, set_license, string::IntoJson,
 };
 use std::io;
 use std::path::{Path, PathBuf};
+use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::SeqCst;
-use std::sync::atomic::{AtomicBool, AtomicUsize};
 use thiserror::Error;
 
+use crate::enterprise::EnterpriseCheckoutStatus;
 use crate::main_thread::{MainThreadAction, MainThreadHandler};
 use crate::progress::ProgressCallback;
 use crate::rc::Ref;
@@ -32,12 +33,9 @@ use std::thread::JoinHandle;
 use std::time::Duration;
 
 static MAIN_THREAD_HANDLE: Mutex<Option<JoinHandle<()>>> = Mutex::new(None);
-/// Prevent two threads from calling init() at the same time
-static INIT_LOCK: Mutex<()> = Mutex::new(());
-/// Used to prevent shutting down Binary Ninja if there are other [`Session`]'s.
+
+/// Used to prevent shutting down Binary Ninja if there is another active [`Session`].
 static SESSION_COUNT: AtomicUsize = AtomicUsize::new(0);
-/// If we checked out a floating license and should release it on shutdown
-static NEED_LICENSE_RELEASE: AtomicBool = AtomicBool::new(false);
 
 #[derive(Error, Debug)]
 pub enum InitializationError {
@@ -49,15 +47,15 @@ pub enum InitializationError {
     InvalidLicense,
     #[error("no license could located, please see `binaryninja::set_license` for details")]
     NoLicenseFound,
-    #[error("could not acquire initialization mutex")]
-    InitMutex,
+    #[error("initialization already managed by ui")]
+    AlreadyManaged,
 }
 
 /// 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`].
+/// The preferred method for core initialization is [`Session`], use that instead of this where possible.
 ///
 /// If you need to customize initialization, use [`init_with_opts`] instead.
 pub fn init() -> Result<(), InitializationError> {
@@ -67,15 +65,12 @@ pub fn init() -> Result<(), InitializationError> {
 
 /// Unloads plugins, stops all worker threads, and closes open logs.
 ///
-/// If the core was initialized using an enterprise license, that will also be freed.
-///
-/// ⚠️ Important! Must be called at the end of scripts. ⚠️
+/// This function does _NOT_ release floating licenses; it is expected that you call [`enterprise::release_license`].
 pub fn shutdown() {
-    match crate::product().to_string().as_str() {
+    match crate::product().as_str() {
         "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" => {
-            if NEED_LICENSE_RELEASE.load(SeqCst) {
-                enterprise::release_license()
-            }
+            // By default, we do not release floating licenses.
+            enterprise::release_license(false)
         }
         _ => {}
     }
@@ -91,9 +86,9 @@ pub fn is_shutdown_requested() -> bool {
 pub struct InitializationOptions {
     /// A license to override with, you can use this to make sure you initialize with a specific license.
     pub license: Option<String>,
-    /// If you need to make sure that you do not check out a license set this to false.
+    /// 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.
+    /// This is really only useful if you have a headless license but are using an enterprise-enabled core.
     pub checkout_license: bool,
     /// Whether to register the default main thread handler.
     ///
@@ -105,11 +100,11 @@ pub struct InitializationOptions {
     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.
+    /// 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.
+    /// Set this to false if your use might be impacted by a repo-installed plugin.
     pub repo_plugins: bool,
 }
 
@@ -129,9 +124,9 @@ impl InitializationOptions {
         self
     }
 
-    /// If you need to make sure that you do not check out a license set this to false.
+    /// 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.
+    /// This is really only useful if you have a headless license but are using an enterprise-enabled core.
     pub fn with_license_checkout(mut self, should_checkout: bool) -> Self {
         self.checkout_license = should_checkout;
         self
@@ -151,13 +146,13 @@ impl InitializationOptions {
         self
     }
 
-    /// Set this to false if your use might be impacted by a user installed plugin.
+    /// 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.
+    /// 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
@@ -181,7 +176,11 @@ impl Default for InitializationOptions {
 
 /// This initializes the core with the given [`InitializationOptions`].
 pub fn init_with_opts(options: InitializationOptions) -> Result<(), InitializationError> {
-    // If we are the main thread that means there is no main thread, we should register a main thread handler.
+    if is_ui_enabled() {
+        return Err(InitializationError::AlreadyManaged);
+    }
+
+    // If we are the main thread, that means there is no main thread, we should register a main thread handler.
     if options.register_main_thread_handler && is_main_thread() {
         let mut main_thread_handle = MAIN_THREAD_HANDLE.lock().unwrap();
         if main_thread_handle.is_none() {
@@ -192,7 +191,7 @@ pub fn init_with_opts(options: InitializationOptions) -> Result<(), Initializati
             let join_handle = std::thread::Builder::new()
                 .name("HeadlessMainThread".to_string())
                 .spawn(move || {
-                    // We must register the main thread within said thread.
+                    // We must register the main thread within the thread.
                     main_thread.register();
                     while let Ok(action) = receiver.recv() {
                         action.execute();
@@ -204,16 +203,13 @@ pub fn init_with_opts(options: InitializationOptions) -> Result<(), Initializati
         }
     }
 
-    match crate::product().as_str() {
-        "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" => {
-            if options.checkout_license {
-                // We are allowed to check out a license, so do it!
-                if enterprise::checkout_license(options.floating_license_duration)? {
-                    NEED_LICENSE_RELEASE.store(true, SeqCst);
-                }
-            }
+    if is_enterprise_product() && options.checkout_license {
+        // We are allowed to check out a license, so do it!
+        let checkout_status = enterprise::checkout_license(options.floating_license_duration)?;
+        if checkout_status == EnterpriseCheckoutStatus::AlreadyManaged {
+            // Should be impossible, but just in case.
+            return Err(InitializationError::AlreadyManaged);
         }
-        _ => {}
     }
 
     if let Some(license) = options.license {
@@ -232,7 +228,7 @@ pub fn init_with_opts(options: InitializationOptions) -> Result<(), Initializati
     }
 
     if !is_license_validated() {
-        // Unfortunately you must have a valid license to use Binary Ninja.
+        // Unfortunately, you must have a valid license to use Binary Ninja.
         Err(InitializationError::InvalidLicense)
     } else {
         Ok(())
@@ -258,20 +254,27 @@ impl MainThreadHandler for HeadlessMainThreadSender {
     }
 }
 
+fn is_enterprise_product() -> bool {
+    match crate::product().as_str() {
+        "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" => true,
+        _ => false,
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
 pub enum LicenseLocation {
     /// The license used when initializing will be the environment variable `BN_LICENSE`.
     EnvironmentVariable,
     /// The license used when initializing will be the file in the Binary Ninja user directory.
     File,
-    /// The license is retrieved using keychain credentials for `BN_ENTERPRISE_USERNAME`.
+    /// The license is retrieved using keychain credentials, this is only available for floating enterprise licenses.
     Keychain,
 }
 
 /// Attempts to identify the license location type, this follows the same order as core initialization.
 ///
 /// This is useful if you want to know whether the core will use your license. If this returns `None`
-/// you should look setting the `BN_LICENSE` environment variable, or calling [`set_license`].
+/// you should look into setting the `BN_LICENSE` environment variable or calling [`set_license`].
 pub fn license_location() -> Option<LicenseLocation> {
     match std::env::var("BN_LICENSE") {
         Ok(_) => Some(LicenseLocation::EnvironmentVariable),
@@ -279,33 +282,26 @@ pub fn license_location() -> Option<LicenseLocation> {
             // Check the license_path to see if a file is there.
             if license_path().exists() {
                 Some(LicenseLocation::File)
-            } else {
-                // Check to see if we might be authorizing with enterprise
-                if crate::product().as_str() == "Binary Ninja Enterprise Client"
-                    || crate::product().as_str() == "Binary Ninja Ultimate"
-                {
-                    // If we can't initialize enterprise, we probably are missing enterprise.server.url
-                    // and our license surely is not valid.
-                    if !enterprise::is_server_initialized() && !enterprise::initialize_server() {
-                        return None;
-                    }
-                    // If Enterprise thinks we are using a floating license, then report it will be in the keychain
-                    if enterprise::is_server_floating_license() {
-                        Some(LicenseLocation::Keychain)
-                    } else {
-                        None
-                    }
-                } else {
-                    None
+            } else if is_enterprise_product() {
+                // If we can't initialize enterprise, we probably are missing enterprise.server.url
+                // and our license surely is not valid.
+                if !enterprise::is_server_initialized() && !enterprise::initialize_server() {
+                    return None;
                 }
+                // If Enterprise thinks we are using a floating license, then report it will be in the keychain
+                enterprise::is_server_floating_license().then_some(LicenseLocation::Keychain)
+            } else {
+                // If we are not using an enterprise license, we can't check the keychain, nowhere else to check.
+                None
             }
         }
     }
 }
 
 /// 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.
+#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
 pub struct Session {
-    index: usize,
+    license_duration: Option<Duration>,
 }
 
 impl Session {
@@ -313,10 +309,8 @@ impl Session {
     ///
     /// 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 {
-            index: previous_count,
-        }
+        let _previous_count = SESSION_COUNT.fetch_add(1, SeqCst);
+        Self::default()
     }
 
     /// Before calling new you must make sure that the license is retrievable, otherwise the core won't be able to initialize.
@@ -326,19 +320,7 @@ impl Session {
     pub fn new() -> Result<Self, InitializationError> {
         if license_location().is_some() {
             // We were able to locate a license, continue with initialization.
-
-            // Grab the lock before initialization to prevent another thread from initializing
-            // and racing the call to BNInitPlugins.
-            let _lock = INIT_LOCK
-                .lock()
-                .map_err(|_| InitializationError::InitMutex)?;
-            let session = Self::registered_session();
-            // Since this whole section is locked, we're guaranteed to be index 0 if we're first.
-            // Only the first thread hitting this should be allowed to call BNInitPlugins
-            if session.index == 0 {
-                init()?;
-            }
-            Ok(session)
+            Self::new_with_opts(InitializationOptions::default())
         } else {
             // There was no license that could be automatically retrieved, you must call [Self::new_with_license].
             Err(InitializationError::NoLicenseFound)
@@ -348,19 +330,10 @@ impl Session {
     /// Initialize with options, the same rules apply as [`Session::new`], see [`InitializationOptions::default`] for the regular options passed.
     ///
     /// This differs from [`Session::new`] in that it does not check to see if there is a license that the core
-    /// can discover by itself, therefor it is expected that you know where your license is when calling this directly.
+    /// can discover by itself, therefore, it is expected that you know where your license is when calling this directly.
     pub fn new_with_opts(options: InitializationOptions) -> Result<Self, InitializationError> {
-        // Grab the lock before initialization to prevent another thread from initializing
-        // and racing the call to BNInitPlugins.
-        let _lock = INIT_LOCK
-            .lock()
-            .map_err(|_| InitializationError::InitMutex)?;
         let session = Self::registered_session();
-        // Since this whole section is locked, we're guaranteed to be index 0 if we're first.
-        // Only the first thread hitting this should be allowed to call BNInitPlugins
-        if session.index == 0 {
-            init_with_opts(options)?;
-        }
+        init_with_opts(options)?;
         Ok(session)
     }
 
@@ -457,7 +430,7 @@ impl Drop for Session {
     fn drop(&mut self) {
         let previous_count = SESSION_COUNT.fetch_sub(1, SeqCst);
         if previous_count == 1 {
-            // We were the last session, therefor we can safely shut down.
+            // We were the last session, therefore, we can safely shut down.
             shutdown();
         }
     }