#include <errno.h>
#include <inttypes.h>
#include <setjmp.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <asm/prctl.h>
#include <fcntl.h>
#include <linux/elf.h>
#include <linux/unistd.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <unistd.h>

#include "selfie.h"

#define PAGESZ 0x1000

enum {
	Fread  = 1,
	Fwrite = 2,
	Fexec  = 4,
	Fpriv  = 8,
	Fvdso  = 16,
	Fstk   = 32,
};

typedef struct {
	uint64_t beg;
	uint64_t end;
	int flags;
} Line;

typedef struct {
	int num;
	struct sigaction act;
} Signal;

static char zero[PAGESZ];
static char maps[4*PAGESZ];
static size_t mapssz;
static uint64_t fsreg;

/* variables used to persist data in
 * the selfie image */
static void *stkptr;
Elf64_Phdr stk;
static jmp_buf env;
static Signal sig[] = {
	{SIGABRT}, {SIGALRM}, {SIGFPE}, {SIGHUP}, {SIGILL}, {SIGINT},
	{SIGKILL}, {SIGPIPE}, {SIGQUIT}, {SIGSEGV}, {SIGTERM}, {SIGUSR1},
	{SIGUSR2}, {SIGCHLD}, {SIGCONT}, {SIGSTOP}, {SIGTSTP}, {SIGTTIN},
	{SIGTTOU}, {SIGVTALRM}, {SIGPROF}, {SIGBUS}, {SIGPOLL}, {SIGSYS},
	{SIGTRAP}, {SIGURG}, {SIGXCPU}, {SIGXFSZ}, {-1}
};

static int
readmaps()
{
	int fd, rd, n;

	fd = open("/proc/self/maps", O_RDONLY);
	if (fd < 0)
		return -1;
	rd = 0;
	do {
		n = read(fd, maps+rd, sizeof maps - rd);
		if (n < 0) {
			if (errno == EINTR)
				continue;
			close(fd);
			return -1;
		}
		rd += n;
	} while (n);
	close(fd);
	if (rd == sizeof maps) {
		errno = ENOMEM;
		return -1;
	}
	mapssz = rd;
	return 0;
}


static void
mapsit(void (*f)(void *, Line), void *data)
{
	char *ptr, *end;
	uint64_t stk;
	int n, s;
	Line l;

	stk = (uint64_t)&stk;
	end = &maps[mapssz];
	ptr = maps;
	while (ptr < end) {
		s = sscanf(ptr, "%"SCNx64"-%"SCNx64" %n", &l.beg, &l.end, &n);
		if (s != 2)
			return;
		ptr += n;
		if (ptr > end-4)
			return;
		l.flags = 0;
		if (*ptr++ != '-')
			l.flags |= Fread;
		if (*ptr++ != '-')
			l.flags |= Fwrite;
		if (*ptr++ != '-')
			l.flags |= Fexec;
		if (*ptr++ != 's')
			l.flags |= Fpriv;
		if (stk - l.beg < l.end - l.beg)
			l.flags |= Fstk;
		ptr = strchr(ptr, '\n');
		if (!ptr)
			return;
		if (ptr >= maps+6)
		if (strncmp(ptr-6, "[vvar]", 6) == 0
		|| strncmp(ptr-6, "[vdso]", 6) == 0)
			l.flags |= Fvdso;
		f(data, l);
		ptr++;
	}
}

static int
noread(int flags)
{
	return !(flags & Fread) || (flags & Fvdso);
}

static void
countmaps(void *pnmaps, Line l)
{
	if (noread(l.flags))
		return;
	++*(int *)pnmaps;
}

static int
dump(int fd, void *data, size_t n)
{
	ssize_t w;
	char *ptr;

	ptr = data;
	while (n != 0) {
		w = write(fd, ptr, n);
		if (w < 0) {
			if (errno == EAGAIN)
				continue;
			return -1;
		}
		n -= w;
		ptr += w;
	}
	return 0;
}

typedef struct {
	int fd;
	int failed;
	size_t offset;
	size_t stksz;
} State;

static void
dumpphdr(void *data, Line l)
{
	Elf64_Phdr phdr;
	size_t size, addr;
	State *s;

	s = data;
	if (s->failed || noread(l.flags))
		return;

	phdr = (Elf64_Phdr){0};
	phdr.p_type = PT_LOAD;
	phdr.p_flags = PF_R;
	if (l.flags & Fwrite)
		phdr.p_flags |= PF_W;
	if (l.flags & Fexec)
		phdr.p_flags |= PF_X;
	s->offset = (s->offset + PAGESZ-1) & -PAGESZ;
	phdr.p_offset = s->offset;
	if ((l.flags & Fstk) && s->stksz) {
		size = s->stksz;
		addr = l.end - size;
	} else {
		size = l.end - l.beg;
		addr = l.beg;
	}
	phdr.p_vaddr = addr;
	phdr.p_paddr = addr;
	phdr.p_filesz = size;
	phdr.p_memsz = size;
	s->offset += size;
	phdr.p_align = PAGESZ;
	if (l.flags & Fstk)
		stk = phdr;
	if (dump(s->fd, &phdr, sizeof phdr) < 0)
		s->failed = 1;
}

