#include "dprintf.h"
#include "stage1.h"
#include "stage2.h"
#include "system.h"
#include "vdl.h"
#include "futex.h"
#include "vdl-alloc.h"
#include "vdl-list.h"
#include "vdl-utils.h"
#include "machine.h"
#include <elf.h>
#include <link.h>


#define READ_LONG(p)				\
  ({long v = *((long*)p);			\
    p+=sizeof(long);				\
    v;})

#define READ_POINTER(p)				\
  ({char * v = *((char**)p);			\
    p+=sizeof(char*);				\
    v;})
static struct Stage2Input
prepare_stage2 (unsigned long entry_point_struct)
{
  struct Stage2Input stage2_input;
  unsigned long tmp = entry_point_struct;
  ElfW(auxv_t) *auxvt, *auxvt_tmp;
  // although the C convention for argc is to be an int (as such, 4 bytes
  // on both i386 and x86_64), the kernel ABI has a long argc (as such,
  // 4 bytes on i386 and 8 bytes on x86_64).
  stage2_input.program_argc = READ_LONG (tmp); // skip argc
  DPRINTF("argc=0x%x\n", stage2_input.program_argc);
  stage2_input.program_argv = (char **)tmp;
  tmp += sizeof(char *)*(stage2_input.program_argc+1); // skip argv
  stage2_input.program_envp = (char **)tmp;
  while (READ_POINTER (tmp) != 0) {} // skip envp
  auxvt = (ElfW(auxv_t) *)tmp; // save aux vector start
  stage2_input.sysinfo = 0;
  stage2_input.program_phdr = 0;
  stage2_input.program_phnum = 0;
  auxvt_tmp = auxvt;
  while (auxvt_tmp->a_type != AT_NULL)
    {
      if (auxvt_tmp->a_type == AT_PHDR)
	{
	  stage2_input.program_phdr = (ElfW(Phdr) *)auxvt_tmp->a_un.a_val;
	}
      else if (auxvt_tmp->a_type == AT_PHNUM)
	{
	  stage2_input.program_phnum = auxvt_tmp->a_un.a_val;
	}
      else if (auxvt_tmp->a_type == AT_SYSINFO)
	{
	  stage2_input.sysinfo = auxvt_tmp->a_un.a_val;
	}
      auxvt_tmp++;
    }
  if (stage2_input.program_phdr == 0 ||
      stage2_input.program_phnum == 0)
    {
      system_exit (-3);
    }
  return stage2_input;
}

static void global_initialize (unsigned long interpreter_load_base)
{
  // after this call to vdl_alloc_initialize is completed,
  // we are allowed to allocate heap memory.
  vdl_alloc_initialize ();

  struct Vdl *vdl = &g_vdl;
  vdl->version = 1;
  vdl->link_map = 0;
  vdl->breakpoint = 0;
  vdl->state = VDL_CONSISTENT;
  vdl->interpreter_load_base = interpreter_load_base;
  vdl->bind_now = 0; // by default, do lazy binding
  vdl->finalized = 0;
  vdl->ldso = 0;
  vdl->contexts = vdl_list_new ();
  vdl->search_dirs = vdl_utils_splitpath (machine_get_system_search_dirs ());
  vdl->tls_gen = 1;
  vdl->tls_static_total_size = 0;
  vdl->tls_static_current_size = 0;
  vdl->tls_static_align = 0;
  vdl->tls_n_dtv = 0;
  vdl->tls_next_index = 1;
  vdl->futex = futex_new ();
  vdl->errors = vdl_list_new ();
  vdl->n_added = 0;
  vdl->n_removed = 0;
}



