Ben Cohen
← back
project

Building a Memory Allocator

What i learned while writing a custom safe memory allocator from scratch

• 11 min read
oscheap exploitation

The final code for this project is available on GitHub.

I wanted to understand how malloc works, so I built my own. Here’s what I learned.

Malloc & Free

Most CS students, including me before this project, treat malloc and free as black boxes. You call malloc to get some memory, and you call free to give it back.

Under the hood, however, the standard library isn’t asking the kernel for memory on every allocation. System calls like mmap and sbrk are expensive, so the allocator will request memory by large pages and manage it itself. This is where the fun begins.

The Chunk’s Structure

At the core of an allocator is the chunk. When memory is allocated, the caller gets a pointer to usable space. But the allocator needs a way to track its freed chunks.

To solve this, before each usable chunk we can store a header with metadata about it:

typedef struct heapchunk
{
    size_t size; // The raw size of the allocated memory
    bool is_inuse; // currently allocated or free?

} heapchunk;

We also need to save freed chunks in a list so we can reuse them later. For that, we can add doubly linked list pointers to the header.

At last, we want to access the usable memory of a chunk easily in the code, so we can add a “pointer”:

Note that zero length arrays are a GCC extension, and are not part of the C standard

typedef struct heapchunk
{
    size_t size;
    bool is_inuse;

    struct heapchunk *next;
    struct heapchunk *prev;

    uint8_t payload[0]; // data label

} heapchunk;

The Union Trick

Technically, we don’t need the free list pointers for allocated chunks, and we don’t need the memory space for free chunks. To save space, we can use a union to share the same memory for both:

this is the final structure of a chunk:

typedef struct heapchunk
{
    size_t canary; // We will talk about this later
    size_t size;
    bool is_inuse;

    union
    {
        struct
        {
            struct heapchunk *next;
            struct heapchunk *prev;
        } list;

        uint8_t payload[0];
    };
} heapchunk;

Organizing the free list: Segregated bins

As we said before, we need to keep track of freed chunks so we can reuse them later. A simple way to do this is to keep a free list of all freed chunks. For this, we can have a global “heap” struct:

typedef struct heapinfo
{
    heapchunk *free_list; // head of the free list
    bool initialized; // has the heap been initialized?
    size_t avail; // total available memory in the heap
} heapinfo;

The Arena

To initialize the heap, we can request 2MB of memory pages (which we will call an arena) from the OS using mmap and carve up the first chunk to be the entire heap. Arena

mmap is a system call that allows us to request memory pages from the OS, and will be explained in more detail later.

We also need to know the arena’s boundaries so we don’t access memory outside of it. For that, we create a top chunk that is always at the end of the arena and has a size of 0. This way, we can verify that we don’t go past the top chunk.

One question that some of you could have is: what happens when the initial 2MB arena gets filled up? When this happens, we can manually request more memory from the OS using mmap and add it to the free list.

Efficiently finding free chunks

Having a free list is great, but keeping all free chunks in a single list means searching for a free block of size N becomes O(n).

To solve this, we can use segregated bins divided into size ranges. For example, we can have a bin for chunks of size 16-32, 32-64, 64-128, etc. This way, when we need to allocate a chunk of size N, we can go directly to the bin for that size class and find a free chunk fast.

This is the final heapinfo struct:

typedef struct heapinfo
{
    heapchunk *bins[NUM_BINS];
    bool initialized;
    size_t avail;

} heapinfo;

Splitting, Coalescing, and Chunk Fragmentation

When the user wants to allocate a 32 byte chunk, giving him the entire 2MB arena is like giving him a whole Ferrari when he just wanted new wheels.

Obviously, we want to split the chunk into a 32 byte chunk, give that to the user, and keep the rest of the memory in the free list for future allocations. This is called splitting. Splitting should be done smartly, so we don’t end up with a lot of small chunks that are too small to be useful.

if (avail_chunk->size < requested_size + HEADER_SIZE + MIN_CHUNK_SIZE)
{
return; // not enough space to split
}

We only split a chunk if the remaining space is enough to hold a new chunk with its header and a minimum payload size.

The problem occurs when a lot of small chunks are left in the free list. Imagine that the user allocates the entire arena with 32 byte chunks. When freeing them, the free list will be filled with 32 byte chunks. If the user now wants to allocate a 64 byte chunk, there won’t be any free chunks big enough to satisfy the request, even though there is enough total free memory. Fragmentation

This problem is called fragmentation. To solve it, we can do the opposite of splitting: coalescing. When a chunk is freed, we can check if its neighbour chunks are free, and if they are, we can merge them into a single larger chunk.

Keep in mind that “neighbour” means physically in memory, not in the free list.

One problem that arises is that we need to find the neighbour chunks in memory. To find the right neighbour, all we need is to add the chunk’s size to its payload address.

Finding the left neighbour is a bit trickier, because we don’t know the previous chunk’s size. We can solve this in some ways, like adding prev_size to the chunk’s header, or adding a footer to the chunk, but i decided to ignore this problem for now and just coalesce with the right neighbour.

Talking to the OS: mmap, page faults and madvise

As we said before, when we initialize the heap we request 2MB of memory from the OS using mmap. How mmap works is that the OS doesn’t actually give us the memory.

