From ee66e4ec111af51095d6a6c98859eea6b9cc0fb9 Mon Sep 17 00:00:00 2001
From: ergrelet <ergrelet@users.noreply.github.com>
Date: Sat, 6 Apr 2024 02:39:16 +0200
Subject: [PATCH] Add support for AARCH64 binaries

---
 CHANGELOG.md            |  4 ++++
 src/commands.cc         | 13 ++++++++++++-
 src/meta_basic_block.cc | 40 +++++++++++++++++++++++++++++++++++-----
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4de6b5b..d404a29 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- Add support for ARM64 binaries
+
 ## [0.1.1] - 2024-04-06
 
 ### Changed
diff --git a/src/commands.cc b/src/commands.cc
index e27a01e..581e935 100644
--- 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 =
@@ -95,6 +100,11 @@ void SimplifyFunctionCommand(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;
   }
 
   // Create `MetaBasicBlock`s from the current Binja function's basic blocks
@@ -130,7 +140,8 @@ bool ValidateSimplifyFunctionCommand(BinaryView* p_view) {
   // Check platform compatibility
   const std::string architecture_name =
       view.GetDefaultArchitecture()->GetName();
-  if (architecture_name != "x86_64" && architecture_name != "x86") {
+  if (architecture_name != "x86_64" && architecture_name != "x86" &&
+      architecture_name != "aarch64") {
     LogError("Unsupported architecture");
     return false;
   }
diff --git a/src/meta_basic_block.cc b/src/meta_basic_block.cc
index 656fbc1..6dd7f56 100644
--- 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);
@@ -41,7 +43,6 @@ std::vector<MetaBasicBlock> ExtractMetaBasicBlocksFromBasicBlock(
                                           max_instr_len, binja_instruction)) {
       continue;
     }
-    LogDebug("0x%p - %zu", (void*)cur_instr_addr, binja_instruction.length);
 
     // Add disassembled instruction to the basic block
     triton::arch::Instruction new_instr(
@@ -55,10 +56,11 @@ std::vector<MetaBasicBlock> ExtractMetaBasicBlocksFromBasicBlock(
       return {};
     }
     triton_bb.add(new_instr);
+    LogDebug("0x%p - %zu - '%s'", (void*)cur_instr_addr,
+             binja_instruction.length, new_instr.getDisassembly().c_str());
 
     // 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 +177,8 @@ 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)) {
+      LogDebug("jump detected: %s", last_instr.getDisassembly().c_str());
       cur_triton_bb.remove(last_instr_index);
     }
   }
@@ -190,6 +192,34 @@ 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:
+      // Match `bl` and `blr`
+      return instr.getDisassembly().find("bl") == 0;
+    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();
+      // Match `b` but not `bl` or `bl.XX`, so we add a space at the end
+      return disassembly.find("b ") == 0;
+    }
+    default:
+      return false;
+  }
+}
+
 // Simplify the given `MetaBasicBlock`s with Triton's dead store elimination
 // pass
 std::vector<MetaBasicBlock> SimplifyMetaBasicBlocks(