diff --git a/lib/src/capstone/context.rs b/lib/src/capstone/context.rs index 3aefd2c..e608939 100644 --- a/lib/src/capstone/context.rs +++ b/lib/src/capstone/context.rs @@ -32,10 +32,10 @@ impl CapstoneContext { } /// Clones an Arc out of the handle field of the given object. - pub fn get(env: &mut JNIEnv, object: &JObject) -> Arc { + pub fn get(env: &mut JNIEnv, object: &JObject) -> JResult> { let guard: MutexGuard> = - unsafe { env.get_rust_field(object, HANDLE_FIELD).unwrap() }; - guard.clone() + unsafe { env.get_rust_field(object, HANDLE_FIELD)? }; + Ok(guard.clone()) } /// Surrenders ownership of the context instance to Java. diff --git a/lib/src/capstone/mod.rs b/lib/src/capstone/mod.rs index 38007f8..5b45442 100644 --- a/lib/src/capstone/mod.rs +++ b/lib/src/capstone/mod.rs @@ -3,6 +3,110 @@ * * Use of this source code is governed by the MIT license found in the LICENSE file. */ +use crate::capstone::context::CapstoneContext; +use crate::capstone::mode::CapstoneMode; +use crate::capstone::output::CapstoneOutput; +use capstone::arch::BuildsCapstone; +use capstone::{arch, Capstone, InsnGroupId, InsnGroupIdInt, InsnId, InsnIdInt, RegId, RegIdInt}; +use jni::objects::{JByteArray, JObject, ReleaseMode}; +use jni::sys::{jint, jlong, jshort}; +use jni::JNIEnv; +use std::ops::BitAnd; + pub mod context; pub mod mode; pub mod output; + +type Result = std::result::Result>; + +pub fn init<'local>( + env: &mut JNIEnv<'local>, + this: JObject<'local>, + mode: JObject<'local>, +) -> Result<()> { + let mode = CapstoneMode::from(env, &mode); + + let capstone = match mode { + Some(CapstoneMode::ARM32) => Capstone::new() + .arm() + .mode(arch::arm::ArchMode::Arm) + .detail(true) + .build(), + Some(CapstoneMode::ARM64) => Capstone::new() + .arm64() + .mode(arch::arm64::ArchMode::Arm) + .detail(true) + .build(), + _ => return Err("invalid argument 'mode'".into()), + } + .map_err(|e| e.to_string())?; + + let instance = CapstoneContext::new(capstone, mode.unwrap()); + + CapstoneContext::surrender_instance(instance, env, &this)?; + Ok(()) +} + +pub fn disassemble<'local>( + env: &mut JNIEnv<'local>, + this: JObject<'local>, + result_object: JObject<'local>, + data: JByteArray<'local>, + count: jint, + address: jlong, +) -> Result<()> { + let code: Vec = { + let result = unsafe { env.get_array_elements(&data, ReleaseMode::NoCopyBack)? }; + result.iter().map(|b| (*b as u8).bitand(0xff)).collect() + }; + + let ctx = CapstoneContext::get(env, &this)?; + let capstone = ctx.capstone.lock().unwrap(); + + let instructions = { + if count == 0 { + capstone.disasm_all(&code, address as u64) + } else { + capstone.disasm_count(&code, address as u64, count as usize) + } + .map_err(|e| e.to_string())? + }; + + let mut output = CapstoneOutput::new(env, &ctx.mode, &capstone, &result_object); + output.copy_instructions(instructions) +} + +pub fn get_insn_name<'local>( + env: &mut JNIEnv<'local>, + this: JObject<'local>, + insn_id: jint, +) -> Result> { + let ctx = CapstoneContext::get(env, &this)?; + let capstone = ctx.capstone.lock().unwrap(); + Ok(capstone.insn_name(InsnId(insn_id as InsnIdInt))) +} + +pub fn get_reg_name<'local>( + env: &mut JNIEnv<'local>, + this: JObject<'local>, + reg_id: jint, +) -> Result> { + let ctx = CapstoneContext::get(env, &this)?; + let capstone = ctx.capstone.lock().unwrap(); + Ok(capstone.reg_name(RegId(reg_id as RegIdInt))) +} + +pub fn get_group_name<'local>( + env: &mut JNIEnv<'local>, + this: JObject<'local>, + group_id: jshort, +) -> Result> { + let ctx = CapstoneContext::get(env, &this)?; + let capstone = ctx.capstone.lock().unwrap(); + Ok(capstone.group_name(InsnGroupId(group_id as InsnGroupIdInt))) +} + +pub fn throw(env: &mut JNIEnv, message: &str) { + env.throw_new("org/native4j/capstone/exception/CapstoneException", message) + .unwrap(); +} diff --git a/lib/src/lib.rs b/lib/src/lib.rs index fb1368a..c930175 100644 --- a/lib/src/lib.rs +++ b/lib/src/lib.rs @@ -3,50 +3,25 @@ * * Use of this source code is governed by the MIT license found in the LICENSE file. */ -use std::ops::BitAnd; - -use ::capstone::arch::BuildsCapstone; -use ::capstone::{arch, Capstone, InsnGroupId, InsnGroupIdInt, InsnId, InsnIdInt, RegId, RegIdInt}; -use jni::objects::{JByteArray, JObject, ReleaseMode}; +use jni::objects::{JByteArray, JObject}; use jni::sys::{jint, jlong, jshort, jstring}; use jni::JNIEnv; use crate::capstone::context::CapstoneContext; -use crate::capstone::mode::CapstoneMode; -use crate::capstone::output::CapstoneOutput; mod capstone; mod obj; mod util; mod writer; + #[no_mangle] pub extern "system" fn Java_org_native4j_capstone_Capstone_init<'local>( mut env: JNIEnv<'local>, this: JObject<'local>, mode: JObject<'local>, ) -> jstring { - let mode = CapstoneMode::from(&mut env, &mode); - - let capstone = match mode { - Some(CapstoneMode::ARM32) => Capstone::new() - .arm() - .mode(arch::arm::ArchMode::Arm) - .detail(true) - .build(), - Some(CapstoneMode::ARM64) => Capstone::new() - .arm64() - .mode(arch::arm64::ArchMode::Arm) - .detail(true) - .build(), - _ => return make_error!(env, "invalid argument 'mode'"), - }; - check_result!(env, capstone); - - let instance = CapstoneContext::new(capstone.unwrap(), mode.unwrap()); - - let result = CapstoneContext::surrender_instance(instance, &mut env, &this); + let result = capstone::init(&mut env, this, mode); check_result!(env, result); - 0 as jstring /* null */ } @@ -57,7 +32,6 @@ pub extern "system" fn Java_org_native4j_capstone_Capstone_shutdown<'local>( ) -> jstring { let result = CapstoneContext::drop_instance(&mut env, &this); check_result!(env, result); - 0 as jstring /* null */ } @@ -70,36 +44,8 @@ pub extern "system" fn Java_org_native4j_capstone_Capstone_disassemble<'local>( count: jint, address: jlong, ) -> jstring { - let code: Vec = { - let result = unsafe { env.get_array_elements(&data, ReleaseMode::NoCopyBack) }; - check_result!(env, result); - result - .unwrap() - .iter() - .map(|b| (*b as u8).bitand(0xff)) - .collect() - }; - - let ctx = CapstoneContext::get(&mut env, &this); - let capstone = ctx.capstone.lock().unwrap(); - - let instructions = { - let result = if count == 0 { - capstone.disasm_all(&code, address as u64) - } else { - capstone.disasm_count(&code, address as u64, count as usize) - }; - if let Err(e) = result { - return make_error!(env, e.to_string()); - } - result.unwrap() - }; - - let mut output = CapstoneOutput::new(&mut env, &ctx.mode, &capstone, &result_object); - - let result = output.copy_instructions(instructions); + let result = capstone::disassemble(&mut env, this, result_object, data, count, address); check_result!(env, result); - 0 as jstring /* null */ } @@ -109,12 +55,15 @@ pub extern "system" fn Java_org_native4j_capstone_Capstone_getInsnName<'local>( this: JObject<'local>, insn_id: jint, ) -> jstring { - let ctx = CapstoneContext::get(&mut env, &this); - let capstone = ctx.capstone.lock().unwrap(); - match capstone.insn_name(InsnId(insn_id as InsnIdInt)) { - Some(str) => make_jstring!(env, str), - None => 0 as jstring, /* null */ + let result = capstone::get_insn_name(&mut env, this, insn_id); + if let Err(e) = result { + capstone::throw(&mut env, &e.to_string()); + return 0 as jstring /* null */; } + result + .unwrap() + .map(|str| make_jstring!(env, str)) + .unwrap_or(0 as jstring /* null */) } #[no_mangle] @@ -123,12 +72,15 @@ pub extern "system" fn Java_org_native4j_capstone_Capstone_getRegName<'local>( this: JObject<'local>, reg_id: jint, ) -> jstring { - let ctx = CapstoneContext::get(&mut env, &this); - let capstone = ctx.capstone.lock().unwrap(); - match capstone.reg_name(RegId(reg_id as RegIdInt)) { - Some(str) => make_jstring!(env, str), - None => 0 as jstring, /* null */ + let result = capstone::get_reg_name(&mut env, this, reg_id); + if let Err(e) = result { + capstone::throw(&mut env, &e.to_string()); + return 0 as jstring /* null */; } + result + .unwrap() + .map(|str| make_jstring!(env, str)) + .unwrap_or(0 as jstring /* null */) } #[no_mangle] @@ -137,10 +89,13 @@ pub extern "system" fn Java_org_native4j_capstone_Capstone_getGroupName<'local>( this: JObject<'local>, group_id: jshort, ) -> jstring { - let ctx = CapstoneContext::get(&mut env, &this); - let capstone = ctx.capstone.lock().unwrap(); - match capstone.group_name(InsnGroupId(group_id as InsnGroupIdInt)) { - Some(str) => make_jstring!(env, str), - None => 0 as jstring, /* null */ + let result = capstone::get_group_name(&mut env, this, group_id); + if let Err(e) = result { + capstone::throw(&mut env, &e.to_string()); + return 0 as jstring /* null */; } + result + .unwrap() + .map(|str| make_jstring!(env, str)) + .unwrap_or(0 as jstring /* null */) } diff --git a/native4j-capstone/src/test/java/org/native4j/capstone/CapstoneTests.java b/native4j-capstone/src/test/java/org/native4j/capstone/CapstoneTests.java index 3c7231d..3aeeafb 100644 --- a/native4j-capstone/src/test/java/org/native4j/capstone/CapstoneTests.java +++ b/native4j-capstone/src/test/java/org/native4j/capstone/CapstoneTests.java @@ -8,11 +8,13 @@ import java.util.ArrayList; import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.Test; import static org.junit.jupiter.api.Assertions.*; +import org.native4j.capstone.exception.CapstoneException; import org.native4j.capstone.insn.CapstoneResult; import org.native4j.capstone.insn.arm.CsInsnArm64; @@ -22,21 +24,37 @@ public class CapstoneTests { (byte) 0xFF, (byte) 0x97, 0x00, 0x00, (byte) 0x80, 0x52, (byte) 0xFD, 0x7B, (byte) 0xC1, (byte) 0xA8, (byte) 0xC0, 0x03, 0x5F, (byte) 0xD6, }; - private void verifyInstructions(CsInsnArm64[] insns) { - assertEquals(insns[0].mnemonic, "stp"); - assertEquals(insns[0].operand, "x29, x30, [sp, #-0x10]!"); - assertEquals(insns[0].address, 0x1000); - assertEquals(insns[0].conditionCodes, 0); - - assertTrue(insns[0].writebackRequired); - assertFalse(insns[0].updatesFlags); - assertNull(insns[0].regsRead); - assertNull(insns[0].regsWrite); - assertNull(insns[0].groups); - - assertEquals(insns[0].operands[0].getReg(), 2); - assertEquals(insns[0].operands[1].getReg(), 3); - assertEquals(insns[0].operands[2].getMem().displacement, -0x10); + private boolean verifyInstructions(Capstone capstone, CsInsnArm64[] insns) { + CsInsnArm64 first = insns[0]; + + assertEquals(first.mnemonic, "stp"); + assertEquals(capstone.getInsnName(first.instructionId), "stp"); + + assertEquals(first.operand, "x29, x30, [sp, #-0x10]!"); + assertEquals(first.address, 0x1000); + assertEquals(first.conditionCodes, 0); + + assertTrue(first.writebackRequired); + assertFalse(first.updatesFlags); + assertNull(first.regsRead); + assertNull(first.regsWrite); + assertNull(first.groups); + + assertEquals(first.operands[0].getReg(), 2); + assertEquals(first.operands[1].getReg(), 3); + assertEquals(first.operands[2].getMem().displacement, -0x10); + + CsInsnArm64 fifth = insns[4]; + assertEquals(fifth.mnemonic, "bl"); + assertEquals(capstone.getInsnName(fifth.instructionId), "bl"); + + assertEquals(fifth.operands[0].getImm(), 3804); + + assertEquals(capstone.getGroupName(fifth.groups[0]), "call"); + assertEquals(capstone.getGroupName(fifth.groups[1]), "jump"); + assertEquals(capstone.getGroupName(fifth.groups[2]), "branch_relative"); + + return true; } @Test @@ -45,12 +63,27 @@ void testARM64Disassembly() { CapstoneResult result = new CapstoneResult(); capstone.disassembleAll(result, code, 0x1000); CsInsnArm64[] insns = result.toArray(CsInsnArm64[].class); - verifyInstructions(insns); + verifyInstructions(capstone, insns); + } + } + + @Test + void testNullHandle() { + { + Capstone capstone = new Capstone(CapstoneMode.ARM64); + capstone.close(); + assertThrows(CapstoneException.class, () -> capstone.getInsnName(0)); + } + { + Capstone capstone = new Capstone(CapstoneMode.ARM64); + capstone.close(); + assertThrows(CapstoneException.class, capstone::close); } } @Test - void testThreading() { + void testThreading() throws InterruptedException { + AtomicBoolean failed = new AtomicBoolean(true); try (Capstone capstone = new Capstone(CapstoneMode.ARM64)) { List threads = new ArrayList<>(); for (int i = 0; i < 4; i++) { @@ -59,19 +92,17 @@ void testThreading() { CapstoneResult result = new CapstoneResult(); capstone.disassembleAll(result, code, 0x1000); CsInsnArm64[] insns = result.toArray(CsInsnArm64[].class); - verifyInstructions(insns); + failed.set(!verifyInstructions(capstone, insns)); } }); th.start(); threads.add(th); } for (Thread th : threads) { - try { - th.join(); - } catch (InterruptedException e) { - fail(e); - } + th.join(); } } + if (failed.get()) + fail("Failed to verify instructions"); } }