Support assemble callback in Rust architecture plugins

2 files changed, 194 insertions, 24 deletions

diff --git a/rust/src/architecture.rs b/rust/src/architecture.rs
index 580454e9..f84ec122 100644
--- a/rust/src/architecture.rs
+++ b/rust/src/architecture.rs
@@ -21,7 +21,7 @@ use binaryninjacore_sys::*;
 use std::{
     borrow::{Borrow, Cow},
     collections::HashMap,
-    ffi::CStr,
+    ffi::{c_char, c_int, CStr, CString},
     hash::Hash,
     mem::zeroed,
     ops, ptr, slice,
@@ -29,6 +29,7 @@ use std::{
 
 use crate::{
     callingconvention::CallingConvention,
+    databuffer::DataBuffer,
     disassembly::InstructionTextToken,
     llil::{
         get_default_flag_cond_llil, get_default_flag_write_llil, FlagWriteOp, LiftedExpr, Lifter,
@@ -487,6 +488,13 @@ pub trait Architecture: 'static + Sized + AsRef<CoreArchitecture> {
         None
     }
 
+    fn can_assemble(&self) -> bool {
+        false
+    }
+    fn assemble(&self, _code: &str, _addr: u64) -> Result<Vec<u8>, String> {
+        Err("Assemble unsupported".into())
+    }
+
     fn handle(&self) -> Self::Handle;
 }
 
@@ -1399,6 +1407,40 @@ impl Architecture for CoreArchitecture {
         Some(CoreIntrinsic(self.0, id))
     }
 
+    fn can_assemble(&self) -> bool {
+        unsafe { BNCanArchitectureAssemble(self.0) }
+    }
+
+    fn assemble(&self, code: &str, addr: u64) -> Result<Vec<u8>, String> {
+        let code = CString::new(code).map_err(|_| "Invalid encoding in code string".to_string())?;
+
+        let result = match DataBuffer::new(&[]) {
+            Ok(result) => result,
+            Err(_) => return Err("Result buffer allocation failed".to_string()),
+        };
+        let mut error_raw: *mut c_char = ptr::null_mut();
+        let res = unsafe {
+            BNAssemble(
+                self.0,
+                code.as_ptr(),
+                addr,
+                result.as_raw(),
+                &mut error_raw as *mut *mut c_char,
+            )
+        };
+
+        let error = raw_to_string(error_raw);
+        unsafe {
+            BNFreeString(error_raw);
+        }
+
+        if res {
+            Ok(result.get_data().to_vec())
+        } else {
+            Err(error.unwrap_or("Assemble failed".into()))
+        }
+    }
+
     fn handle(&self) -> CoreArchitecture {
         *self
     }
@@ -2320,22 +2362,48 @@ where
         }
     }
 
