about summary refs log tree commit diff
path: root/frida_mode/src/instrument/instrument_x64_cache.c
blob: ef10e13333ed6ee8fe4e53d8db0a094ade6e2d15 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
#include <sys/mman.h>
#include <sys/resource.h>

#include "instrument.h"
#include "util.h"

#if defined(__x86_64__)

  #define INVALID 1
  #define DEFAULT_CACHE_SIZE (256ULL << 20)

gboolean         instrument_cache_enabled = TRUE;
gsize            instrument_cache_size = DEFAULT_CACHE_SIZE;
static gpointer *map_base = MAP_FAILED;

void instrument_cache_config(void) {

  instrument_cache_enabled = (getenv("AFL_FRIDA_INST_NO_CACHE") == NULL);

  if (getenv("AFL_FRIDA_INST_CACHE_SIZE") != NULL) {

    if (!instrument_cache_enabled) {

      FFATAL(
          "AFL_FRIDA_INST_CACHE_SIZE incomatible with "
          "AFL_FRIDA_INST_NO_CACHE");

    }

    instrument_cache_size =
        util_read_address("AFL_FRIDA_INST_CACHE_SIZE", DEFAULT_CACHE_SIZE);
    util_log2(instrument_cache_size);

  }

}

void instrument_cache_init(void) {

  FOKF(cBLU "Instrumentation" cRST " - " cGRN "cache:" cYEL " [%c]",
       instrument_cache_enabled ? 'X' : ' ');
  if (!instrument_cache_enabled) { return; }

  FOKF(cBLU "Instrumentation" cRST " - " cGRN "cache size:" cYEL " [0x%016lX]",
       instrument_cache_size);

  const struct rlimit data_limit = {.rlim_cur = RLIM_INFINITY,
                                    .rlim_max = RLIM_INFINITY};

  if (setrlimit(RLIMIT_AS, &data_limit) != 0) {

    FFATAL("Failed to setrlimit: %d", errno);

  }

  map_base =
      gum_memory_allocate(NULL, instrument_cache_size, instrument_cache_size,
                          GUM_PAGE_READ | GUM_PAGE_WRITE);
  if (map_base == MAP_FAILED) { FFATAL("Failed to map segment: %d", errno); }

  FOKF(cBLU "Instrumentation" cRST " - " cGRN "cache addr:" cYEL " [0x%016lX]",
       GUM_ADDRESS(map_base));

}

static gpointer *instrument_cache_get_addr(gpointer addr) {

  gsize mask = (instrument_cache_size / sizeof(gpointer)) - 1;
  return &map_base[GPOINTER_TO_SIZE(addr) & mask];

}

void instrument_cache_insert(gpointer real_address, gpointer code_address) {

  if (!instrument_cache_enabled) { return; }

  gpointer *target = instrument_cache_get_addr(real_address);
  if (*target == code_address) {

    return;

  } else if (*target == NULL) {

    *target = code_address;

  } else {

    *target = GSIZE_TO_POINTER(INVALID);

  }

}

static gboolean instrument_cache_relocate(GumAddress old_pc, GumAddress new_pc,
                                          gint32  old_offset,
                                          gint32 *new_offset) {

  guint64 old_target = old_pc + old_offset;
  gint64  relocated = old_target - new_pc;

  if (relocated > G_MAXINT32 || relocated < G_MININT32) { return FALSE; }

  *new_offset = relocated;
  return TRUE;

}

