/* See LICENSE file for copyright and license details. */
#include "memcheck.h"

#include <sys/syscall.h>
#include <sys/mman.h>
#include <errno.h>
#include <stdarg.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#define UNW_LOCAL_ONLY
#include <libunwind.h>
#include <elfutils/libdwfl.h>


#if defined(__GNUC__)
# pragma GCC diagnostic ignored "-Wpadded"
#endif
#if defined(__clang__)
# pragma clang diagnostic ignored "-Wunsafe-buffer-usage"
# pragma clang diagnostic ignored "-Wcomma"
# pragma clang diagnostic ignored "-Wunsafe-buffer-usage" /* clang is currently exceptionally stupid */
#endif


enum alloctype {
	ALLOCTYPE_REALLOC,
	ALLOCTYPE_MMAP
};

#ifdef PRINT_BACKTRACES
struct backtrace {
	size_t alloc_size;
	size_t n;
	uintptr_t rips[];
};
#endif

struct allocinfo {
	void *real;
	void *ptr;
	size_t real_size;
	enum alloctype type;
	int flags;
	int dont_track;
#ifdef PRINT_BACKTRACES
	struct backtrace *backtrace;
#endif
};


static struct allocinfo *allocs = NULL;
static size_t nallocs = 0;
static size_t allocs_size = 0;
static size_t alloc_fail_exclusion = 0;
static size_t memleak_exclusion = 1;
static size_t alloc_fail_in = 0;
static int alloc_fail_all = 0;
static int have_custom_malloc = 0;


static void *
memcheck_mmap(void *addr, size_t len, int prot, int flags, int fd, off_t off)
{
#ifdef SYS_mmap2
	unsigned long pgoff = (unsigned long)(off / 4096);
	long int r = syscall(SYS_mmap2, addr, len, (unsigned long)prot, (unsigned long)flags, fd, pgoff);
#else
	long int r = syscall(SYS_mmap, addr, len, (unsigned long)prot, (unsigned long)flags, fd, off);
#endif
	return (void *)r;
}


static void *
memcheck_mremap(void *old, size_t old_size, size_t new_size, int flags, void *new_address)
{
	long int r = syscall(SYS_mremap, old, old_size, new_size, flags, new_address);
	return (void *)r;
}


static int
memcheck_munmap(void *ptr, size_t n)
{
	return (int)syscall(SYS_munmap, ptr, n);
}


