zlib can use m_malloc/m_free too

--HG--
branch : fuzz

3 files changed, 15 insertions, 7 deletions

diff --git a/common-kex.c b/common-kex.c
index a95182c..1f3d51b 100644
--- a/common-kex.c
+++ b/common-kex.c
@@ -391,6 +391,14 @@ int is_compress_recv() {
 			&& ses.keys->recv.algo_comp == DROPBEAR_COMP_ZLIB_DELAY);
 }
 
+static void* dropbear_zalloc(void* UNUSED(opaque), uInt items, uInt size) {
+	return m_calloc(items, size);
+}
+
+static void dropbear_zfree(void* UNUSED(opaque), void* ptr) {
+	m_free(ptr);
+}
+
 /* Set up new zlib compression streams, close the old ones. Only
  * called from gen_new_keys() */
 static void gen_new_zstream_recv() {
@@ -399,11 +407,10 @@ static void gen_new_zstream_recv() {
 	if (ses.newkeys->recv.algo_comp == DROPBEAR_COMP_ZLIB
 			|| ses.newkeys->recv.algo_comp == DROPBEAR_COMP_ZLIB_DELAY) {
 		ses.newkeys->recv.zstream = (z_streamp)m_malloc(sizeof(z_stream));
-		ses.newkeys->recv.zstream->zalloc = Z_NULL;
-		ses.newkeys->recv.zstream->zfree = Z_NULL;
+		ses.newkeys->recv.zstream->zalloc = dropbear_zalloc;
+		ses.newkeys->recv.zstream->zfree = dropbear_zfree;
 		
 		if (inflateInit(ses.newkeys->recv.zstream) != Z_OK) {
-			m_free(ses.newkeys->recv.zstream);
 			dropbear_exit("zlib error");
 		}
 	} else {
@@ -424,8 +431,8 @@ static void gen_new_zstream_trans() {
 	if (ses.newkeys->trans.algo_comp == DROPBEAR_COMP_ZLIB
 			|| ses.newkeys->trans.algo_comp == DROPBEAR_COMP_ZLIB_DELAY) {
 		ses.newkeys->trans.zstream = (z_streamp)m_malloc(sizeof(z_stream));
-		ses.newkeys->trans.zstream->zalloc = Z_NULL;
-		ses.newkeys->trans.zstream->zfree = Z_NULL;
+		ses.newkeys->trans.zstream->zalloc = dropbear_zalloc;
+		ses.newkeys->trans.zstream->zfree = dropbear_zfree;
 	
 		if (deflateInit2(ses.newkeys->trans.zstream, Z_DEFAULT_COMPRESSION,
 					Z_DEFLATED, DROPBEAR_ZLIB_WINDOW_BITS, 
diff --git a/dbmalloc.c b/dbmalloc.c
index e62e020..0542bf9 100644
--- a/dbmalloc.c
+++ b/dbmalloc.c
@@ -77,7 +77,9 @@ void * m_malloc(size_t size) {
 }
 
 void * m_calloc(size_t nmemb, size_t size) {
-    assert(nmemb <= 1000 && size <= 10000);
+    if (SIZE_T_MAX / nmemb < size) {
+        dropbear_exit("m_calloc failed");
+    }
     return m_malloc(nmemb*size);
 }
 
diff --git a/dbmalloc.h b/dbmalloc.h
index f6cccbb..f27ab39 100644
--- a/dbmalloc.h
+++ b/dbmalloc.h
@@ -4,7 +4,6 @@
 #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);