/* -*- c-file-style: "GNU" -*- */
/*
 * Copyright (C) CNRS, INRIA, Université Bordeaux 1, Télécom SudParis
 * See COPYING in top-level directory.
 *
 *
 * hijack.c
 *
 * Check the instructions that we are going to overwrite with a trampoline
 *
 */

#ifndef __x86_64__
#error "This file is for x86_64 only"
#endif

#ifndef ENABLE_BINARY_INSTRUMENTATION
#error "This file is for binary instrumentation only"
#endif

#include <eztrace-core/eztrace_config.h>
#include "opcodes.h"
#include "eztrace-instrumentation/tracing.h"
#include "eztrace-instrumentation/errors.h"

#include <bfd.h>
#include <dis-asm.h>

int dummy_print(void* stream __attribute__((unused)), const char* fmt __attribute__((unused)), ...) {
  return 0;
}

/* Check the instructions that we are going to overwrite with a trampoline
 * Return 1 if the trampoline can be inserted without any issue.
 * return 0 if the trampoline may break the program
 */
int check_instructions(void* bin, pid_t child, word_uint sym_addr, word_uint reloc_addr,
                       size_t length) {
  /* copy the child memory locally so that we can modify it */
  uint8_t buffer[1024];
  trace_read(child, sym_addr, buffer, length);

  bfd* abfd = (bfd*)bin;

  /* set to 1 if it is possible to modify the child program */
  int can_change = 1;

  asection* section;
  section = bfd_get_section_by_name(abfd, ".text");
  if (section == NULL) {
    return 0;
  }

  pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                "Check instruction %d bytes starting at %p. relocating to %p\n", length, sym_addr, reloc_addr);
  /* reloc_addr and sym_addr are separated by more than 2GB
   * -> we can't use 'small jumps'
   */
  if (reloc_addr - sym_addr > 1 << 30) {
    can_change = 0;
  }

  struct disassemble_info i;
#if HAVE_BINUTILS_2_38_50_OR_HIGHER
  init_disassemble_info(&i, NULL, dummy_print, NULL);
#else
  init_disassemble_info(&i, NULL, dummy_print);
#endif

  i.arch = bfd_get_arch(abfd);
  i.mach = bfd_get_mach(abfd);

#if HAVE_BINUTILS_2_28_OR_OLDER
  disassembler_ftype disassemble = disassembler(abfd);
#else
  disassembler_ftype disassemble = disassembler(i.arch, bfd_big_endian(abfd), i.mach, abfd);
#endif

  if (disassemble == NULL) {
    return 0;
  }

  i.buffer_vma = section->vma;
  i.buffer_length = section->size;
  i.section = section;

  bfd_malloc_and_get_section(abfd,
                             section,
                             &i.buffer);
  disassemble_init_for_target(&i);

  unsigned long count, pc;
  pc = sym_addr;
  count = 0;
  unsigned start_addr = 0;
  int need_to_push_changes = 0;

  /* browse the instructions and search for problematic opcodes. We may have a
   * problem in the following situations:
   * - one of the instructions is a jump
   * - one of the instruction uses RIP-relative addressing
   *   -> in this case we can try to change the offset
   */
  do {
    pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                  "0x%x  ", pc);
    count = disassemble(pc, &i);
    pc += count;
    pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                  "\n\t");

    /* TODO: check for other problematic instructions (ie. jump)
     */
    if (buffer[start_addr] == 0x8b &&
        buffer[start_addr + 1] == 0x5) {
      /* the instruction is a 'mov offset(%rip), %register'
       * we need to change the value of offset
       */
      if (can_change) {
        /* it is possible to change the offset */
        //uint32_t *p_addr = (uint32_t*)&buffer[start_addr+2];
        uint32_t addr = buffer[start_addr + 2] | buffer[start_addr + 3] << 8 | buffer[start_addr + 4] << 16;
        pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                      "\t\told offset: %lx\n", addr);
        addr = addr - (reloc_addr - sym_addr);
        pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                      "\t\tnew offset: %lx\n", addr);

        buffer[start_addr + 2] = addr & 0xff;
        buffer[start_addr + 3] = (addr & 0xff00) >> 8;
        buffer[start_addr + 4] = (addr & 0xff0000) >> 16;

        addr = buffer[start_addr + 2] | buffer[start_addr + 3] << 8 | buffer[start_addr + 4] << 16;

        need_to_push_changes = 1;
      } else {
        /* it's not possible to change the offset. We should not instrument this
	 * function !
	 */
        return 0;
      }
    }
    start_addr += count;

  } while (count > 0 &&
           pc < bfd_get_start_address(abfd) + section->size &&
           pc <= sym_addr + length);

  if (need_to_push_changes && can_change) {
    /* some instructions were modified. We need to push the modifications */
    pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                  "\tWriting the modified instructions.\n");
    trace_write(child, sym_addr, buffer, length);
  }

  return 1;
}
