#include "google/coredumper.h"

#include <elf.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/sysctl.h>
#include <sys/time.h>
#include <sys/types.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

# pragma weak __libc_stack_end
extern void* __libc_stack_end;
extern void *__bss_start;
extern void *end;

#ifdef __x86_64__
  #define ELF_CLASS ELFCLASS64
  #define ELF_ARCH  EM_X86_64
  #define Ehdr      Elf64_Ehdr
  #define Phdr      Elf64_Phdr
  #define Shdr      Elf64_Shdr
  #define Nhdr      Elf64_Nhdr
#else
  #define ELF_CLASS ELFCLASS32
  #define ELF_ARCH  EM_386
  #define Ehdr      Elf32_Ehdr
  #define Phdr      Elf32_Phdr
  #define Shdr      Elf32_Shdr
  #define Nhdr      Elf32_Nhdr
#endif

typedef struct i386_regs {      /* Normal (non-FPU) CPU registers            */
#ifdef __x86_64__
  #define BP rbp
  #define SP rsp
  #define IP rip
  uint64_t  r15,r14,r13,r12,rbp,rbx,r11,r10;
  uint64_t  r9,r8,rax,rcx,rdx,rsi,rdi,orig_rax;
  uint64_t  rip,cs,eflags;
  uint64_t  rsp,ss;
  uint64_t  fs_base, gs_base;
  uint64_t  ds,es,fs,gs;
#else
  #define BP ebp
  #define SP esp
  #define IP eip
  uint32_t  ebx, ecx, edx, esi, edi, ebp, eax;
  uint16_t  ds, __ds, es, __es;
  uint16_t  fs, __fs, gs, __gs;
  uint32_t  orig_eax, eip;
  uint16_t  cs, __cs;
  uint32_t  eflags, esp;
  uint16_t  ss, __ss;
#endif
} i386_regs;

void *get_stack_begin() {
    int lb[2];
    char b;
    void *p = &b;
    pipe(lb);
    while (1) {
        if (write(lb[1], p, 1) < 0 || read(lb[0], &b, 1) < 0) break;
        p++;
    }
    errno = 0;
    close(lb[0]);
    close(lb[1]);
    return p;
}

void restore(const char* core) {
    int i;
    Ehdr ehdr;
    Phdr phdr;
    Nhdr nhdr;
    int num_mappings;
    int fd = open(core, O_RDONLY);
    int *mappings;
    i386_regs regs;
    void* regsp;
    int mem_offset = 0;
    char *stack;
    int stacksize;
    char *stackaddr;

    read(fd, &ehdr, sizeof(ehdr));

    read(fd, &phdr, sizeof(phdr));
    num_mappings = (phdr.p_offset - sizeof(Ehdr)) / sizeof(Phdr) - 1;
    mappings = (int*)alloca(2*num_mappings * 4);

    for (i = 0; i < num_mappings; i++) {
        int start, size;
        read(fd, &phdr, sizeof(phdr));
        start = phdr.p_vaddr;
        size = phdr.p_filesz;
        mem_offset += size;
        mappings[i*2] = start;
        mappings[i*2+1] = size;
    }

    while (1) {
        read(fd, &nhdr, sizeof(nhdr));
        lseek(fd, 4, SEEK_CUR);
        if (nhdr.n_type == NT_PRXREG) {
            read(fd, &regs, sizeof(regs));
            break;
        }
        else {
            lseek(fd, nhdr.n_descsz, SEEK_CUR);
        }
    }

    {
        char* orig;
        int gap;
        int zerocnt = 0;
        void **stp, **tmpst;

        stacksize = mappings[num_mappings*2-1];
        stackaddr = get_stack_begin();

        orig = (char*)mappings[num_mappings*2-2];
        orig += stacksize;

        tmpst = (void**)malloc(stacksize);
        lseek(fd, -stacksize, SEEK_END);
        read(fd, tmpst, stacksize);
        for (stp = tmpst; !*stp; stp++) zerocnt++;
        stacksize -= sizeof(void*)*zerocnt;
        stack = (char*)alloca(stacksize);
        memcpy(stack, stp, stacksize);
        free(tmpst);
        stack += stacksize;

        gap = stackaddr - orig;
        if (gap) {
            void **ebp;
            int c = stack - stackaddr;
            regs.esp += gap;
            regs.ebp += gap;

            ebp = (void**)(regs.ebp+c);
            while (*ebp) {
                char **p = (char**)ebp;
                *p += gap;
                ebp = (void**)((char*)*p+c);
            }
        }
        regs.eip += 16;
    }

    regsp = alloca(sizeof(regs));
    memcpy(regsp, &regs, sizeof(regs));

    /* is it correct? */
    lseek(fd, -mem_offset, SEEK_END);

    for (i = 0; i < num_mappings-1; i++) {
        void *start = (void*)mappings[i*2];
        int size = mappings[i*2+1];
        if (size == 0) continue;
        if (start < sbrk(0)) {
            read(fd, (void*)start, size);
        }
        else {
            lseek(fd, size, SEEK_CUR);
        }
    }

    close(fd);

    /* let's copy the stack! */
    __asm__("  movl %0, %%eax         \n\t"
            "  movl %1, %%ebx         \n\t"
            "  movl %2, %%ecx         \n\t"
            "  movl %3, %%esi         \n\t"
            "loop:                    \n\t"
            "  subl $4, %%ebx         \n\t"
            "  subl $4, %%ecx         \n\t"
            "  movl (%%ecx), %%edx    \n\t"
            "  movl %%edx, (%%ebx)    \n\t"
            "  subl $4, %%eax         \n\t"
            "  cmpl $0, %%eax         \n\t"
            "  jne loop               \n\t"
            "  movl %%esi, %%eax      \n\t"
            "  movl (%%eax), %%ebx    \n\t"
            "  movl 4(%%eax), %%ecx   \n\t"
            "  movl 8(%%eax), %%edx   \n\t"
            "  movl 12(%%eax), %%esi  \n\t"
            "  movl 16(%%eax), %%edi  \n\t"
            "  movl 20(%%eax), %%ebp  \n\t"
            "  movl 60(%%eax), %%esp  \n\t"
//            "  int $3                 \n\t"
//            "  movl 24(%%eax), %%eax  \n\t"
            "  jmpl *48(%%eax)        \n\t"
            : : "m"(stacksize), "m"(stackaddr), "m"(stack), "m"(regsp));
}

int sum;

void run() {
    int i;
    int *fact = malloc(sizeof(int));
    sum = 0;
    *fact = 1;

    for (i = 0; i < 3; i++) {
        char buf[256];
        sprintf(buf, "core%d", i);
        WriteCoreDump(buf);

        sum += i;
        *fact *= i+1;
        printf("%d %d %d\n", i, sum, *fact);
    }
}

int main(int argc, char* argv[]) {
    if (argc == 2) {
        restore(argv[1]);
    }

    run();

    return 0;
}