/* -*- c-file-style: "GNU" -*- */
/*
 * Copyright (C) CNRS, INRIA, Université Bordeaux 1, Télécom SudParis
 * See COPYING in top-level directory.
 */

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

#include "binary.h"
#include "eztrace-instrumentation/errors.h"
#include "eztrace-instrumentation/tracing.h"
#include "isize.h"
#include "memory.h"
#include "opcodes.h"

#include <assert.h>
#include <stdlib.h>

#ifndef __x86_64__
#error unsupported CPU architecture
#endif

int check_instructions(void* bin, pid_t child, word_uint sym_addr, word_uint reloc_addr,
                       size_t length);
#ifndef HAVE_CHECK_INSTRUCTIONS
/* check_instructions is not defined.
 * define it as a dummy function that always allow the insertion of a trampoline
 */
int check_instructions(void* bin, pid_t child, word_uint sym_addr, word_uint reloc_addr,
                       size_t length) {
  return 1;
}
#endif

#define HIJACK_CODE_FAILED(test) \
  if (test) {                    \
    goto out_failed;             \
  }

static ssize_t _install_jump(pid_t child,
			     ssize_t max_size,
			     word_uint src_addr,
			     word_uint dest_addr) {
  // create the trampoline back to the original address
  uint8_t* jump = (uint8_t*)malloc(max_size);
  ssize_t jump_size = generate_trampoline(jump,
                                          max_size,
                                          src_addr,
                                          dest_addr);
  HIJACK_CODE_FAILED(jump_size < 0);

  // install the jump
  pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                "(hijack)\tinserting a jump from (0x%lx) to (0x%lx)\n",
                (unsigned long)(src_addr), (unsigned long)(dest_addr));
  trace_write(child, src_addr, jump, jump_size);

out:
  assert(jump_size <= max_size);
  free(jump);
  return jump_size;
out_failed:
  jump_size = -1;
  goto out;
}

/* copies the overwritten opcodes, creates the back jump from
 * reloc_addr+overwrite_size to sym_addr+overwrite_size and installs it
 * return the size of the created back jump (or -1 in case of an error)
 */
static ssize_t _install_back_jump(pid_t child,
				  word_uint reloc_addr,
				  word_uint sym_addr,
				  ssize_t overwrite_size) {

  // copy the overwritten code to reloc_addr
  pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                "(hijack)\tcopying overwritten code (%d bytes) from original code (0x%lx) to relocated position (0x%lx)\n",
                overwrite_size, (unsigned long)sym_addr, (unsigned long)(reloc_addr));
  trace_copy(child, sym_addr, reloc_addr, overwrite_size);

  ssize_t back_jump_size = _install_jump(child,
					 MAX_TRAMPOLINE_SIZE,
					 reloc_addr + overwrite_size,
					 sym_addr + overwrite_size);
  HIJACK_CODE_FAILED(back_jump_size < 0);

out:
  return back_jump_size;

out_failed:
  back_jump_size = -1;
  goto out;
}

/* creates the first jump from sym_addr to repl_addr
 *
 * return the size that will be overwritten when the jump is installed
 * or -1 in case of an error.
 */
static ssize_t _prepare_first_jump(void* bin,
				   pid_t child,
				   word_uint sym_addr,
				   word_uint sym_size,
				   word_uint repl_addr,
				   uint8_t** first_jump) {

  *first_jump = (uint8_t*)malloc(sym_size);

  // Build the first jump from sym_addr to repl_addr
  ssize_t first_jump_size = generate_trampoline(*first_jump, sym_size, sym_addr,
                                                repl_addr);
  HIJACK_CODE_FAILED(first_jump_size < 0);

  // get the size to override
  ssize_t over = get_overridden_size(bin, child, sym_addr, first_jump_size);
  HIJACK_CODE_FAILED(over < first_jump_size || over > sym_size);

out:
  return over;
out_failed:
  /* no need to free first_jump since it will be freed by the calling function */
  over = -1;
  goto out;
}

/* install the first jump from sym_addr to repl_addr and free *first_jump
 */
static void _finalize_first_jump(pid_t child,
				 word_uint sym_addr,
				 word_uint repl_addr,
				 uint8_t* first_jump,
				 ssize_t overwrite_size) {
  // install first_jump
  pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                "(hijack)\tinserting first jump from 0x%lx to 0x%lx\n",
                (unsigned long)sym_addr, (unsigned long)repl_addr);
  trace_write(child, sym_addr, first_jump, overwrite_size); // Insert the long jump

  free(first_jump);
}

/* set the callback address to reloc_addr
 */
