More Memsqueezing fixes.

4 files changed, 70 insertions, 27 deletions

diff --git a/fitz/filt_basic.c b/fitz/filt_basic.c
index 3d211645..7e5f2f21 100644
--- a/fitz/filt_basic.c
+++ b/fitz/filt_basic.c
@@ -39,13 +39,24 @@ fz_stream *
 fz_open_null(fz_stream *chain, int len)
 {
 	struct null_filter *state;
+	fz_stream *stream;
+	fz_context *ctx = chain->ctx;
 
 	assert(chain);
-	state = fz_malloc(chain->ctx, sizeof(struct null_filter));
+	state = fz_malloc(ctx, sizeof(struct null_filter));
 	state->chain = chain;
 	state->remain = len;
 
-	return fz_new_stream(chain->ctx, state, read_null, close_null);
+	fz_try(ctx)
+	{
+		stream = fz_new_stream(ctx, state, read_null, close_null);
+	}
+	fz_catch(ctx)
+	{
+		fz_free(ctx, state);
+		fz_rethrow(ctx);
+	}
+	return stream;
 }
 
 /* ASCII Hex Decode */
diff --git a/fitz/filt_flate.c b/fitz/filt_flate.c
index d7077bc8..9b3370c4 100644
--- a/fitz/filt_flate.c
+++ b/fitz/filt_flate.c
@@ -12,7 +12,7 @@ struct fz_flate_s
 
 static void *zalloc(void *opaque, unsigned int items, unsigned int size)
 {
-	return fz_malloc_array(opaque, items, size);
+	return fz_malloc_array_no_throw(opaque, items, size);
 }
 
 static void zfree(void *opaque, void *ptr)
@@ -84,20 +84,34 @@ fz_stream *
 fz_open_flated(fz_stream *chain)
 {
 	fz_flate *state;
-	int code;
+	int code = Z_OK;
+	fz_context *ctx = chain->ctx;
+	fz_stream *stream;
+
+	fz_var(code);
 
-	state = fz_malloc(chain->ctx, sizeof(fz_flate));
+	state = fz_malloc(ctx, sizeof(fz_flate));
 	state->chain = chain;
 
 	state->z.zalloc = zalloc;
 	state->z.zfree = zfree;
-	state->z.opaque = chain->ctx;
+	state->z.opaque = ctx;
 	state->z.next_in = NULL;
 	state->z.avail_in = 0;
 
-	code = inflateInit(&state->z);
-	if (code != Z_OK)
-		fz_warn(chain->ctx, "zlib error: inflateInit: %s", state->z.msg);
-
-	return fz_new_stream(chain->ctx, state, read_flated, close_flated);
+	fz_try(ctx)
+	{
+		code = inflateInit(&state->z);
+		if (code != Z_OK)
+			fz_throw(ctx, "zlib error: inflateInit: %s", state->z.msg);
+		stream = fz_new_stream(chain->ctx, state, read_flated, close_flated);
+	}
+	fz_catch(ctx)
+	{
+		if (code == Z_OK)
+			inflateEnd(&state->z);
+		fz_free(ctx, state);
+		fz_rethrow(ctx);
+	}
+	return stream;
 }
diff --git a/fitz/stm_buffer.c b/fitz/stm_buffer.c
index 2722be1a..0f846987 100644
--- a/fitz/stm_buffer.c
+++ b/fitz/stm_buffer.c
@@ -9,7 +9,15 @@ fz_new_buffer(fz_context *ctx, int size)
 
 	b = fz_malloc(ctx, sizeof(fz_buffer));
 	b->refs = 1;
-	b->data = fz_malloc(ctx, size);
+	fz_try(ctx)
+	{
+		b->data = fz_malloc(ctx, size);
+	}
+	fz_catch(ctx)
+	{
+		fz_free(ctx, b);
+		fz_rethrow(ctx);
+	}
 	b->cap = size;
 	b->len = 0;
 
diff --git a/fitz/stm_read.c b/fitz/stm_read.c
index 9f9fc302..3d76d51a 100644
--- a/fitz/stm_read.c
+++ b/fitz/stm_read.c
@@ -94,27 +94,37 @@ fz_read_all(fz_stream *stm, int initial)
 	int n;
 	fz_context *ctx = stm->ctx;
 
-	if (initial < 1024)
-		initial = 1024;
+	fz_var(buf);
 
-	buf = fz_new_buffer(ctx, initial);
-
-	while (1)
+	fz_try(ctx)
 	{
-		if (buf->len == buf->cap)
-			fz_grow_buffer(ctx, buf);
+		if (initial < 1024)
+			initial = 1024;
+
+		buf = fz_new_buffer(ctx, initial);
 
-		if (buf->len / 200 > initial)
+		while (1)
 		{
-			fz_drop_buffer(ctx, buf);
-			fz_throw(ctx, "compression bomb detected");
-		}
+			if (buf->len == buf->cap)
+				fz_grow_buffer(ctx, buf);
 
-		n = fz_read(stm, buf->data + buf->len, buf->cap - buf->len);
-		if (n == 0)
-			break;
+			if (buf->len / 200 > initial)
+			{
+				fz_drop_buffer(ctx, buf);
+				fz_throw(ctx, "compression bomb detected");
+			}
+
+			n = fz_read(stm, buf->data + buf->len, buf->cap - buf->len);
+			if (n == 0)
+				break;
 
-		buf->len += n;
+			buf->len += n;
+		}
+	}
+	fz_catch(ctx)
+	{
+		fz_drop_buffer(ctx, buf);
+		fz_rethrow(ctx);
 	}
 
 	return buf;