#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <unistd.h>
#include <string.h>
#include <fcntl.h>
#include <errno.h>
#include <elf.h>
#include <time.h>

#include <sys/time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>

#include "zdtmtst.h"

const char *test_doc = "Check if we can use vDSO using direct vDSO calls\n";
const char *test_author = "Cyrill Gorcunov <gorcunov@openvz.org";

#ifdef __i386__

#define Ehdr_t Elf32_Ehdr
#define Sym_t  Elf32_Sym
#define Phdr_t Elf32_Phdr
#define Word_t Elf32_Word
#define Dyn_t  Elf32_Dyn

#define ELF_ST_TYPE ELF32_ST_TYPE
#define ELF_ST_BIND ELF32_ST_BIND

const char elf_ident[] = {
	0x7f, 0x45, 0x4c, 0x46, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
};

#else /* non-i386 */

#define Ehdr_t Elf64_Ehdr
#define Sym_t  Elf64_Sym
#define Phdr_t Elf64_Phdr
#define Word_t Elf64_Word
#define Dyn_t  Elf64_Dyn

#ifndef ELF_ST_TYPE
#define ELF_ST_TYPE ELF64_ST_TYPE
#endif
#ifndef ELF_ST_BIND
#define ELF_ST_BIND ELF64_ST_BIND
#endif

/*
 * See Elf specification for this magic values.
 */
const char elf_ident[] = {
	0x7f, 0x45, 0x4c, 0x46, 0x02, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
};

#endif

typedef int(__vdso_clock_gettime_t)(clockid_t clock, struct timespec *ts);
typedef long(__vdso_getcpu_t)(unsigned *cpu, unsigned *node, void *unused);
typedef int(__vdso_gettimeofday_t)(struct timeval *tv, struct timezone *tz);
typedef time_t(__vdso_time_t)(time_t *t);

#define TIME_DELTA_SEC (3)

#define BUILD_BUG_ON(condition) ((void)sizeof(char[1 - 2 * !!(condition)]))

#define VDSO_BAD_ADDR (-1ul)

struct vdso_symbol {
	char name[32];
	unsigned long offset;
};

#define VDSO_SYMBOL_INIT                 \
	{                                \
		.offset = VDSO_BAD_ADDR, \
	}

/* Check if symbol present in symtable */
static inline bool vdso_symbol_empty(struct vdso_symbol *s)
{
	return s->offset == VDSO_BAD_ADDR && s->name[0] == '\0';
}

enum {
	VDSO_SYMBOL_CLOCK_GETTIME,
	VDSO_SYMBOL_GETCPU,
	VDSO_SYMBOL_GETTIMEOFDAY,
	VDSO_SYMBOL_TIME,

	VDSO_SYMBOL_MAX
};

const char *vdso_symbols[VDSO_SYMBOL_MAX] = {
	[VDSO_SYMBOL_CLOCK_GETTIME] = "__vdso_clock_gettime",
	[VDSO_SYMBOL_GETCPU] = "__vdso_getcpu",
	[VDSO_SYMBOL_GETTIMEOFDAY] = "__vdso_gettimeofday",
	[VDSO_SYMBOL_TIME] = "__vdso_time",
};

struct vdso_symtable {
	unsigned long vma_start;
	unsigned long vma_end;
	struct vdso_symbol symbols[VDSO_SYMBOL_MAX];
};

#define VDSO_SYMTABLE_INIT                                                                  \
	{                                                                                   \
		.vma_start = VDSO_BAD_ADDR, .vma_end = VDSO_BAD_ADDR,                       \
		.symbols = {                                                                \
			[0 ... VDSO_SYMBOL_MAX - 1] = (struct vdso_symbol)VDSO_SYMBOL_INIT, \
		},                                                                          \
	}

static bool __ptr_oob(void *ptr, void *start, size_t size)
{
	void *end = (void *)((unsigned long)start + size);
	return ptr > end || ptr < start;
}

static unsigned long elf_hash(const unsigned char *name)
{
	unsigned long h = 0, g;

	while (*name) {
		h = (h << 4) + *name++;
		g = h & 0xf0000000ul;
		if (g)
			h ^= g >> 24;
		h &= ~g;
	}
	return h;
}