The mmap syscall asks the kernel to reserve a range of virtual memory addresses for our process, but it doesnt actually allocate physical memory for it.

When we access a page of memory for the first time, the kernel will raise a page fault, and only then will it allocate physical memory for that page. This is called demand paging. This is a great optimization, because it means that if we never access a page, the kernel will never allocate physical memory for it. Virtual memory illustration

picture from https://www.geeksforgeeks.org/

Madvise Syscall

When we free a large chunk, we don’t currently need the memory pages that fit inside of it. We can tell the kernel that we don’t need the physical memory for that chunk anymore, so that it can be reused for other processes. This is done using the madvise syscall.

When we call madvise with the MADV_FREE flag, the kernel will mark the pages as free, and if another process needs memory, it can reuse those pages. If we access those pages again, the kernel will raise a page fault and allocate physical memory for them again.

Madvise is still a system call, so it is expensive. We don’t want to call it on every free, so we only call it when we free a large chunk that multiple pages fit inside of it. This way we optimize our memory usage and avoid wasting physical memory for chunks that are not currently in use.

Realloc: changing the size of an allocated chunk

Reallocating an allocated chunk is a bit trickier than regular allocation. Yes, we can just allocate a new chunk, copy the data over and free the old chunk, but this is not very efficient. So we can separate the process into 3 cases:

Shrinking

When the requested size is smaller than the current, we can just shrink the chunk’s size and split the remaining space into a new free chunk (if enough space is left). This is the easiest case. shrinking a 48 byte chunk to 16 bytes

Growing in place

When the requested size is larger than the current, we can check if the right neighbour chunk is free and has enough space. If it does, we can merge the two and grow the chunk in place. This is the most efficient case, because we don’t need to allocate a new chunk or copy any data. growing a 16 byte chunk in place

Growing with a new allocation

Sometimes the right neighbour chunk is not free or doesn’t have enough space. In this case, we can resort to allocating a new chunk, copying the data over and freeing the old chunk.

Defending Against Heap Corruption

At this point of the project, I was pretty done with the basic functions of the allocator. But I decided to take it a step further and add some security features to defend against heap corruption.

Heap Canaries

Most heap corruptions involve overwriting the metadata of a chunk, which can lead to arbitrary code execution.

One feature i decided to add after researching are heap canaries. A heap canary is a random value that is stored in the chunk’s header. When the chunk is freed, the allocator checks if the canary value has been modified. If it has, it means that the chunk has been corrupted and the program will abort.

A good heap canary should be a unique value that is hard to guess or leak. For this, I decided to generate a random value at the start of the program, then XOR it with the chunk’s address to get a unique canary for each chunk. This way, even if an attacker knows the canary value of one chunk, they won’t be able to calculate the canary value of another.

void init_canary() // called at the start of the program
{
    FILE *urandom = fopen("/dev/urandom", "r");
    if (urandom != NULL)
    {
        fread(&global_cookie, sizeof(size_t), 1, urandom);
        fclose(urandom);
    }
}

inline size_t calculate_canary(heapchunk *chunk)
{
    return global_cookie ^ (size_t)chunk;
}

The updated chunk structure with the canary looks like this:

typedef struct heapchunk
{
    size_t canary;
    size_t size;
    ...
} heapchunk;

Now when you try to overwrite the metadata of a chunk with an overflow, the canary will be modified and the program will abort.

Safe Unlinking

Heap canaries are great, but can only be detected at free(), and only work if the attacker has an overflow.

Another common heap corruption technique is to overwrite the free list pointers of a chunk, and then free it. This will cause the allocator to write to an arbitrary address, and with manipulation can lead again to arbitrary code execution.

There are great pwn challenges that exploit this technique, like “unlink” from pwnable.kr, and others from the pwn.college yellow belt.

A common way to defend against this is to use safe unlinking. This technique checks if the free list pointers of a chunk are valid before unlinking it from the free list. If they are not valid, the program will abort.

// disconnecting the previous chunk
if (chunk->list.prev->list.next != chunk)
{
    fprintf(stderr, "detected a corrupted doubly-linked list, abort.\n");
    abort();
}

// disconnecting the next chunk
if (chunk->list.next->list.prev != chunk)
{
    fprintf(stderr, "detected a corrupted doubly-linked list, abort.\n");
    abort();
}

Stopping Use-After-Free bugs and exploits

UAF is a common vulnerability that occurs when a program continues to use a pointer after freeing the memory. This can lead to arbitrary code execution if the freed memory is reallocated and used in a way that allows an attacker to control its contents.

To try and prevent this programming error as much as possible, I decided to implement “heap poisoning”. When a chunk is freed, we can overwrite its payload with a dummy value (0xDEADBEEF), so that if the program tries to use it after freeing, it will hopefully crash.

This is not a perfect solution, but it can help catch some UAF bugs and exploits. UAF is a hard problem to solve because it is more of a programming error.

Conclusion

Building a memory allocator from scratch was a great learning experience, which I highly recommend to anyone interested in low-level programming, operating systems, or security. It gave me a deeper understanding of how memory management works, and the challenges involved in protecting against heap corruption and exploitation.

You can review the complete code, suggest improvements, or try to exploit it yourself over on GitHub. Hope you enjoyed reading this post, and if you have any questions or suggestions, feel free to reach out via the contact form.