Initial attempt at a minimal allocator based on dlmalloc

1 files changed, 202 insertions, 7 deletions

diff --git a/preload.c b/preload.c
index cae2233..5054567 100644
--- a/preload.c
+++ b/preload.c
@@ -1,6 +1,7 @@
 #include <stdio.h>
 #include <unistd.h>
 #include <stdlib.h>
+#include <string.h>
 
 #include <sys/ioctl.h>
 #include <sys/mman.h>
@@ -22,11 +23,53 @@ static inline int memfd_create(const char *name, unsigned int flags)
 }
 
 
-
-#define PAGE_SIZE 4096
+/* 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) {
@@ -35,31 +78,183 @@ static void check(int cond, const char *str)
 	}
 }
 
+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, PAGE_SIZE);
+	ret = ftruncate(fd, SEG_SIZE);
 	check(ret < 0, "ftruncate");
 
-	p = mmap(NULL, PAGE_SIZE, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
+	p = mmap(NULL, SEG_SIZE, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
 	check(p == MAP_FAILED, "mmap");
 
-	secure_page = p;
+	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 < PAGE_SIZE)
-		return secure_page;
+	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);
+	}
+}