hardened_malloc/malloc.c
2018-08-23 17:53:10 -04:00

873 lines
23 KiB
C

#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, 85, 64, 51, 42, 36, 64,
/* 32 */ 51, 64, 54, 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;
}
}
}
static void pre_fork(void) {
pthread_mutex_lock(&regions_lock);
for (unsigned i = 0; i < N_SIZE_CLASSES; i++) {
pthread_mutex_lock(&size_class_metadata[i].mutex);
}
}
static void post_fork_parent(void) {
pthread_mutex_unlock(&regions_lock);
for (unsigned i = 0; i < N_SIZE_CLASSES; i++) {
pthread_mutex_unlock(&size_class_metadata[i].mutex);
}
}
static void post_fork_child(void) {
if (pthread_mutex_init(&regions_lock, NULL)) {
fatal_error("mutex initialization failed");
}
for (unsigned i = 0; i < N_SIZE_CLASSES; i++) {
if (pthread_mutex_init(&size_class_metadata[i].mutex, NULL)) {
fatal_error("mutex initialization failed");
}
}
}
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;
}
pthread_atfork(pre_fork, post_fork_parent, post_fork_child);
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;
}