initial commit

.gitignore

@@ -0,0 +1,2 @@
+*.o
+*.so

CREDITS

@@ -0,0 +1,34 @@
+get_random_size_uniform:
+
+    Copyright (c) 2008, Damien Miller <djm@openbsd.org>
+
+    Permission to use, copy, modify, and distribute this software for any
+    purpose with or without fee is hereby granted, provided that the above
+    copyright notice and this permission notice appear in all copies.
+
+    THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+    WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+    MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+    ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+    WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+    ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+    OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+open-addressed hash table:
+
+    Copyright (c) 2008, 2010, 2011, 2016 Otto Moerbeek <otto@drijf.net>
+    Copyright (c) 2012 Matthew Dempsky <matthew@openbsd.org>
+    Copyright (c) 2008 Damien Miller <djm@openbsd.org>
+    Copyright (c) 2000 Poul-Henning Kamp <phk@FreeBSD.org>
+
+    Permission to use, copy, modify, and distribute this software for any
+    purpose with or without fee is hereby granted, provided that the above
+    copyright notice and this permission notice appear in all copies.
+
+    THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+    WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+    MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+    ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+    WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+    ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+    OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

LICENSE

@@ -0,0 +1 @@
+Copyright (c) 2018 Daniel Micay

Makefile

@@ -0,0 +1,15 @@
+CPPFLAGS := -D_GNU_SOURCE
+CFLAGS := -std=c11 -Wall -Wextra -O2 -flto -fPIC -fvisibility=hidden -pedantic
+LDFLAGS := -Wl,--as-needed
+LDLIBS := -lpthread
+OBJECTS := malloc.o random.o util.o
+
+hardened_malloc.so: $(OBJECTS)
+	$(CC) $(CFLAGS) $(LDFLAGS) -shared $^ $(LDLIBS) -o $@
+
+malloc.o: malloc.c malloc.h random.h util.h
+random.o: random.c random.h util.h
+util.o: util.c util.h
+
+clean:
+	rm -f hardened_malloc.so $(OBJECTS)

calculate_waste.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+
+size_classes = [
+    16, 32, 48, 64, 80, 96, 112, 128,
+    160, 192, 224, 256,
+    320, 384, 448, 512,
+    640, 768, 896, 1024,
+    1280, 1536, 1792, 2048,
+    2560, 3072, 3584, 4096,
+    5120, 6144, 7168, 8192,
+    10240, 12288, 14336, 16384
+]
+
+def page_align(size):
+    return (size + 4095) & ~4095
+
+max_bits = 256
+max_page_span = 16
+
+print("maximum bitmap size is {}-bit".format(max_bits))
+print("maximum page span size is {} ({})".format(max_page_span, max_page_span * 4096))
+
+for size_class in size_classes:
+    choices = []
+    for bits in range(1, max_bits + 1):
+        used = size_class * bits
+        real = page_align(used)
+        if real > 65536:
+            continue
+        pages = real / 4096
+        efficiency = used / real * 100
+        choices.append((bits, used, real, pages, efficiency))
+
+    choices.sort(key=lambda x: x[4], reverse=True)
+
+    print()
+    print("size_class:", size_class)
+    for choice in choices[:10]:
+        print(choice)

malloc.c