// relocate entries in DT_REL
static void
relocate_dt_rel (ElfW(Dyn) *dynamic, unsigned long load_base)
{
  ElfW(Dyn) *tmp = dynamic;
  ElfW(Rel) *dt_rel = 0;
  unsigned long dt_relsz = 0;
  unsigned long dt_relent = 0;
  ElfW(Rela) *dt_rela = 0;
  unsigned long dt_relasz = 0;
  unsigned long dt_relaent = 0;
  // search DT_REL, DT_RELSZ, DT_RELENT, DT_RELA, DT_RELASZ, DT_RELAENT
  while (tmp->d_tag != DT_NULL && 
	 (dt_rel == 0 || dt_relsz == 0 || dt_relent == 0 ||
	  dt_rela == 0 || dt_relasz == 0 || dt_relaent == 0))
    {
      //DEBUG_HEX(tmp->d_tag);
      if (tmp->d_tag == DT_REL)
	{
	  dt_rel = (ElfW(Rel) *)(load_base + tmp->d_un.d_ptr);
	}
      else if (tmp->d_tag == DT_RELSZ)
	{
	  dt_relsz = tmp->d_un.d_val;
	}
      else if (tmp->d_tag == DT_RELENT)
	{
	  dt_relent = tmp->d_un.d_val;
	}
      else if (tmp->d_tag == DT_RELA)
	{
	  dt_rela = (ElfW(Rela) *)(load_base + tmp->d_un.d_ptr);
	}
      else if (tmp->d_tag == DT_RELASZ)
	{
	  dt_relasz = tmp->d_un.d_val;
	}
      else if (tmp->d_tag == DT_RELAENT)
	{
	  dt_relaent = tmp->d_un.d_val;
	}
      tmp++;
    }
  DPRINTF ("dt_rel=0x%x, dt_relsz=%d, dt_relent=%d, dt_rela=0x%x, dt_relasz=%d, dt_relaent=%d\n", 
	   dt_rel, dt_relsz, dt_relent, dt_rela, dt_relasz, dt_relaent);

  // relocate entries in dt_rel and dt_rela. 
  // since we are relocating the dynamic loader itself here,
  // the entries will always be of type R_XXX_RELATIVE.
  uint32_t i;
  for (i = 0; i < dt_relsz; i+=dt_relent)
    {
      ElfW(Rel) *tmp = (ElfW(Rel)*)(((uint8_t*)dt_rel) + i);
      ElfW(Addr) *reloc_addr = (void *)(load_base + tmp->r_offset);
      if (!machine_reloc_is_relative (ELFW_R_TYPE (tmp->r_info)))
	{
	  DPRINTF ("Invalid reloc entry type: %s\n",
		   machine_reloc_type_to_str (ELFW_R_TYPE (tmp->r_info)));
	  goto error;
	}
      *reloc_addr += (ElfW(Addr))load_base;
    }
  for (i = 0; i < dt_relasz; i+=dt_relaent)
    {
      ElfW(Rela) *tmp = (ElfW(Rela)*)(((uint8_t*)dt_rela) + i);
      ElfW(Addr) *reloc_addr = (void *)(load_base + tmp->r_offset);
      if (!machine_reloc_is_relative (ELFW_R_TYPE (tmp->r_info)))
	{
	  DPRINTF ("Invalid reloc entry type: %s\n",
		   machine_reloc_type_to_str (ELFW_R_TYPE (tmp->r_info)));
	  goto error;
	}
      *reloc_addr = (ElfW(Addr))load_base + tmp->r_addend;
    }

  // Note that, technically, we could also relocate DT_JMPREL entries but
  // this would be fairly complex so, it's easier to just make sure that
  // our generated ldso binary does not contain any.
error:
  return;
}

void stage1_finalize (void)
{
  stage2_finalize ();
  g_vdl.finalized = 1;
}

void stage1_freeres (void)
{
  // We are called by libc_freeres_interceptor which is our wrapper
  // around __libc_freeres.
  // If stage1_freeres is called while g_vdl.finalized is set to 1, it
  // means that stage1_finalize has already run which means that
  // we are called by valgrind. If this is so, we do perform final shutdown
  // of everything here. We are allowed to do so because 
  // this function will return to vgpreload_core and the process
  // will terminate immediately after.
  if (!g_vdl.finalized)
    {
      return;
    }
  stage2_freeres ();
  vdl_utils_str_list_delete (g_vdl.search_dirs);
  vdl_list_delete (g_vdl.contexts);
  futex_delete (g_vdl.futex);
  {
    void **i;
    for (i = vdl_list_begin (g_vdl.errors);
	 i != vdl_list_end (g_vdl.errors);
	 i = vdl_list_next (i))
      {
	struct VdlError *error = *i;
	vdl_alloc_free (error->prev_error);
	vdl_alloc_free (error->error);
	vdl_alloc_delete (error);
      }
    vdl_list_delete (g_vdl.errors);
  }

  // After this call, we can't do any malloc/free anymore.
  vdl_alloc_destroy ();

  g_vdl.search_dirs = 0;
  g_vdl.contexts = 0;
  g_vdl.futex = 0;
  g_vdl.errors = 0;
}

// Called from stage0 entry point asm code.
void stage1 (struct Stage1InputOutput *input_output)
{
  // The linker defines the symbol _DYNAMIC to give you the offset from
  // the load base to the start of the PT_DYNAMIC area which has been 
  // mapped by the OS loader as part of the rw PT_LOAD segment.
  void *dynamic = _DYNAMIC;
  dynamic += input_output->load_base;
  relocate_dt_rel (dynamic, input_output->load_base);

  // Now that access to global variables is possible, we initialize 
  // our main global variable. After this function call completes,
  // we are allowed to do memory allocations.
  global_initialize (input_output->load_base);

  struct Stage2Input stage2_input = prepare_stage2 (input_output->entry_point_struct);
  stage2_input.interpreter_load_base = input_output->load_base;

  // Now that we have relocated this binary, we can access global variables
  // so, we switch to stage2 to complete the loader initialization.
  struct Stage2Output stage2_output = stage2_initialize (stage2_input);

  // We are all done, so we update the caller's data structure to be able
  // jump in the program's entry point.
  input_output->entry_point = stage2_output.entry_point;
  input_output->dl_fini = 0;
  input_output->dl_fini = (unsigned long) stage1_finalize;
}