Skip to content

Commit

Permalink
relocator: 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.

Co-authored-by: Håvard Sørbø <[email protected]>
  • Loading branch information
@oleavr @hsorbo
oleavr and hsorbo committed Sep 19, 2023
1 parent 43b8d2c commit 3efb405
Showing 1 changed file with 130 additions and 22 deletions.
152 changes: 130 additions & 22 deletions gum/arch-arm64/gumarm64relocator.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,14 @@

#include "gummemory.h"

#include <string.h>

#define GUM_MAX_INPUT_INSN_COUNT (100)

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

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

struct _GumRegState
{
GumRegAccess access;
GumAddress address;
};

struct _GumRegMap
{
csh capstone;
GumRegState states[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 +63,11 @@ 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 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 +374,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 +394,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 +424,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 +438,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 +516,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 +535,8 @@ gum_arm64_relocator_can_relocate (gpointer address,
{
current_code = NULL;
}

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

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

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

x16_used = FALSE;
x17_used = FALSE;
*available_scratch_reg = ARM64_REG_INVALID;

for (insn_index = 0; insn_index != n / 4; insn_index++)
for (i = 0; i != G_N_ELEMENTS (reg_map.states); i++)
{
const cs_insn * insn = rl.input_insns[insn_index];
const cs_arm64 * info = &insn->detail->arm64;
uint8_t op_index;
const GumRegState * state = &reg_map.states[i];

for (op_index = 0; op_index != info->op_count; op_index++)
if (state->access == GUM_REG_ACCESS_CLOBBERED &&
state->address >= rl.input_pc)
{
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;
}
*available_scratch_reg = ARM64_REG_X0 + i;
break;
}
}

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 (*available_scratch_reg == ARM64_REG_INVALID)
{
const GumRegState * x16 = &reg_map.states[16];
const GumRegState * x17 = &reg_map.states[17];

if (x16->access == GUM_REG_ACCESS_UNKNOWN ||
x16->address >= rl.input_pc)
{
*available_scratch_reg = ARM64_REG_X16;
}
else if (x17->access == GUM_REG_ACCESS_UNKNOWN ||
x17->address >= rl.input_pc)
{
*available_scratch_reg = ARM64_REG_X17;
}
}
}

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

return TRUE;
}

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

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;
GumRegState * state;

reg_index = gum_reg_map_resolve_register (regs_read[i]);
if (reg_index == -1)
continue;
state = &self->states[reg_index];

if (state->access == GUM_REG_ACCESS_UNKNOWN)
{
state->access = GUM_REG_ACCESS_READ;
state->address = insn->address;
}
}

for (i = 0; i != write_count; i++)
{
gint reg_index;
GumRegState * state;

reg_index = gum_reg_map_resolve_register (regs_write[i]);
if (reg_index == -1)
continue;
state = &self->states[reg_index];

if (state->access == GUM_REG_ACCESS_UNKNOWN)
{
state->access = GUM_REG_ACCESS_CLOBBERED;
state->address = insn->address;
}
}
}

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;
}

0 comments on commit 3efb405

Please sign in to comment.