#include <sys/cdefs.h>
// #define extern __hidden__

#define PIEDEBUG

extern char __executable_start;

/* Warning: this code sets up the machine registers and segments needed
 * for -fstack-protector to work.  If you compile this function with
 * -fstack-protector, it will reference those registers before they are
 * set up properly, causing a segmentation fault.  Ubuntu adds
 * -fstack-protector to their gcc default options, so this breaks.  I
 * added a workaround to the Makefile to make sure this code is always
 * compiled with -fno-stack-protector for this reason.  Or, as a friend
 * put it: yo dawg. I herd u liek stack protektion. :-)
 */
#include <unistd.h>
#include <fcntl.h>
#include <alloca.h>
#include <sys/time.h>
#include <sys/tls.h>
#include <endian.h>
#include <elf.h>
#include <stdlib.h>
#include "dietfeatures.h"
#include <sys/auxv.h>
#include <string.h>

#ifdef IN_LDSO
#undef WANT_TLS
#undef WANT_SSP
#endif

#ifdef WANT_GNU_STARTUP_BLOAT
char* program_invocation_name;
char* program_invocation_short_name;
#endif

void* __vdso;

extern int main(int argc,char* argv[],char* envp[]);

#if defined(WANT_SSP)
extern __hidden__ unsigned long __guard;
#endif

#if defined(WANT_VALGRIND_SUPPORT)
int __valgrind=1;
#endif

#ifdef __i386__
int __modern_linux;
#endif

#ifdef WANT_TLS
/* __tdatasize is the size of the initialized thread local data section
 * __tmemsize is the size of the complete thread local data section
 *   (including uninitialized data)
 * __tdataptr is a pointer to the initialized thread local data section
 * __tmemsize is already rounded up to meet alignment
 * the final memory layout is [tdata] [tbss (zero)] [tcb] */
size_t __tdatasize, __tmemsize;
void* __tdataptr;

static void findtlsdata(long* auxvec) {
#if (__WORDSIZE == 64)
  Elf64_Phdr* x=0;
#else
  Elf32_Phdr* x=0;
#endif
  size_t i,n=0;
  while (*auxvec) {
    if (auxvec[0]==3) {	/* AT_PHDR */
      x=(void*)auxvec[1];
      if (n) break;
    } else if (auxvec[0]==5) { /* AT_PHNUM */
      n=auxvec[1];
      if (x) break;
    }
    auxvec+=2;
  } /* if we don't find the entry, the kernel let us down */
  if (!x || !n) return;	/* a kernel this old does not support thread local storage anyway */
  for (i=0; i<n; ++i)
    if (x[i].p_type==PT_TLS) {
      __tdataptr=(void*)x[i].p_vaddr;
      __tdatasize=x[i].p_filesz;
      __tmemsize=x[i].p_memsz;
      if (__tmemsize&15)
	__tmemsize=(__tmemsize+15)&~15;
      break;
    }
  /* if there is no PT_TLS section, there is no thread-local data, and
   * we just leave the __t* variables zero */
}
#endif

#if defined(WANT_SSP) || defined(WANT_TLS)
tcbhead_t* __tcb_mainthread;

void __setup_tls(tcbhead_t*);

void __setup_tls(tcbhead_t* mainthread) {
  mainthread->tcb=mainthread;
  mainthread->dtv=0;
  mainthread->self=0;
  mainthread->multiple_threads=0;
#if defined(WANT_SSP)
  mainthread->stack_guard=__guard;
#endif

#if defined(__x86_64__)

  arch_prctl(ARCH_SET_FS, mainthread);

#elif defined(__i386__)

  unsigned int sd[4];
  sd[0]=-1;
  sd[1]=(unsigned long int)mainthread;
  sd[2]=0xfffff; /* 4 GB limit */
  sd[3]=0x51; /* bitfield, see struct user_desc in asm-i386/ldt.h */
  if (__modern_linux>=0) {
    if (set_thread_area((struct user_desc*)(void*)&sd)==0) {
      asm volatile ("movw %w0, %%gs" :: "q" (sd[0]*8+3));
      __modern_linux=1;
    } else
      __modern_linux=-1;
  }

#elif defined(__alpha__) || defined(__s390__)
  __builtin_set_thread_pointer(mainthread);
#elif defined(__mips__)
  set_thread_area((char*)(void *)mainthread);
#elif defined(__aarch64__)
  asm volatile ("msr tpidr_el0, %0" :: "r"(mainthread));
#elif defined(__arm__)
  extern void __arm_set_tls(void*);
  __arm_set_tls(mainthread);
#elif defined(__hppa__)
  /* control register 27 is used as thread pointer on PA-RISC Linux,
   * but it can only be set from Ring0. The Linux kernel provides
   * privileged code to set this register, so call that. (cf. syscalls,
   * which branch to 0x100(%%sr2, %%r0), instead.) PA-RISC has
   * 1-instruction delayed branching. The register may be read by any
   * code however (using mfctl %cr27, %rXX). r26 is used as input for
   * the kernel code, r31 is the return address pointer set by the
   * branch instruction, so clobber both. */
  asm volatile ("ble 0xe0(%%sr2, %%r0)\n\t"
                "copy %0, %%r26"
                :: "r" (mainthread) : "r26", "r31");
#elif defined(__ABI_TLS_REGISTER)
  register tcbhead_t* __thread_self __asm__(__ABI_TLS_REGISTER);
  __thread_self=mainthread;
  __asm__ __volatile__("" : : "r"(__thread_self) : "memory");
#else
#warning "no idea how to enable TLS on this platform, edit lib/stackgap.c"
#endif
}
#endif

