Merge branch 'fix-vec-leak' of github.com:rbran/binaryninja-api into dev

2 files changed, 54 insertions, 67 deletions

diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs
index cb71a9cd..db887e4f 100644
--- a/rust/src/architecture.rs
+++ b/rust/src/architecture.rs
@@ -1162,10 +1162,11 @@ impl Architecture for CoreArchitecture {
                 &mut result as *mut _,
                 &mut count as *mut _,
             ) {
-                let vec = Vec::<BNInstructionTextToken>::from_raw_parts(result, count, count)
+                let vec = slice::from_raw_parts(result, count)
                     .iter()
-                    .map(|x| InstructionTextToken::from_raw(x))
+                    .map(|x| InstructionTextToken::from_raw(x).to_owned())
                     .collect();
+                BNFreeInstructionText(result, count);
                 Some((consumed, vec))
             } else {
                 None
@@ -1810,27 +1811,25 @@ where
         let data = unsafe { slice::from_raw_parts(data, *len) };
         let result = unsafe { &mut *result };
 
-        match custom_arch.instruction_text(data, addr) {
-            Some((res_size, mut res_tokens)) => {
-                unsafe {
-                    // TODO: Can't use into_raw_parts as it's unstable so we do this instead...
-                    let r_ptr = res_tokens.as_mut_ptr();
-                    let r_count = res_tokens.len();
-                    mem::forget(res_tokens);
+        let Some((res_size, res_tokens)) = custom_arch.instruction_text(data, addr) else {
+            return false;
+        };
 
-                    *result = &mut (*r_ptr).0;
-                    *count = r_count;
-                    *len = res_size;
-                }
-                true
-            }
-            None => false,
+        let res_tokens: Box<[_]> = res_tokens.into_boxed_slice();
+        unsafe {
+            let res_tokens = Box::leak(res_tokens);
+            let r_ptr = res_tokens.as_mut_ptr();
+            let r_count = res_tokens.len();
+
+            *result = &mut (*r_ptr).0;
+            *count = r_count;
+            *len = res_size;
         }
+        true
     }
 
     extern "C" fn cb_free_instruction_text(tokens: *mut BNInstructionTextToken, count: usize) {
-        let _tokens =
-            unsafe { Vec::from_raw_parts(tokens as *mut InstructionTextToken, count, count) };
+        let _tokens = unsafe { Box::from_raw(ptr::slice_from_raw_parts_mut(tokens, count)) };
     }
 
     extern "C" fn cb_instruction_llil<A>(
@@ -1930,15 +1929,7 @@ where
         if len == 0 {
             ptr::null_mut()
         } else {
-            let mut res = Vec::with_capacity(len + 1);
-
-            res.push(len as u32);
-
-            for i in items {
-                res.push(i);
-            }
-
-            assert!(res.len() == len + 1);
+            let mut res: Box<[_]> = [len as u32].into_iter().chain(items).collect();
 
             let raw = res.as_mut_ptr();
             mem::forget(res);
@@ -2279,7 +2270,8 @@ where
         unsafe {
             let actual_start = regs.offset(-1);
             let len = *actual_start + 1;
-            let _regs = Vec::from_raw_parts(actual_start, len as usize, len as usize);
+            let regs_ptr = ptr::slice_from_raw_parts_mut(actual_start, len.try_into().unwrap());
+            let _regs = Box::from_raw(regs_ptr);
         }
     }
 
@@ -2419,28 +2411,25 @@ where
     {
         let custom_arch = unsafe { &*(ctxt as *mut A) };
 
-        if let Some(intrinsic) = custom_arch.intrinsic_from_id(intrinsic) {
-            let inputs = intrinsic.inputs();
-            let mut res = Vec::with_capacity(inputs.len());
-            for input in inputs {
-                res.push(unsafe { Ref::into_raw(input) }.into_raw());
-            }
-
-            unsafe {
-                *count = res.len();
-                if res.is_empty() {
-                    ptr::null_mut()
-                } else {
-                    let raw = res.as_mut_ptr();
-                    mem::forget(res);
-                    raw
-                }
-            }
-        } else {
+        let Some(intrinsic) = custom_arch.intrinsic_from_id(intrinsic) else {
             unsafe {
                 *count = 0;
             }
-            ptr::null_mut()
+            return ptr::null_mut();
+        };
+
+        let inputs = intrinsic.inputs();
+        let mut res: Box<[_]> = inputs.into_iter().map(|input| input.into_raw()).collect();
+
+        unsafe {
+            *count = res.len();
+            if res.is_empty() {
+                ptr::null_mut()
+            } else {
+                let raw = res.as_mut_ptr();
+                mem::forget(res);
+                raw
+            }
         }
     }
 
@@ -2452,8 +2441,8 @@ where
 
         if !nt.is_null() {
             unsafe {
-                let list = Vec::from_raw_parts(nt, count, count);
-                for nt in list {
+                let name_and_types = Box::from_raw(ptr::slice_from_raw_parts_mut(nt, count));
+                for nt in name_and_types.into_iter() {
                     BnString::from_raw(nt.name);
                 }
             }
@@ -2472,10 +2461,7 @@ where
 
         if let Some(intrinsic) = custom_arch.intrinsic_from_id(intrinsic) {
             let inputs = intrinsic.outputs();
-            let mut res = Vec::with_capacity(inputs.len());
-            for input in inputs.iter() {
-                res.push(input.as_ref().into());
-            }
+            let mut res: Box<[_]> = inputs.iter().map(|input| input.as_ref().into()).collect();
 
             unsafe {
                 *count = res.len();
@@ -2504,9 +2490,7 @@ where
     {
         let _custom_arch = unsafe { &*(ctxt as *mut A) };
         if !tl.is_null() {
-            unsafe {
-                let _list = Vec::from_raw_parts(tl, count, count);
-            }
+            let _type_list = unsafe { Box::from_raw(ptr::slice_from_raw_parts_mut(tl, count)) };
         }
     }
 
diff --git a/rust/src/disassembly.rs b/rust/src/disassembly.rs
index 8b3a4cb5..855807b1 100644
--- a/rust/src/disassembly.rs
+++ b/rust/src/disassembly.rs
@@ -73,7 +73,7 @@ pub type InstructionTextTokenContext = BNInstructionTextTokenContext;
 // IndirectImportToken = 69,
 // ExternalSymbolToken = 70,
 
-#[repr(C)]
+#[repr(transparent)]
 pub struct InstructionTextToken(pub(crate) BNInstructionTextToken);
 
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
@@ -99,8 +99,8 @@ pub enum InstructionTextTokenContents {
 }
 
 impl InstructionTextToken {
-    pub(crate) unsafe fn from_raw(raw: &BNInstructionTextToken) -> Self {
-        Self(*raw)
+    pub(crate) unsafe fn from_raw(raw: &BNInstructionTextToken) -> &Self {
+        mem::transmute(raw)
     }
 
     pub fn new(text: &str, contents: InstructionTextTokenContents) -> Self {
@@ -254,13 +254,16 @@ impl Clone for InstructionTextToken {
     }
 }
 
-// TODO : There is almost certainly a memory leak here - in the case where
-//  `impl CoreOwnedArrayProvider for InstructionTextToken` doesn't get triggered
-// impl Drop for InstructionTextToken {
-//     fn drop(&mut self) {
-//         let _owned = unsafe { BnString::from_raw(self.0.text) };
-//     }
-// }
+impl Drop for InstructionTextToken {
+    fn drop(&mut self) {
+        if !self.0.text.is_null() {
+            let _owned = unsafe { BnString::from_raw(self.0.text) };
+        }
+        if !self.0.typeNames.is_null() && self.0.namesCount != 0 {
+            unsafe { BNFreeStringList(self.0.typeNames, self.0.namesCount) }
+        }
+    }
+}
 
 pub struct DisassemblyTextLine(pub(crate) BNDisassemblyTextLine);
 
@@ -290,7 +293,7 @@ impl DisassemblyTextLine {
         unsafe {
             std::slice::from_raw_parts::<BNInstructionTextToken>(self.0.tokens, self.0.count)
                 .iter()
-                .map(|&x| InstructionTextToken::from_raw(&x))
+                .map(|x| InstructionTextToken::from_raw(x).clone())
                 .collect()
         }
     }