preload.c

#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 NOSECRET 1
#ifndef NOSECRET
#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)
#else
#define MFD_SECRET 0
#endif

#define ASSERT(x) do { if (!(x)) { printf("ASSERTION failed at line %d\n", __LINE__); exit(1); } } while (0)

/* 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

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

static int debug = 0;

#define DEBUG(...) do { if (debug) printf(__VA_ARGS__); } while(0)
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 {
	char *base;
	/* no size because they're always SEG_SIZE */
	struct segptr *next;
};

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

static 	struct malloc_state *m;

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 struct segptr *chunk_to_segment(struct malloc_chunk *c)
{
	struct segptr *sp;

	for (sp = &m->seg; sp != NULL; sp = sp->next)
		if (sp->base <= (char *)c &&
		    (char *)c < sp->base + SEG_SIZE)
			return sp;
	return NULL;
}

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)
{
	struct segptr *seg = chunk_to_segment(c);
	void *n = (char*)c + chunk_size(c);

	ASSERT((char *)n > seg->base &&
	       (char *)n <= seg->base + SEG_SIZE);
	if (n == seg->base + SEG_SIZE)
		return NULL;

	return n;
}



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

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 show_segment(void)
{
	struct malloc_chunk *c;

	if (!debug)
		return;

	printf("SHOW SEGMENT\n");
	for (c = (struct malloc_chunk *)m->seg.base; c != NULL; c = next_chunk(c)) {
		printf("%p:%d:%d:%04x:%04x", c, in_use(c), prev_in_use(c),
		       c->prev_foot, chunk_size(c));
		if (!in_use(c))
			printf(":%p:%p", c->fd, c->bk);
		printf("\n");
	}

	printf("SHOW SEGMENT END\n");
}

void __attribute__ ((constructor)) preload_setup(void)
{
	int fd;
	int ret;
	void *p;
	struct malloc_chunk *c;
	const size_t msize = pad_request(sizeof(*m));

	if (getenv("MALLOC_DEBUG") != NULL)
		debug = 1;

	fd = memfd_create("secure", MFD_CLOEXEC|MFD_SECRET);
	check(fd < 0, "memfd_create");

#ifndef NOSECRET
	ret = ioctl(fd, MFD_SECRET_EXCLUSIVE);
	check(ret < 0, "ioctl");
#endif

	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");

	DEBUG("initial malloc at 0x%p-0x%p\n", p, p+SEG_SIZE);
	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;
	DEBUG("next chunk at 0x%p, head=%x, prev=%x\n", c, chunk_size(c), c->prev_foot);
	m->free.fd = &m->free;
	m->free.bk = &m->free;

	link_free_chunk(c);
}

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

	for (c = m->free.fd; 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 split_free_chunk(struct malloc_chunk *c, size_t size)
{
	struct malloc_chunk *new_c, *n;
	size_t csize = chunk_size(c);

	if (csize == size) {
		/* nothing to split, just give everything up */
		unlink_free_chunk(c);
		c->head |= CINUSE_BIT;
		return;
	}

	/* 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) | PINUSE_BIT;
	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;
	/* and adjust the next prev_foot if there is one */
	n = next_chunk(new_c);
	if (n)
		n->prev_foot = chunk_size(new_c);
}

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

	size = pad_request(size);
	c = find_free(size);
	if (c == NULL) {
		printf("failed to find free\n");
		/* FIXME ADD MORE */
		return NULL;
	}
	DEBUG("found chunk 0x%p\n", c);

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

void *CRYPTO_malloc(size_t size, const char *file, int line)
{
	void *ret = NULL;
	if (size < SEG_SIZE)
		ret = dlmalloc(size);
	DEBUG("in crypto malloc from %s:%d =%p\n", file, line, ret);
	show_segment();
	return ret;
}

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

	DEBUG("in crypto free from %s:%d\n", file, line);
	if (ptr == NULL)
		return;

	c = mem2chunk(ptr);
	/* shred the data */
	memset(ptr, 0, chunk_size(c) - CHUNK_SIZE);

	n = next_chunk(c);
	DEBUG("free c=%p, n=%p, next_in_use=%d, prev_in_use=%d\n",
	       c,n,in_use(n),prev_in_use(c));

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

		p->head += chunk_size(c);
		if (n)
			n->prev_foot = chunk_size(p);

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

	/* and finally consolidation with next */
	if (n && !in_use(n)) {
		unlink_free_chunk(n);
		c->head += chunk_size(n);
		n = next_chunk(c);
		if (n)
			n->prev_foot = chunk_size(c);
	} else if (n) {
		ASSERT(prev_in_use(n));
		n->head &= ~PINUSE_BIT;
	}
	c->head &= ~CINUSE_BIT;
	show_segment();
}