static void* find_in_auxvec(long* x,long what) {
  while (*x) {
    if (*x==what)
      return (void*)x[1];
    x+=2;
  }
  return NULL;
}

static long* _auxvec;

unsigned long getauxval(unsigned long type) {
  return (long)find_in_auxvec(_auxvec,type);
}



#ifdef __PIE__

#if (__WORDSIZE == 64)

#define phdr Elf64_Phdr
#define ehdr Elf64_Ehdr
#define shdr Elf64_Shdr
#define sym Elf64_Sym
#define dyn Elf64_Dyn
#define rela Elf64_Rela
#define R_SYM ELF64_R_SYM
#define R_TYPE ELF64_R_TYPE

#else

#define phdr Elf32_Phdr
#define ehdr Elf32_Ehdr
#define shdr Elf32_Shdr
#define sym Elf32_Sym
#define dyn Elf32_Dyn
#define rela Elf32_Rela
#define R_SYM ELF32_R_SYM
#define R_TYPE ELF32_R_TYPE

#endif

#ifdef PIEDEBUG
static size_t Strlen(const char* s) {
  size_t i;
  for (i=0; s[i]; ++i) ;
  return i;
}

static void Write(int fd,const char* s,size_t len) {
  asm("syscall\n" : : "a" (1), "D" (fd), "S" (s), "d" (len));
}

static void _puts(const char* s) {
  Write(1,s,Strlen(s));
}

static void _putn(size_t n) {
  char buf[100];
  size_t j=sizeof(buf);
  if (!n) buf[--j]='0';
  for (; j>0 && n; ) {
    buf[--j]=(n%10)+'0';
    n/=10;
  }
  Write(1,buf+j,sizeof(buf)-j);
}

static void _putx(size_t n) {
  char buf[100];
  size_t j=sizeof(buf);
  if (!n) buf[--j]='0';
  for (; j>0 && n; ) {
    char x = n%16;
    buf[--j]=x<10 ? x+'0' : x+'a'-10;
    n/=16;
  }
  Write(1,buf+j,sizeof(buf)-j);
}

struct lookup { size_t n; const char s[20]; };

static void _putl(size_t n,struct lookup* l) {
  size_t i;
  for (i=0; l[i].n!=0; ++i)
    if (l[i].n==n) {
      _puts(l[i].s);
      return;
    }
  Write(1,"[",1);
  _putx(n);
  Write(1,"]",1);
}

static void dump_maps() {
  int fd=open("/proc/self/maps",O_RDONLY);
  char buf[2048];
  if (fd!=-1) {
    int r=read(fd,buf,sizeof(buf));
    if (r>0) write(1,buf,r);
    close(fd);
  }
}

