From cf3604fcefe45adf43425300310c45b1cde2fefc Mon Sep 17 00:00:00 2001
From: s1341 <github@shmarya.net>
Date: Tue, 13 Feb 2024 09:22:55 +0200
Subject: [PATCH] stalker-arm: add another test and fix previous one

---
 tests/core/arch-arm/stalker-arm.c | 47 ++++++++++++++++++++++++++++++-
 1 file changed, 46 insertions(+), 1 deletion(-)

diff --git a/tests/core/arch-arm/stalker-arm.c b/tests/core/arch-arm/stalker-arm.c
index c455750cf8..dd70d09a22 100644
--- a/tests/core/arch-arm/stalker-arm.c
+++ b/tests/core/arch-arm/stalker-arm.c
@@ -99,6 +99,7 @@ TESTLIST_BEGIN (stalker)
 
   TESTENTRY (custom_transformer)
   TESTENTRY (arm_transformer_should_be_able_to_replace_call_with_callout)
+  TESTENTRY (arm_transformer_should_be_able_to_replace_jumpout_with_callout)
   TESTENTRY (arm_callout)
   TESTENTRY (thumb_callout)
   TESTENTRY (unfollow_should_be_allowed_before_first_transform)
@@ -141,6 +142,8 @@ static void duplicate_adds (GumStalkerIterator * iterator,
     GumStalkerOutput * output, gpointer user_data);
 static void replace_call_with_callout (GumStalkerIterator * iterator,
     GumStalkerOutput * output, gpointer user_data);
+static void replace_jumpout_with_callout (GumStalkerIterator * iterator,
+    GumStalkerOutput * output, gpointer user_data);
 static void callout_set_cool (GumCpuContext * cpu_context, gpointer user_data);
 static void transform_arm_return_value (GumStalkerIterator * iterator,
     GumStalkerOutput * output, gpointer user_data);
@@ -3238,7 +3241,7 @@ TESTCASE (arm_transformer_should_be_able_to_replace_call_with_callout)
   memcpy (code, arm_simple_call, CODE_SIZE (arm_simple_call));
 
   fixture->transformer = gum_stalker_transformer_make_from_callback (
-      insert_callout_after_cmp, NULL, NULL);
+      replace_call_with_callout, NULL, NULL);
 
   INVOKE_ARM_EXPECTING (GUM_EXEC, code, 0xc001);
 }
@@ -3264,6 +3267,48 @@ replace_call_with_callout (GumStalkerIterator * iterator,
   }
 }
 
+TESTCODE (arm_simple_jumpout,
+    0x14, 0x05, 0x00, 0xe3, /* mov r0, 1300    */
+    0xff, 0xff, 0xff, 0xea, /* b bump_number */
+    /* bump_number:                            */
+    0x25, 0x00, 0x80, 0xe2, /* add r0, 37      */
+    0x1e, 0xff, 0x2f, 0xe1, /* bx lr           */
+);
+
+TESTCASE (arm_transformer_should_be_able_to_replace_jumpout_with_callout)
+{
+  guint32 code[CODE_SIZE (Arm_simple_jumpout) / sizeof (guint32)], val;
+
+  memcpy (code, arm_simple_jumpout, CODE_SIZE (arm_simple_jumpout));
+
+  fixture->transformer = gum_stalker_transformer_make_from_callback (
+      replace_jumpout_with_callout, NULL, NULL);
+
+  INVOKE_ARM_EXPECTING (GUM_EXEC, code, 0xc001);
+}
+
+static void
+replace_jumpout_with_callout (GumStalkerIterator * iterator,
+                           GumStalkerOutput * output,
+                           gpointer user_data)
+{
+  gint * num_cmp_callouts = user_data;
+  GumMemoryAccess access;
+  const cs_insn * insn;
+
+  while (gum_stalker_iterator_next (iterator, &insn))
+  {
+    if (insn->id == ARM_INS_B)
+    {
+      gum_stalker_iterator_put_callout (iterator, callout_set_cool,
+          NULL, NULL);
+      gum_stalker_iterator_put_chaining_return (iterator);
+      continue;
+    }
+    gum_stalker_iterator_keep (iterator);
+  }
+}
+
 static void
 callout_set_cool (GumCpuContext * cpu_context,
                   gpointer user_data)