static void _assign_callback_address(pid_t child,
				     word_uint reloc_addr,
				     word_uint orig_addr) {
  pptrace_debug(PPTRACE_DEBUG_LEVEL_ALL,
                "(hijack)\tassigning callback (located at addr 0x%lx) to 0x%lx\n",
                (unsigned long)orig_addr, (unsigned long)reloc_addr);

  /* the struct ezt_instrumented_function structure contains a callback field that points to
   * a callback to reloc_addr.
   * we need to do something like:
   * *f->callback = reloc_addr
   */

  /* get the value of f->callback */
  word_uint callback_addr = orig_addr; /* &f->callback */
  word_uint buffer;		       /* f->callback */
  trace_read(child, callback_addr, (uint8_t*)&buffer,  sizeof(callback_addr));

  /* now write *f->callback = reloc_addr; */
  trace_write(child, buffer, (uint8_t*)&reloc_addr,  sizeof(callback_addr));
}

static ssize_t hijack_code_small_jump(void* bin, pid_t child, word_uint sym_addr,
                                      word_uint sym_size, word_uint reloc_addr,
                                      word_uint orig_addr, word_uint repl_addr) {

  // If the size is not available, then we suppose we have enough room
  if (sym_size <= 0) {
    sym_size = MAX_TRAMPOLINE_SIZE;
  }
  uint8_t* small_jump = NULL;

  /* sym_addr: address of the original function (in which we install a call to eztrace)
   * reloc_addr: address of the block of data that jumps to eztrace
   * orig_addr: address of the eztrace version of the function
   * repl_addr: address of the callback in eztrace that permits to call the user-defined function
   */

  /* the goal here is to change the child process so that when the application calls the sym_addr function,
   * the following behavior occurs:
   * call sym_addr: (small) jump to reloc_addr (<-- small_jump)
   * reloc_addr:  (long) jump to repl_addr (<-- long_jump)
   * orig_addr: perform eztrace stuff (recording events, etc.)
   * orig_addr: call *repl_addr (=reloc_addr+over)
   *   repl_addr: replay the overwritten opcodes
   *   repl_addr: (long) jump to sym_addr+x (<-- back_jump)
   *   sym_addr: process stuff (it is the application function)
   *   sym_addr: ret -> go back to repl_addr
   * orig_addr: perform eztrace stuff (recording events, etc.)
   * orig_addr: ret -> go back to the application
   */

  /* We need to :
     - install a small (ie. relative) jump from sym_addr to reloc_addr
     - install long (ie. absolute) jump from reloc_addr to orig_addr
     - copy the overwritten opcodes from sym_addr to reloc_addr+sizeof(long_jump)
     - install a jump from reloc_addr+sizeof(long_jump) to sym_addr
     - set the callback value to reloc_addr+sizeof(long_jump)
   */

  // Prepare the first (small) jump from sym_addr to reloc_addr
  ssize_t over = _prepare_first_jump(bin, child, sym_addr, sym_size, reloc_addr, &small_jump);
  HIJACK_CODE_FAILED((!check_instructions(bin, child, sym_addr, reloc_addr, over)));

  // Install the (long) jump from reloc_addr to orig_addr
  ssize_t long_jump_size = _install_jump(child, MAX_TRAMPOLINE_SIZE, reloc_addr, orig_addr);
  HIJACK_CODE_FAILED(long_jump_size < 0);

  // copy the overwritten opcodes and installs the jump back to sym_addr
  word_uint back_jump_addr = reloc_addr + long_jump_size;
  ssize_t back_jump_size = _install_back_jump(child, back_jump_addr, sym_addr, over);
  HIJACK_CODE_FAILED(back_jump_size < 0);

  // Insert the trampoline into the dest process
  _finalize_first_jump(child, sym_addr, reloc_addr, small_jump, over);

  // set the callback to back_jump_addr
  _assign_callback_address(child, back_jump_addr, repl_addr);

  // Return the size of the inserted code
  ssize_t retval = long_jump_size + back_jump_size + over;
out:
  return retval;
out_failed:
  free(small_jump);
  retval = -1;
  goto out;
}