-    // TODO : I have no idea what I'm doing and this is likely wrong!
-    extern "C" fn cb_can_assemble(_ctxt: *mut c_void) -> bool {
-        false
+    extern "C" fn cb_can_assemble<A>(ctxt: *mut c_void) -> bool
+    where
+        A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync,
+    {
+        let custom_arch = unsafe { &*(ctxt as *mut A) };
+        custom_arch.can_assemble()
     }
 
-    extern "C" fn cb_assemble(
-        _ctxt: *mut c_void,
-        _code: *const c_char,
-        _addr: u64,
-        _result: *mut BNDataBuffer,
+    extern "C" fn cb_assemble<A>(
+        ctxt: *mut c_void,
+        code: *const c_char,
+        addr: u64,
+        buffer: *mut BNDataBuffer,
         errors: *mut *mut c_char,
-    ) -> bool {
-        unsafe {
-            *errors = ptr::null_mut();
-        }
-        false
+    ) -> bool
+    where
+        A: 'static + Architecture<Handle = CustomArchitectureHandle<A>> + Send + Sync,
+    {
+        let custom_arch = unsafe { &*(ctxt as *mut A) };
+        let code = raw_to_string(code).unwrap_or("".into());
+        let mut buffer = DataBuffer::from_raw(buffer);
+
+        let result = match custom_arch.assemble(&code, addr) {
+            Ok(result) => {
+                buffer.set_data(&result);
+                unsafe {
+                    *errors = BnString::new("").into_raw();
+                }
+                true
+            }
+            Err(result) => {
+                unsafe {
+                    *errors = BnString::new(result).into_raw();
+                }
+                false
+            }
+        };
+
+        // Caller owns the data buffer, don't free it
+        mem::forget(buffer);
+
+        result
     }
 
     extern "C" fn cb_patch_unavailable(
@@ -2438,8 +2506,8 @@ where
         getIntrinsicOutputs: Some(cb_intrinsic_outputs::<A>),
         freeTypeList: Some(cb_free_type_list::<A>),
 
-        canAssemble: Some(cb_can_assemble),
-        assemble: Some(cb_assemble),
+        canAssemble: Some(cb_can_assemble::<A>),
+        assemble: Some(cb_assemble::<A>),
 
         isNeverBranchPatchAvailable: Some(cb_patch_unavailable),
         isAlwaysBranchPatchAvailable: Some(cb_patch_unavailable),
@@ -2502,3 +2570,91 @@ where
         unsafe { &*self.handle }
     }
 }
+
+#[repr(i32)]
+pub enum LlvmServicesDialect {
+    Unspecified = 0,
+    Att = 1,
+    Intel = 2,
+}
+
+#[repr(i32)]
+pub enum LlvmServicesCodeModel {
+    Default = 0,
+    Small = 1,
+    Kernel = 2,
+    Medium = 3,
+    Large = 4,
+}
+
+#[repr(i32)]
+pub enum LlvmServicesRelocMode {
+    Static = 0,
+    PIC = 1,
+    DynamicNoPIC = 2,
+}
+
+pub fn llvm_assemble(
+    code: &str,
+    dialect: LlvmServicesDialect,
+    arch_triple: &str,
+    code_model: LlvmServicesCodeModel,
+    reloc_mode: LlvmServicesRelocMode,
+) -> Result<Vec<u8>, String> {
+    let code = CString::new(code).map_err(|_| "Invalid encoding in code string".to_string())?;
+    let arch_triple = CString::new(arch_triple)
+        .map_err(|_| "Invalid encoding in architecture triple string".to_string())?;
+    let mut out_bytes: *mut c_char = ptr::null_mut();
+    let mut out_bytes_len: c_int = 0;
+    let mut err_bytes: *mut c_char = ptr::null_mut();
+    let mut err_len: c_int = 0;
+
+    unsafe {
+        BNLlvmServicesInit();
+    }
+
+    let result = unsafe {
+        BNLlvmServicesAssemble(
+            code.as_ptr(),
+            dialect as i32,
+            arch_triple.as_ptr(),
+            code_model as i32,
+            reloc_mode as i32,
+            &mut out_bytes as *mut *mut c_char,
+            &mut out_bytes_len as *mut c_int,
+            &mut err_bytes as *mut *mut c_char,
+            &mut err_len as *mut c_int,
+        )
+    };
+
+    let out = if out_bytes_len == 0 {
+        Vec::new()
+    } else {
+        unsafe {
+            slice::from_raw_parts(
+                out_bytes as *const c_char as *const u8,
+                out_bytes_len as usize,
+            )
+        }
+        .to_vec()
+    };
+
+    let errors = if err_len == 0 {
+        "".into()
+    } else {
+        String::from_utf8_lossy(unsafe {
+            slice::from_raw_parts(err_bytes as *const c_char as *const u8, err_len as usize)
+        })
+        .into_owned()
+    };
+
+    unsafe {
+        BNLlvmServicesAssembleFree(out_bytes, err_bytes);
+    }
+
+    if result == 0 {
+        Ok(out)
+    } else {
+        Err(errors)
+    }
+}
diff --git a/rust/src/databuffer.rs b/rust/src/databuffer.rs
index e8bf32c9..f14b71a0 100644
--- a/rust/src/databuffer.rs
+++ b/rust/src/databuffer.rs
@@ -16,6 +16,7 @@
 
 use binaryninjacore_sys::*;
 
+use std::ffi::c_void;
 use std::ptr;
 use std::slice;
 
@@ -25,6 +26,9 @@ impl DataBuffer {
     pub(crate) fn from_raw(raw: *mut BNDataBuffer) -> Self {
         DataBuffer(raw)
     }
+    pub(crate) fn as_raw(&self) -> *mut BNDataBuffer {
+        self.0
+    }
 
     pub fn get_data(&self) -> &[u8] {
         if self.0.is_null() {
@@ -39,6 +43,16 @@ impl DataBuffer {
         }
     }
 
+    pub fn set_data(&mut self, data: &[u8]) {
+        unsafe {
+            BNSetDataBufferContents(
+                self.0,
+                data.as_ptr() as *const c_void as *mut c_void,
+                data.len(),
+            );
+        }
+    }
+
     pub fn len(&self) -> usize {
         unsafe { BNGetDataBufferLength(self.0) }
     }
@@ -47,14 +61,14 @@ impl DataBuffer {
         unsafe { BNGetDataBufferLength(self.0) == 0 }
     }
 
-    // pub fn new(data: ?, len: usize) -> Result<Self> {
-    //   let read_buffer = unsafe { BNCreateDataBuffer(data, len) };
-    //   if read_buffer.is_null() {
-    //     Err(())
-    //   } else {
-    //     Ok(DataBuffer::from_raw(read_buffer))
-    //   }
-    // }
+    pub fn new(data: &[u8]) -> Result<Self, ()> {
+        let buffer = unsafe { BNCreateDataBuffer(data.as_ptr() as *const c_void, data.len()) };
+        if buffer.is_null() {
+            Err(())
+        } else {
+            Ok(DataBuffer::from_raw(buffer))
+        }
+    }
 }
 
 // TODO : delete this