4 files changed, 44 insertions, 9 deletions
diff --git a/frida_mode/src/instrument/instrument.c b/frida_mode/src/instrument/instrument.c
index 414dc84c..8ee21f5b 100644
--- a/frida_mode/src/instrument/instrument.c
+++ b/frida_mode/src/instrument/instrument.c
@@ -193,7 +193,20 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
       instrument_debug_start(instr->address, output);
       instrument_coverage_start(instr->address);
 
+#if defined(__arm__)
+      if (output->encoding == GUM_INSTRUCTION_SPECIAL) {
+
+        prefetch_write(GSIZE_TO_POINTER(instr->address + 1));
+
+      } else {
+
+        prefetch_write(GSIZE_TO_POINTER(instr->address));
+
+      }
+
+#else
       prefetch_write(GSIZE_TO_POINTER(instr->address));
+#endif
 
       if (likely(!excluded)) {
 
@@ -213,7 +226,7 @@ static void instrument_basic_block(GumStalkerIterator *iterator,
 
     }
 
-    instrument_debug_instruction(instr->address, instr->size);
+    instrument_debug_instruction(instr->address, instr->size, output);
 
     if (likely(!excluded)) {
 
diff --git a/frida_mode/src/instrument/instrument_arm32.c b/frida_mode/src/instrument/instrument_arm32.c
index fa8b0bd2..16e8eaab 100644
--- a/frida_mode/src/instrument/instrument_arm32.c
+++ b/frida_mode/src/instrument/instrument_arm32.c
@@ -28,7 +28,15 @@ void instrument_coverage_optimize_init(void) {
 
 void instrument_flush(GumStalkerOutput *output) {
 
-  gum_arm_writer_flush(output->writer.arm);
+  if (output->encoding == GUM_INSTRUCTION_SPECIAL) {
+
+    gum_thumb_writer_flush(output->writer.thumb);
+
+  } else {
+
+    gum_arm_writer_flush(output->writer.arm);
+
+  }
 
 }
 
diff --git a/frida_mode/src/instrument/instrument_debug.c b/frida_mode/src/instrument/instrument_debug.c
index a175b585..9c95857f 100644
--- a/frida_mode/src/instrument/instrument_debug.c
+++ b/frida_mode/src/instrument/instrument_debug.c
@@ -32,18 +32,27 @@ static void instrument_debug(char *format, ...) {
 
 }
 
-static void instrument_disasm(guint8 *start, guint8 *end) {
+static void instrument_disasm(guint8 *start, guint8 *end,
+                              GumStalkerOutput *output) {
 
   csh      capstone;
   cs_err   err;
+  cs_mode  mode;
   uint16_t size;
   cs_insn *insn;
   size_t   count = 0;
   size_t   i;
   uint16_t len;
 
+  mode = GUM_DEFAULT_CS_MODE | GUM_DEFAULT_CS_ENDIAN;
+
+#if defined(__arm__)
+  if (output->encoding == GUM_INSTRUCTION_SPECIAL) { mode |= CS_MODE_THUMB; }
+#endif
+
   err = cs_open(GUM_DEFAULT_CS_ARCH,
-                GUM_DEFAULT_CS_MODE | GUM_DEFAULT_CS_ENDIAN, &capstone);
+                CS_MODE_THUMB | GUM_DEFAULT_CS_MODE | GUM_DEFAULT_CS_ENDIAN,
+                &capstone);
   g_assert(err == CS_ERR_OK);
 
   size = GPOINTER_TO_SIZE(end) - GPOINTER_TO_SIZE(start);
@@ -121,11 +130,12 @@ void instrument_debug_start(uint64_t address, GumStalkerOutput *output) {
 
 }
 
-void instrument_debug_instruction(uint64_t address, uint16_t size) {
+void instrument_debug_instruction(uint64_t address, uint16_t size,
+                                  GumStalkerOutput *output) {
 
   if (likely(debugging_fd < 0)) { return; }
   uint8_t *start = (uint8_t *)GSIZE_TO_POINTER(address);
-  instrument_disasm(start, start + size);
+  instrument_disasm(start, start + size, output);
 
 }
 
@@ -136,7 +146,7 @@ void instrument_debug_end(GumStalkerOutput *output) {
 
   instrument_debug("\nGenerated block %p-%p\n", instrument_gen_start,
                    instrument_gen_end);
-  instrument_disasm(instrument_gen_start, instrument_gen_end);
+  instrument_disasm(instrument_gen_start, instrument_gen_end, output);
 
 }
 
diff --git a/frida_mode/src/main.c b/frida_mode/src/main.c
index 913e3a46..1be63bc4 100644
--- a/frida_mode/src/main.c
+++ b/frida_mode/src/main.c
@@ -219,6 +219,8 @@ __attribute__((visibility("default"))) void afl_frida_start(void) {
 
 static int on_main(int argc, char **argv, char **envp) {
 
+  int ret;
+
   on_main_os(argc, argv, envp);
 
   intercept_unhook_self();
@@ -227,14 +229,16 @@ static int on_main(int argc, char **argv, char **envp) {
 
   if (js_main_hook != NULL) {
 
-    return js_main_hook(argc, argv, envp);
+    ret = js_main_hook(argc, argv, envp);
 
   } else {
 
-    return main_fn(argc, argv, envp);
+    ret = main_fn(argc, argv, envp);
 
   }
 
+  return ret;
+
 }
 
 #if defined(EMBEDDED)