static struct lookup atype[] = {
  { 1, "AT_IGNORE" },
  { 2, "AT_EXECFD" },
  { 3, "AT_PHDR" },
  { 4, "AT_PHENT" },
  { 5, "AT_PHNUM" },
  { 6, "AT_PAGESZ" },
  { 7, "AT_BASE" },
  { 8, "AT_FLAGS" },
  { 9, "AT_ENTRY" },
  { 10, "AT_NOTELF" },
  { 11, "AT_UID" },
  { 12, "AT_EUID" },
  { 13, "AT_GID" },
  { 14, "AT_EGID" },
  { 17, "AT_CLKTCK" },
  { 15, "AT_PLATFORM" },
  { 16, "AT_HWCAP" },
  { 18, "AT_FPUCW" },
  { 19, "AT_DCACHEBSIZE" },
  { 20, "AT_ICACHEBSIZE" },
  { 21, "AT_UCACHEBSIZE" },
  { 22, "AT_IGNOREPPC" },
  { 23, "AT_SECURE" },
  { 24, "AT_BASE_PLATFORM" },
  { 25, "AT_RANDOM" },
  { 26, "AT_HWCAP2" },
  { 31, "AT_EXECFN" },
  { 32, "AT_SYSINFO" },
  { 33, "AT_SYSINFO_EHDR" },
  { 0 },
};

static struct lookup ptype[] = {
  { 1, "PT_LOAD" },
  { 2, "PT_DYNAMIC" },
  { 3, "PT_INTERP" },
  { 4, "PT_NOTE" },
  { 5, "PT_SHLIB" },
  { 6, "PT_PHDR" },
  { 7, "PT_TLS" },
  { 8, "PT_NUM" },
  { 0x60000000, "PT_LOOS" },
  { 0x6474e550, "PT_GNU_EH_FRAME" },
  { 0x6474e551, "PT_GNU_STACK" },
  { 0x6ffffffa, "PT_LOSUNW" },
  { 0x6ffffffa, "PT_SUNWBSS" },
  { 0x6ffffffb, "PT_SUNWSTACK" },
  { 0x6fffffff, "PT_HISUNW" },
  { 0x6fffffff, "PT_HIOS" },
  { 0x70000000, "PT_LOPROC" },
  { 0x7fffffff, "PT_HIPROC" },
  { 0 } };

static struct lookup dtag[] = {
  { 1, "DT_NEEDED" },
  { 2, "DT_PLTRELSZ" },
  { 3, "DT_PLTGOT" },
  { 4, "DT_HASH" },
  { 5, "DT_STRTAB" },
  { 6, "DT_SYMTAB" },
  { 7, "DT_RELA" },
  { 8, "DT_RELASZ" },
  { 9, "DT_RELAENT" },
  { 10, "DT_STRSZ" },
  { 11, "DT_SYMENT" },
  { 12, "DT_INIT" },
  { 13, "DT_FINI" },
  { 14, "DT_SONAME" },
  { 15, "DT_RPATH" },
  { 16, "DT_SYMBOLIC" },
  { 17, "DT_REL" },
  { 18, "DT_RELSZ" },
  { 19, "DT_RELENT" },
  { 20, "DT_PLTREL" },
  { 21, "DT_DEBUG" },
  { 22, "DT_TEXTREL" },
  { 23, "DT_JMPREL" },
  { 24, "DT_BIND_NOW" },
  { 25, "DT_INIT_ARRAY" },
  { 26, "DT_FINI_ARRAY" },
  { 27, "DT_INIT_ARRAYSZ" },
  { 28, "DT_FINI_ARRAYSZ" },
  { 29, "DT_RUNPATH" },
  { 30, "DT_FLAGS" },
  { 32, "DT_ENCODING" },
  { 32, "DT_PREINIT_ARRAY" },
  { 33, "DT_PREINIT_ARRAYSZ" },
  { 34, "DT_NUM" },
  { 0x6000000d, "DT_LOOS" },
  { 0x6ffff000, "DT_HIOS" },
  { 0x6ffffff9, "DT_RELACOUNT" },
  { 0x6ffffffa, "DT_RELCOUNT" },
  { 0x70000000, "DT_LOPROC" },
  { 0x7fffffff, "DT_HIPROC" },
  { 0 }
};
#endif

static void callback() {
#ifdef PIEDEBUG
  _puts("callback called!\n");
#endif
  _exit(111);
}

#endif	/* __PIE__ */

// extern size_t _GLOBAL_OFFSET_TABLE_;






__hidden__ char _DYNAMIC;

