From 02c1436e9f274258827670b52a5782dfd1e69b9b Mon Sep 17 00:00:00 2001
From: Sebastian Rasmussen <sebras@gmail.com>
Date: Tue, 25 Sep 2018 14:38:08 +0800
Subject: [PATCH] [mupdf] Add custom allocator to avoid having fuzzer kill
 process (#1830) (#1832)

This fixes oss-fuzz #5679 and oss-fuzz #7803 for the mupdf project.
---
 projects/mupdf/pdf_fuzzer.cc | 90 +++++++++++++++++++++++++++++++++++-
 1 file changed, 89 insertions(+), 1 deletion(-)

diff --git a/projects/mupdf/pdf_fuzzer.cc b/projects/mupdf/pdf_fuzzer.cc
index 05c87f3ed..fd8ad7faf 100644
--- a/projects/mupdf/pdf_fuzzer.cc
+++ b/projects/mupdf/pdf_fuzzer.cc
@@ -17,11 +17,98 @@
 */
 
 #include <cstdint>
+#include <stdlib.h>
+#include <string.h>
+#include <inttypes.h>
 
 #include <mupdf/fitz.h>
 
+#define ALIGNMENT 16
+#define MAX_ALLOCATION (1024 * 1024 * 1024)
+
+static uint64_t total = 0;
+
+static void *
+fz_malloc_ossfuzz(void *opaque, size_t size)
+{
+	char *ptr = NULL;
+
+	if (size == 0)
+		return NULL;
+	if (size > SIZE_MAX - ALIGNMENT)
+		return NULL;
+
+	if (size > MAX_ALLOCATION - ALIGNMENT - total)
+		return NULL;
+
+	ptr = (char *) malloc(size + ALIGNMENT);
+	if (ptr == NULL)
+		return NULL;
+
+	memcpy(ptr, &size, sizeof(size));
+	total += size + ALIGNMENT;
+
+	return ptr + ALIGNMENT;
+}
+
+static void
+fz_free_ossfuzz(void *opaque, void *ptr)
+{
+	size_t size;
+
+	if (ptr == NULL)
+		return;
+
+	ptr = ((char *) ptr) - ALIGNMENT;
+
+	memcpy(&size, ptr, sizeof(size));
+	total -= size - ALIGNMENT;
+	free(ptr);
+}
+
+static void *
+fz_realloc_ossfuzz(void *opaque, void *old, size_t size)
+{
+	size_t oldsize;
+	char *ptr;
+
+	if (old == NULL)
+		return fz_malloc_ossfuzz(opaque, size);
+	if (size == 0)
+	{
+		fz_free_ossfuzz(opaque, old);
+		return NULL;
+	}
+	if (size > SIZE_MAX - ALIGNMENT)
+		return NULL;
+
+	old = ((char *) old) - ALIGNMENT;
+	memcpy(&oldsize, old, sizeof(oldsize));
+
+	if (size > MAX_ALLOCATION - total + oldsize)
+		return NULL;
+
+	ptr = (char *) realloc(old, size + ALIGNMENT);
+	if (ptr == NULL)
+		return NULL;
+
+	total -= oldsize + ALIGNMENT;
+	memcpy(ptr, &size, sizeof(size));
+	total += size + ALIGNMENT;
+
+	return ptr + ALIGNMENT;
+}
+
+static fz_alloc_context fz_alloc_ossfuzz =
+{
+	NULL,
+	fz_malloc_ossfuzz,
+	fz_realloc_ossfuzz,
+	fz_free_ossfuzz
+};
+
 extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
-  fz_context *ctx = fz_new_context(nullptr, nullptr, FZ_STORE_DEFAULT);
+  fz_context *ctx = fz_new_context(&fz_alloc_ossfuzz, nullptr, FZ_STORE_DEFAULT);
 
   fz_stream *stream = NULL;
   fz_document *doc = NULL;
@@ -35,6 +122,7 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
     fz_register_document_handlers(ctx);
     stream = fz_open_memory(ctx, data, size);
     doc = fz_open_document_with_stream(ctx, "pdf", stream);
+
     for (int i = 0; i < fz_count_pages(ctx, doc); i++) {
       pix = fz_new_pixmap_from_page_number(ctx, doc, i, fz_identity, fz_device_rgb(ctx), 0);
       fz_drop_pixmap(ctx, pix);