Skip to content

Commit

Permalink
[WIP] interceptor: Improve scratch register strategy on arm64
Browse files Browse the repository at this point in the history
Such that if the first basic block writes to a register before reading
it, we pick that register instead of gambling on either X16 or X17
being safe to clobber.
  • Loading branch information
oleavr committed Sep 18, 2023
1 parent 23ffc79 commit 53cc84c
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 26 deletions.
144 changes: 122 additions & 22 deletions gum/arch-arm64/gumarm64relocator.c
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,13 @@

#include "gummemory.h"

#include <string.h>

#define GUM_MAX_INPUT_INSN_COUNT (100)

typedef struct _GumCodeGenCtx GumCodeGenCtx;
typedef struct _GumRegMap GumRegMap;
typedef guint GumRegAccess;

struct _GumCodeGenCtx
{
Expand All @@ -22,6 +26,19 @@ struct _GumCodeGenCtx
GumArm64Writer * output;
};

struct _GumRegMap
{
csh capstone;
GumRegAccess accesses[18];
};

enum _GumRegAccess
{
GUM_REG_ACCESS_UNKNOWN,
GUM_REG_ACCESS_READ,
GUM_REG_ACCESS_CLOBBERED,
};

static gboolean gum_arm64_branch_is_unconditional (const cs_insn * insn);