int stackgap(int argc,char* argv[],char* envp[]);
int stackgap(int argc,char* argv[],char* envp[]) {
  long* auxvec=(long*)envp;
#if defined(WANT_STACKGAP) || defined(WANT_SSP) || defined(WANT_TLS)
  char* rand;
  char* tlsdata;
#endif
  while (*auxvec) ++auxvec;			/* skip envp to get to auxvec */
  ++auxvec;

  _auxvec=auxvec;
#ifdef WANT_STACKGAP
  unsigned short s;
#endif
#ifdef __PIE__
  /* This code is meant for static PIE executables. */
  /* Since we do not have an interpreter, we'll have to set up the GOT ourselves. */
  {
    size_t i;
    const char* base=0;
    rela* r=0;
    size_t relasz=0,relaent=0;
    size_t phent=0, phnum=0;
    size_t saddr=0;
    const char* elfimage=0;
    for (i=0; auxvec[i]; i+=2) {
      switch (auxvec[i]) {
      case AT_PHDR: elfimage=(const char*)auxvec[i+1]; break;
      case AT_PHENT: phent=auxvec[i+1]; break;
      case AT_PHNUM: phnum=auxvec[i+1]; break;
      case AT_ENTRY: saddr=auxvec[i+1]; break;
      }
    }
    /* All four must be there, or we are hosed. Not checking. */

#ifdef PIEDEBUG
    dump_maps();
#endif

#ifdef PIEDEBUG
    _puts("auxvec:\n");
    for (i=0; auxvec[i]; i+=2) {
      _putl(auxvec[i],atype);
      _puts(" - ");
      _putx(auxvec[i+1]);
      _puts("\n");
    }
    _puts("\nme=");
    _putx((size_t)&stackgap);
    _puts("\n");
#endif

    if (elfimage) {

      size_t* pltgot=0;
      size_t pltrelsz=0,pltrel=0;

#ifdef PIEDEBUG
      _puts("\nphdr:\n");
      for (i=0; i<phnum; ++i) {
	const phdr* ph=(const phdr*)(elfimage + i*phent);
	_puts("type ");
	_putl(ph->p_type,ptype); _puts("\tofs ");
	_putx(ph->p_offset); _puts("\tvaddr ");
	_putx(ph->p_vaddr); _puts("\tpaddr ");
	_putx(ph->p_paddr); _puts("\tfilesz ");
	_putx(ph->p_filesz); _puts("\tmemsz ");
	_putx(ph->p_memsz); _puts("\tflags ");
	_putx(ph->p_flags); _puts("\talign ");
	_putx(ph->p_align); _puts("\n");
      }
      _puts("\n\n");
#endif

      for (i=0; i<phnum; ++i) {
	const phdr* ph=(const phdr*)(elfimage + i*phent);
	if (ph->p_type==PT_DYNAMIC) {	/* found .dynamic */
	  const dyn* dh;
	  size_t j;
	  
	  base = &_DYNAMIC - ph->p_vaddr;
#ifdef PIEDEBUG
	  if (((ehdr*)base)->e_entry + (uintptr_t)base != saddr) {
	    _puts("fail: ");
	    _putx(((ehdr*)base)->e_entry + (uintptr_t)base);
	    _puts(" != ");
	    _putx(saddr);
	    _puts("!\n\n");
	  }
#endif

	  for (j=0; j<ph->p_memsz; j+=sizeof(dyn)) {
	    dh=(const dyn*)(base + ph->p_vaddr + j);
	    if (!dh->d_tag) break;
#ifdef PIEDEBUG
	    _putl(dh->d_tag,dtag); _puts(" - "); _putx(dh->d_un.d_val); _puts("\n");
#endif
	    switch (dh->d_tag) {
	    case DT_PLTGOT: pltgot=(size_t*)(base+dh->d_un.d_ptr); break;
	    case DT_PLTRELSZ: pltrelsz=dh->d_un.d_val; break;
	    case DT_PLTREL: pltrel=dh->d_un.d_val; break;

	    case DT_RELA: r=(rela*)(base+dh->d_un.d_ptr); break;
	    case DT_RELASZ: relasz=dh->d_un.d_val; break;
	    case DT_RELAENT: relaent=dh->d_un.d_val; break;
	    }
	  }
	  break;
	}
      }

#ifdef PIEDEBUG
      _puts("got pltgot "); _putx((uintptr_t)pltgot); _puts(", pltrelsz "); _putn(pltrelsz); _puts(", pltrel "); _putn(pltrel); _puts("\n");
#endif
      if (r && relasz && relaent) {
#ifdef PIEDEBUG
	_puts("got rela ");
	_putx((uintptr_t)r);
	_puts(", relasz ");
	_putx(relasz);
	_puts(", relaent ");
	_putx(relaent);
	_puts("\n");

	_puts("\nrela dump:\n");
#endif

	for (i=relasz/relaent; i; --i) {

#ifdef PIEDEBUG
	  _puts("@ "); _putx((uintptr_t)base + r->r_offset);
	  _puts(", ofs "); _putx(r->r_offset);
	  _puts(", info "); _putx(r->r_info);
	  _puts(", addend "); _putx(r->r_addend);
	  _puts("\n");
#endif

#if __WORDSIZE == 64
	  unsigned int type = ELF64_R_TYPE(r->r_info);
#else
	  unsigned int type = ELF32_R_TYPE(r->r_info);
#endif
	  size_t* x=(size_t*)((char*)(base+r->r_offset));
	  switch (type) {
#if defined(__x86_64__)
	  case R_X86_64_RELATIVE: *x = (uintptr_t)base + r->r_addend; break;
#else
#error architecture not supported yet (edit stackgap.c)
#endif
#ifdef PIEDEBUG
	  default: _puts("unsupported relocation type "); _putx(type); break;
#endif
	  }
	  ++r;
	}
	if (pltgot)
	  pltgot[2]=(uintptr_t)callback;
      }
    }
  }
#endif

#if defined(WANT_STACKGAP) || defined(WANT_SSP) || defined(WANT_TLS)
#if defined(WANT_STACKGAP) || defined(WANT_SSP)
  volatile char* gap;
  rand=find_in_auxvec(auxvec,25);
  if (!rand) {
    rand=alloca(10);
    int fd=open("/dev/urandom",O_RDONLY);	// If this fails, there is not much we can do. Limp on.
    read(fd,rand,10);
    close(fd);
  }
#endif
#ifdef WANT_STACKGAP
#ifdef __UNALIGNED_MEMORY_ACCESS_OK
  s=*(unsigned short*)(rand+8);
#else
  s=rand[8]+(rand[9]<<8);
#endif	// __UNALIGNED_MEMORY_ACCESS_OK
#endif	// WANT_STACKGAP
#ifdef WANT_SSP
#ifdef __UNALIGNED_MEMORY_ACCESS_OK
  __guard=*(unsigned long*)rand;
#else
  memcpy(&__guard,rand,sizeof(__guard));
#endif	// __UNALIGNED_MEMORY_ACCESS_OK
#endif	// WANT_SSP
#ifdef WANT_STACKGAP
  gap=alloca(s);
#endif
#endif

  __vdso=find_in_auxvec(auxvec,33);	// AT_SYSINFO_EHDR -> vdso start address
#ifdef __x86_64__
  if (!__vdso) __vdso=(char*)0xffffffffff600000;
#endif

#ifdef WANT_TLS
  findtlsdata(auxvec);
  if (__unlikely(__tmemsize+sizeof(tcbhead_t)<sizeof(tcbhead_t)) ||
      __unlikely(__tmemsize>512*1024*1024) ||
      __unlikely(__tmemsize<__tdatasize))
    return 111;
  /* Ugly temporary workaround until a proper fix is available:

     The TCB is allocated on the stack in stackgap (called by the
     startup code). However, exit() will be called by the startup
     code after the end of main, and some atexit functions might
     still want to access the current TCB.

     The workaround for now is to allocate 512 bytes more so that
     a small stack within atexit() functions will not overwrite the
     TCB. This is not a permanent solution, and we should probably
     mmap() the real TCB - but for that we need to have an mmap
     implementation that doesn't set errno, because that might be a
     TLS variable.
   */
  tlsdata=alloca(__tmemsize+sizeof(tcbhead_t)+512);
  memcpy(tlsdata,__tdataptr,__tdatasize);
  memset(tlsdata+__tdatasize,0,__tmemsize-__tdatasize);
  __setup_tls(__tcb_mainthread=(tcbhead_t*)(tlsdata+__tmemsize));
  __tcb_mainthread->sysinfo=(uintptr_t)find_in_auxvec(auxvec,32);
#elif defined(WANT_SSP)
  tlsdata=alloca(sizeof(tcbhead_t));
  __setup_tls(__tcb_mainthread=(tcbhead_t*)(tlsdata));
#endif
#if defined(WANT_VALGRIND_SUPPORT)
  {
    const char* v=getenv("LD_PRELOAD");
    __valgrind=(v && strstr(v,"valgrind"));
  }
#endif
#ifdef WANT_GNU_STARTUP_BLOAT
  program_invocation_name=argv[0];
  {
    char* c;
    for (c=program_invocation_short_name=program_invocation_name; *c; ++c)
      if (*c=='/') program_invocation_short_name=c+1;
  }
#endif
  return main(argc,argv,envp);
}