static void instrument_cache_rewrite_branch_insn(const cs_insn    *instr,
                                                 GumStalkerOutput *output) {

  GumX86Writer *cw = output->writer.x86;
  cs_x86       *x86 = &instr->detail->x86;
  guint8        modified[sizeof(instr->bytes)] = {0};
  guint8        offset = 0;
  guint8        skip = 0;

  g_assert(sizeof(x86->prefix) == 4);
  g_assert(sizeof(x86->opcode) == 4);

  /*
   * If the target is simply RAX, we can skip writing the code to load the
   * RIP
   */
  if (x86->operands[0].type == X86_OP_REG ||
      x86->operands[0].reg == X86_REG_RAX) {

    return;

  }

  /* Write the prefix */
  for (gsize i = 0; i < sizeof(x86->prefix); i++) {

    if (x86->prefix[i] != 0) {

      if (x86->prefix[i] == 0xf2) {

        skip++;

      } else {

        modified[offset++] = x86->prefix[i];
        skip++;

      }

    }

  }

  /* Write the REX */
  if (x86->rex == 0) {

    /*
     * CALL (near) and JMP (near) default to 64-bit operands, MOV does not,
     * write REX.W
     */
    modified[offset++] = 0x48;

  } else {

    if ((x86->rex & 0xF8) != 0x40) {

      FATAL("Unexpected REX byte: 0x%02x", x86->rex);

    }

    modified[offset++] = x86->rex | 0x08;
    skip++;

  }

  /*
   * CALL is FF /2, JMP is FF /4. The remaining op-code fields should thus be
   * unused
   */

  if (x86->opcode[0] != 0xFF || x86->opcode[1] != 0x00 ||
      x86->opcode[2] != 0x00 || x86->opcode[3] != 0x00) {

    FFATAL("Unexpected Op-code: 0x%02x 0x%02x 0x%02x 0x%02x", x86->opcode[0],
           x86->opcode[1], x86->opcode[2], x86->opcode[3]);

  }

  /* The reg field of the ModRM should be set to 2 for CALL and 4 for JMP */
  guint8 reg = (x86->modrm >> 3) & 7;
  if (reg != 0x4 && reg != 0x2) {

    FFATAL("Unexpected Reg: 0x%02x, ModRM: 0x%02x", reg, x86->modrm);

  }

  /* MOV */
  modified[offset++] = 0x8b;
  skip++;

  /* Clear the reg field (RAX) */
  modified[offset++] = x86->modrm & 0xc7;
  skip++;

  /* Operands */
  guint8 op_len = instr->size - skip;

  /* If our branch was RIP relative, we'll need to fix-up the offset */
  if (x86->operands[0].type == X86_OP_MEM &&
      x86->operands[0].mem.base == X86_REG_RIP) {

    /* RIP relative offsets should be 32-bits */
    if (op_len != sizeof(gint32)) {

      FFATAL("Unexpected operand length: %d\n", op_len);

    }

    gint32 old_offset = *(gint32 *)&instr->bytes[skip];
    gint32 new_offset = 0;
    if (instrument_cache_relocate(instr->address, cw->pc, old_offset,
                                  &new_offset)) {

      gint32 *output = (gint32 *)&modified[offset];
      *output = new_offset;
      offset += sizeof(gint32);

    } else {

      GumAddress target = instr->address + old_offset;
      gum_x86_writer_put_mov_reg_address(cw, GUM_X86_RAX, target);
      gum_x86_writer_put_mov_reg_reg_ptr(cw, GUM_X86_RAX, GUM_X86_RAX);
      return;

    }

  } else {

    for (int i = 0; i < op_len; i++) {

      guint8 val = instr->bytes[i + skip];
      modified[offset++] = val;

    }

  }

  gum_x86_writer_put_bytes(cw, modified, offset);

}

static void instrument_cache_write_push_frame(GumX86Writer *cw) {

  gum_x86_writer_put_mov_reg_offset_ptr_reg(
      cw, GUM_X86_XSP, -(GUM_RED_ZONE_SIZE + (1 * sizeof(gpointer))),
      GUM_X86_XAX);
  gum_x86_writer_put_lahf(cw);
  gum_x86_writer_put_mov_reg_offset_ptr_reg(
      cw, GUM_X86_XSP, -(GUM_RED_ZONE_SIZE + (2 * sizeof(gpointer))),
      GUM_X86_XAX);
  gum_x86_writer_put_mov_reg_offset_ptr_reg(
      cw, GUM_X86_XSP, -(GUM_RED_ZONE_SIZE + (3 * sizeof(gpointer))),
      GUM_X86_XBX);

}

static void instrument_cache_write_pop_frame(GumX86Writer *cw) {

  gum_x86_writer_put_mov_reg_reg_offset_ptr(
      cw, GUM_X86_XBX, GUM_X86_XSP,
      -(GUM_RED_ZONE_SIZE + (3 * sizeof(gpointer))));
  gum_x86_writer_put_mov_reg_reg_offset_ptr(
      cw, GUM_X86_XAX, GUM_X86_XSP,
      -(GUM_RED_ZONE_SIZE + (2 * sizeof(gpointer))));
  gum_x86_writer_put_sahf(cw);
  gum_x86_writer_put_mov_reg_reg_offset_ptr(
      cw, GUM_X86_XAX, GUM_X86_XSP,
      -(GUM_RED_ZONE_SIZE + (1 * sizeof(gpointer))));

}

static void instrument_cache_write_lookup(GumX86Writer *cw) {

  /* &map_base[GPOINTER_TO_SIZE(addr) & MAP_MASK]; */

  gsize mask = (instrument_cache_size / sizeof(gpointer)) - 1;
  gum_x86_writer_put_mov_reg_u64(cw, GUM_X86_XBX, mask);
  gum_x86_writer_put_and_reg_reg(cw, GUM_X86_XAX, GUM_X86_XBX);
  gum_x86_writer_put_shl_reg_u8(cw, GUM_X86_XAX, util_log2(sizeof(gpointer)));
  gum_x86_writer_put_mov_reg_u64(cw, GUM_X86_XBX, GPOINTER_TO_SIZE(map_base));
  gum_x86_writer_put_add_reg_reg(cw, GUM_X86_XAX, GUM_X86_XBX);

  /* Read the return address lookup */
  gum_x86_writer_put_mov_reg_reg_ptr(cw, GUM_X86_XAX, GUM_X86_XAX);

}

