From d4dc797e4a65871134f6a244cb8421a33fe91360 Mon Sep 17 00:00:00 2001 From: Sebastian Rasmussen Date: Sun, 15 Sep 2019 15:52:26 +0200 Subject: jbig2dec: Add command line option to limit memory usage. --- jbig2dec/jbig2dec.c | 159 +++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 152 insertions(+), 7 deletions(-) (limited to 'jbig2dec') 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", ¶ms->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 memory limit expressed in bytes\n" " -o \n" " --output send decoded output to \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, ¶ms); @@ -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; -- cgit v1.2.1