interceptor: Relocate tiny targets on arm64

Co-authored-by: Håvard Sørbø <[email protected]>

gum/arch-arm64/gumarm64relocator.c

@@ -1,5 +1,6 @@
 /*
  * Copyright (C) 2014-2023 Ole André Vadla Ravnås <[email protected]>
+ * Copyright (C) 2023 Håvard Sørbø <[email protected]>
  *
  * Licence: wxWindows Library Licence, Version 3.1
  */
@@ -370,6 +371,7 @@ gum_arm64_relocator_can_relocate (gpointer address,
                                   arm64_reg * available_scratch_reg)
 {
   guint n = 0;
+  gboolean fully_relocated = FALSE;
   guint8 * buf;
   GumArm64Writer cw;
   GumArm64Relocator rl;
@@ -425,6 +427,7 @@ gum_arm64_relocator_can_relocate (gpointer address,
     GHashTable * checked_targets, * targets_to_check;
     csh capstone;
     guint basic_block_index;
+    gsize first_basic_block_size;
     cs_insn * insn;
     const guint8 * current_code;
     uint64_t current_address;
@@ -439,6 +442,7 @@ gum_arm64_relocator_can_relocate (gpointer address,
     cs_option (capstone, CS_OPT_DETAIL, CS_OPT_ON);
 
     basic_block_index = 0;
+    first_basic_block_size = 0;
     insn = cs_malloc (capstone);
     current_code = rl.input_cur;
     current_address = rl.input_pc;
@@ -521,6 +525,12 @@ gum_arm64_relocator_can_relocate (gpointer address,
           gum_reg_map_apply_instruction (&reg_map, insn);
       }
 
+      if (basic_block_index == 0)
+      {
+        first_basic_block_size =
+            (insn->address + insn->size) - GPOINTER_TO_SIZE (address);
+      }
+
       g_hash_table_iter_init (&iter, targets_to_check);
       if (g_hash_table_iter_next (&iter, &target, NULL))
       {
@@ -540,15 +550,25 @@ gum_arm64_relocator_can_relocate (gpointer address,
     }
     while (current_code != NULL);
 
-    g_hash_table_iter_init (&iter, checked_targets);
-    while (g_hash_table_iter_next (&iter, &target, NULL))
+    fully_relocated =
+        g_hash_table_size (checked_targets) == 1 &&
+        first_basic_block_size <= 128;
+    if (fully_relocated)
+    {
+      n = first_basic_block_size;
+    }
+    else
     {
-      gssize offset = (gssize) target - (gssize) address;
-      if (offset > 0 && offset < (gssize) n)
+      g_hash_table_iter_init (&iter, checked_targets);
+      while (g_hash_table_iter_next (&iter, &target, NULL))
       {
-        n = offset;
-        if (n == 4)
-          break;
+        gssize offset = (gssize) target - (gssize) address;
+        if (offset > 0 && offset < (gssize) n)
+        {
+          n = offset;
+          if (n == 4)
+            break;
+        }
       }
     }
 
@@ -571,7 +591,7 @@ gum_arm64_relocator_can_relocate (gpointer address,
       const GumRegState * state = &reg_map.states[i];
 
       if (state->access == GUM_REG_ACCESS_CLOBBERED &&
-          state->address >= rl.input_pc)
+          (fully_relocated || state->address >= rl.input_pc))
       {
         *available_scratch_reg = ARM64_REG_X0 + i;
         break;
@@ -584,16 +604,19 @@ gum_arm64_relocator_can_relocate (gpointer address,
       const GumRegState * x17 = &reg_map.states[17];
 
       if (x16->access == GUM_REG_ACCESS_UNKNOWN ||
-          x16->address >= rl.input_pc)
+          (fully_relocated || x16->address >= rl.input_pc))
       {
         *available_scratch_reg = ARM64_REG_X16;
       }
       else if (x17->access == GUM_REG_ACCESS_UNKNOWN ||
-          x17->address >= rl.input_pc)
+          (fully_relocated || x17->address >= rl.input_pc))
       {
         *available_scratch_reg = ARM64_REG_X17;
       }
     }
+
+    if (*available_scratch_reg == ARM64_REG_INVALID && fully_relocated)
+      *available_scratch_reg = ARM64_REG_X16;
   }
 
   gum_arm64_relocator_clear (&rl);

gum/backend-arm64/guminterceptor-arm64.c

@@ -1,6 +1,7 @@
 /*
- * Copyright (C) 2014-2022 Ole André Vadla Ravnås <[email protected]>
+ * Copyright (C) 2014-2023 Ole André Vadla Ravnås <[email protected]>
  * Copyright (C) 2022 Francesco Tamagni <[email protected]>
+ * Copyright (C) 2023 Håvard Sørbø <[email protected]>
  *
  * Licence: wxWindows Library Licence, Version 3.1
  */
@@ -53,6 +54,7 @@ struct _GumInterceptorBackend
 struct _GumArm64FunctionContextData
 {
   guint redirect_code_size;
+  guint reloc_code_size;
   arm64_reg scratch_reg;
 };
 
@@ -630,6 +632,7 @@ gum_interceptor_backend_prepare_trampoline (GumInterceptorBackend * self,
       GUM_SCENARIO_ONLINE, &redirect_limit, &data->scratch_reg))
   {
     data->redirect_code_size = 16;
+    data->reloc_code_size = redirect_limit;
 
     ctx->trampoline_slice = gum_code_allocator_alloc_slice (self->allocator);
   }
@@ -661,6 +664,8 @@ gum_interceptor_backend_prepare_trampoline (GumInterceptorBackend * self,
       return FALSE;
     }
 
+    data->reloc_code_size = data->redirect_code_size;
+
     ctx->trampoline_slice = gum_code_allocator_try_alloc_slice_near (
         self->allocator, &spec, alignment);
     if (ctx->trampoline_slice == NULL)
@@ -695,7 +700,7 @@ _gum_interceptor_backend_create_trampoline (GumInterceptorBackend * self,
   gpointer deflector_target;
   GString * signature;
   gboolean is_eligible_for_lr_rewriting;
-  guint reloc_bytes;
+  guint reloc_bytes, overwritten_bytes;
 
   if (!gum_interceptor_backend_prepare_trampoline (self, ctx, &need_deflector))
     return FALSE;
@@ -761,18 +766,25 @@ _gum_interceptor_backend_create_trampoline (GumInterceptorBackend * self,
 
   signature = g_string_sized_new (16);
 
+  overwritten_bytes = 0;
   do
   {
     const cs_insn * insn;
 
     reloc_bytes = gum_arm64_relocator_read_one (ar, &insn);
     g_assert (reloc_bytes != 0);
 
+    if (overwritten_bytes == 0 &&
+        reloc_bytes >= data->redirect_code_size)
+    {
+      overwritten_bytes = reloc_bytes;
+    }
+
     if (signature->len != 0)
       g_string_append_c (signature, ';');
     g_string_append (signature, insn->mnemonic);
   }
-  while (reloc_bytes < data->redirect_code_size);
+  while (reloc_bytes < data->reloc_code_size);
 
   /*
    * Try to deal with minimal thunks that determine their caller and pass
@@ -865,8 +877,8 @@ _gum_interceptor_backend_create_trampoline (GumInterceptorBackend * self,
   gum_arm64_writer_flush (aw);
   g_assert (gum_arm64_writer_offset (aw) <= ctx->trampoline_slice->size);
 
-  ctx->overwritten_prologue_len = reloc_bytes;
-  gum_memcpy (ctx->overwritten_prologue, function_address, reloc_bytes);
+  ctx->overwritten_prologue_len = overwritten_bytes;
+  gum_memcpy (ctx->overwritten_prologue, function_address, overwritten_bytes);
 
   return TRUE;
 }