static int vdso_fill_symtable(char *mem, size_t size, struct vdso_symtable *t)
{
	Phdr_t *dynamic = NULL, *load = NULL;
	Ehdr_t *ehdr = (void *)mem;
	Dyn_t *dyn_strtab = NULL;
	Dyn_t *dyn_symtab = NULL;
	Dyn_t *dyn_strsz = NULL;
	Dyn_t *dyn_syment = NULL;
	Dyn_t *dyn_hash = NULL;
	Word_t *hash = NULL;
	Phdr_t *phdr;
	Dyn_t *d;

	Word_t *bucket, *chain;
	Word_t nbucket, nchain;

	char *dynsymbol_names;
	unsigned int i, j, k;

	BUILD_BUG_ON(sizeof(elf_ident) != sizeof(ehdr->e_ident));

	test_msg("Parsing at %lx %lx\n", (long)mem, (long)mem + (long)size);

	/*
	 * Make sure it's a file we support.
	 */
	if (memcmp(ehdr->e_ident, elf_ident, sizeof(elf_ident))) {
		pr_perror("Elf header magic mismatch");
		return -EINVAL;
	}

	/*
	 * We need PT_LOAD and PT_DYNAMIC here. Each once.
	 */
	phdr = (void *)&mem[ehdr->e_phoff];
	for (i = 0; i < ehdr->e_phnum; i++, phdr++) {
		if (__ptr_oob(phdr, mem, size))
			goto err_oob;
		switch (phdr->p_type) {
		case PT_DYNAMIC:
			if (dynamic) {
				pr_perror("Second PT_DYNAMIC header");
				return -EINVAL;
			}
			dynamic = phdr;
			break;
		case PT_LOAD:
			if (load) {
				pr_perror("Second PT_LOAD header");
				return -EINVAL;
			}
			load = phdr;
			break;
		}
	}

	if (!load || !dynamic) {
		pr_perror("One of obligated program headers is missed");
		return -EINVAL;
	}

	test_msg("PT_LOAD p_vaddr: %lx\n", (unsigned long)load->p_vaddr);

	/*
	 * Dynamic section tags should provide us the rest of information
	 * needed. Note that we're interested in a small set of tags.
	 */
	d = (void *)&mem[dynamic->p_offset];
	for (i = 0; i < dynamic->p_filesz / sizeof(*d); i++, d++) {
		if (__ptr_oob(d, mem, size))
			goto err_oob;

		if (d->d_tag == DT_NULL) {
			break;
		} else if (d->d_tag == DT_STRTAB) {
			dyn_strtab = d;
		} else if (d->d_tag == DT_SYMTAB) {
			dyn_symtab = d;
		} else if (d->d_tag == DT_STRSZ) {
			dyn_strsz = d;
		} else if (d->d_tag == DT_SYMENT) {
			dyn_syment = d;
		} else if (d->d_tag == DT_HASH) {
			dyn_hash = d;
		}
	}

	if (!dyn_strtab || !dyn_symtab || !dyn_strsz || !dyn_syment || !dyn_hash) {
		pr_perror("Not all dynamic entries are present");
		return -EINVAL;
	}

	dynsymbol_names = &mem[dyn_strtab->d_un.d_val - load->p_vaddr];
	if (__ptr_oob(dynsymbol_names, mem, size))
		goto err_oob;

	hash = (void *)&mem[(unsigned long)dyn_hash->d_un.d_ptr - (unsigned long)load->p_vaddr];
	if (__ptr_oob(hash, mem, size))
		goto err_oob;

	nbucket = hash[0];
	nchain = hash[1];
	bucket = &hash[2];
	chain = &hash[nbucket + 2];

	test_msg("nbucket %lu nchain %lu bucket %p chain %p\n", (long)nbucket, (long)nchain, bucket, chain);

	for (i = 0; i < ARRAY_SIZE(vdso_symbols); i++) {
		k = elf_hash((const unsigned char *)vdso_symbols[i]);

		for (j = bucket[k % nbucket]; j < nchain && j != STN_UNDEF; j = chain[j]) {
			Sym_t *sym = (void *)&mem[dyn_symtab->d_un.d_ptr - load->p_vaddr];
			char *name;

			sym = &sym[j];
			if (__ptr_oob(sym, mem, size))
				continue;

			if (ELF_ST_TYPE(sym->st_info) != STT_FUNC && ELF_ST_BIND(sym->st_info) != STB_GLOBAL)
				continue;

			name = &dynsymbol_names[sym->st_name];
			if (__ptr_oob(name, mem, size))
				continue;

			if (strcmp(name, vdso_symbols[i]))
				continue;

			memcpy(t->symbols[i].name, name, sizeof(t->symbols[i].name));
			t->symbols[i].offset = (unsigned long)sym->st_value - load->p_vaddr;
			test_msg("symbol %s offset %lx\n", t->symbols[i].name, t->symbols[i].offset);
			break;
		}
	}

	return 0;

err_oob:
	pr_perror("Corrupted Elf data");
	return -EFAULT;
}

