cbz, tbz, tbnz support for aarch64 cmplog (#298)

* add support for cbz/tbz

* remove unecessary print

* implemented support for tbz

* add support for tbnz

* fix an error in the emitted code for both tbz/tbnz

Co-authored-by: Omree <Omree10@gmail.com>

libafl_frida/src/cmplog_rt.rs

@@ -9,6 +9,8 @@ extern "C" {
 
 pub struct CmpLogRuntime {
     ops_save_register_and_blr_to_populate: Option<Box<[u8]>>,
+    ops_handle_tbz_masking: Option<Box<[u8]>>,
+    ops_handle_tbnz_masking: Option<Box<[u8]>>,
 }
 
 impl CmpLogRuntime {
@@ -16,6 +18,8 @@ impl CmpLogRuntime {
     pub fn new() -> CmpLogRuntime {
         Self {
             ops_save_register_and_blr_to_populate: None,
+            ops_handle_tbz_masking: None,
+            ops_handle_tbnz_masking: None,
         }
     }
 
@@ -72,10 +76,52 @@ impl CmpLogRuntime {
             );};
         }
 
+        macro_rules! tbz_masking {
+            ($ops:ident) => {dynasm!($ops
+                ; .arch aarch64
+                ; mov x5, #1
+                ; lsl x5, x5, x1
+                ; eor x5, x5, #255
+                ; orr x1, x0, x5
+            );};
+        }
+
+        macro_rules! tbnz_masking {
+            ($ops:ident) => {dynasm!($ops
+                ; .arch aarch64
+                ; mov x5, #1
+                ; lsl x5, x5, x1
+                ; orr x1, x0, x5
+            );};
+
+        }
+
+        let mut ops_handle_tbz_masking =
+            dynasmrt::VecAssembler::<dynasmrt::aarch64::Aarch64Relocation>::new(0);
+        tbz_masking!(ops_handle_tbz_masking);
+
+        let mut ops_handle_tbnz_masking =
+            dynasmrt::VecAssembler::<dynasmrt::aarch64::Aarch64Relocation>::new(0);
+        tbnz_masking!(ops_handle_tbnz_masking);
+
         let mut ops_save_register_and_blr_to_populate =
             dynasmrt::VecAssembler::<dynasmrt::aarch64::Aarch64Relocation>::new(0);
         blr_to_populate!(ops_save_register_and_blr_to_populate);
 
+        self.ops_handle_tbz_masking = Some(
+            ops_handle_tbz_masking
+                .finalize()
+                .unwrap()
+                .into_boxed_slice(),
+        );
+
+        self.ops_handle_tbnz_masking = Some(
+            ops_handle_tbnz_masking
+                .finalize()
+                .unwrap()
+                .into_boxed_slice(),
+        );
+
         self.ops_save_register_and_blr_to_populate = Some(
             ops_save_register_and_blr_to_populate
                 .finalize()
@@ -93,6 +139,20 @@ impl CmpLogRuntime {
     pub fn ops_save_register_and_blr_to_populate(&self) -> &[u8] {
         self.ops_save_register_and_blr_to_populate.as_ref().unwrap()
     }
+
+    /// Get the blob which handles the tbz opcode masking
+    #[inline]
+    #[must_use]
+    pub fn ops_handle_tbz_masking(&self) -> &[u8] {
+        self.ops_handle_tbz_masking.as_ref().unwrap()
+    }
+
+    /// Get the blob which handles the tbnz opcode masking
+    #[inline]
+    #[must_use]
+    pub fn ops_handle_tbnz_masking(&self) -> &[u8] {
+        self.ops_handle_tbnz_masking.as_ref().unwrap()
+    }
 }
 
 impl Default for CmpLogRuntime {

libafl_frida/src/helper.rs

@@ -46,6 +46,11 @@ enum CmplogOperandType {
     Mem(capstone::RegId, capstone::RegId, i32, u32),
 }
 
+enum SpecialCmpLogCase {
+    Tbz,
+    Tbnz,
+}
+
 #[cfg(target_vendor = "apple")]
 const ANONYMOUS_FLAG: MapFlags = MapFlags::MAP_ANON;
 #[cfg(not(target_vendor = "apple"))]
@@ -373,11 +378,17 @@ impl<'a> FridaInstrumentationHelper<'a> {
                             todo!("Implement cmplog for non-aarch64 targets");
                             #[cfg(all(feature = "cmplog", target_arch = "aarch64"))]
                             // check if this instruction is a compare instruction and if so save the registers values
-                            if let Ok((op1, op2)) =
+                            if let Ok((op1, op2, special_case)) =
                                 helper.cmplog_is_interesting_instruction(address, instr)
                             {
                                 //emit code that saves the relevant data in runtime(passes it to x0, x1)
-                                helper.emit_comparison_handling(address, &output, op1, op2);
+                                helper.emit_comparison_handling(
+                                    address,
+                                    &output,
+                                    op1,
+                                    op2,
+                                    special_case,
+                                );
                             }
                         }
 
@@ -423,6 +434,7 @@ impl<'a> FridaInstrumentationHelper<'a> {
         output: &StalkerOutput,
         op1: CmplogOperandType,
         op2: CmplogOperandType,
+        special_case: Option<SpecialCmpLogCase>,
     ) {
         let writer = output.writer();
 
@@ -504,6 +516,17 @@ impl<'a> FridaInstrumentationHelper<'a> {
         match op2 {
             CmplogOperandType::Imm(value) | CmplogOperandType::Cimm(value) => {
                 writer.put_ldr_reg_u64(Aarch64Register::X1, value);
+                match special_case {
+                    Some(inst) => match inst {
+                        SpecialCmpLogCase::Tbz => {
+                            writer.put_bytes(&self.cmplog_runtime.ops_handle_tbz_masking());
+                        }
+                        SpecialCmpLogCase::Tbnz => {
+                            writer.put_bytes(&self.cmplog_runtime.ops_handle_tbnz_masking());
+                        }
+                    },
+                    None => (),
+                }
             }
             CmplogOperandType::Regid(reg) => {
                 let reg = self.writer_register(reg);
@@ -1054,10 +1077,18 @@ impl<'a> FridaInstrumentationHelper<'a> {
         &self,
         _address: u64,
         instr: &Insn,
-    ) -> Result<(CmplogOperandType, CmplogOperandType), ()> {
+    ) -> Result<
+        (
+            CmplogOperandType,
+            CmplogOperandType,
+            Option<SpecialCmpLogCase>,
+        ),
+        (),
+    > {
         // We only care for compare instrunctions - aka instructions which set the flags
         match instr.mnemonic().unwrap() {
-            "cmp" | "ands" | "subs" | "adds" | "negs" | "ngcs" | "sbcs" | "bics" | "cls" => (),
+            "cmp" | "ands" | "subs" | "adds" | "negs" | "ngcs" | "sbcs" | "bics" | "cls"
+            | "cbz" | "tbz" | "tbnz" => (),
             _ => return Err(()),
         }
         let operands = self
@@ -1066,9 +1097,14 @@ impl<'a> FridaInstrumentationHelper<'a> {
             .unwrap()
             .arch_detail()
             .operands();
-        if operands.len() != 2 {
+
+        // cbz - 1 operand, tbz - 3 operands
+        let special_case = ["cbz", "tbz", "tbnz"].contains(&instr.mnemonic().unwrap());
+        if operands.len() != 2 || !special_case {
             return Err(());
         }
+        // cbz marked as special since there is only 1 operand
+        let special_case = instr.mnemonic().unwrap() == "cbz";
 
         let operand1 = if let Arm64Operand(arm64operand) = operands.first().unwrap() {
             match arm64operand.op_type {
@@ -1087,7 +1123,10 @@ impl<'a> FridaInstrumentationHelper<'a> {
             None
         };
 
-        let operand2 = if let Arm64Operand(arm64operand2) = operands.last().unwrap() {
+        let operand2 = match special_case {
+            true => Some(CmplogOperandType::Imm(0)),
+            false => {
+                if let Arm64Operand(arm64operand2) = &operands[1] {
                     match arm64operand2.op_type {
                         Arm64OperandType::Reg(regid) => Some(CmplogOperandType::Regid(regid)),
                         Arm64OperandType::Imm(val) => Some(CmplogOperandType::Imm(val as u64)),
@@ -1102,10 +1141,19 @@ impl<'a> FridaInstrumentationHelper<'a> {
                     }
                 } else {
                     None
+                }
+            }
+        };
+
+        // tbz will need to have special handling at emit time(masking operand1 value with operand2)
+        let special_case = match instr.mnemonic().unwrap() {
+            "tbz" => Some(SpecialCmpLogCase::Tbz),
+            "tbnz" => Some(SpecialCmpLogCase::Tbnz),
+            _ => None,
         };
 
         if operand1.is_some() && operand2.is_some() {
-            Ok((operand1.unwrap(), operand2.unwrap()))
+            Ok((operand1.unwrap(), operand2.unwrap(), special_case))
         } else {
             Err(())
         }