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

#include <opcodes.h>
#include <tracing.h>
#include <isize.h>
#include <memory.h>
#include <binary.h>
#include <errors.h>
#include <stdlib.h>
#include <assert.h>

#if ENABLE_BINARY_INSTRUMENTATION

#ifdef __x86_64__
#ifdef HAVE_LIBOPCODE
#include "arch/x86_64/hijack.c"
#endif	/* HAVE_LIBOPCODES */
#endif

#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);
  trace_write(child, orig_addr, (uint8_t*) (&reloc_addr), sizeof(reloc_addr));
}



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
   * 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: (small) jump to reloc_addr (<-- small_jump)
   * reloc_addr:  (long) jump to repl_addr (<-- long_jump)
   * repl_addr: perform eztrace stuff (recording events, etc.)
   * repl_addr: call *orig_addr (=reloc_addr+over)
   *   orig_addr: replay the overwritten opcodes
   *   orig_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 small (ie. relative) jump from sym_addr to reloc_addr
     - install long (ie. absolute) jump from reloc_addr to repl_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 repl_addr
  ssize_t long_jump_size = __install_jump(child, MAX_TRAMPOLINE_SIZE,reloc_addr, repl_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, orig_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;
}


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 symbol %s (0x%lx) with %s (0x%lx) using %s (0x%lx) as function pointer...\n",
      toHijack->symbol_name, (unsigned long) addr, repl->symbol_name,
      (unsigned long) repl_addr, orig->symbol_name, (unsigned long) orig_addr);
  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;
}
#endif	/* ENABLE_BINARY_INSTRUMENTATION */
