[Rust] Enter more session scoped tracing spans for debug info and binary view callbacks

5 files changed, 87 insertions, 47 deletions

diff --git a/rust/src/custom_binary_view.rs b/rust/src/custom_binary_view.rs
index 8e070b15..13a124e9 100644
--- a/rust/src/custom_binary_view.rs
+++ b/rust/src/custom_binary_view.rs
@@ -48,12 +48,10 @@ where
     where
         T: CustomBinaryViewType,
     {
-        ffi_wrap!("BinaryViewTypeBase::is_valid_for", unsafe {
-            let view_type = &*(ctxt as *mut T);
-            let data = BinaryView::ref_from_raw(BNNewViewReference(data));
-
-            view_type.is_valid_for(&data)
-        })
+        let view_type = unsafe { &*(ctxt as *mut T) };
+        let data = unsafe { BinaryView::ref_from_raw(BNNewViewReference(data)) };
+        let _span = ffi_span!("BinaryViewTypeBase::is_valid_for", data);
+        view_type.is_valid_for(&data)
     }
 
     extern "C" fn cb_deprecated<T>(ctxt: *mut c_void) -> bool
@@ -89,6 +87,7 @@ where
                 actual_parent: &data,
             };
 
+            let _span = ffi_span!("BinaryViewTypeBase::create", data);
             match view_type.create_custom_view(&data, builder) {
                 Ok(bv) => {
                     // force a leak of the Ref; failure to do this would result
@@ -114,6 +113,7 @@ where
                 actual_parent: &data,
             };
 
+            let _span = ffi_span!("BinaryViewTypeBase::parse", data);
             match view_type.parse_custom_view(&data, builder) {
                 Ok(bv) => {
                     // force a leak of the Ref; failure to do this would result
@@ -134,6 +134,7 @@ where
             let view_type = &*(ctxt as *mut T);
             let data = BinaryView::ref_from_raw(BNNewViewReference(data));
 
+            let _span = ffi_span!("BinaryViewTypeBase::load_settings", data);
             match view_type.load_settings_for_data(&data) {
                 Some(settings) => Ref::into_raw(settings).handle,
                 None => ptr::null_mut() as *mut _,
diff --git a/rust/src/debuginfo.rs b/rust/src/debuginfo.rs
index 5879f40d..f4cee4ff 100644
--- a/rust/src/debuginfo.rs
+++ b/rust/src/debuginfo.rs
@@ -216,12 +216,10 @@ impl DebugInfoParser {
         where
             C: CustomDebugInfoParser,
         {
-            ffi_wrap!("CustomDebugInfoParser::is_valid", unsafe {
-                let cmd = &*(ctxt as *const C);
-                let view = BinaryView::ref_from_raw(view);
-
-                cmd.is_valid(&view)
-            })
+            let cmd = unsafe { &*(ctxt as *const C) };
+            let view = unsafe { BinaryView::ref_from_raw(view) };
+            let _span = ffi_span!("CustomDebugInfoParser::is_valid", view);
+            cmd.is_valid(&view)
         }
 
         extern "C" fn cb_parse_info<C>(
@@ -235,28 +233,27 @@ impl DebugInfoParser {
         where
             C: CustomDebugInfoParser,
         {
-            ffi_wrap!("CustomDebugInfoParser::parse_info", unsafe {
-                let cmd = &*(ctxt as *const C);
-                let view = BinaryView::ref_from_raw(view);
-                let debug_file = BinaryView::ref_from_raw(debug_file);
-                let mut debug_info = DebugInfo::ref_from_raw(debug_info);
+            let cmd = unsafe { &*(ctxt as *const C) };
+            let view = unsafe { BinaryView::ref_from_raw(view) };
+            let debug_file = unsafe { BinaryView::ref_from_raw(debug_file) };
+            let mut debug_info = unsafe { DebugInfo::ref_from_raw(debug_info) };
 
-                cmd.parse_info(
-                    &mut debug_info,
-                    &view,
-                    &debug_file,
-                    Box::new(move |cur: usize, max: usize| match progress {
-                        Some(func) => {
-                            if func(progress_ctxt, cur, max) {
-                                Ok(())
-                            } else {
-                                Err(())
-                            }
+            let _span = ffi_span!("CustomDebugInfoParser::parse_info", view);
+            cmd.parse_info(
+                &mut debug_info,
+                &view,
+                &debug_file,
+                Box::new(move |cur: usize, max: usize| match progress {
+                    Some(func) => unsafe {
+                        if func(progress_ctxt, cur, max) {
+                            Ok(())
+                        } else {
+                            Err(())
                         }
-                        _ => Ok(()),
-                    }),
-                )
-            })
+                    },
+                    _ => Ok(()),
+                }),
+            )
         }
 
         let name = name.to_cstr();
diff --git a/rust/src/ffi.rs b/rust/src/ffi.rs
index b6124dd4..d1c63216 100644
--- a/rust/src/ffi.rs
+++ b/rust/src/ffi.rs
@@ -20,7 +20,7 @@ macro_rules! ffi_wrap {
         use std::process;
 
         panic::catch_unwind(|| $b).unwrap_or_else(|_| {
-            tracing::error!("ffi callback caught panic: {}", $n);
+            ::tracing::error!("ffi callback caught panic: {}", $n);
             process::abort()
         })
     }};
@@ -37,10 +37,10 @@ macro_rules! ffi_span {
         use $crate::binary_view::BinaryViewExt;
         #[allow(unused_imports)]
         use $crate::file_metadata::FileMetadata;
-        tracing::info_span!($name, session_id = $bv.file().session_id().0).entered()
+        ::tracing::info_span!($name, session_id = $bv.file().session_id().0).entered()
     }};
     ($name:expr) => {
-        tracing::info_span!($name).entered()
+        ::tracing::info_span!($name).entered()
     };
 }
 
diff --git a/rust/src/logger.rs b/rust/src/logger.rs
index f77cb52c..31696d25 100644
--- a/rust/src/logger.rs
+++ b/rust/src/logger.rs
@@ -137,6 +137,7 @@ unsafe impl Sync for Logger {}
 /// This is typically used in headless usage. It can also be used to temporarily log core
 /// messages to something like a file while some analysis is occurring, once the [`LogGuard`] is
 /// dropped, the listener will be unregistered.
+#[must_use]
 pub fn register_log_listener<L: LogListener>(listener: L) -> LogGuard<L> {
     use binaryninjacore_sys::BNRegisterLogListener;
 
diff --git a/rust/src/tracing.rs b/rust/src/tracing.rs
index 2afdcd5e..539b697f 100644
--- a/rust/src/tracing.rs
+++ b/rust/src/tracing.rs
@@ -2,14 +2,11 @@
 
 use crate::file_metadata::SessionId;
 use crate::logger::{
-    bn_log_with_session, BnLogLevel, LogContext, LogListener, LOGGER_DEFAULT_SESSION_ID,
+    bn_log_with_session, BnLogLevel, LogContext, LogGuard, LogListener, LOGGER_DEFAULT_SESSION_ID,
 };
 use tracing::{Event, Id, Level, Subscriber};
-use tracing_subscriber::prelude::*;
-
-// Re-export specific things to make it easy for the user
-pub use tracing::{debug, error, info, trace, warn};
 use tracing_subscriber::layer::Context;
+use tracing_subscriber::prelude::*;
 use tracing_subscriber::registry::LookupSpan;
 use tracing_subscriber::Layer;
 
@@ -26,7 +23,7 @@ use tracing_subscriber::Layer;
 /// #[unsafe(no_mangle)]
 /// pub unsafe extern "C" fn CorePluginInit() -> bool {
 ///     binaryninja::tracing_init!("MyPlugin");
-///     binaryninja::tracing::info!("Core plugin initialized");
+///     tracing::info!("Core plugin initialized");
 ///     true
 /// }
 /// ```
@@ -264,7 +261,7 @@ impl tracing::field::Visit for BnFieldVisitor {
 ///     // Register our tracing subscriber, this will send tracing events to stdout.
 ///     tracing_subscriber::fmt::init();
 ///     // Register our log listener, this will send logs from the core to our tracing subscriber.
-///     let _listener = register_log_listener(TracingLogListener::new(BnLogLevel::DebugLog));
+///     let _listener = TracingLogListener::new(BnLogLevel::DebugLog).register();
 ///     // Should see logs from the core in regard to initialization show up.
 ///     let _session = Session::new().expect("Failed to create session");
 ///     bn_log("Test", BnLogLevel::DebugLog, "Hello, world!");
@@ -275,9 +272,29 @@ pub struct TracingLogListener {
 }
 
 impl TracingLogListener {
-    pub fn new(minimum_level: BnLogLevel) -> Self {
+    /// Create a [`TracingLogListener`] with the minimum log level set to [`BnLogLevel::InfoLog`].
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn new_with_lvl(minimum_level: BnLogLevel) -> Self {
         Self { minimum_level }
     }
+
+    /// Register the [`TracingLogListener`] and send logs to the registered tracing subscriber until
+    /// the [`LogGuard`] is dropped, make sure to register your tracing subscriber before registering.
+    #[must_use]
+    pub fn register(self) -> LogGuard<Self> {
+        crate::logger::register_log_listener(self)
+    }
+}
+
+impl Default for TracingLogListener {
+    fn default() -> Self {
+        Self {
+            minimum_level: BnLogLevel::InfoLog,
+        }
+    }
 }
 
 impl LogListener for TracingLogListener {
@@ -285,16 +302,40 @@ impl LogListener for TracingLogListener {
         let session = ctx.session_id.map(|s| s.0);
         match level {
             BnLogLevel::ErrorLog | BnLogLevel::AlertLog => {
-                error!(session_id = session, target = %ctx.logger_name, "{}", message)
+                tracing::error!(
+                    target: "binaryninja",
+                    session_id = session,
+                    logger = %ctx.logger_name,
+                    "{}",
+                    message
+                )
             }
             BnLogLevel::WarningLog => {
-                warn!(session_id = session, target = %ctx.logger_name, "{}", message)
+                tracing::warn!(
+                    target: "binaryninja",
+                    session_id = session,
+                    logger = %ctx.logger_name,
+                    "{}",
+                    message
+                )
             }
             BnLogLevel::InfoLog => {
-                info!(session_id = session, target = %ctx.logger_name, "{}", message)
+                tracing::info!(
+                    target: "binaryninja",
+                    session_id = session,
+                    logger = %ctx.logger_name,
+                    "{}",
+                    message
+                )
             }
             BnLogLevel::DebugLog => {
-                debug!(session_id = session, target = %ctx.logger_name, "{}", message)
+                tracing::debug!(
+                    target: "binaryninja",
+                    session_id = session,
+                    logger = %ctx.logger_name,
+                    "{}",
+                    message
+                )
             }
         };
     }