static int vdso_fill_self_symtable(struct vdso_symtable *s)
{
	char buf[512];
	int ret = -1;
	FILE *maps;

	*s = (struct vdso_symtable)VDSO_SYMTABLE_INIT;

	maps = fopen("/proc/self/maps", "r");
	if (!maps) {
		pr_perror("Can't open self-vma");
		return -1;
	}

	while (fgets(buf, sizeof(buf), maps)) {
		unsigned long start, end;

		if (!strstr(buf, "[vdso]"))
			continue;

		ret = sscanf(buf, "%lx-%lx", &start, &end);
		if (ret != 2) {
			ret = -1;
			pr_perror("Can't find vDSO bounds");
			goto err;
		}

		s->vma_start = start;
		s->vma_end = end;

		ret = vdso_fill_symtable((void *)start, end - start, s);
		break;
	}

	test_msg("[vdso] %lx-%lx\n", s->vma_start, s->vma_end);
err:
	fclose(maps);
	return ret;
}

static int vdso_clock_gettime_handler(void *func)
{
	__vdso_clock_gettime_t *vdso_clock_gettime = func;
	struct timespec ts1, ts2;

	clock_gettime(CLOCK_REALTIME, &ts1);
	vdso_clock_gettime(CLOCK_REALTIME, &ts2);

	test_msg("clock_gettime: tv_sec %li vdso_clock_gettime: tv_sec %li\n", ts1.tv_sec, ts2.tv_sec);

	if (labs(ts1.tv_sec - ts2.tv_sec) > TIME_DELTA_SEC) {
		pr_perror("Delta is too big");
		return -1;
	}

	return 0;
}

static int vdso_getcpu_handler(void *func)
{
	__vdso_getcpu_t *vdso_getcpu = func;
	unsigned cpu, node;

	vdso_getcpu(&cpu, &node, NULL);
	test_msg("vdso_getcpu: cpu %d node %d\n", cpu, node);

	return 0;
}

static int vdso_gettimeofday_handler(void *func)
{
	__vdso_gettimeofday_t *vdso_gettimeofday = func;
	struct timeval tv1, tv2;
	struct timezone tz;

	gettimeofday(&tv1, &tz);
	vdso_gettimeofday(&tv2, &tz);

	test_msg("gettimeofday: tv_sec %li vdso_gettimeofday: tv_sec %li\n", tv1.tv_sec, tv2.tv_sec);

	if (labs(tv1.tv_sec - tv2.tv_sec) > TIME_DELTA_SEC) {
		pr_perror("Delta is too big");
		return -1;
	}

	return 0;
}

static int vdso_time_handler(void *func)
{
	__vdso_time_t *vdso_time = func;
	time_t t1, t2;

	t1 = time(NULL);
	t2 = vdso_time(NULL);

	test_msg("time: %li vdso_time: %li\n", (long)t1, (long)t1);

	if (labs(t1 - t2) > TIME_DELTA_SEC) {
		pr_perror("Delta is too big");
		return -1;
	}

	return 0;
}

static int call_handlers(struct vdso_symtable *symtable)
{
	typedef int(handler_t)(void *func);
	handler_t *handlers[VDSO_SYMBOL_MAX] = {
		[VDSO_SYMBOL_CLOCK_GETTIME] = vdso_clock_gettime_handler,
		[VDSO_SYMBOL_GETCPU] = vdso_getcpu_handler,
		[VDSO_SYMBOL_GETTIMEOFDAY] = vdso_gettimeofday_handler,
		[VDSO_SYMBOL_TIME] = vdso_time_handler,
	};
	size_t i;

	for (i = 0; i < ARRAY_SIZE(symtable->symbols); i++) {
		struct vdso_symbol *s = &symtable->symbols[i];
		handler_t *func;

		if (vdso_symbol_empty(s) || i > ARRAY_SIZE(handlers))
			continue;
		func = handlers[i];

		if (func((void *)(s->offset + symtable->vma_start))) {
			pr_perror("Handler error");
			return -1;
		}
	}

	return 0;
}

int main(int argc, char *argv[])
{
	struct vdso_symtable symtable;

	test_init(argc, argv);

	if (vdso_fill_self_symtable(&symtable)) {
		pr_perror("Failed to parse vdso");
		return -1;
	}

	if (call_handlers(&symtable))
		return -1;

	test_daemon();
	test_waitsig();

	/*
	 * After restore the vDSO must remain in old place.
	 */
	if (call_handlers(&symtable)) {
		fail("Failed to call vdso handlers from symtable after C/R");
		return -1;
	}

	pass();

	return 0;
}