examples/jpeg: allow to specify maximum limit for total number of pixels

1 files changed, 15 insertions, 10 deletions

diff --git a/examples/libjpeg/persistent-jpeg.c b/examples/libjpeg/persistent-jpeg.c
index 49952242..076d2600 100644
--- a/examples/libjpeg/persistent-jpeg.c
+++ b/examples/libjpeg/persistent-jpeg.c
@@ -26,19 +26,18 @@ struct jpegErrorManager {
 
 struct jpegErrorManager jerr;
 
-void jpegErrorExit(j_common_ptr cinfo)
-{
+void jpegErrorExit(j_common_ptr cinfo) {
     struct jpegErrorManager* myerr = (struct jpegErrorManager*)cinfo->err;
     longjmp(myerr->setjmp_buffer, 1);
 }
 
 static const char* const cdjpeg_message_table[] = {
 #include "cderror.h"
-    NULL
+	NULL
 };
 
-int LLVMFuzzerInitialize(int* argc, char*** argv)
-{
+static uint64_t max_total_pixels = 100000000ULL; /* 100M */
+int LLVMFuzzerInitialize(int* argc, char*** argv) {
     null_fd = open("/dev/null", O_WRONLY);
 
     cinfo.err = jpeg_std_error(&jerr.pub);
@@ -49,11 +48,15 @@ int LLVMFuzzerInitialize(int* argc, char*** argv)
     jerr.pub.last_addon_message = JMSG_LASTADDONCODE;
 
     jpeg_create_decompress(&cinfo);
+
+	/* If there are any arguments provided, limit width*heigth to this value */
+    if (*argc > 1) {
+        max_total_pixels = strtoull((*argv)[1], NULL, 0);
+    }
     return 0;
 }
 
-int LLVMFuzzerTestOneInput(const uint8_t* buf, size_t len)
-{
+int LLVMFuzzerTestOneInput(const uint8_t* buf, size_t len) {
     if (setjmp(jerr.setjmp_buffer)) {
         goto out;
     }
@@ -64,12 +67,14 @@ int LLVMFuzzerTestOneInput(const uint8_t* buf, size_t len)
         goto out;
     }
 
-    if (cinfo.output_height > 10000 || cinfo.output_width > 10000) {
+    /* Limit total number of pixels to decode to 50M */
+    uint64_t total_pix = (uint64_t)cinfo.output_height * (uint32_t)cinfo.output_width;
+    if (total_pix > max_total_pixels) {
         goto out;
     }
 
-    cinfo.mem->max_memory_to_use = (1024 * 1024 * 1024);
-    cinfo.mem->max_alloc_chunk = (1024 * 128 * 256);
+    cinfo.mem->max_memory_to_use = (1024ULL * 1024ULL * 1024ULL);
+    cinfo.mem->max_alloc_chunk = (1024ULL * 1024ULL * 1024ULL);
 
     jpeg_start_decompress(&cinfo);