OrionLibC
Barry Fixing malloc heap exhaustion bug d0db5bd (3 years, 2 months ago)
#include <stdint.h>
#include <stdlib.h>
#include <stddef.h>
#include <stdio.h>
#include <string.h>
#include <sys/mman.h>
#include <errno.h>
#define BLOCK_SIZE 16
typedef struct Header Header;
typedef struct Arena Arena;
extern char _end[];
/* Structure for a Memory Header */
struct Header {
Header *prev, *next;
size_t size;
char magic[4];
};
/* Structure for a Memory Arena */
struct Arena {
Arena *next;
void *start, *end;
size_t size;
};
Arena primaryArena = {
.next = NULL,
.start = NULL,
.end = NULL,
.size = 0
};
/* Allocate a region of memory in an arena */
static void *
arena_malloc(Arena *arena, size_t size)
{
size_t blockSize, gapSize;
intptr_t blockEnd;
Header *prev, *head, *next;
head = prev = arena->start;
next = NULL;
/* Minimum allocation */
if (size % BLOCK_SIZE)
size += BLOCK_SIZE - (size % BLOCK_SIZE);
/* Block does not exist, create heap */
if (head->prev != head || memcmp(head->magic, "HEAP", 4)) {
head->prev = head;
head->next = NULL;
head->size = size;
memcpy(head->magic, "HEAP", 4);
memset((void *) (head + 1), 0, size);
return (void *) (head + 1);
}
/* Find gap */
while (head->next) {
next = head->next;
blockSize = sizeof(Header) + head->size;
blockEnd = (uintptr_t) head + blockSize;
gapSize = (size_t) next - blockEnd;
prev = head;
/* Fit in gap */
if (gapSize >= size + sizeof(Header)) {
head = (void *) blockEnd;
head->next = next;
head->prev = prev;
prev->next = head;
next->prev = head;
head->size = size;
memcpy(head->magic, "HEAP", 4);
memset((void *) (head + 1), 0, size);
return (void *) (head + 1);
}
head = head->next;
}
/* Add to end */
blockSize = sizeof(Header) + head->size;
blockEnd = (uintptr_t) head + blockSize;
/* Ensure more arenas can be allocated */
gapSize = (size_t) arena->end - blockEnd - sizeof(Arena);
/* Fit in gap */
if (gapSize >= size + sizeof(Header)) {
prev = head;
head = (void *) blockEnd;
head->next = NULL;
head->prev = prev;
prev->next = head;
head->size = size;
memcpy(head->magic, "HEAP", 4);
memset((void *) (head + 1), 0, size);
return (void *) (head + 1);
}
return NULL;
}
/* Allocate a region of memory */
void *
malloc(size_t size)
{
void *addr = NULL;
Arena *arena;
size_t arenaSize = 64 * 1024;
while (arenaSize < size + sizeof(Arena))
arenaSize += 64 * 1024;
for (arena = &primaryArena; !addr; arena = arena->next) {
/* Allocate a new arena */
if (!arena->size) {
arena->size = arenaSize;
arena->start = mmap(NULL, arena->size,
PROT_READ | PROT_WRITE,
MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
if (arena->start == MAP_FAILED) {
errno = ENOMEM;
return NULL;
}
arena->end = arena->start + arena->size;
}
/* Attempt to use arena */
addr = arena_malloc(arena, size);
if (!addr && !arena->next)
arena->next = arena->start + arena->size
- sizeof(Arena);
}
return addr;
}
/* Free an allocated region of memory */
void
free(void *addr)
{
Header *prev, *head, *next;
head = (Header *) addr - 1;
if (memcmp(head->magic, "HEAP", 4)) {
printf("free(): invalid pointer\n");
abort();
}
prev = head->prev;
next = head->next;
memset(head, 0, sizeof(Header));
if (prev != head)
prev->next = next;
if (next)
next->prev = prev;
}
/* Re-allocate a region of memory */
void *
realloc(void *addr, size_t size)
{
void *ptr = malloc(size);
if (addr) {
Header *old = (Header *) addr - 1;
if (memcmp(old->magic, "HEAP", 4)) {
printf("realloc(): invalid pointer\n");
abort();
}
memcpy(ptr, addr, old->size);
free(addr);
}
return ptr;
}