static void
dumpmap(void *data, Line l)
{
	State *s;
	size_t align, n;

	s = data;
	if (s->failed || noread(l.flags))
		return;

	align = -s->offset % PAGESZ;
	if ((l.flags & Fstk) && s->stksz)
		align += s->stksz - (l.end - l.beg);
	s->offset += align + (l.end - l.beg);
	while (align > 0) {
		n = align > PAGESZ ? PAGESZ : align;
		if (dump(s->fd, zero, n) < 0) {
			s->failed = 1;
			return;
		}
		align -= n;
	}
	if (dump(s->fd, (void *)l.beg, l.end - l.beg) < 0)
		s->failed = 1;
}

static uint64_t entry(void);

static int
dumpelf(int fd, size_t stksz)
{
	Elf64_Ehdr hdr;
	State s;
	int nmaps, doff;

	s = (State){.fd=fd, .stksz=stksz};
	nmaps = 0;
	mapsit(countmaps, &nmaps);

	hdr = (Elf64_Ehdr){0};
	memcpy(hdr.e_ident, ELFMAG"\x02\x01\x01\x00", 8);
	hdr.e_type = ET_EXEC;
	hdr.e_machine = EM_X86_64;
	hdr.e_version = EV_CURRENT;
	hdr.e_entry = entry();
	hdr.e_phoff = sizeof(Elf64_Ehdr);
	hdr.e_ehsize = sizeof(Elf64_Ehdr);
	hdr.e_phentsize = sizeof(Elf64_Phdr);
	hdr.e_phnum = nmaps;
	hdr.e_shentsize = sizeof(Elf64_Shdr);

	if (dump(fd, &hdr, sizeof hdr) < 0)
		return -1;

	doff = sizeof hdr + nmaps * hdr.e_phentsize;
	s.offset = doff;
	mapsit(dumpphdr, &s);
	s.offset = doff;
	mapsit(dumpmap, &s);
	return s.failed ? -1 : 0;
}

static void
stkrestore()
{
	uint64_t sz, off;
	ssize_t n;
	int fd;
	char *ptr;

	sz = stk.p_memsz - ((uint64_t)stkptr - stk.p_vaddr);
	off = stk.p_offset + (stk.p_memsz - sz);
	fd = open("/proc/self/exe", O_RDONLY);
	if (fd < 0)
		return;
	if (lseek(fd, off, SEEK_SET) < 0)
		sz = 0;
	for (ptr = stkptr; sz > 0;) {
		n = sz > 0x4000 ? 0x4000 : sz;
		n = read(fd, ptr, n);
		if (n == -1) {
			if (errno == EAGAIN)
				continue;
			break;
		}
		sz -= n;
		ptr += n;
	}
	close(fd);
}

static int
resume()
{
	Signal *s;

	/* fs is used for dynamic loading, restore it
	 * before calling any function */
	__asm__(
		"syscall"
		:: "D"(ARCH_SET_FS), "S"(fsreg), "a"(__NR_arch_prctl)
		: "rcx", "r11", "memory");

	/* the brk pointer is used by memory allocation
	 * routines (e.g., malloc())
	 * glibc has a __curbrk cache for the current
	 * brk; after the resume it is invalid (because
	 * of the restored dynlinked libs and stack)
	 * we call brk(0) to refresh its contents */
	brk(0);

	for (s=sig; s->num>=0; s++)
		sigaction(s->num, &s->act, 0);

	/* restore the stack up to stkptr, this is in
	 * case ASLR was disabled and the kernel
	 * clobbered the base of the stack */
	stkrestore();

	/* pop back into selfie() */
	longjmp(env, 1);
}

static uint64_t
entry()
{
	uint64_t addr;

	/* call this function to get the address of
	 * the entry point code; this fine control
	 * prevents the compiler from ruining the
	 * party */
	__asm__(
		"leaq 42f(%%rip), %0\n\t"
		"jmp 51f\n"
		"42:\n\t"
		"movq stkptr(%%rip), %%rsp\n\t"
		"subq $1024, %%rsp\n\t" /* slack space */
		"jmp resume\n"
		"51:\n\t"
		: "=r"(addr)
		/* extra arguments given to prevent
		 * elision of these static objects */
		: "r"(stkptr), "r"(resume));
	return addr;
}

int
selfie(int fd)
{
	/* amazingly, arch_prctl(2) recommends declaring
	 * the prototype!... */
	extern int arch_prctl(int, uint64_t *);
	struct rlimit stklim;
	Signal *s;

	if (readmaps() < 0)
		return -1;

	/* setjmp/longjmp will take care of saving
	 * and restoring the callee-save registers */
	if (setjmp(env) == 1)
		return 1;

	for (s=sig; s->num>=0; s++)
		sigaction(s->num, 0, &s->act);
	stkptr = __builtin_frame_address(0);
	if (arch_prctl(ARCH_GET_FS, &fsreg) < 0)
		return -1;
	if (getrlimit(RLIMIT_STACK, &stklim) < 0)
		stklim.rlim_cur = 0;
	stklim.rlim_cur &= -PAGESZ;

	return dumpelf(fd, stklim.rlim_cur);
}
