From 26ad6853d20f0845697e7477b2224d1d5d1546cb Mon Sep 17 00:00:00 2001
From: Matt Johnston <matt@ucc.asn.au>
Date: Thu, 1 Mar 2018 23:46:24 +0800
Subject: [PATCH] Only use malloc wrapper if fuzzing

---
 dbmalloc.c   | 80 ++++++++++++++++++++++++++++++++++++++--------------
 dbmalloc.h   | 11 +++++++-
 sysoptions.h |  2 ++
 3 files changed, 71 insertions(+), 22 deletions(-)

diff --git a/dbmalloc.c b/dbmalloc.c
index e887d36..8c6701b 100644
--- a/dbmalloc.c
+++ b/dbmalloc.c
@@ -1,6 +1,64 @@
 #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;
@@ -90,13 +148,6 @@ void * m_malloc(size_t size) {
     return &mem[sizeof(struct dbmalloc_header)];
 }
 
-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_realloc(void* ptr, size_t size) {
     char* mem = NULL;
     struct dbmalloc_header* header = NULL;
@@ -128,17 +179,4 @@ void m_free_direct(void* ptr) {
     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;
-}
-
-
+#endif /* DROPBEAR_TRACKING_MALLOC */
diff --git a/dbmalloc.h b/dbmalloc.h
index f05f8f5..a5a13ad 100644
--- a/dbmalloc.h
+++ b/dbmalloc.h
@@ -7,10 +7,19 @@ void * m_malloc(size_t size);
 void * m_calloc(size_t nmemb, size_t size);
 void * m_strdup(const char * str);
 void * m_realloc(void* ptr, size_t size);
+
+#if DROPBEAR_TRACKING_MALLOC
 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, int dofree);
 
+#else
+/* plain wrapper */
+
+#define m_free(X) do {free(X); (X) = NULL;} while (0)
+
+#endif
+
+
 #endif /* DBMALLOC_H_ */
diff --git a/sysoptions.h b/sysoptions.h
index 1666a1e..0028199 100644
--- a/sysoptions.h
+++ b/sysoptions.h
@@ -316,4 +316,6 @@ If you test it please contact the Dropbear author */
 #define DROPBEAR_CLIENT_TCP_FAST_OPEN 0
 #endif
 
+#define DROPBEAR_TRACKING_MALLOC (DROPBEAR_FUZZ)
+
 /* no include guard for this file */