[jbig2dec] Fixes for bugs in the custom allocator (#3409)

* [jbig2dec] Fix issues in library glue logic.

This addresses the out of memory condition exhibited in
project-jbig2dec bug #17168.

* [jbig2dec] Print peak memory usage.

This would be helpful when we're approaching OOM conditions.

1 files changed, 29 insertions, 4 deletions

diff --git a/projects/jbig2dec/jbig2_fuzzer.cc b/projects/jbig2dec/jbig2_fuzzer.cc
index fda32f67..c4a713d5 100644
--- a/projects/jbig2dec/jbig2_fuzzer.cc
+++ b/projects/jbig2dec/jbig2_fuzzer.cc
@@ -24,9 +24,12 @@
 #include "jbig2.h"
 
 #define ALIGNMENT 16
-#define MAX_ALLOCATION (1024 * 1024 * 1024)
+#define MBYTE (1024 * 1024)
+#define GBYTE (1024 * MBYTE)
+#define MAX_ALLOCATION (3 * GBYTE)
 
 static uint64_t total = 0;
+static uint64_t peak = 0;
 
 static void *jbig2_alloc(Jbig2Allocator *allocator, size_t size)
 {
@@ -34,11 +37,21 @@ static void *jbig2_alloc(Jbig2Allocator *allocator, size_t size)
 
   if (size == 0)
     return NULL;
+  if (size > MAX_ALLOCATION - ALIGNMENT - total)
+    return NULL;
 
   ptr = malloc(size + ALIGNMENT);
+  if (ptr == NULL)
+    return NULL;
+
   memcpy(ptr, &size, sizeof(size));
   total += size + ALIGNMENT;
 
+  if (peak == 0 || total / MBYTE > peak / MBYTE) {
+	  peak = total;
+	  fprintf(stderr, "memory: limit: %u Mbyte peak usage: %u Mbyte\n", MAX_ALLOCATION, peak);
+  }
+
   return (unsigned char *) ptr + ALIGNMENT;
 }
 
@@ -61,15 +74,16 @@ static void *jbig2_realloc(Jbig2Allocator *allocator, void *p, size_t size)
   if (size > SIZE_MAX - ALIGNMENT)
     return NULL;
 
-  if (size > MAX_ALLOCATION - ALIGNMENT - total)
-    return NULL;
-
   if (oldp == NULL)
   {
     if (size == 0)
       return NULL;
+    if (size > MAX_ALLOCATION - ALIGNMENT - total)
+      return NULL;
 
     p = malloc(size + ALIGNMENT);
+    if (p == NULL)
+      return NULL;
   }
   else
   {
@@ -83,6 +97,9 @@ static void *jbig2_realloc(Jbig2Allocator *allocator, void *p, size_t size)
       return NULL;
     }
 
+    if (size > MAX_ALLOCATION - total + oldsize)
+      return NULL;
+
     p = realloc(oldp, size + ALIGNMENT);
     if (p == NULL)
       return NULL;
@@ -92,6 +109,12 @@ static void *jbig2_realloc(Jbig2Allocator *allocator, void *p, size_t size)
 
   memcpy(p, &size, sizeof(size));
   total += size + ALIGNMENT;
+
+  if (peak == 0 || total / MBYTE > peak / MBYTE) {
+	  peak = total;
+	  fprintf(stderr, "memory: limit: %u Mbyte peak usage: %u Mbyte\n", MAX_ALLOCATION, peak);
+  }
+
   return (unsigned char *) p + ALIGNMENT;
 }
 
@@ -121,5 +144,7 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
   }
   jbig2_ctx_free(ctx);
 
+  fprintf(stderr, "memory: limit: %u Mbyte peak usage: %u Mbyte\n", MAX_ALLOCATION, peak);
+
   return 0;
 }