jbig2dec: Add command line option to limit memory usage.

1 files changed, 152 insertions, 7 deletions

diff --git a/jbig2dec/jbig2dec.c b/jbig2dec/jbig2dec.c
index 6763421ad..dad786f3c 100644
--- a/jbig2dec/jbig2dec.c
+++ b/jbig2dec/jbig2dec.c
@@ -65,6 +65,7 @@ typedef struct {
     SHA1_CTX *hash_ctx;
     char *output_filename;
     jbig2dec_format output_format;
+    size_t memory_limit;
 } jbig2dec_params_t;
 
 typedef struct {
@@ -75,9 +76,118 @@ typedef struct {
     long repeats;
 } jbig2dec_error_callback_state_t;
 
+typedef struct {
+    Jbig2Allocator super;
+    Jbig2Ctx *ctx;
+    size_t memory_limit;
+    size_t memory_used;
+    size_t memory_peak;
+} jbig2dec_allocator_t;
+
 static int print_version(void);
 static int print_usage(void);
 
+#define ALIGNMENT 16
+#define KBYTE 1024
+#define MBYTE (1024 * KBYTE)
+
+static void *jbig2dec_alloc(Jbig2Allocator *allocator_, size_t size)
+{
+    jbig2dec_allocator_t *allocator = (jbig2dec_allocator_t *) allocator_;
+    void *ptr;
+
+    if (size == 0)
+        return NULL;
+    if (size > allocator->memory_limit - ALIGNMENT - allocator->memory_used)
+        return NULL;
+
+    ptr = malloc(size + ALIGNMENT);
+    if (ptr == NULL)
+        return NULL;
+    memcpy(ptr, &size, sizeof(size));
+    allocator->memory_used += size + ALIGNMENT;
+
+    if (allocator->memory_used > allocator->memory_peak) {
+        allocator->memory_peak = allocator->memory_used;
+
+        if (allocator->ctx) {
+            size_t limit_mb = allocator->memory_limit / MBYTE;
+            size_t peak_mb = allocator->memory_peak / MBYTE;
+            jbig2_error(allocator->ctx, JBIG2_SEVERITY_DEBUG, -1, "memory: limit: %lu Mbyte peak usage: %lu Mbyte", limit_mb, peak_mb);
+        }
+    }
+
+    return (unsigned char *) ptr + ALIGNMENT;
+}
+
+static void jbig2dec_free(Jbig2Allocator *allocator_, void *p)
+{
+    jbig2dec_allocator_t *allocator = (jbig2dec_allocator_t *) allocator_;
+    size_t size;
+
+    if (p == NULL)
+        return;
+
+    memcpy(&size, (unsigned char *) p - ALIGNMENT, sizeof(size));
+    allocator->memory_used -= size + ALIGNMENT;
+    free((unsigned char *) p - ALIGNMENT);
+}
+
+static void *jbig2dec_realloc(Jbig2Allocator *allocator_, void *p, size_t size)
+{
+    jbig2dec_allocator_t *allocator = (jbig2dec_allocator_t *) allocator_;
+    unsigned char *oldp = p ? (unsigned char *) p - ALIGNMENT : NULL;
+
+    if (size > SIZE_MAX - ALIGNMENT)
+        return NULL;
+
+    if (oldp == NULL)
+    {
+        if (size == 0)
+            return NULL;
+        if (size > allocator->memory_limit - ALIGNMENT - allocator->memory_used)
+            return NULL;
+
+        p = malloc(size + ALIGNMENT);
+    }
+    else
+    {
+        size_t oldsize;
+        memcpy(&oldsize, oldp, sizeof(oldsize));
+
+        if (size == 0)
+        {
+            allocator->memory_used -= oldsize + ALIGNMENT;
+            free(oldp);
+            return NULL;
+        }
+
+        if (size > allocator->memory_limit - allocator->memory_used + oldsize)
+            return NULL;
+
+        p = realloc(oldp, size + ALIGNMENT);
+        if (p == NULL)
+            return NULL;
+
+        allocator->memory_used -= oldsize + ALIGNMENT;
+    }
+
+    memcpy(p, &size, sizeof(size));
+    allocator->memory_used += size + ALIGNMENT;
+
+    if (allocator->memory_used > allocator->memory_peak) {
+        allocator->memory_peak = allocator->memory_used;
+
+        if (allocator->ctx) {
+            size_t limit_mb = allocator->memory_limit / MBYTE;
+            size_t peak_mb = allocator->memory_peak / MBYTE;
+            jbig2_error(allocator->ctx, JBIG2_SEVERITY_DEBUG, -1, "memory: limit: %lu Mbyte peak usage: %lu Mbyte", limit_mb, peak_mb);
+        }
+    }
+
+    return (unsigned char *) p + ALIGNMENT;
+}
+
 /* page hashing functions */
 static void
 hash_init(jbig2dec_params_t *params)
@@ -155,9 +265,10 @@ parse_options(int argc, char *argv[], jbig2dec_params_t *params)
     };
     int option_idx = 1;
     int option;