static gboolean gum_arm64_relocator_rewrite_ldr (GumArm64Relocator * self,
Expand All @@ -39,6 +56,12 @@ static gboolean gum_arm64_relocator_rewrite_cbz (GumArm64Relocator * self,
static gboolean gum_arm64_relocator_rewrite_tbz (GumArm64Relocator * self,
GumCodeGenCtx * ctx);

static void gum_reg_map_init (GumRegMap * map, csh capstone);
static void gum_reg_map_apply_instruction (GumRegMap * self,
const cs_insn * insn);
static arm64_reg gum_reg_map_find_available_register (const GumRegMap * self);
static gint gum_reg_map_resolve_register (arm64_reg reg);

GumArm64Relocator *
gum_arm64_relocator_new (gconstpointer input_code,
GumArm64Writer * output)
Expand Down Expand Up @@ -345,12 +368,15 @@ gum_arm64_relocator_can_relocate (gpointer address,
GumArm64Writer cw;
GumArm64Relocator rl;
guint reloc_bytes;
GumRegMap reg_map;

buf = g_alloca (3 * min_bytes);
gum_arm64_writer_init (&cw, buf);

gum_arm64_relocator_init (&rl, address, &cw);

gum_reg_map_init (&reg_map, rl.capstone);

do
{
const cs_insn * insn;
Expand All @@ -362,6 +388,8 @@ gum_arm64_relocator_can_relocate (gpointer address,

n = reloc_bytes;

gum_reg_map_apply_instruction (&reg_map, insn);

if (scenario == GUM_SCENARIO_ONLINE)
{
switch (insn->id)
Expand Down Expand Up @@ -390,6 +418,7 @@ gum_arm64_relocator_can_relocate (gpointer address,
{
GHashTable * checked_targets, * targets_to_check;
csh capstone;
guint basic_block_index;
cs_insn * insn;
const guint8 * current_code;
uint64_t current_address;
Expand All @@ -403,6 +432,7 @@ gum_arm64_relocator_can_relocate (gpointer address,
cs_open (CS_ARCH_ARM64, GUM_DEFAULT_CS_ENDIAN, &capstone);
cs_option (capstone, CS_OPT_DETAIL, CS_OPT_ON);

basic_block_index = 0;
insn = cs_malloc (capstone);
current_code = rl.input_cur;
current_address = rl.input_pc;
Expand Down Expand Up @@ -480,6 +510,9 @@ gum_arm64_relocator_can_relocate (gpointer address,
default:
break;
}

if (basic_block_index == 0)
gum_reg_map_apply_instruction (&reg_map, insn);
}

g_hash_table_iter_init (&iter, targets_to_check);
Expand All @@ -496,6 +529,8 @@ gum_arm64_relocator_can_relocate (gpointer address,
{
current_code = NULL;
}

basic_block_index++;
}
while (current_code != NULL);

Expand All @@ -521,36 +556,38 @@ gum_arm64_relocator_can_relocate (gpointer address,

if (available_scratch_reg != NULL)
{
gboolean x16_used, x17_used;
guint insn_index;

x16_used = FALSE;
x17_used = FALSE;

for (insn_index = 0; insn_index != n / 4; insn_index++)
*available_scratch_reg = gum_reg_map_find_available_register (&reg_map);
if (*available_scratch_reg == ARM64_REG_INVALID)
{
const cs_insn * insn = rl.input_insns[insn_index];
const cs_arm64 * info = &insn->detail->arm64;
uint8_t op_index;
gboolean x16_used, x17_used;
guint insn_index;

x16_used = FALSE;
x17_used = FALSE;

for (op_index = 0; op_index != info->op_count; op_index++)
for (insn_index = 0; insn_index != n / 4; insn_index++)
{
const cs_arm64_op * op = &info->operands[op_index];
const cs_insn * insn = rl.input_insns[insn_index];
const cs_arm64 * info = &insn->detail->arm64;
uint8_t op_index;

if (op->type == ARM64_OP_REG)
for (op_index = 0; op_index != info->op_count; op_index++)
{
x16_used |= op->reg == ARM64_REG_X16;
x17_used |= op->reg == ARM64_REG_X17;
const cs_arm64_op * op = &info->operands[op_index];

if (op->type == ARM64_OP_REG)
{
x16_used |= op->reg == ARM64_REG_X16;
x17_used |= op->reg == ARM64_REG_X17;
}
}
}
}

if (!x16_used)
*available_scratch_reg = ARM64_REG_X16;
else if (!x17_used)
*available_scratch_reg = ARM64_REG_X17;
else
*available_scratch_reg = ARM64_REG_INVALID;
if (!x16_used)
*available_scratch_reg = ARM64_REG_X16;
else if (!x17_used)
*available_scratch_reg = ARM64_REG_X17;
}
}

gum_arm64_relocator_clear (&rl);
Expand Down Expand Up @@ -778,3 +815,66 @@ gum_arm64_relocator_rewrite_tbz (GumArm64Relocator * self,

return TRUE;
}

static void
gum_reg_map_init (GumRegMap * map,
csh capstone)
{
map->capstone = capstone;
memset (map->accesses, 0, sizeof (map->accesses));
}

static void
gum_reg_map_apply_instruction (GumRegMap * self,
const cs_insn * insn)
{
cs_regs regs_read, regs_write;
uint8_t read_count, write_count, i;

cs_regs_access (self->capstone, insn, regs_read, &read_count,
regs_write, &write_count);

for (i = 0; i != read_count; i++)
{
gint reg_index = gum_reg_map_resolve_register (regs_read[i]);
if (reg_index == -1)
continue;
if (self->accesses[reg_index] == GUM_REG_ACCESS_UNKNOWN)
self->accesses[reg_index] = GUM_REG_ACCESS_READ;
}

for (i = 0; i != write_count; i++)
{
gint reg_index = gum_reg_map_resolve_register (regs_write[i]);
if (reg_index == -1)
continue;
if (self->accesses[reg_index] == GUM_REG_ACCESS_UNKNOWN)
self->accesses[reg_index] = GUM_REG_ACCESS_CLOBBERED;
}
}

static arm64_reg
gum_reg_map_find_available_register (const GumRegMap * self)
{
guint i;

for (i = 0; i != G_N_ELEMENTS (self->accesses); i++)
{
if (self->accesses[i] == GUM_REG_ACCESS_CLOBBERED)
return ARM64_REG_X0 + i;
}

return ARM64_REG_INVALID;
}

static gint
gum_reg_map_resolve_register (arm64_reg reg)
{
if (reg >= ARM64_REG_X0 && reg <= ARM64_REG_X17)
return reg - ARM64_REG_X0;

if (reg >= ARM64_REG_W0 && reg <= ARM64_REG_W17)
return reg - ARM64_REG_W0;

return -1;
}
8 changes: 4 additions & 4 deletions gum/backend-arm64/guminterceptor-arm64.c
Original file line number Diff line number Diff line change
Expand Up @@ -950,12 +950,12 @@ _gum_interceptor_backend_activate_trampoline (GumInterceptorBackend * self,
gum_arm64_writer_put_b_imm (aw, on_enter);
break;
case 8:
gum_arm64_writer_put_adrp_reg_address (aw, ARM64_REG_X16, on_enter);
gum_arm64_writer_put_br_reg_no_auth (aw, ARM64_REG_X16);
gum_arm64_writer_put_adrp_reg_address (aw, data->scratch_reg, on_enter);
gum_arm64_writer_put_br_reg_no_auth (aw, data->scratch_reg);
break;
case 16:
gum_arm64_writer_put_ldr_reg_address (aw, ARM64_REG_X16, on_enter);
gum_arm64_writer_put_br_reg (aw, ARM64_REG_X16);
gum_arm64_writer_put_ldr_reg_address (aw, data->scratch_reg, on_enter);
gum_arm64_writer_put_br_reg (aw, data->scratch_reg);
break;
default:
g_assert_not_reached ();
Expand Down

0 comments on commit 53cc84c

Please sign in to comment.