4 files changed, 80 insertions, 0 deletions
diff --git a/tests/common/SConscript b/tests/common/SConscript
index f1dee0e2..4581bea1 100644
--- a/tests/common/SConscript
+++ b/tests/common/SConscript
@@ -15,3 +15,7 @@ strict.Append(CFLAGS = strict['CORECFLAGS'])
 strict.Object("pb_decode.o", "$NANOPB/pb_decode.c")
 strict.Object("pb_encode.o", "$NANOPB/pb_encode.c")
 strict.Object("pb_common.o", "$NANOPB/pb_common.c")
+
+mw = env.Object("malloc_wrappers.o", "malloc_wrappers.c")
+Depends(mw, ["malloc_wrappers_syshdr.h"])
+
diff --git a/tests/common/malloc_wrappers.c b/tests/common/malloc_wrappers.c
new file mode 100644
index 00000000..ad69f1ce
--- /dev/null
+++ b/tests/common/malloc_wrappers.c
@@ -0,0 +1,54 @@
+#include "malloc_wrappers.h"
+#include <stdint.h>
+#include <assert.h>
+#include <string.h>
+
+static size_t alloc_count = 0;
+
+/* Allocate memory and place check values before and after. */
+void* malloc_with_check(size_t size)
+{
+    size_t size32 = (size + 3) / 4 + 3;
+    uint32_t *buf = malloc(size32 * sizeof(uint32_t));
+    buf[0] = size32;
+    buf[1] = 0xDEADBEEF;
+    buf[size32 - 1] = 0xBADBAD;
+    return buf + 2;
+}
+
+/* Free memory allocated with malloc_with_check() and do the checks. */
+void free_with_check(void *mem)
+{
+    uint32_t *buf = (uint32_t*)mem - 2;
+    assert(buf[1] == 0xDEADBEEF);
+    assert(buf[buf[0] - 1] == 0xBADBAD);
+    free(buf);
+}
+
+/* Track memory usage */
+void* counting_realloc(void *ptr, size_t size)
+{
+    /* Don't allocate crazy amounts of RAM when fuzzing */
+    if (size > 1000000)
+        return NULL;
+
+    if (!ptr && size)
+        alloc_count++;
+    
+    return realloc(ptr, size);
+}
+
+void counting_free(void *ptr)
+{
+    if (ptr)
+    {
+        assert(alloc_count > 0);
+        alloc_count--;
+        free(ptr);
+    }
+}
+
+size_t get_alloc_count()
+{
+    return alloc_count;
+}
diff --git a/tests/common/malloc_wrappers.h b/tests/common/malloc_wrappers.h
new file mode 100644
index 00000000..7eec7952
--- /dev/null
+++ b/tests/common/malloc_wrappers.h
@@ -0,0 +1,7 @@
+#include <stdlib.h>
+
+void* malloc_with_check(size_t size);
+void free_with_check(void *mem);
+void* counting_realloc(void *ptr, size_t size);
+void counting_free(void *ptr);
+size_t get_alloc_count();
diff --git a/tests/common/malloc_wrappers_syshdr.h b/tests/common/malloc_wrappers_syshdr.h
new file mode 100644
index 00000000..d295d9ed
--- /dev/null
+++ b/tests/common/malloc_wrappers_syshdr.h
@@ -0,0 +1,15 @@
+/* This is just a wrapper in order to get our own malloc wrappers into nanopb core. */
+
+#define pb_realloc(ptr,size) counting_realloc(ptr,size)
+#define pb_free(ptr) counting_free(ptr)
+
+#ifdef PB_OLD_SYSHDR
+#include PB_OLD_SYSHDR
+#else
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+#include <string.h>
+#endif
+
+#include <malloc_wrappers.h>