#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>

#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/syscall.h>
#include <sys/types.h>

#include <linux/memfd.h>

/* bits to get memfd_create to work */
#define MFD_SECRET 0x0008U
#define MFD_SECRET_IOCTL '-'
#define MFD_SECRET_EXCLUSIVE	_IOW(MFD_SECRET_IOCTL, 0x13, unsigned long)
#define MFD_SECRET_UNCACHED	_IOW(MFD_SECRET_IOCTL, 0x14, unsigned long)

/* glibc should have defined this by now, sigh */
static inline int memfd_create(const char *name, unsigned int flags)
{
	return syscall(__NR_memfd_create, name, flags);
}


/* segment size.  Matches hugepage size */
#define SEG_SIZE 2*1024*1024

static void *secure_page;

#define CHUNK_SIZE (2 * sizeof(size_t))
#define CHUNK_ALIGNMENT 0xf
#define MIN_FREE_CHUNK 256

static size_t pad_request(size_t s)
{
	return (s + CHUNK_SIZE + CHUNK_ALIGNMENT) & ~CHUNK_ALIGNMENT;
}

#define PINUSE_BIT	0x01
#define CINUSE_BIT	0x02

#define FLAG_BITS	(CINUSE_BIT | PINUSE_BIT)

struct malloc_chunk {
  size_t               prev_foot;  /* Size of previous chunk (if free).  */
  size_t               head;       /* Size and inuse bits. */
  struct malloc_chunk* fd;         /* double links -- used only if free. */
  struct malloc_chunk* bk;
};

struct segptr {
	void *base;
	/* no size because they're always SEG_SIZE */
	struct segptr *next;
};

struct malloc_state {
	struct segptr seg;
	struct malloc_chunk *free;
};

static int in_use(struct malloc_chunk *c)
{
	return c->head & CINUSE_BIT ? 1 : 0;
}

static int prev_in_use(struct malloc_chunk *c)
{
	return c->head & PINUSE_BIT ? 1 : 0;
}

static void check(int cond, const char *str)
{
	if (cond) {
		perror(str);
		exit(1);
	}
}

static void *chunk2mem(struct malloc_chunk *c)
{
	return (char *)c + CHUNK_SIZE;
}

static struct malloc_chunk *mem2chunk(void *p)
{
	return (struct malloc_chunk *)((char *)p - CHUNK_SIZE);
}

static size_t chunk_size(struct malloc_chunk *c)
{
	return c->head & ~FLAG_BITS;
}

static struct malloc_chunk *next_chunk(struct malloc_chunk *c)
{
	return (struct malloc_chunk *)((char*)c + chunk_size(c));
}


static struct malloc_chunk *prev_chunk(struct malloc_chunk *c)
{
	return (struct malloc_chunk *)((char*)c - c->prev_foot);
}

static 	struct malloc_state *m;

void __attribute__ ((constructor)) preload_setup(void)
{
	int fd = memfd_create("secure", MFD_CLOEXEC|MFD_SECRET);
	int ret;
	void *p;
	struct malloc_chunk *c;
	const size_t msize = pad_request(sizeof(*m));

	check(fd < 0, "memfd_create");

	ret = ioctl(fd, MFD_SECRET_EXCLUSIVE);
	check(ret < 0, "ioctl");

	ret = ftruncate(fd, SEG_SIZE);
	check(ret < 0, "ftruncate");

	p = mmap(NULL, SEG_SIZE, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
	check(p == MAP_FAILED, "mmap");

	c = p;
	m = chunk2mem(c);
	memset(m, 0, sizeof(*m));
	c->head = msize | CINUSE_BIT | PINUSE_BIT;
	m->seg.base = p;
	c = next_chunk(c);
	c->head = (size_t)(((char *)p + SEG_SIZE) - (char *)c) | PINUSE_BIT;
	c->prev_foot = msize;
	m->free = c;
	c->bk = c->fd = c;
}

static struct malloc_chunk *find_free(size_t size)
{
	struct malloc_chunk *c, *found = NULL;

	if (m->free == m->free->fd)
		return NULL;

	for (c = m->free; c != m->free; c = c->fd) {
		if (chunk_size(c) < size)
			continue;
		if (found && chunk_size(found) < chunk_size(c))
			continue;
		found = c;
	}
	return found;
}

static void link_free_chunk(struct malloc_chunk *c)
{
	struct malloc_chunk *f = m->free;
	struct malloc_chunk *b = f->bk;

	c->fd = f;
	f->bk = c;
		
	b->fd = c;
	c->bk = b;
}

static void unlink_free_chunk(struct malloc_chunk *c)
{
	struct malloc_chunk *f = c->fd;
	struct malloc_chunk *b = c->bk;

	b->fd = f;
	f->bk = b;
}

static void split_free_chunk(struct malloc_chunk *c, size_t size)
{
	struct malloc_chunk *new_c;
	size_t csize = chunk_size(c);

	if (csize < size + MIN_FREE_CHUNK) {
		unlink_free_chunk(c);
		return;
	}

	/* here we need to split the chunk, so pad up the size to the min */

	if (size < MIN_FREE_CHUNK)
		size = MIN_FREE_CHUNK;

	/* set the old chunk to the size */
	c->head = size | CINUSE_BIT | PINUSE_BIT;
	/* get the new part of the split */
	new_c = next_chunk(c);
	new_c->head = csize - size;
	new_c->prev_foot  = size;
	/* now replace the new chunk with the old chunk */
	new_c->fd = c->fd;
	new_c->bk = c->bk;
	c->bk->fd = new_c;
	c->fd->bk = new_c;
}

static void *alloc(size_t size)
{
	struct malloc_chunk *c;

	size = pad_request(size);
	c = find_free(size);
	if (c == NULL)
		/* FIXME ADD MORE */
		return NULL;

	split_free_chunk(c, size);
	return chunk2mem(c);
}

void *CRYPTO_malloc(size_t size, const char *file, int line)
{
	printf("in crypto malloc from %s:%d\n", file, line);
	if (size < SEG_SIZE)
		return alloc(size);
	else
		return NULL;
}

void *CRYPTO_free(void *ptr, const char *file, int line)
{
	struct malloc_chunk *c, *n;

	printf("in crypto frss from %s:%d\n", file, line);
	c = mem2chunk(ptr);
	/* shred the data */
	memset(ptr, 0, chunk_size(c) - CHUNK_SIZE);

	n = next_chunk(c);

	/* now check for consolidation with previous */
	if (!prev_in_use(c)) {
		struct malloc_chunk *p = prev_chunk(c);

		p->head += chunk_size(c);

		/* the new consolidated chunk becomes our current
		 * chunk for the next free check below.  The previous
		 * chunk was already linked */
		c = p;
	} else {
		link_free_chunk(c);
	}

	/* and finally consolidation with next */
	if (!in_use(n)) {
		unlink_free_chunk(n);
		c->head += chunk_size(n);
	}
}