Add support for AARCH64 binaries

ergrelet · Apr 6, 2024 · 90173e3 · 90173e3
1 parent 088c92a
commit 90173e3
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 4 deletions.
diff --git a/src/commands.cc b/src/commands.cc
@@ -44,6 +44,11 @@ void SimplifyBasicBlockCommand(BinaryNinja::BinaryView* p_view) {
     triton.setArchitecture(triton::arch::ARCH_X86_64);
   } else if (architecture_name == "x86") {
     triton.setArchitecture(triton::arch::ARCH_X86);
+  } else if (architecture_name == "aarch64") {
+    triton.setArchitecture(triton::arch::ARCH_AARCH64);
+  } else {
+    LogError("Unsupported architecture '%s'", architecture_name.c_str());
+    return;
   }
 
   auto meta_basic_blocks =

diff --git a/src/meta_basic_block.cc b/src/meta_basic_block.cc
@@ -6,6 +6,8 @@ namespace triton_bn {
 
 using namespace BinaryNinja;
 
+static bool IsCallInstruction(const triton::arch::Instruction&);
+static bool IsJumpInstruction(const triton::arch::Instruction& instr);
 static void MergeLinkedBasicBlocks(const BasicBlockEdge& edge,
                                    MetaBasicBlock& root_bb,
                                    MetaBasicBlock& target_bb);
@@ -57,8 +59,7 @@ std::vector<MetaBasicBlock> ExtractMetaBasicBlocksFromBasicBlock(
     triton_bb.add(new_instr);
 
     // Split basic blocks on `call` instructions to make them simplifiable
-    // TODO: Fix for ARM64
-    if (new_instr.getDisassembly().find("call") == 0) {
+    if (IsCallInstruction(new_instr)) {
       LogDebug("call detected: %s", new_instr.getDisassembly().c_str());
       // Add basic block to the result
       result.emplace_back(MetaBasicBlock(triton_bb, basic_block));
@@ -175,8 +176,7 @@ static void MergeLinkedBasicBlocks(const BasicBlockEdge& edge,
     const triton::arch::Instruction last_instr =
         cur_triton_bb.getInstructions()[last_instr_index];
     // Remove last instruction if it's a `jmp`
-    // TODO: Fix for ARM64
-    if (last_instr.getDisassembly().find("jmp") == 0) {
+    if (IsJumpInstruction(last_instr)) {
       cur_triton_bb.remove(last_instr_index);
     }
   }
@@ -190,6 +190,33 @@ static void MergeLinkedBasicBlocks(const BasicBlockEdge& edge,
   root_bb.AddOutgoingEdges(target_bb.outgoing_edges());
 }
 
+static bool IsCallInstruction(const triton::arch::Instruction& instr) {
+  switch (instr.getArchitecture()) {
+    case triton::arch::ARCH_X86_64:
+    case triton::arch::ARCH_X86:
+      return instr.getDisassembly().find("call") == 0;
+    case triton::arch::ARCH_AARCH64:
+      return instr.getDisassembly().find("bl") == 0;
+      return false;
+    default:
+      return false;
+  }
+}
+
+static bool IsJumpInstruction(const triton::arch::Instruction& instr) {
+  switch (instr.getArchitecture()) {
+    case triton::arch::ARCH_X86_64:
+    case triton::arch::ARCH_X86:
+      return instr.getDisassembly().find("jmp") == 0;
+    case triton::arch::ARCH_AARCH64: {
+      const std::string disassembly = instr.getDisassembly();
+      return disassembly.find("b") == 0 && disassembly.find("bl") != 0;
+    }
+    default:
+      return false;
+  }
+}
+
 // Simplify the given `MetaBasicBlock`s with Triton's dead store elimination
 // pass
 std::vector<MetaBasicBlock> SimplifyMetaBasicBlocks(