basic_block.c

#include "basic_block.h"
#include "instr.h"

#ifdef BASIC_BLOCK_TRACE
#include <stdio.h>
#endif

#if GUM_NATIVE_CPU == GUM_CPU_AMD64

static const guint8 afl_log_code[] = {
    0x9c,                                      /* pushfq */
    0x50,                                      /* push rax */
    0x51,                                      /* push rcx */
    0x52,                                      /* push rdx */

    0x48, 0x8d, 0x05, 0x27, 0x00, 0x00, 0x00,  /* lea rax, sym._afl_area_ptr_ptr */
    0x48, 0x8b, 0x00,                          /* mov rax, qword [rax] */
    0x48, 0x8b, 0x00,                          /* mov rax, qword [rax] */
    0x48, 0x8d, 0x0d, 0x22, 0x00, 0x00, 0x00,  /* lea rcx, sym.___afl_prev_loc */
    0x48, 0x8b, 0x11,                          /* mov rdx, qword [rcx] */
    0x48, 0x8b, 0x12,                          /* mov rdx, qword [rdx] */
    0x48, 0x31, 0xfa,                          /* xor rdx, rdi */
    0xfe, 0x04, 0x10,                          /* inc byte [rax + rdx] */
    0x48, 0xd1, 0xef,                          /* shr rdi, 1 */
    0x48, 0x8b, 0x01,                          /* mov rax, qword [rcx] */
    0x48, 0x89, 0x38,                          /* mov qword [rax], rdi */

    0x5a,                                      /* pop rdx */
    0x59,                                      /* pop rcx */
    0x58,                                      /* pop rax */
    0x9d,                                      /* popfq */

    0xc3,                                      /* ret */

    /* Read-only data goes here: */
        /* uint8_t** afl_area_ptr_ptr */
        /* uint64_t* afl_prev_loc_ptr */
};

/*
 * Generated by Clang from:
 *
 *     #include <stdint.h>
 *     #include <stdlib.h>
 *
 *     uint8_t** afl_area_ptr_ptr;
 *     uint64_t* __afl_prev_loc;
 *
 *     void afl_log(size_t offset) {
 *         (*afl_area_ptr_ptr)[offset ^ *__afl_prev_loc]++;
 *         *__afl_prev_loc = offset >> 1;
 *     }
 *
 * With four push/pop pairs manually added to the prolog/epilog.
 */

#else

static void on_basic_block(GumCpuContext* context, gpointer user_data) {
    guint64 current_pc = GUM_ADDRESS (user_data);
#ifdef BASIC_BLOCK_TRACE
    printf("Entered BB @ 0x%llx\n", current_pc);
#endif
    afl_maybe_log(current_pc);
}

#endif

void instr_basic_block(GumStalkerIterator* iterator, GumStalkerOutput* output, gpointer user_data) {
    range_t* range = (range_t*) user_data;

    const cs_insn* instr;
    gboolean begin = TRUE;
    while (gum_stalker_iterator_next(iterator, &instr)) {
        if (begin) {
            guint64 current_pc = instr->address - range->base_address;
            if (range->code_start <= current_pc && range->code_end >= current_pc) {
#ifdef BASIC_BLOCK_TRACE
              printf("Transforming BB @ 0x%llx\n", current_pc);
#endif
#if GUM_NATIVE_CPU == GUM_CPU_AMD64
              guint64 area_offset = (current_pc >> 4) ^ (current_pc << 8);
              area_offset &= MAP_SIZE - 1;
              if (area_offset < afl_instr_rms) {
                  GumX86Writer* cw = output->writer.x86;

                  if (range->current_log_impl == 0 ||
                          !gum_x86_writer_can_branch_directly_between(cw->pc, range->current_log_impl) ||
                          !gum_x86_writer_can_branch_directly_between(cw->pc + 128, range->current_log_impl)) {
                      gconstpointer after_log_impl = cw->code + 1;

                      gum_x86_writer_put_jmp_near_label(cw, after_log_impl);

                      range->current_log_impl = cw->pc;
                      gum_x86_writer_put_bytes(cw, afl_log_code, sizeof(afl_log_code));

                      uint8_t** afl_area_ptr_ptr = &afl_area_ptr;
                      uint64_t* afl_prev_loc_ptr = &range->afl_prev_loc;
                      gum_x86_writer_put_bytes(cw, (const guint8 *) &afl_area_ptr_ptr, sizeof(afl_area_ptr_ptr));
                      gum_x86_writer_put_bytes(cw, (const guint8 *) &afl_prev_loc_ptr, sizeof(afl_prev_loc_ptr));

                      gum_x86_writer_put_label(cw, after_log_impl);
                  }

                  gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_RSP, GUM_REG_RSP, -GUM_RED_ZONE_SIZE);
                  gum_x86_writer_put_push_reg(cw, GUM_REG_RDI);
                  gum_x86_writer_put_mov_reg_address(cw, GUM_REG_RDI, area_offset);
                  gum_x86_writer_put_call_address(cw, range->current_log_impl);
                  gum_x86_writer_put_pop_reg(cw, GUM_REG_RDI);
                  gum_x86_writer_put_lea_reg_reg_offset(cw, GUM_REG_RSP, GUM_REG_RSP, GUM_RED_ZONE_SIZE);
              }
#else
              gum_stalker_iterator_put_callout(iterator, on_basic_block, GSIZE_TO_POINTER (current_pc), NULL);
#endif
              begin = FALSE;
            }
        }

        gum_stalker_iterator_keep(iterator);
    }
}