+    int ret;
 
     while (1) {
-        option = getopt_long(argc, argv, "Vh?qv:do:t:e", long_options, &option_idx);
+        option = getopt_long(argc, argv, "Vh?qv:do:t:eM:", long_options, &option_idx);
         if (option == -1)
             break;
 
@@ -200,9 +311,14 @@ parse_options(int argc, char *argv[], jbig2dec_params_t *params)
         case 'e':
             params->embedded = 1;
             break;
+        case 'M':
+            ret = sscanf(optarg, "%zu", &params->memory_limit);
+            if (ret != 1)
+                fprintf(stderr, "could not parse memory limit argument\n");
+            break;
         default:
             if (!params->verbose)
-                fprintf(stdout, "unrecognized option: -%c\n", option);
+                fprintf(stderr, "unrecognized option: -%c\n", option);
             break;
         }
     }
@@ -238,6 +354,7 @@ print_usage(void)
             "    -V --version    program name and version information\n"
             "    -m --hash       print a hash of the decoded document\n"
             "    -e --embedded   expect embedded bit stream without file header\n"
+            "    -M <limit>      memory limit expressed in bytes\n"
             "    -o <file>\n"
             "    --output <file> send decoded output to <file>\n"
             "                    Defaults to the the input with a different\n"
@@ -428,11 +545,13 @@ write_document_hash(jbig2dec_params_t *params)
 int
 main(int argc, char **argv)
 {
-    FILE *f = NULL, *f_page = NULL;
-    Jbig2Ctx *ctx = NULL;
-    uint8_t buf[4096];
     jbig2dec_params_t params;
     jbig2dec_error_callback_state_t error_callback_state;
+    jbig2dec_allocator_t allocator_;
+    jbig2dec_allocator_t *allocator = &allocator_;
+    Jbig2Ctx *ctx = NULL;
+    FILE *f = NULL, *f_page = NULL;
+    uint8_t buf[4096];
     int filearg;
     int result = 1;
     int code;
@@ -444,6 +563,7 @@ main(int argc, char **argv)
     params.output_filename = NULL;
     params.output_format = jbig2dec_format_none;
     params.embedded = 0;
+    params.memory_limit = 0;
 
     filearg = parse_options(argc, argv, &params);
 
@@ -496,7 +616,20 @@ main(int argc, char **argv)
         error_callback_state.last_message = NULL;
         error_callback_state.repeats = 0;
 
-        ctx = jbig2_ctx_new(NULL, (Jbig2Options)(f_page != NULL || params.embedded ? JBIG2_OPTIONS_EMBEDDED : 0), NULL, error_callback, &error_callback_state);
+        if (params.memory_limit == 0)
+            allocator = NULL;
+        else
+        {
+            allocator->super.alloc = jbig2dec_alloc;
+            allocator->super.free = jbig2dec_free;
+            allocator->super.realloc = jbig2dec_realloc;
+            allocator->ctx = NULL;
+            allocator->memory_limit = params.memory_limit;
+            allocator->memory_used = 0;
+            allocator->memory_peak = 0;
+        }
+
+        ctx = jbig2_ctx_new((Jbig2Allocator *) allocator, (Jbig2Options)(f_page != NULL || params.embedded ? JBIG2_OPTIONS_EMBEDDED : 0), NULL, error_callback, &error_callback_state);
         if (ctx == NULL) {
             fclose(f);
             if (f_page)
@@ -504,6 +637,9 @@ main(int argc, char **argv)
             goto cleanup;
         }
 
+        if (allocator != NULL)
+            allocator->ctx = ctx;
+
         /* pull the whole file/global stream into memory */
         for (;;) {
             int n_bytes = fread(buf, 1, sizeof(buf), f);
@@ -519,8 +655,11 @@ main(int argc, char **argv)
         if (f_page != NULL) {
             Jbig2GlobalCtx *global_ctx = jbig2_make_global_ctx(ctx);
 
-            ctx = jbig2_ctx_new(NULL, JBIG2_OPTIONS_EMBEDDED, global_ctx, error_callback, &error_callback_state);
+            ctx = jbig2_ctx_new((Jbig2Allocator *) allocator, JBIG2_OPTIONS_EMBEDDED, global_ctx, error_callback, &error_callback_state);
             if (ctx != NULL) {
+                if (allocator != NULL)
+                    allocator->ctx = ctx;
+
                 for (;;) {
                     int n_bytes = fread(buf, 1, sizeof(buf), f_page);
 
@@ -597,6 +736,12 @@ main(int argc, char **argv)
 
     }                           /* end params.mode switch */
 
+    if (allocator != NULL && allocator->ctx != NULL) {
+        size_t limit_mb = allocator->memory_limit / MBYTE;
+        size_t peak_mb = allocator->memory_peak / MBYTE;
+        jbig2_error(allocator->ctx, JBIG2_SEVERITY_DEBUG, -1, "memory: limit: %lu Mbyte peak usage: %lu Mbyte", limit_mb, peak_mb);
+    }
+
     /* fin */
     result = 0;