/* The wrapper functions in this file work like regular malloc() and free(),

 * but store check values before and after the allocation. This helps to catch

 * any buffer overrun errors in the test cases.

 *

 * This is done to help testing on embedded targets, where regular valgrind

 * and other systems are not necessarily available.

 */

#include "malloc_wrappers.h"

#include <stdint.h>

#include <stdio.h>

#include <assert.h>

#include <string.h>

#define GUARD_SIZE (sizeof(size_t)*3)

#define PREFIX_SIZE (sizeof(size_t)*2)

#define CHECK1 ((size_t)0xDEADBEEF)

#define CHECK2 ((size_t)0x600DCAFE)

#ifndef MAX_ALLOC_BYTES

#define MAX_ALLOC_BYTES 16*1024*1024

#endif

#ifndef DEBUG_MALLOC

#define DEBUG_MALLOC 0

#endif

static size_t g_alloc_count = 0;

static size_t g_alloc_bytes = 0;

static size_t g_max_alloc_bytes = MAX_ALLOC_BYTES;

#ifdef LLVMFUZZER

/* LLVM libsanitizer has a realloc() implementation that always copies

 * the whole memory block, even if there would be space to expand it in

 * place. This gets pretty slow when fuzzing, so this wrapper limits the

 * realloc() calls by rounding allocation size upwards. Real world

 * realloc() implementations are hopefully smarter. */

static size_t round_blocksize(size_t size)

{

    if (size < 256)

    {

        return size;

    }

    else

    {

        return (size + 1023) / 1024 * 1024;

    }

}

#else

static size_t round_blocksize(size_t size)

{

    return size;

}

#endif

/* Allocate memory and place check values before and after. */

void* malloc_with_check(size_t size)

{

    const size_t check2 = CHECK2;

    char *buf = NULL;

    if (size <= g_max_alloc_bytes - g_alloc_bytes)

    {

        buf = malloc(round_blocksize(size + GUARD_SIZE));

    }

    if (buf)

    {

        ((size_t*)buf)[0] = size;

        ((size_t*)buf)[1] = CHECK1;

        memcpy(buf + PREFIX_SIZE + size, &check2, sizeof(check2));

        g_alloc_count++;

        g_alloc_bytes += size;

        if (DEBUG_MALLOC) fprintf(stderr, "Alloc 0x%04x/%u\n", (unsigned)(uintptr_t)(buf + PREFIX_SIZE), (unsigned)size);

        return buf + PREFIX_SIZE;

    }

    else

    {

        if (DEBUG_MALLOC) fprintf(stderr, "malloc(%u) failed\n", (unsigned)size);

        return NULL;

    }

}

/* Free memory allocated with malloc_with_check() and do the checks. */

void free_with_check(void *mem)

{

    if (mem)

    {

        char *buf = (char*)mem - PREFIX_SIZE;

        size_t check2 = 0;

        size_t size = ((size_t*)buf)[0];

        if (DEBUG_MALLOC) fprintf(stderr, "Release 0x%04x/%u\n", (unsigned)(uintptr_t)mem, (unsigned)size);

        assert(((size_t*)buf)[1] == CHECK1);

        memcpy(&check2, buf + PREFIX_SIZE + size, sizeof(check2));

        assert(check2 == CHECK2);

        assert(g_alloc_count > 0);

        assert(g_alloc_bytes >= size);

        ((size_t*)buf)[1] = 0;

        memset(buf + PREFIX_SIZE + size, 0, sizeof(check2));

        g_alloc_count--;

        g_alloc_bytes -= size;

        free(buf);

    }

}

/* Reallocate block and check / write guard values */

void* realloc_with_check(void *ptr, size_t size)

{

    if (!ptr && size)

    {

        /* Allocate new block and write guard values */

        return malloc_with_check(size);

    }

    else if (ptr && size)

    {

        /* Change block size */

        char *buf = (char*)ptr - PREFIX_SIZE;

        size_t check2 = 0;

        size_t oldsize = ((size_t*)buf)[0];

        assert(((size_t*)buf)[1] == CHECK1);

        memcpy(&check2, buf + PREFIX_SIZE + oldsize, sizeof(check2));

        assert(check2 == CHECK2);

        assert(g_alloc_count > 0);

        assert(g_alloc_bytes >= oldsize);

        if (size <= g_max_alloc_bytes - (g_alloc_bytes - oldsize))

        {

            size_t new_rounded = round_blocksize(size + GUARD_SIZE);

            size_t old_rounded = round_blocksize(oldsize + GUARD_SIZE);

            if (new_rounded != old_rounded)

            {

                buf = realloc(buf, new_rounded);

            }

        }

        else

        {

            buf = NULL;

        }

        if (!buf)

        {

            if (DEBUG_MALLOC) fprintf(stderr, "Realloc 0x%04x/%u to %u failed\n", (unsigned)(uintptr_t)ptr, (unsigned)oldsize, (unsigned)size);

            return NULL;

        }

        ((size_t*)buf)[0] = size;

        ((size_t*)buf)[1] = CHECK1;

        memcpy(buf + PREFIX_SIZE + size, &check2, sizeof(check2));

        g_alloc_bytes -= oldsize;

        g_alloc_bytes += size;

        if (DEBUG_MALLOC) fprintf(stderr, "Realloc 0x%04x/%u to 0x%04x/%u\n", (unsigned)(uintptr_t)ptr, (unsigned)oldsize, (unsigned)(uintptr_t)(buf + PREFIX_SIZE), (unsigned)size);

        return buf + PREFIX_SIZE;

    }

    else if (ptr && !size)

    {

        /* Deallocate */

        free_with_check(ptr);

        return NULL;

    }

    else

    {

        /* No action */

        return NULL;

    }

}

/* Return total number of allocations not yet released */

size_t get_alloc_count()

{

    return g_alloc_count;

}

/* Return allocated size for a pointer returned from malloc(). */

size_t get_allocation_size(const void *mem)

{

    char *buf = (char*)mem - PREFIX_SIZE;

    return ((size_t*)buf)[0];

}

/* Get total number of allocated bytes */

size_t get_alloc_bytes()

{

    return g_alloc_bytes;

}

/* Set limit for allocation size */

void set_max_alloc_bytes(size_t max_bytes)

{

    g_max_alloc_bytes = max_bytes;

}

size_t get_max_alloc_bytes()

{

    return g_max_alloc_bytes;

}