static void *
memcheck_mmap_ram(size_t n)
{
	return memcheck_mmap(NULL, n, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
}


#ifdef PRINT_BACKTRACES
static struct backtrace *
memcheck_create_backtrace(void)
{
	struct backtrace *ret = NULL, *new;
	unw_word_t rip;
	unw_cursor_t cursor;
	unw_context_t context;
	size_t n = 0, size = 0;
	int saved_errno = errno;

	memleak_exclusion += 1;

	if (unw_getcontext(&context))
		goto fail;
	if (unw_init_local(&cursor, &context))
		goto fail;

	n = 8;
	size = offsetof(struct backtrace, rips) + n * sizeof(*ret->rips);
	ret = memcheck_mmap_ram(size);
	if (!ret)
		goto fail;
	ret->alloc_size = size;

	ret->n = 0;
	while (unw_step(&cursor) > 0) {
		if (unw_get_reg(&cursor, UNW_REG_IP, &rip))
			goto fail;
		if (ret->n == n) {
			n += 8;
			size = offsetof(struct backtrace, rips) + n * sizeof(*ret->rips);
			new = memcheck_mmap_ram(size);
			if (!new)
				goto fail;
			new->alloc_size = size;
			new->n = ret->n;
			memcpy(new->rips, ret->rips, n * sizeof(*ret->rips));
			memcheck_munmap(ret, ret->alloc_size);
			ret = new;
		}
		ret->rips[ret->n++] = (uintptr_t)rip;
	}

	errno = saved_errno;
	memleak_exclusion -= 1;
	return ret;

fail:
	if (ret)
		memcheck_munmap(ret, ret->alloc_size);
	errno = saved_errno;
	memleak_exclusion -= 1;
	return NULL;
}
#endif


#ifdef PRINT_BACKTRACES
static void
memcheck_print_backtrace(struct backtrace *backtrace, const char *indent)
{
	int saved_errno = errno;
	int lineno = 0; /* initialised to silence compiler */
	char *debuginfo_path = NULL;
	Dwarf_Addr ip;
	Dwfl_Callbacks callbacks;
	Dwfl *dwfl = NULL;
	Dwfl_Line *line = NULL;
	Dwfl_Module *module = NULL;
	const char *filename = NULL;
	const char *funcname = NULL;
	size_t i;

	if (!backtrace)
		return;

	memleak_exclusion += 1;

	memset(&callbacks, 0, sizeof(callbacks));

	callbacks.find_elf       = dwfl_linux_proc_find_elf;
	callbacks.find_debuginfo = dwfl_standard_find_debuginfo;
	callbacks.debuginfo_path = &debuginfo_path;

	dwfl = dwfl_begin(&callbacks);

	if (dwfl) {
		if (dwfl_linux_proc_report(dwfl, getpid()) ||
		    dwfl_report_end(dwfl, NULL, NULL)) {
			perror("");
			dwfl_end(dwfl);
			dwfl = NULL;
		}
	}

	for (i = 0; i < backtrace->n; i++) {
		ip = (Dwarf_Addr)backtrace->rips[i];

		if (dwfl) {
			module   = dwfl_addrmodule(dwfl, ip);
			funcname = module ? dwfl_module_addrname(module, ip) : NULL;
			line     = dwfl_getsrc(dwfl, ip);
			if (line)
				filename = dwfl_lineinfo(line, &(Dwarf_Addr){0}, &lineno, NULL, NULL, NULL);
		}

		dprintf(STDERR_FILENO, "%s%s 0x%016"PRIxPTR": %s",
		        indent, !i ? "at" : "by",
			(uintptr_t)ip,
		        funcname ? funcname : "???");
		if (line)
			dprintf(STDERR_FILENO, " (%s:%i)\n", filename, lineno);
		else
			dprintf(STDERR_FILENO, "\n");
	}

	if (dwfl)
		dwfl_end(dwfl);

	errno = saved_errno;
	memleak_exclusion -= 1;
}
#endif


#if defined(__clang__)
__attribute__((__format__(__printf__, 1, 2)))
#elif defined(__GNUC__)
__attribute__((__format__(__gnu_printf__, 1, 2)))
#endif
static _Noreturn void
memcheck_errorf(const char *fmt, ...)
{
	va_list ap;
	memleak_exclusion = 1;
	va_start(ap, fmt);
	vdprintf(STDERR_FILENO, fmt, ap);
	dprintf(STDERR_FILENO, ", aborting...\n");
	va_end(ap);
#ifdef PRINT_BACKTRACES
	memcheck_print_backtrace(memcheck_create_backtrace(), "\t");
#endif
	abort();
}


#if defined(__GNUC__)
__attribute__((__pure__))
#endif
static struct allocinfo *
memcheck_getalloc(void *ptr)
{
	size_t i;
	for (i = 0; i < nallocs; i++)
		if (allocs[i].ptr == ptr)
			return &allocs[i];
	return NULL;
}


static int
memcheck_fake_fail(void)
{
	have_custom_malloc = 1;

	if (!alloc_fail_exclusion) {
	        if (alloc_fail_all || (alloc_fail_in && !--alloc_fail_in)) {
		        alloc_fail_all = 1;
	                errno = ENOMEM;
                        return 1;
		}
	}

	return 0;
}


static void
memcheck_putalloc(struct allocinfo *info)
{
	if (nallocs == allocs_size) {
		size_t old_allocs_size = allocs_size;
		void *new;
		if (allocs_size > SIZE_MAX / sizeof(*allocs) - 64)
			memcheck_errorf("size of memory bookkeeping will exceed SIZE_MAX");
		allocs_size += 64;
		new = memcheck_mmap_ram(allocs_size * sizeof(*allocs));
		if (!new)
			memcheck_errorf("failed to allocate enought memory for bookkeeping");
		if (allocs) {
			memcpy(new, allocs, nallocs * sizeof(*allocs));
			memcheck_munmap(allocs, old_allocs_size * sizeof(*allocs));
		}
		allocs = new;
	}
	allocs[nallocs++] = *info;
}


static int
memcheck_free(void *ptr, size_t size, enum alloctype type)
{
	int ret;
	struct allocinfo *alloc;
	alloc = memcheck_getalloc(ptr);
	if (!alloc)
		memcheck_errorf("attempting to deallocate unknown pointer %p", ptr);
	if (alloc->type != type)
		memcheck_errorf("attempting to deallocating with incompatible function");
	if (type == ALLOCTYPE_MMAP && size != alloc->real_size)
		memcheck_errorf("attempting to munmap(2) pointer %p with wrong size", ptr);
	ret = memcheck_munmap(alloc->real, alloc->real_size);
#ifdef PRINT_BACKTRACES
	if (alloc->backtrace)
		memcheck_munmap(alloc->backtrace, alloc->backtrace->alloc_size);
#endif
	*alloc = allocs[--nallocs];
	if (!nallocs) {
		memcheck_munmap(allocs, allocs_size * sizeof(*allocs));
		allocs = NULL;
		nallocs = 0;
		allocs_size = 0;
	}
	return ret;
}


static void *
memcheck_realloc(void *old, size_t n, size_t m, size_t alignment, int initbyte)
{
	struct allocinfo info;
	uintptr_t addr;
	void *ret;

	if (memcheck_fake_fail())
		return NULL;

	if (!n || !m)
		memcheck_errorf("attempting to allocate with zero size which is implementation-defined behaviour");
	if (n > SIZE_MAX / m)
		memcheck_errorf("attempting allocate more than SIZE_MAX");

	if (!alignment)
		alignment = _Alignof(max_align_t);

	n *= m;
	if (n > SIZE_MAX - (alignment - 1U))
		memcheck_errorf("attempting allocate more than SIZE_MAX after alignment padding");

#ifdef PRINT_BACKTRACES
	info.backtrace = memleak_exclusion ? NULL : memcheck_create_backtrace();
#endif
	info.dont_track = memleak_exclusion > 0;
	info.flags = 0;
	info.type = ALLOCTYPE_REALLOC;
	info.real_size = n + (alignment - 1U);
	info.real = memcheck_mmap_ram(info.real_size);
	if (!info.real)
		memcheck_errorf("failed to allocate enough memory");
	memset(info.real, initbyte, info.real_size);

	addr = (uintptr_t)info.real;
	if (addr % alignment)
		addr += alignment - addr % alignment;
	ret = info.ptr = (void *)addr;

	memcheck_putalloc(&info);

	if (old) {
		struct allocinfo *oldinfo;
		oldinfo = memcheck_getalloc(old);
		if (!oldinfo)
			memcheck_errorf("attempting to reallocate unknown pointer %p", old);
		memcpy(info.real, oldinfo->real, info.real_size < oldinfo->real_size ? info.real_size : oldinfo->real_size);
		memcheck_free(old, 0, ALLOCTYPE_REALLOC);
	}

	return ret;
}



/* Implementation of checked functions */


#if defined(__GNUC__)
# pragma GCC diagnostic push
# pragma GCC diagnostic ignored "-Wmissing-prototypes"
#endif


void *
mmap(void *addr, size_t len, int prot, int flags, int fd, off_t off)
{
	struct allocinfo info;
	void *ret;

	if (off % 4096)
		memcheck_errorf("attempting to mmap(2) with offset %ju which is not a multiple of 4096", (uintmax_t)off);

#ifdef PRINT_BACKTRACES
	info.backtrace = memleak_exclusion ? NULL : memcheck_create_backtrace();
#endif
	info.dont_track = memleak_exclusion > 0;
	info.flags = flags;
	info.type = ALLOCTYPE_MMAP;
	info.real_size = len;
	info.real = memcheck_mmap(addr, len, prot, flags, fd, off);
	if (!info.real)
		memcheck_errorf("failed to mmap(2)");
	ret = info.ptr = info.real;

	memcheck_putalloc(&info);
	return ret;
}


int
munmap(void *ptr, size_t n)
{
	return memcheck_free(ptr, n, ALLOCTYPE_MMAP);
}


void *
mremap(void *old, size_t old_size, size_t new_size, int flags, ...)
{
	struct allocinfo *oldinfo;
	void *new_address = NULL;
	void *ret;

	if (memcheck_fake_fail())
		return NULL;

	oldinfo = memcheck_getalloc(old);
	if (!oldinfo)
		memcheck_errorf("attempting to mremap(2) unknown pointer %p", old);
	if (oldinfo->type != ALLOCTYPE_MMAP)
		memcheck_errorf("attempting to mremap(2) pointer %p that was not allocated with mmap(2) or mremap(2)", old);

	if (flags & MREMAP_FIXED) {
		va_list ap;
		va_start(ap, flags);
		new_address = va_arg(ap, void *);
		va_end(ap);
	}

	ret = memcheck_mremap(old, old_size, new_size, flags, new_address);
	if (!ret)
		return NULL;

	if (ret == old || !(flags & MREMAP_DONTUNMAP)) {
		oldinfo->ptr = oldinfo->real = ret;
		oldinfo->real_size = new_size;
	} else {
		struct allocinfo info;
		info = *oldinfo;
		info.ptr = info.real = ret;
		info.real_size = new_size;
		memcheck_putalloc(&info);
	}

	return ret;
}


void
free(void *ptr)
{
	if (!ptr)
		return;
	memcheck_free(ptr, 0, ALLOCTYPE_REALLOC);
}


void *
malloc(size_t n)
{
	return memcheck_realloc(NULL, 1, n, 0, 'x');
}


void *
calloc(size_t n, size_t m)
{
	return memcheck_realloc(NULL, n, m, 0, 0);
}


void *
realloc(void *old, size_t n)
{
	return memcheck_realloc(old, 1, n, 0, 'x');
}


void *
reallocarray(void *old, size_t n, size_t m)
{
	return memcheck_realloc(old, n, m, 0, 'x');
}


void *
memdup(const void *s, size_t n)
{
	char *r = memcheck_realloc(NULL, 1, n, 0, 'x');
	if (r)
		memcpy(r, s, n);
	return r;
}


char *
strdup(const char *s)
{
	size_t n = strlen(s) + 1U;
	char *r = memcheck_realloc(NULL, 1, n, 0, 'x');
	if (r)
		memcpy(r, s, n);
	return r;
}


char *
strndup(const char *s, size_t max)
{
	size_t n = strnlen(s, max) + 1U;
	char *r = memcheck_realloc(NULL, 1, n, 0, 'x');
	if (n > max)
		n = max;
	if (r)
		memcpy(r, s, n);
	return r;
}


int
posix_memalign(void **ret_out, size_t align, size_t n)
{
	int saved_errno, r;
	saved_errno = errno, errno = 0;
	if (!align || (align & (align - 1U)) || align % sizeof(void *))
		memcheck_errorf("invalid alignment for posix_memalign: %zu", align);
	*ret_out = memcheck_realloc(NULL, 1, n, align, 'x');
	return r = errno, errno = saved_errno, r;
}


void *
memalign(size_t align, size_t n)
{
	if (!align)
		memcheck_errorf("invalid alignment for memalign: %zu", align);
	return memcheck_realloc(NULL, 1, n, align, 'x');
}


void *
aligned_alloc(size_t align, size_t n)
{
	if (!align || (align & (align - 1U)))
		memcheck_errorf("invalid alignment for aligned_alloc: %zu", align);
	return memcheck_realloc(NULL, 1, n, align, 'x');
}


void *
valloc(size_t n)
{
	return memcheck_realloc(NULL, 1, n, 4096U, 'x');
}


void *
pvalloc(size_t n)
{
	if (n % 4096U) {
		n |= 4096U - 1U;
		if (n == SIZE_MAX)
			memcheck_errorf("attempting to allocate beyond SIZE_MAX");
		n += 1U;
	}
	return memcheck_realloc(NULL, 1, n, 4096U, 'x');
}


#if defined(__GNUC__)
# pragma GCC diagnostic pop
#endif



/* Test functions */


int
memcheck_have_custom_malloc(void)
{
	static void *volatile test_have_custom_malloc_ptr = NULL;
        int saved_errno = errno;
        alloc_fail_exclusion++;
        memleak_exclusion++;
        test_have_custom_malloc_ptr = malloc(1);
        free(test_have_custom_malloc_ptr);
        memleak_exclusion--;
        alloc_fail_exclusion--;
	errno = saved_errno;
	return have_custom_malloc;
}

size_t
memcheck_alloc_fail_in(size_t n)
{
	size_t ret = alloc_fail_in;
	alloc_fail_in = n;
	alloc_fail_all = 0;
	return ret;
}


int
memcheck_check_memleaks(void)
{
	struct allocinfo *leak;
	size_t i, count = 0;
	int first = 1;
	memleak_exclusion = 1;
	for (i = 0; i < nallocs; i++) {
		leak = &allocs[i];
		if (leak->dont_track)
			continue;
		count += 1;
	}
	if (!count)
               	return 1;
	dprintf(STDERR_FILENO, "%zu %s!\n", count, count == 1 ? "memory leak" : "memory leaks");
#ifdef PRINT_BACKTRACES
	for (i = 0; i < nallocs; i++) {
		leak = &allocs[i];
		if (leak->dont_track)
			continue;
		if (first)
			first = 0;
		else
			dprintf(STDERR_FILENO, "\n");
		memcheck_print_backtrace(leak->backtrace, "\t");
	}
#endif
	return 0;
}

void memcheck_alloc_fail_exclusion_begin(void) { alloc_fail_exclusion++; }
void memcheck_alloc_fail_exclusion_end(void) { alloc_fail_exclusion--; }
void memcheck_exclusion_begin(void) { memleak_exclusion++; }
void memcheck_exclusion_end(void) { memleak_exclusion--; }
void memcheck_begin(void) { memleak_exclusion = 0; }