1 files changed, 182 insertions, 0 deletions
diff --git a/dbmalloc.c b/dbmalloc.c
new file mode 100644
index 0000000..8c6701b
--- /dev/null
+++ b/dbmalloc.c
@@ -0,0 +1,182 @@
+#include "dbmalloc.h"
+#include "dbutil.h"
+
+
+void * m_calloc(size_t nmemb, size_t size) {
+    if (SIZE_T_MAX / nmemb < size) {
+        dropbear_exit("m_calloc failed");
+    }
+    return m_malloc(nmemb*size);
+}
+
+void * m_strdup(const char * str) {
+    char* ret;
+    unsigned int len;
+    len = strlen(str);
+
+    ret = m_malloc(len+1);
+    if (ret == NULL) {
+        dropbear_exit("m_strdup failed");
+    }
+    memcpy(ret, str, len+1);
+    return ret;
+}
+
+#if !DROPBEAR_TRACKING_MALLOC
+
+/* Simple wrappers around malloc etc */
+void * m_malloc(size_t size) {
+
+	void* ret;
+
+	if (size == 0) {
+		dropbear_exit("m_malloc failed");
+	}
+	ret = calloc(1, size);
+	if (ret == NULL) {
+		dropbear_exit("m_malloc failed");
+	}
+	return ret;
+
+}
+
+void * m_realloc(void* ptr, size_t size) {
+
+	void *ret;
+
+	if (size == 0) {
+		dropbear_exit("m_realloc failed");
+	}
+	ret = realloc(ptr, size);
+	if (ret == NULL) {
+		dropbear_exit("m_realloc failed");
+	}
+	return ret;
+}
+
+
+#else
+
+/* For fuzzing */
+
+struct dbmalloc_header {
+    unsigned int epoch;
+    struct dbmalloc_header *prev;
+    struct dbmalloc_header *next;
+};
+
+static void put_alloc(struct dbmalloc_header *header);
+static void remove_alloc(struct dbmalloc_header *header);
+
+/* end of the linked list */
+static struct dbmalloc_header* staple;
+
+unsigned int current_epoch = 0;
+
+void m_malloc_set_epoch(unsigned int epoch) {
+    current_epoch = epoch;
+}
+
+void m_malloc_free_epoch(unsigned int epoch, int dofree) {
+    struct dbmalloc_header* header;
+    struct dbmalloc_header* nextheader = NULL;
+    struct dbmalloc_header* oldstaple = staple;
+    staple = NULL;
+    /* free allocations from this epoch, create a new staple-anchored list from
+    the remainder */
+    for (header = oldstaple; header; header = nextheader)
+    {
+        nextheader = header->next;
+        if (header->epoch == epoch) {
+            if (dofree) {
+                free(header);
+            }
+        } else {
+            header->prev = NULL;
+            header->next = NULL;
+            put_alloc(header);
+        }
+    }
+}
+
+static void put_alloc(struct dbmalloc_header *header) {
+    assert(header->next == NULL);
+    assert(header->prev == NULL);
+    if (staple) {
+        staple->prev = header;
+    }
+    header->next = staple;
+    staple = header;
+}
+
+static void remove_alloc(struct dbmalloc_header *header) {
+    if (header->prev) {
+        header->prev->next = header->next;
+    }
+    if (header->next) {
+        header->next->prev = header->prev;
+    }
+    if (staple == header) {
+        staple = header->next;
+    }
+    header->prev = NULL;
+    header->next = NULL;
+}
+
+static struct dbmalloc_header* get_header(void* ptr) {
+    char* bptr = ptr;
+    return (struct dbmalloc_header*)&bptr[-sizeof(struct dbmalloc_header)];
+}
+
+void * m_malloc(size_t size) {
+    char* mem = NULL;
+    struct dbmalloc_header* header = NULL;
+
+    if (size == 0 || size > 1e9) {
+        dropbear_exit("m_malloc failed");
+    }
+
+    size = size + sizeof(struct dbmalloc_header);
+
+    mem = calloc(1, size);
+    if (mem == NULL) {
+        dropbear_exit("m_malloc failed");
+    }
+    header = (struct dbmalloc_header*)mem;
+    put_alloc(header);
+    header->epoch = current_epoch;
+    return &mem[sizeof(struct dbmalloc_header)];
+}
+
+void * m_realloc(void* ptr, size_t size) {
+    char* mem = NULL;
+    struct dbmalloc_header* header = NULL;
+    if (size == 0 || size > 1e9) {
+        dropbear_exit("m_realloc failed");
+    }
+
+    header = get_header(ptr);
+    remove_alloc(header);
+
+    size = size + sizeof(struct dbmalloc_header);
+    mem = realloc(header, size);
+    if (mem == NULL) {
+        dropbear_exit("m_realloc failed");
+    }
+
+    header = (struct dbmalloc_header*)mem;
+    put_alloc(header);
+    return &mem[sizeof(struct dbmalloc_header)];
+}
+
+void m_free_direct(void* ptr) {
+    struct dbmalloc_header* header = NULL;
+    if (!ptr) {
+        return;
+    }
+    header = get_header(ptr);
+    remove_alloc(header);
+    free(header);
+}
+
+#endif /* DROPBEAR_TRACKING_MALLOC */