void instrument_cache_jmp_call(const cs_insn *instr, GumStalkerOutput *output) {

  GumX86Writer *cw = output->writer.x86;
  cs_x86       *x86 = &instr->detail->x86;

  if (x86->op_count != 1) { FFATAL("Unexpected operand count"); }

  if (x86->operands[0].type == X86_OP_IMM) { return; }

  gconstpointer null = cw->code;

  instrument_cache_write_push_frame(cw);

  /*
   * We are about to re-write the CALL or JMP instruction, but replace the
   * op-code with that for a MOV into RAX. Since we are keeping the operand from
   * the JMP exactly the same, it is imperative that the target register state
   * be exactly the same as how the target left it. Since `LAHF` spoils `RAX` we
   * must restore it from the stack. We also must avoid adjusting `RSP`, so we
   * use `MOV` instructions to store our context into the stack beyond the
   * red-zone.
   */
  gum_x86_writer_put_mov_reg_reg_offset_ptr(
      cw, GUM_X86_XAX, GUM_X86_XSP,
      -(GUM_RED_ZONE_SIZE + (1 * sizeof(gpointer))));

  instrument_cache_rewrite_branch_insn(instr, output);

  instrument_cache_write_lookup(cw);

  /* Test if its set*/
  gum_x86_writer_put_cmp_reg_i32(cw, GUM_X86_XAX, INVALID);
  gum_x86_writer_put_jcc_short_label(cw, X86_INS_JLE, null, GUM_UNLIKELY);

  /* If it's set, then stash the address beyond the red-zone */
  gum_x86_writer_put_mov_reg_offset_ptr_reg(
      cw, GUM_X86_XSP, -(GUM_RED_ZONE_SIZE + (4 * sizeof(gpointer))),
      GUM_X86_XAX);

  if (instr->id == X86_INS_JMP) {

    instrument_cache_write_pop_frame(cw);
    gum_x86_writer_put_jmp_reg_offset_ptr(
        cw, GUM_X86_XSP, -(GUM_RED_ZONE_SIZE + (4 * sizeof(gpointer))));

  } else {

    gum_x86_writer_put_mov_reg_address(
        cw, GUM_X86_XAX, GUM_ADDRESS(instr->address + instr->size));
    gum_x86_writer_put_mov_reg_offset_ptr_reg(cw, GUM_X86_XSP,
                                              -sizeof(gpointer), GUM_X86_XAX);

    instrument_cache_write_pop_frame(cw);

    gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_X86_XSP, GUM_X86_XSP,
                                          -sizeof(gpointer));
    gum_x86_writer_put_jmp_reg_offset_ptr(
        cw, GUM_X86_XSP, -(GUM_RED_ZONE_SIZE + ((4 - 1) * sizeof(gpointer))));

  }

  /* Tidy up our mess and let FRIDA handle it */
  gum_x86_writer_put_label(cw, null);
  instrument_cache_write_pop_frame(cw);

}

void instrument_cache_ret(const cs_insn *instr, GumStalkerOutput *output) {

  GumX86Writer *cw = output->writer.x86;
  cs_x86       *x86 = &instr->detail->x86;
  guint16       n = 0;

  if (x86->op_count != 0) {

    if (x86->operands[0].type != X86_OP_IMM) {

      FFATAL("Unexpected operand type");

    }

    n = x86->operands[0].imm;

  }

  gconstpointer null = cw->code;

  instrument_cache_write_push_frame(cw);

  gum_x86_writer_put_mov_reg_reg_ptr(cw, GUM_X86_XAX, GUM_X86_XSP);

  instrument_cache_write_lookup(cw);

  /* Test if its set*/
  gum_x86_writer_put_cmp_reg_i32(cw, GUM_X86_XAX, INVALID);
  gum_x86_writer_put_jcc_short_label(cw, X86_INS_JLE, null, GUM_UNLIKELY);

  /* If it's set, then overwrite our return address and return */
  gum_x86_writer_put_mov_reg_ptr_reg(cw, GUM_X86_XSP, GUM_X86_XAX);
  instrument_cache_write_pop_frame(cw);

  if (n == 0) {

    gum_x86_writer_put_ret(cw);

  } else {

    gum_x86_writer_put_ret_imm(cw, n);

  }

  /* Tidy up our mess and let FRIDA handle it */
  gum_x86_writer_put_label(cw, null);
  instrument_cache_write_pop_frame(cw);

}

void instrument_cache(const cs_insn *instr, GumStalkerOutput *output) {

  if (!instrument_cache_enabled) { return; }

  switch (instr->id) {

    case X86_INS_RET:
      instrument_cache_ret(instr, output);
      break;

    case X86_INS_CALL:
    case X86_INS_JMP:
      instrument_cache_jmp_call(instr, output);
      break;

    default:
      return;

  }

}

#endif