@@ -0,0 +1,845 @@
+#include <assert.h>
+#include <errno.h>
+#include <stdatomic.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <unistd.h>
+#include <pthread.h>
+#include <sys/mman.h>
+
+#include <malloc.h>
+
+#include "malloc.h"
+#include "random.h"
+#include "util.h"
+
+static_assert(sizeof(void *) == 8, "64-bit only");
+
+#define PAGE_SHIFT 12
+#define PAGE_SIZE ((size_t)1 << PAGE_SHIFT)
+#define PAGE_MASK ((size_t)(PAGE_SIZE - 1))
+#define PAGE_CEILING(s) (((s) + PAGE_MASK) & ~PAGE_MASK)
+
+#define MIN_ALIGN 16
+#define ALIGNMENT_CEILING(s, alignment)	(((s) + (alignment - 1)) & ((~(alignment)) + 1))
+
+static const size_t guard_size = PAGE_SIZE;
+
+// TODO: can be removed once the work is further along
+COLD static noreturn void unimplemented(void) {
+    fatal_error("unimplemented");
+}
+
+static void *memory_map(size_t size) {
+    void *p = mmap(NULL, size, PROT_NONE, MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
+    if (p == MAP_FAILED) {
+        return NULL;
+    }
+    return p;
+}
+
+static int memory_unmap(void *ptr, size_t size) {
+    int ret = munmap(ptr, size);
+    if (ret && errno != ENOMEM) {
+        fatal_error("non-ENOMEM munmap failure");
+    }
+    return ret;
+}
+
+static void *allocate_pages(size_t usable_size, bool unprotect) {
+    usable_size = PAGE_CEILING(usable_size);
+
+    size_t real_size;
+    if (__builtin_add_overflow(usable_size, guard_size * 2, &real_size)) {
+        return NULL;
+    }
+    void *real = memory_map(real_size);
+    if (real == NULL) {
+        return NULL;
+    }
+    void *usable = (char *)real + guard_size;
+    if (unprotect && mprotect(usable, usable_size, PROT_READ|PROT_WRITE)) {
+        memory_unmap(real, real_size);
+        return NULL;
+    }
+    return usable;
+}
+
+static void deallocate_pages(void *usable, size_t usable_size) {
+    usable_size = PAGE_CEILING(usable_size);
+
+    memory_unmap((char *)usable - guard_size, usable_size + guard_size * 2);
+}
+
+static void *allocate_pages_aligned(size_t usable_size, size_t alignment) {
+    usable_size = PAGE_CEILING(usable_size);
+
+    size_t alloc_size;
+    if (__builtin_add_overflow(usable_size, alignment - PAGE_SIZE, &alloc_size)) {
+        return NULL;
+    }
+
+    size_t real_alloc_size;
+    if (__builtin_add_overflow(alloc_size, guard_size * 2, &real_alloc_size)) {
+        return NULL;
+    }
+
+    void *real = memory_map(real_alloc_size);
+    if (real == NULL) {
+        return NULL;
+    }
+
+    void *usable = (char *)real + guard_size;
+
+    size_t lead_size = ALIGNMENT_CEILING((uintptr_t)usable, alignment) - (uintptr_t)usable;
+    size_t trail_size = alloc_size - lead_size - usable_size;
+    void *base = (char *)usable + lead_size;
+
+    if (mprotect(base, usable_size, PROT_READ|PROT_WRITE)) {
+        memory_unmap(real, real_alloc_size);
+        return NULL;
+    }
+
+    if (lead_size) {
+        if (memory_unmap(real, lead_size)) {
+            memory_unmap(real, real_alloc_size);
+            return NULL;
+        }
+    }
+
+    if (trail_size) {
+        if (memory_unmap((char *)base + usable_size + guard_size, trail_size)) {
+            memory_unmap(real, real_alloc_size);
+            return NULL;
+        }
+    }
+
+    return base;
+}
+
+static union {
+    struct {
+        void *slab_region_start;
+        void *slab_region_end;
+        atomic_bool initialized;
+    };
+    char padding[PAGE_SIZE];
+} ro __attribute__((aligned(PAGE_SIZE))) = {
+    .initialized = ATOMIC_VAR_INIT(false)
+};
+
+struct slab_metadata {
+    uint64_t bitmap;
+    struct slab_metadata *next;
+    struct slab_metadata *prev;
+};
+
+static const size_t max_slab_size_class = 16384;
+
+static const uint16_t size_classes[] = {
+    /* 16 */ 16, 32, 48, 64, 80, 96, 112, 128,
+    /* 32 */ 160, 192, 224, 256,
+    /* 64 */ 320, 384, 448, 512,
+    /* 128 */ 640, 768, 896, 1024,
+    /* 256 */ 1280, 1536, 1792, 2048,
+    /* 512 */ 2560, 3072, 3584, 4096,
+    /* 1024 */ 5120, 6144, 7168, 8192,
+    /* 2048 */ 10240, 12288, 14336, 16384
+};
+
+static const uint16_t size_class_slots[] = {
+    /* 16 */ 256, 128, 128, 64, 51, 42, 36, 64,
+    /* 32 */ 51, 64, 64, 64,
+    /* 64 */ 64, 64, 64, 64,
+    /* 128 */ 64, 64, 64, 64,
+    /* 256 */ 16, 16, 16, 16,
+    /* 512 */ 8, 8, 8, 8,
+    /* 1024 */ 8, 8, 8, 8,
+    /* 2048 */ 5, 6, 4, 4
+};
+
+#define N_SIZE_CLASSES (sizeof(size_classes) / sizeof(size_classes[0]))
+
+struct size_info {
+    size_t size;
+    size_t class;
+};
+
+static struct size_info get_size_info(size_t size) {
+    for (size_t i = 0; i < N_SIZE_CLASSES; i++) {
+        size_t real_size = size_classes[i];
+        if (size <= real_size) {
+            return (struct size_info){real_size, i};
+        }
+    }
+    fatal_error("invalid size for slabs");
+}
+
+static size_t get_slab_size(size_t slots, size_t size) {
+    return PAGE_CEILING(slots * size);
+}
+
+static struct size_class {
+    void *class_region_start;
+    size_t metadata_allocated;
+    size_t metadata_count;
+    struct slab_metadata *partial_slabs;
+    struct slab_metadata *free_slabs;
+    struct slab_metadata *slab_info;
+    pthread_mutex_t mutex;
+} size_class_metadata[N_SIZE_CLASSES];
+
+static const size_t class_region_size = 128ULL * 1024 * 1024 * 1024;
+static const size_t real_class_region_size = class_region_size * 2;
+static const size_t slab_region_size = real_class_region_size * N_SIZE_CLASSES;
+static_assert(PAGE_SIZE == 4096, "bitmap handling will need adjustment for other page sizes");
+
+static size_t get_metadata_max(size_t slab_size) {
+    return class_region_size / slab_size;
+}
+
+static struct slab_metadata *alloc_metadata(struct size_class *c, size_t slab_size) {
+    if (c->metadata_count == c->metadata_allocated) {
+        size_t metadata_max = get_metadata_max(slab_size);
+        if (c->metadata_count == metadata_max) {
+            return NULL;
+        }
+        size_t allocate = c->metadata_allocated * 2;
+        if (allocate > metadata_max) {
+            allocate = metadata_max;
+        }
+        if (mprotect(c->slab_info, allocate * sizeof(struct slab_metadata), PROT_READ|PROT_WRITE)) {
+            return NULL;
+        }
+        c->metadata_allocated = allocate;
+    }
+
+    struct slab_metadata *metadata = c->slab_info + c->metadata_count;
+    c->metadata_count++;
+    return metadata;
+}
+
+static void check_index(size_t index) {
+    if (index >= 64) {
+        fatal_error("invalid index");
+    }
+}
+
+static void set_slot(struct slab_metadata *metadata, size_t index) {
+    check_index(index);
+    metadata->bitmap |= 1UL << index;
+}
+
+static void clear_slot(struct slab_metadata *metadata, size_t index) {
+    check_index(index);
+    metadata->bitmap &= ~(1UL << index);
+}
+
+static bool get_slot(struct slab_metadata *metadata, size_t index) {
+    check_index(index);
+    return (metadata->bitmap >> index) & 1UL;
+}
+
+static uint64_t get_mask(size_t slots) {
+    if (slots > 64) return 0; // TODO: implement multi-word bitmaps
+    return slots < 64 ? ~0UL << slots : 0;
+}
+
+static size_t first_free_slot(size_t slots, struct slab_metadata *metadata) {
+    size_t masked = metadata->bitmap | get_mask(slots);
+    if (masked == ~0UL) {
+        fatal_error("no zero bits");
+    }
+    return __builtin_ffsl(~masked) - 1;
+}
+
+static bool has_free_slots(size_t slots, struct slab_metadata *metadata) {
+    size_t masked = metadata->bitmap | get_mask(slots);
+    return masked != ~0UL;
+}
+
+static bool is_free_slab(struct slab_metadata *metadata) {
+    return !metadata->bitmap;
+}
+
+static void *get_slab(struct size_class *c, size_t slab_size, struct slab_metadata *metadata) {
+    size_t index = metadata - c->slab_info;
+    return (char *)c->class_region_start + (index * slab_size);
+}
+
+static struct slab_metadata *get_metadata(struct size_class *c, size_t slab_size, void *p) {
+    size_t offset = (char *)p - (char *)c->class_region_start;
+    size_t index = offset / slab_size;
+    return c->slab_info + index;
+}
+
+static void *slab_allocate(size_t requested_size) {
+    struct size_info info = get_size_info(requested_size);
+    size_t size = info.size;
+    struct size_class *c = &size_class_metadata[info.class];
+    size_t slots = size_class_slots[info.class];
+    size_t slab_size = get_slab_size(slots, size);
+
+    pthread_mutex_lock(&c->mutex);
+
+    if (c->partial_slabs == NULL) {
+        if (c->free_slabs != NULL) {
+            struct slab_metadata *metadata = c->free_slabs;
+            c->free_slabs = c->free_slabs->next;
+            if (c->free_slabs) {
+                c->free_slabs->prev = NULL;
+            }
+
+            metadata->next = c->partial_slabs;
+            metadata->prev = NULL;
+
+            if (c->partial_slabs) {
+                c->partial_slabs->prev = metadata;
+            }
+            c->partial_slabs = metadata;
+
+            void *slab = get_slab(c, slab_size, metadata);
+            set_slot(metadata, 0);
+            pthread_mutex_unlock(&c->mutex);
+            return slab;
+        }
+
+        struct slab_metadata *metadata = alloc_metadata(c, slab_size);
+        if (metadata == NULL) {
+            pthread_mutex_unlock(&c->mutex);
+            return NULL;
+        }
+
+        void *slab = get_slab(c, slab_size, metadata);
+        if (mprotect(slab, slab_size, PROT_READ|PROT_WRITE)) {
+            metadata->next = c->free_slabs;
+            if (c->free_slabs) {
+                c->free_slabs->prev = metadata;
+            }
+            c->free_slabs = metadata;
+
+            // TODO: implement memory protected free slabs
+            unimplemented();
+
+            pthread_mutex_unlock(&c->mutex);
+            return NULL;
+        }
+
+        c->partial_slabs = metadata;
+        set_slot(metadata, 0);
+
+        pthread_mutex_unlock(&c->mutex);
+        return slab;
+    }
+
+    struct slab_metadata *metadata = c->partial_slabs;
+    size_t slot = first_free_slot(slots, metadata);
+    set_slot(metadata, slot);
+
+    if (!has_free_slots(slots, metadata)) {
+        c->partial_slabs = c->partial_slabs->next;
+        if (c->partial_slabs) {
+            c->partial_slabs->prev = NULL;
+        }
+    }
+
+    void *slab = get_slab(c, slab_size, metadata);
+    void *p = (char *)slab + slot * size;
+
+    pthread_mutex_unlock(&c->mutex);
+    return p;
+}
+
+static size_t slab_size_class(void *p) {
+    size_t offset = (char *)p - (char *)ro.slab_region_start;
+    return offset / class_region_size;
+}
+
+static size_t slab_usable_size(void *p) {
+    return size_classes[slab_size_class(p)];
+}
+
+static void slab_free(void *p) {
+    size_t class = slab_size_class(p);
+
+    struct size_class *c = &size_class_metadata[class];
+    size_t size = size_classes[class];
+    size_t slots = size_class_slots[class];
+    size_t slab_size = get_slab_size(slots, size);
+
+    pthread_mutex_lock(&c->mutex);
+
+    struct slab_metadata *metadata = get_metadata(c, slab_size, p);
+    if (!has_free_slots(slots, metadata)) {
+        metadata->next = c->partial_slabs;
+        metadata->prev = NULL;
+
+        if (c->partial_slabs) {
+            c->partial_slabs->prev = metadata;
+        }
+        c->partial_slabs = metadata;
+    }
+
+    void *slab = get_slab(c, slab_size, metadata);
+    size_t slot = ((char *)p - (char *)slab) / size;
+    if (!get_slot(metadata, slot)) {
+        fatal_error("double free");
+    }
+    clear_slot(metadata, slot);
+    memset(p, 0, size);
+
+    if (is_free_slab(metadata)) {
+        if (metadata->prev) {
+            metadata->prev->next = metadata->next;
+        } else {
+            if (c->partial_slabs != metadata) {
+                fatal_error("not good");
+            }
+            c->partial_slabs = metadata->next;
+        }
+        if (metadata->next) {
+            metadata->next->prev = metadata->prev;
+        }
+
+        metadata->next = c->free_slabs;
+        metadata->prev = NULL;
+
+        if (c->free_slabs) {
+            c->free_slabs->prev = metadata;
+        }
+        c->free_slabs = metadata;
+    }
+
+    pthread_mutex_unlock(&c->mutex);
+}
+
+struct region_info {
+    void *p;
+    size_t size;
+};
+
+static const size_t initial_region_table_size = 256;
+
+static struct region_info *regions;
+static size_t regions_total = initial_region_table_size;
+static size_t regions_free = initial_region_table_size;
+static pthread_mutex_t regions_lock = PTHREAD_MUTEX_INITIALIZER;
+
+static size_t hash_page(void *p) {
+    uintptr_t u = (uintptr_t)p >> PAGE_SHIFT;
+    size_t sum = u;
+    sum = (sum << 7) - sum + (u >> 16);
+    sum = (sum << 7) - sum + (u >> 32);
+    sum = (sum << 7) - sum + (u >> 48);
+    return sum;
+}
+
+static int regions_grow(void) {
+    if (regions_total > SIZE_MAX / sizeof(struct region_info) / 2) {
+        return 1;
+    }
+
+    size_t newtotal = regions_total * 2;
+    size_t newsize = newtotal * sizeof(struct region_info);
+    size_t mask = newtotal - 1;
+
+    struct region_info *p = allocate_pages(newsize, true);
+    if (p == NULL) {
+        return 1;
+    }
+
+    for (size_t i = 0; i < regions_total; i++) {
+        void *q = regions[i].p;
+        if (q != NULL) {
+            size_t index = hash_page(q) & mask;
+            while (p[index].p != NULL) {
+                index = (index - 1) & mask;
+            }
+            p[index] = regions[i];
+        }
+    }
+
+    deallocate_pages(regions, regions_total * sizeof(struct region_info));
+    regions_free = regions_free + regions_total;
+    regions_total = newtotal;
+    regions = p;
+    return 0;
+}
+
+static int regions_insert(void *p, size_t size) {
+    if (regions_free * 4 < regions_total) {
+        if (regions_grow()) {
+            return 1;
+        }
+    }
+
+    size_t mask = regions_total - 1;
+    size_t index = hash_page(p) & mask;
+    void *q = regions[index].p;
+    while (q != NULL) {
+        index = (index - 1) & mask;
+        q = regions[index].p;
+    }
+    regions[index].p = p;
+    regions[index].size = size;
+    regions_free--;
+    return 0;
+}
+
+static struct region_info *regions_find(void *p) {
+    size_t mask = regions_total - 1;
+    size_t index = hash_page(p) & mask;
+    void *r = regions[index].p;
+    while (r != p && r != NULL) {
+        index = (index - 1) & mask;
+        r = regions[index].p;
+    }
+    return (r == p && r != NULL) ? &regions[index] : NULL;
+}
+
+static void regions_delete(struct region_info *region) {
+    size_t mask = regions_total - 1;
+
+    regions_free++;
+
+    size_t i = region - regions;
+    for (;;) {
+        regions[i].p = NULL;
+        regions[i].size = 0;
+        size_t j = i;
+        for (;;) {
+            i = (i - 1) & mask;
+            if (regions[i].p == NULL) {
+                return;
+            }
+            size_t r = hash_page(regions[i].p) & mask;
+            if ((i <= r && r < j) || (r < j && j < i) || (j < i && i <= r)) {
+                continue;
+            }
+            regions[j] = regions[i];
+            break;
+        }
+    }
+}
+
+COLD static void init_slow_path(void) {
+    static pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
+
+    pthread_mutex_lock(&mutex);
+
+    if (atomic_load_explicit(&ro.initialized, memory_order_acquire)) {
+        pthread_mutex_unlock(&mutex);
+        return;
+    }
+
+    struct random_state rng;
+    random_state_init(&rng);
+
+    regions = allocate_pages(regions_total * sizeof(struct region_info), true);
+    if (regions == NULL) {
+        fatal_error("failed to set up allocator");
+    }
+
+    ro.slab_region_start = memory_map(slab_region_size);
+    if (ro.slab_region_start == NULL) {
+        fatal_error("failed to allocate slab region");
+    }
+    ro.slab_region_end = (char *)ro.slab_region_start + slab_region_size;
+
+    for (unsigned i = 0; i < N_SIZE_CLASSES; i++) {
+        struct size_class *c = &size_class_metadata[i];
+
+        if (pthread_mutex_init(&c->mutex, NULL)) {
+            fatal_error("mutex initialization failed");
+        }
+
+        size_t gap = (get_random_size_uniform(&rng, (real_class_region_size - class_region_size) / PAGE_SIZE) + 1) * PAGE_SIZE;
+        c->class_region_start = (char *)ro.slab_region_start + class_region_size * i + gap;
+
+        size_t size = size_classes[i];
+        size_t slots = size_class_slots[i];
+        size_t metadata_max = get_metadata_max(get_slab_size(slots, size));
+        c->slab_info = allocate_pages(metadata_max * sizeof(struct slab_metadata), false);
+        if (c->slab_info == NULL) {
+            fatal_error("failed to allocate slab metadata");
+        }
+        c->metadata_allocated = 32;
+        if (mprotect(c->slab_info, c->metadata_allocated * sizeof(struct slab_metadata), PROT_READ|PROT_WRITE)) {
+            fatal_error("failed to allocate initial slab info");
+        }
+    }
+
+    atomic_store_explicit(&ro.initialized, true, memory_order_release);
+
+    if (mprotect(&ro, sizeof(ro), PROT_READ)) {
+        fatal_error("failed to protect allocator data");
+    }
+
+    pthread_mutex_unlock(&mutex);
+}
+
+static void init(void) {
+    if (likely(atomic_load_explicit(&ro.initialized, memory_order_acquire))) {
+        return;
+    }
+
+    init_slow_path();
+}
+
+static void enforce_init(void) {
+    if (!atomic_load_explicit(&ro.initialized, memory_order_acquire)) {
+        fatal_error("invalid uninitialized allocator usage");
+    }
+}
+
+static void *allocate(size_t size) {
+    if (size <= max_slab_size_class) {
+        return slab_allocate(size);
+    }
+
+    void *p = allocate_pages(size, true);
+    if (p == NULL) {
+        return NULL;
+    }
+
+    pthread_mutex_lock(&regions_lock);
+    if (regions_insert(p, size)) {
+        pthread_mutex_unlock(&regions_lock);
+        deallocate_pages(p, size);
+        return NULL;
+    }
+    pthread_mutex_unlock(&regions_lock);
+
+    return p;
+}
+
+static void deallocate(void *p) {
+    if (p >= ro.slab_region_start && p < ro.slab_region_end) {
+        slab_free(p);
+        return;
+    }
+
+    pthread_mutex_lock(&regions_lock);
+    struct region_info *region = regions_find(p);
+    if (region == NULL) {
+        fatal_error("invalid free");
+    }
+    size_t size = region->size;
+    regions_delete(region);
+    pthread_mutex_unlock(&regions_lock);
+
+    deallocate_pages(p, size);
+}
+
+EXPORT void *h_malloc(size_t size) {
+    init();
+    return allocate(size);
+}
+
+EXPORT void *h_calloc(size_t nmemb, size_t size) {
+    size_t total_size;
+    if (__builtin_mul_overflow(nmemb, size, &total_size)) {
+        errno = ENOMEM;
+        return NULL;
+    }
+    init();
+    return allocate(total_size);
+}
+
+EXPORT void *h_realloc(void *old, size_t size) {
+    if (old == NULL) {
+        init();
+        return allocate(size);
+    }
+
+    enforce_init();
+
+    if (size == 0) {
+        deallocate(old);
+        return allocate(size);
+    }
+
+    size_t old_size;
+    if (old >= ro.slab_region_start && old < ro.slab_region_end) {
+        old_size = slab_usable_size(old);
+        if (size <= max_slab_size_class && get_size_info(size).size == old_size) {
+            return old;
+        }
+    } else {
+        pthread_mutex_lock(&regions_lock);
+        struct region_info *region = regions_find(old);
+        if (region == NULL) {
+            fatal_error("invalid realloc");
+        }
+        old_size = region->size;
+        if (PAGE_CEILING(old_size) == PAGE_CEILING(size)) {
+            region->size = size;
+            pthread_mutex_unlock(&regions_lock);
+            return old;
+        }
+        pthread_mutex_unlock(&regions_lock);
+    }
+
+    void *new = allocate(size);
+    if (new == NULL) {
+        return NULL;
+    }
+    size_t copy_size = size < old_size ? size : old_size;
+    memcpy(new, old, copy_size);
+    deallocate(old);
+    return new;
+}
+
+static int alloc_aligned(void **memptr, size_t alignment, size_t size, size_t min_alignment) {
+    if ((alignment - 1) & alignment || alignment < min_alignment) {
+        return EINVAL;
+    }
+
+    if (alignment <= PAGE_SIZE) {
+        if (size < alignment) {
+            size = alignment;
+        }
+
+        void *p = allocate(size);
+        if (p == NULL) {
+            return ENOMEM;
+        }
+        *memptr = p;
+        return 0;
+    }
+
+    void *p = allocate_pages_aligned(size, alignment);
+    if (p == NULL) {
+        return ENOMEM;
+    }
+    if (regions_insert(p, size)) {
+        deallocate_pages(p, size);
+        return ENOMEM;
+    }
+    *memptr = p;
+    return 0;
+}
+
+static void *alloc_aligned_simple(size_t alignment, size_t size) {
+    void *ptr;
+    int ret = alloc_aligned(&ptr, alignment, size, 1);
+    if (ret) {
+        errno = ret;
+        return NULL;
+    }
+    return ptr;
+}
+
+EXPORT int h_posix_memalign(void **memptr, size_t alignment, size_t size) {
+    init();
+    return alloc_aligned(memptr, alignment, size, sizeof(void *));
+}
+
+EXPORT void *h_aligned_alloc(size_t alignment, size_t size) {
+    if (size % alignment) {
+        errno = EINVAL;
+        return NULL;
+    }
+    init();
+    return alloc_aligned_simple(alignment, size);
+}
+
+EXPORT void *h_memalign(size_t alignment, size_t size) {
+    init();
+    return alloc_aligned_simple(alignment, size);
+}
+
+EXPORT void *h_valloc(size_t size) {
+    init();
+    return alloc_aligned_simple(PAGE_SIZE, size);
+}
+
+EXPORT void *h_pvalloc(size_t size) {
+    size_t rounded = PAGE_CEILING(size);
+    if (!rounded) {
+        errno = ENOMEM;
+        return NULL;
+    }
+    init();
+    return alloc_aligned_simple(PAGE_SIZE, rounded);
+}
+
+EXPORT void h_free(void *p) {
+    if (p == NULL) {
+        return;
+    }
+
+    enforce_init();
+    deallocate(p);
+}
+
+EXPORT void h_cfree(void *ptr) __attribute__((alias("free")));
+
+EXPORT size_t h_malloc_usable_size(void *p) {
+    if (p == NULL) {
+        return 0;
+    }
+
+    enforce_init();
+
+    if (p >= ro.slab_region_start && p < ro.slab_region_end) {
+        return slab_usable_size(p);
+    }
+
+    pthread_mutex_lock(&regions_lock);
+    struct region_info *region = regions_find(p);
+    if (p == NULL) {
+        fatal_error("invalid malloc_usable_size");
+    }
+    size_t size = region->size;
+    pthread_mutex_unlock(&regions_lock);
+
+    return size;
+}
+
+EXPORT int h_mallopt(UNUSED int param, UNUSED int value) {
+    return 0;
+}
+
+static const size_t pad_threshold = 16 * 1024 * 1024;
+
+EXPORT int h_malloc_trim(size_t pad) {
+    if (pad > pad_threshold) {
+        return 0;
+    }
+
+    if (!atomic_load_explicit(&ro.initialized, memory_order_acquire)) {
+        return 0;
+    }
+
+    for (unsigned i = 0; i < N_SIZE_CLASSES; i++) {
+        struct size_class *c = &size_class_metadata[i];
+        pthread_mutex_lock(&c->mutex);
+        // TODO: purge and mprotect all free slabs
+        pthread_mutex_unlock(&c->mutex);
+    }
+
+    return 0;
+}
+
+EXPORT void h_malloc_stats(void) {}
+
+EXPORT struct mallinfo h_mallinfo(void) {
+    return (struct mallinfo){0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+}
+
+EXPORT int h_malloc_info(UNUSED int options, UNUSED FILE *fp) {
+    errno = ENOSYS;
+    return -1;
+}
+
+COLD EXPORT void *h_malloc_get_state(void) {
+    return NULL;
+}
+
+COLD EXPORT int h_malloc_set_state(UNUSED void *state) {
+    return -2;
+}

malloc.h

@@ -0,0 +1,51 @@
+#ifndef ALLOCATOR_H
+#define ALLOCATOR_H
+
+#ifndef H_MALLOC_PREFIX
+#define h_malloc malloc
+#define h_calloc calloc
+#define h_realloc realloc
+#define h_posix_memalign posix_memalign
+#define h_aligned_alloc aligned_alloc
+#define h_memalign memalign
+#define h_valloc valloc
+#define h_pvalloc pvalloc
+#define h_free free
+#define h_malloc_usable_size malloc_usable_size
+#define h_mallopt mallopt
+#define h_malloc_trim malloc_trim
+#define h_malloc_stats malloc_stats
+#define h_mallinfo mallinfo
+#define h_malloc_info malloc_info
+#define h_malloc_get_state malloc_get_state
+#define h_malloc_set_state malloc_set_state
+#define h_cfree cfree
+#endif
+
+// C standard
+void *h_malloc(size_t size);
+void *h_calloc(size_t nmemb, size_t size);
+void *h_realloc(void *ptr, size_t size);
+void *h_aligned_alloc(size_t alignment, size_t size);
+void h_free(void *ptr);
+
+// POSIX
+int h_posix_memalign(void **memptr, size_t alignment, size_t size);
+
+// glibc extensions
+size_t h_malloc_usable_size(void *ptr);
+int h_mallopt(int param, int value);
+int h_malloc_trim(size_t pad);
+void h_malloc_stats(void);
+struct mallinfo h_mallinfo(void);
+int h_malloc_info(int options, FILE *fp);
+void *h_malloc_get_state(void);
+int h_malloc_set_state(void *state);
+
+// obsolete glibc extensions
+void *h_memalign(size_t alignment, size_t size);
+void *h_valloc(size_t size);
+void *h_pvalloc(size_t size);
+void h_cfree(void *ptr);
+
+#endif

preload.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+dir=$(cd "$(dirname ${BASH_SOURCE[0]})" && pwd)
+export LD_PRELOAD+=" $dir/hardened_malloc.so"
+exec $@

random.c

@@ -0,0 +1,53 @@
+#include <errno.h>
+
+#include <sys/random.h>
+
+#include "random.h"
+#include "util.h"
+
+void get_random_seed(void *buf, size_t size) {
+    while (size > 0) {
+        ssize_t r;
+
+        do {
+            r = getrandom(buf, size, 0);
+        } while (r == -1 && errno == EINTR);
+
+        if (r <= 0) {
+            fatal_error("getrandom failed");
+        }
+
+        buf = (char *)buf + r;
+        size -= r;
+    }
+}
+
+void random_state_init(UNUSED struct random_state *state) {
+}
+
+// TODO: add ChaCha20-based CSPRNG, for now avoid using this other than during initialization...
+void get_random_bytes(UNUSED struct random_state *state, void *buf, size_t size) {
+    get_random_seed(buf, size);
+}
+
+size_t get_random_size(struct random_state *state) {
+    size_t size;
+    get_random_bytes(state, &size, sizeof(size));
+    return size;
+}
+
+// based on OpenBSD arc4random_uniform
+size_t get_random_size_uniform(struct random_state *state, size_t bound) {
+    if (bound < 2) {
+        return 0;
+    }
+
+    size_t min = -bound % bound;
+
+    size_t r;
+    do {
+        r = get_random_size(state);
+    } while (r < min);
+
+    return r % bound;
+}

random.h

@@ -0,0 +1,14 @@
+#ifndef RANDOM_H
+#define RANDOM_H
+
+struct random_state {
+    char unused;
+};
+
+void get_random_seed(void *buf, size_t size);
+void random_state_init(struct random_state *state);
+void get_random_bytes(struct random_state *state, void *buf, size_t size);
+size_t get_random_size(struct random_state *state);
+size_t get_random_size_uniform(struct random_state *state, size_t bound);
+
+#endif

util.c

@@ -0,0 +1,12 @@
+#include <stdlib.h>
+#include <string.h>
+
+#include <unistd.h>
+
+#include "util.h"
+
+COLD noreturn void fatal_error(const char *s) {
+    write(STDERR_FILENO, s, strlen(s));
+    write(STDERR_FILENO, "\n", 1);
+    abort();
+}

util.h

@@ -0,0 +1,15 @@
+#ifndef UTIL_H
+#define UTIL_H
+
+#include <stdnoreturn.h>
+
+#define likely(x) __builtin_expect((x), 1)
+#define unlikely(x) __builtin_expect((x), 0)
+
+#define COLD __attribute__((cold))
+#define UNUSED __attribute__((unused))
+#define EXPORT __attribute__((visibility("default")))
+
+COLD noreturn void fatal_error(const char *s);
+
+#endif