static ssize_t hijack_code_long_jump(void* bin, pid_t child, word_uint sym_addr,
                                     word_uint sym_size, word_uint reloc_addr,
                                     word_uint orig_addr, word_uint repl_addr) {

  // If the size is not available, then we assume we have enough room
  if (sym_size <= 0) {
    sym_size = MAX_TRAMPOLINE_SIZE;
  }

  /* sym_addr: address of the original function (in which we install a call to eztrace)
   * reloc_addr: address of the block of data that jumps back to the application
   * repl_addr: address of the eztrace version of the function
   * orig_addr: address of the callback in eztrace that permits to call the user-defined function
   */

  /* the goal here is to change the child process so that when the application calls the sym_addr function,
   * the following behavior occurs:
   * call sym_addr: jump to repl_addr (<-- first_jump)
   * repl_addr: perform eztrace stuff (recording events, etc.)
   * repl_addr: call *orig_addr (=reloc_addr)
   *   reloc_addr: replay the overwritten opcodes
   *   reloc_addr: (long) jump to sym_addr+x (<-- back_jump)
   *   sym_addr: process stuff (it is the application function)
   *   sym_addr: ret -> go back to repl_addr
   * repl_addr: perform eztrace stuff (recording events, etc.)
   * repl_addr: ret -> go back to the application
   */

  /* We need to :
     - install a jump from sym_addr to repl_addr
     - copy the overwritten opcodes from sym_addr to reloc_addr
     - install a jump from reloc_addr to sym_addr
     - set the callback value to reloc_addr
   */

  // Prepare the first jump from sym_addr to repl_addr
  uint8_t* first_jump = NULL;
  ssize_t overwrite_size = _prepare_first_jump(bin, child, sym_addr, sym_size, repl_addr, &first_jump);
  HIJACK_CODE_FAILED(overwrite_size < 0);

  // Copy the opcodes that are going to be overwritten by first_jump
  // and install the jump back to sym_addr

  /* BUG: we should check that the relocated opcodes to not depend on their address
   * For instance, mov (%rip+x), %eax fails
   */
  ssize_t back_jump_size = _install_back_jump(child, reloc_addr, sym_addr, overwrite_size);
  HIJACK_CODE_FAILED(back_jump_size < 0);

  // Now we can install the first jump (and thus overwrite the first opcode)
  _finalize_first_jump(child, sym_addr, repl_addr, first_jump, overwrite_size);

  // set the callback to reloc_addr
#ifdef __arm__
  /* we need to switch to Thumb momde */
  reloc_addr |= 0x01;
#endif
  _assign_callback_address(child, reloc_addr, orig_addr);

  // Return the size of the inserted code
  ssize_t retval = back_jump_size + overwrite_size;
out:
  return retval;

out_failed:
  free(first_jump);
  retval = 1;
  goto out;
}

ssize_t hijack_code(void* bin, pid_t child, word_uint sym_addr,
                    word_uint sym_size, word_uint reloc_addr,
                    word_uint orig_addr, word_uint repl_addr) {
#if __arm__
  return hijack_code_long_jump(bin, child, sym_addr, sym_size, reloc_addr, orig_addr, repl_addr);
#else
  return hijack_code_small_jump(bin, child, sym_addr, sym_size, reloc_addr,
                                orig_addr, repl_addr);
#endif
}

ssize_t hijack(void* bin, pid_t child, zzt_symbol* toHijack, zzt_symbol* orig,
               zzt_symbol* repl) {
  // Compute the addresses from the symbols

  word_uint addr =
      (word_uint)(toHijack->symbol_offset + toHijack->section_addr);
  word_uint sym_size = (word_uint)(toHijack->symbol_size);
  word_uint orig_addr = (word_uint)(orig->symbol_offset + orig->section_addr);
  word_uint repl_addr = (word_uint)(repl->symbol_offset + repl->section_addr);

  pptrace_debug(PPTRACE_DEBUG_LEVEL_DEBUG,
		"Hijacking application symbol %s (0x%lx)\n",
		toHijack->symbol_name, (unsigned long)addr); /* application symbol */
  
  pptrace_debug(PPTRACE_DEBUG_LEVEL_DEBUG,
		"\tthe library function is %s (0x%lx)\n",
		orig->symbol_name, (unsigned long)orig_addr); /* library symbol */

  pptrace_debug(PPTRACE_DEBUG_LEVEL_DEBUG,
		"\tThe replay address will be stored in a callback in %s (0x%lx).\n",
		repl->symbol_name, (unsigned long)repl_addr); /* callback */

  pptrace_debug(PPTRACE_DEBUG_LEVEL_DEBUG,
                "Allocating buffer for relocating bytes... ");
  word_uint reloc_addr = allocate_buffer(child, MAX_TRAMPOLINE_SIZE);
  if (reloc_addr != 0) {
    pptrace_debug(PPTRACE_DEBUG_LEVEL_DEBUG, "ok (0x%lx)\n",
                  (unsigned long)reloc_addr);
    ssize_t hijack_size = hijack_code(bin, child, addr, sym_size, reloc_addr,
                                      orig_addr, repl_addr);

    if (hijack_size > 0 && hijack_size != MAX_TRAMPOLINE_SIZE) {
      correct_buffer_allocation(child, MAX_TRAMPOLINE_SIZE, hijack_size);
    }
    pptrace_debug(PPTRACE_DEBUG_LEVEL_DEBUG, "Symbol hijacked...\n");
    return hijack_size;
  }
  pptrace_debug(PPTRACE_DEBUG_LEVEL_DEBUG, "failed!\n");
  return -1;
}
