Rust: Better support for floating licenses

Author: CouleeApps

rust/src/enterprise.rs

@@ -20,10 +20,12 @@ pub enum EnterpriseCheckoutError {
 }
 
 /// Initialize the enterprise server connection to check out a floating license.
-pub fn checkout_license(duration: Duration) -> Result<(), EnterpriseCheckoutError> {
+/// 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> {
     if crate::is_ui_enabled() {
         // We only need to check out a license if running headlessly.
-        return Ok(());
+        return Ok(false);
     }
 
     // The disparate core functions we call here might already have mutexes to guard.
@@ -72,9 +74,10 @@ pub fn checkout_license(duration: Duration) -> Result<(), EnterpriseCheckoutErro
                 last_error,
             ));
         }
+        Ok(true)
+    } else {
+        Ok(false)
     }
-
-    Ok(())
 }
 
 pub fn release_license() {

rust/src/headless.rs

@@ -12,17 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::{
-    binary_view, bundled_plugin_directory, enterprise, is_license_validated, is_main_thread,
-    license_path, set_bundled_plugin_directory, set_license, string::IntoJson,
-};
+use crate::{binary_view, bundled_plugin_directory, enterprise, is_license_validated, is_main_thread, license_path, product, 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 thiserror::Error;
 
-use crate::enterprise::release_license;
 use crate::main_thread::{MainThreadAction, MainThreadHandler};
 use crate::progress::ProgressCallback;
 use crate::rc::Ref;
@@ -32,11 +28,14 @@ use std::sync::Mutex;
 use std::thread::JoinHandle;
 use std::time::Duration;
 
+static INIT_LOCK: Mutex<()> = Mutex::new(());
 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);
 
+static NEED_LICENSE_RELEASE: Mutex<bool> = Mutex::new(false);
+
 #[derive(Error, Debug)]
 pub enum InitializationError {
     #[error("main thread could not be started: {0}")]
@@ -47,6 +46,8 @@ pub enum InitializationError {
     InvalidLicense,
     #[error("no license could located, please see `binaryninja::set_license` for details")]
     NoLicenseFound,
+    #[error("could not acquire initialization mutex")]
+    Mutex,
 }
 
 /// Loads plugins, core architecture, platform, etc.
@@ -68,11 +69,16 @@ pub fn init() -> Result<(), InitializationError> {
 /// ⚠️ Important! Must be called at the end of scripts. ⚠️
 pub fn shutdown() {
     match crate::product().as_str() {
-        "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" => enterprise::release_license(),
+        "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" => {
+            if let Ok(need_release) = NEED_LICENSE_RELEASE.lock() {
+                if *need_release {
+                    enterprise::release_license()
+                }
+            }
+        }
         _ => {}
     }
     unsafe { binaryninjacore_sys::BNShutdown() };
-    release_license();
     // TODO: We might want to drop the main thread here, however that requires getting the handler ctx to drop the sender.
 }
 
@@ -201,7 +207,9 @@ pub fn init_with_opts(options: InitializationOptions) -> Result<(), Initializati
         "Binary Ninja Enterprise Client" | "Binary Ninja Ultimate" => {
             if options.checkout_license {
                 // We are allowed to check out a license, so do it!
-                enterprise::checkout_license(options.floating_license_duration)?;
+                if enterprise::checkout_license(options.floating_license_duration)? {
+                    *(NEED_LICENSE_RELEASE.lock().map_err(|_| InitializationError::InvalidLicense)?) = true;
+                }
             }
         }
         _ => {}
@@ -271,26 +279,39 @@ pub fn license_location() -> Option<LicenseLocation> {
             if license_path().exists() {
                 Some(LicenseLocation::File)
             } else {
-                // Check to see if we might be authorizing with keychain credentials.
-                match std::env::var("BN_ENTERPRISE_USERNAME") {
-                    Ok(_) => Some(LicenseLocation::Keychain),
-                    Err(_) => None,
+                // Check to see if we might be authorizing with enterprise
+                if product().as_str() == "Binary Ninja Enterprise Client" || product().as_str() == "Binary Ninja Ultimate" {
+                    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
                 }
             }
         }
     }
 }
 
 /// 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 {}
+pub struct Session {
+    index: usize,
+}
 
 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 {}
+        let previous_count = SESSION_COUNT.fetch_add(1, SeqCst);
+        Self {
+            index: previous_count,
+        }
     }
 
     /// Before calling new you must make sure that the license is retrievable, otherwise the core won't be able to initialize.
@@ -302,8 +323,12 @@ impl Session {
             // We were able to locate a license, continue with initialization.
             // Grab the session before initialization to prevent another thread from initializing
             // and shutting down before this thread can increment the SESSION_COUNT.
+            let lock = INIT_LOCK.lock().map_err(|_| InitializationError::Mutex)?;
             let session = Self::registered_session();
-            init()?;
+            if session.index == 0 {
+                init()?;
+            }
+            drop(lock);
             Ok(session)
         } else {
             // There was no license that could be automatically retrieved, you must call [Self::new_with_license].
@@ -316,8 +341,13 @@ impl Session {
     /// 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.
     pub fn new_with_opts(options: InitializationOptions) -> Result<Self, InitializationError> {
-        init_with_opts(options)?;
-        Ok(Self::registered_session())
+        let lock = INIT_LOCK.lock().map_err(|_| InitializationError::Mutex)?;
+        let session = Self::registered_session();
+        if session.index == 0 {
+            init_with_opts(options)?;
+        }
+        drop(lock);
+        Ok(session)
     }
 
     /// ```no_run