From 538565e8b89ed5e2b642c2c27833d9c75a39eb41 Mon Sep 17 00:00:00 2001 From: Mason Reed Date: Tue, 8 Apr 2025 15:05:19 -0400 Subject: Fix workflow Rust API not returning refcounted objects IDK why this was not done prior, leaking the returned core activity and workflow object. --- rust/examples/workflow.rs | 2 +- rust/src/workflow.rs | 42 +++++++++++++++++++++--------------------- rust/tests/workflow.rs | 8 ++++---- 3 files changed, 26 insertions(+), 26 deletions(-) (limited to 'rust') diff --git a/rust/examples/workflow.rs b/rust/examples/workflow.rs index 59678cad..2d8f0054 100644 --- a/rust/examples/workflow.rs +++ b/rust/examples/workflow.rs @@ -53,7 +53,7 @@ pub fn main() { println!("Registering workflow..."); let old_meta_workflow = Workflow::instance("core.function.metaAnalysis"); - let meta_workflow = old_meta_workflow.clone("core.function.metaAnalysis"); + let meta_workflow = old_meta_workflow.clone_to("core.function.metaAnalysis"); let activity = Activity::new_with_action(RUST_ACTIVITY_CONFIG, example_activity); meta_workflow.register_activity(&activity).unwrap(); meta_workflow.insert("core.function.runFunctionRecognizers", [RUST_ACTIVITY_NAME]); diff --git a/rust/src/workflow.rs b/rust/src/workflow.rs index dcc388d9..90d6db84 100644 --- a/rust/src/workflow.rs +++ b/rust/src/workflow.rs @@ -164,16 +164,16 @@ pub struct Activity { } impl Activity { + #[allow(unused)] pub(crate) unsafe fn from_raw(handle: NonNull) -> Self { Self { handle } } - - #[allow(unused)] + pub(crate) unsafe fn ref_from_raw(handle: NonNull) -> Ref { Ref::new(Self { handle }) } - pub fn new(config: S) -> Self { + pub fn new(config: S) -> Ref { unsafe extern "C" fn cb_action_nop(_: *mut c_void, _: *mut BNAnalysisContext) {} let config = config.into_bytes_with_nul(); let result = unsafe { @@ -183,10 +183,10 @@ impl Activity { Some(cb_action_nop), ) }; - unsafe { Activity::from_raw(NonNull::new(result).unwrap()) } + unsafe { Activity::ref_from_raw(NonNull::new(result).unwrap()) } } - pub fn new_with_action(config: S, mut action: F) -> Self + pub fn new_with_action(config: S, mut action: F) -> Ref where S: BnStrCompatible, F: FnMut(&AnalysisContext), @@ -208,7 +208,7 @@ impl Activity { Some(cb_action::), ) }; - unsafe { Activity::from_raw(NonNull::new(result).unwrap()) } + unsafe { Activity::ref_from_raw(NonNull::new(result).unwrap()) } } pub fn name(&self) -> BnString { @@ -256,19 +256,19 @@ impl Workflow { /// Create a new unregistered [Workflow] with no activities. /// - /// To get a copy of an existing registered [Workflow] use [Workflow::clone]. - pub fn new(name: S) -> Self { + /// To get a copy of an existing registered [Workflow] use [Workflow::clone_to]. + pub fn new(name: S) -> Ref { let name = name.into_bytes_with_nul(); let result = unsafe { BNCreateWorkflow(name.as_ref().as_ptr() as *const c_char) }; - unsafe { Workflow::from_raw(NonNull::new(result).unwrap()) } + unsafe { Workflow::ref_from_raw(NonNull::new(result).unwrap()) } } /// Make a new unregistered [Workflow], copying all activities and the execution strategy. /// /// * `name` - the name for the new [Workflow] #[must_use] - pub fn clone(&self, name: S) -> Workflow { - self.clone_with_root(name, "") + pub fn clone_to(&self, name: S) -> Ref { + self.clone_to_with_root(name, "") } /// Make a new unregistered [Workflow], copying all activities, within `root_activity`, and the execution strategy. @@ -276,15 +276,15 @@ impl Workflow { /// * `name` - the name for the new [Workflow] /// * `root_activity` - perform the clone operation with this activity as the root #[must_use] - pub fn clone_with_root( + pub fn clone_to_with_root( &self, name: S, root_activity: A, - ) -> Workflow { + ) -> Ref { let raw_name = name.into_bytes_with_nul(); let activity = root_activity.into_bytes_with_nul(); unsafe { - Self::from_raw( + Self::ref_from_raw( NonNull::new(BNWorkflowClone( self.handle.as_ptr(), raw_name.as_ref().as_ptr() as *const c_char, @@ -295,11 +295,11 @@ impl Workflow { } } - pub fn instance(name: S) -> Workflow { + pub fn instance(name: S) -> Ref { let result = unsafe { BNWorkflowInstance(name.into_bytes_with_nul().as_ref().as_ptr() as *const c_char) }; - unsafe { Workflow::from_raw(NonNull::new(result).unwrap()) } + unsafe { Workflow::ref_from_raw(NonNull::new(result).unwrap()) } } /// List of all registered [Workflow]'s @@ -341,7 +341,7 @@ impl Workflow { /// Register an [Activity] with this Workflow. /// /// * `activity` - the [Activity] to register - pub fn register_activity(&self, activity: &Activity) -> Result { + pub fn register_activity(&self, activity: &Activity) -> Result, ()> { self.register_activity_with_subactivities::>(activity, vec![]) } @@ -353,7 +353,7 @@ impl Workflow { &self, activity: &Activity, subactivities: I, - ) -> Result + ) -> Result, ()> where I: IntoIterator, I::Item: BnStrCompatible, @@ -375,7 +375,7 @@ impl Workflow { ) }; let activity_ptr = NonNull::new(result).ok_or(())?; - unsafe { Ok(Activity::from_raw(activity_ptr)) } + unsafe { Ok(Activity::ref_from_raw(activity_ptr)) } } /// Determine if an Activity exists in this [Workflow]. @@ -418,7 +418,7 @@ impl Workflow { } /// Retrieve the Activity object for the specified `name`. - pub fn activity(&self, name: A) -> Option { + pub fn activity(&self, name: A) -> Option> { let name = name.into_bytes_with_nul(); let result = unsafe { BNWorkflowGetActivity( @@ -426,7 +426,7 @@ impl Workflow { name.as_ref().as_ptr() as *const c_char, ) }; - NonNull::new(result).map(|a| unsafe { Activity::from_raw(a) }) + NonNull::new(result).map(|a| unsafe { Activity::ref_from_raw(a) }) } /// Retrieve the list of activity roots for the [Workflow], or if diff --git a/rust/tests/workflow.rs b/rust/tests/workflow.rs index 148d66fd..bf8818c8 100644 --- a/rust/tests/workflow.rs +++ b/rust/tests/workflow.rs @@ -9,7 +9,7 @@ use binaryninja::workflow::Workflow; fn test_workflow_clone() { let _session = Session::new().expect("Failed to initialize session"); let original_workflow = Workflow::new("core.function.baseAnalysis"); - let mut cloned_workflow = original_workflow.clone("clone_workflow"); + let mut cloned_workflow = original_workflow.clone_to("clone_workflow"); assert_eq!(cloned_workflow.name().as_str(), "clone_workflow"); assert_eq!( @@ -17,7 +17,7 @@ fn test_workflow_clone() { original_workflow.configuration() ); - cloned_workflow = original_workflow.clone(""); + cloned_workflow = original_workflow.clone_to(""); assert_eq!( cloned_workflow.name().as_str(), "core.function.baseAnalysis" @@ -36,7 +36,7 @@ fn test_workflow_registration() { .expect_err("Re-registration of null is allowed"); // Validate new workflows can be registered - let test_workflow = Workflow::instance("core.function.baseAnalysis").clone("test_workflow"); + let test_workflow = Workflow::instance("core.function.baseAnalysis").clone_to("test_workflow"); assert_eq!(test_workflow.name().as_str(), "test_workflow"); assert!(!test_workflow.registered()); @@ -66,7 +66,7 @@ fn test_workflow_registration() { assert!(!Workflow::instance("core.function.baseAnalysis").clear()); // Validate re-registration of baseAnalysis is not allowed - let base_workflow_clone = Workflow::instance("core.function.baseAnalysis").clone(""); + let base_workflow_clone = Workflow::instance("core.function.baseAnalysis").clone_to(""); assert_eq!( base_workflow_clone.name().as_str(), -- cgit v1.3.1