--- a/Makefile.in	Sun May 21 10:48:18 2017 +0800
+++ b/Makefile.in	Sun May 21 10:54:11 2017 +0800
@@ -34,6 +34,7 @@
 		queue.o \
 		atomicio.o compat.o fake-rfc2553.o \
 		ltc_prng.o ecc.o ecdsa.o crypto_desc.o \
+		dbmalloc.o \
 		gensignkey.o gendss.o genrsa.o 
 
 SVROBJS=svr-kex.o svr-auth.o sshpty.o \

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/dbmalloc.c	Sun May 21 10:54:11 2017 +0800
@@ -0,0 +1,128 @@
+#include "dbmalloc.h"
+#include "dbutil.h"
+
+#define LIST_SIZE 1000
+
+struct dbmalloc_header {
+    unsigned int index;
+    unsigned int epoch;
+};
+
+static struct dbmalloc_header* dbmalloc_list[LIST_SIZE];
+
+unsigned int current_epoch = 0;
+
+void m_malloc_set_epoch(unsigned int epoch) {
+    current_epoch = epoch;
+}
+
+void m_malloc_free_epoch(unsigned int epoch) {
+    unsigned int i;
+    unsigned int freed = 0;
+    for (i = 0; i < LIST_SIZE; i++) {
+        if (dbmalloc_list[i] != NULL) {
+            assert(dbmalloc_list[i]->index == i);
+            if (dbmalloc_list[i]->epoch == epoch) {
+                free(dbmalloc_list[i]);
+                dbmalloc_list[i] = NULL;
+                freed++;
+            }
+        }
+    }
+    TRACE(("free_epoch freed %d", freed))
+}
+
+static void put_alloc(struct dbmalloc_header *header) {
+    unsigned int i;
+    for (i = 0; i < LIST_SIZE; i++) {
+        if (dbmalloc_list[i] == NULL) {
+            dbmalloc_list[i] = header;
+            header->index = i;
+            return;
+        }
+    }
+    dropbear_exit("ran out of dbmalloc entries");
+}
+
+static void remove_alloc(struct dbmalloc_header *header) {
+    assert(header->index < LIST_SIZE);
+    assert(dbmalloc_list[header->index] == header);
+    assert(header->epoch == current_epoch);
+    dbmalloc_list[header->index] = 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_calloc(size_t nmemb, size_t size) {
+    assert(nmemb <= 1000 && size <= 10000);
+    return m_malloc(nmemb*size);
+}
+
+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);
+}
+
+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;
+}
+
+

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/dbmalloc.h	Sun May 21 10:54:11 2017 +0800
@@ -0,0 +1,17 @@
+#ifndef DBMALLOC_H_
+#define DBMALLOC_H_
+
+#include "includes.h"
+
+void * m_malloc(size_t size);
+/* m_calloc is limited in size, enough for libtomcrypt */
+void * m_calloc(size_t nmemb, size_t size);
+void * m_strdup(const char * str);
+void * m_realloc(void* ptr, size_t size);
+void m_free_direct(void* ptr);
+#define m_free(X) do {m_free_direct(X); (X) = NULL;} while (0)
+
+void m_malloc_set_epoch(unsigned int epoch);
+void m_malloc_free_epoch(unsigned int epoch);
+
+#endif /* DBMALLOC_H_ */

--- a/dbutil.c	Sun May 21 10:48:18 2017 +0800
+++ b/dbutil.c	Sun May 21 10:54:11 2017 +0800
@@ -520,45 +520,6 @@
 	}
 }
 	
-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_strdup(const char * str) {
-	char* ret;
-
-	ret = strdup(str);
-	if (ret == NULL) {
-		dropbear_exit("m_strdup 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;
-}
-
 void setnonblocking(int fd) {
 
 	TRACE(("setnonblocking: %d", fd))

--- a/dbutil.h	Sun May 21 10:48:18 2017 +0800
+++ b/dbutil.h	Sun May 21 10:54:11 2017 +0800
@@ -30,6 +30,7 @@
 #include "buffer.h"
 #include "queue.h"
 #include "dbhelpers.h"
+#include "dbmalloc.h"
 
 #ifndef DISABLE_SYSLOG
 void startsyslog(const char *ident);
@@ -66,10 +67,6 @@
 int buf_getline(buffer * line, FILE * authfile);
 
 void m_close(int fd);
-void * m_malloc(size_t size);
-void * m_strdup(const char * str);
-void * m_realloc(void* ptr, size_t size);
-#define m_free(X) do {free(X); (X) = NULL;} while (0)
 void setnonblocking(int fd);
 void disallow_core(void);
 int m_str_to_uint(const char* str, unsigned int *val);

--- a/fuzzer-preauth.c	Sun May 21 10:48:18 2017 +0800
+++ b/fuzzer-preauth.c	Sun May 21 10:54:11 2017 +0800
@@ -23,9 +23,11 @@
 	int fakesock = 1;
 	wrapfd_add(fakesock, fuzz.input, PLAIN);
 
+	m_malloc_set_epoch(1);
 	if (setjmp(fuzz.jmp) == 0) {
 		svr_session(fakesock, fakesock);
 	} else {
+		m_malloc_free_epoch(1);
 		TRACE(("dropbear_exit longjmped"))
 		// dropbear_exit jumped here
 	}

--- a/libtomcrypt/src/headers/tomcrypt_custom.h	Sun May 21 10:48:18 2017 +0800
+++ b/libtomcrypt/src/headers/tomcrypt_custom.h	Sun May 21 10:54:11 2017 +0800
@@ -4,6 +4,17 @@
 /* compile options depend on Dropbear options.h */
 #include "options.h"
 
+void * m_malloc(size_t size);
+/* m_calloc is limited in size, enough for libtomcrypt */
+void * m_calloc(size_t nmemb, size_t size);
+void * m_realloc(void* ptr, size_t size);
+void m_free_direct(void* ptr);
+
+#define XMALLOC m_malloc
+#define XFREE m_free_direct
+#define XREALLOC m_realloc
+#define XCALLOC m_calloc
+
 /* macros for various libc functions you can change for embedded targets */
 #ifndef XMALLOC
    #ifdef malloc 

--- a/libtommath/tommath_class.h	Sun May 21 10:48:18 2017 +0800
+++ b/libtommath/tommath_class.h	Sun May 21 10:54:11 2017 +0800
@@ -1000,6 +1000,17 @@
 #undef BN_MP_TOOM_MUL_C
 #undef BN_MP_TOOM_SQR_C
 
+void * m_malloc(size_t size);
+/* m_calloc is limited in size, enough for libtomcrypt */
+void * m_calloc(size_t nmemb, size_t size);
+void * m_realloc(void* ptr, size_t size);
+void m_free_direct(void* ptr);
+
+#define XMALLOC m_malloc
+#define XFREE m_free_direct
+#define XREALLOC m_realloc
+#define XCALLOC m_calloc
+
 /* $Source: /cvs/libtom/libtommath/tommath_class.h,v $ */
 /* $Revision: 1.3 $ */
 /* $Date: 2005/07/28 11:59:32 $ */