Use custom allocator everywhere zlib is used.

4 files changed, 44 insertions, 7 deletions

diff --git a/source/fitz/load-png.c b/source/fitz/load-png.c
index 1df6559e..daa0814a 100644
--- a/source/fitz/load-png.c
+++ b/source/fitz/load-png.c
@@ -61,12 +61,12 @@ static const unsigned char png_signature[8] =
 	137, 80, 78, 71, 13, 10, 26, 10
 };
 
-static void *zalloc(void *opaque, unsigned int items, unsigned int size)
+static void *zalloc_png(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 *address)
+static void zfree_png(void *opaque, void *address)
 {
 	fz_free(opaque, address);
 }
@@ -445,8 +445,8 @@ png_read_image(fz_context *ctx, struct info *info, const unsigned char *p, size_
 
 		info->samples = fz_malloc(ctx, info->size);
 
-		stm.zalloc = zalloc;
-		stm.zfree = zfree;
+		stm.zalloc = zalloc_png;
+		stm.zfree = zfree_png;
 		stm.opaque = ctx;
 
 		stm.next_out = info->samples;
diff --git a/source/fitz/output-png.c b/source/fitz/output-png.c
index 59dbba94..15fb41f0 100644
--- a/source/fitz/output-png.c
+++ b/source/fitz/output-png.c
@@ -4,6 +4,16 @@
 
 #include <zlib.h>
 
+static void *zalloc_outpng(void *opaque, unsigned int items, unsigned int size)
+{
+	return fz_malloc_array_no_throw(opaque, items, size);
+}
+
+static void zfree_outpng(void *opaque, void *address)
+{
+	fz_free(opaque, address);
+}
+
 static inline void big32(unsigned char *buf, unsigned int v)
 {
 	buf[0] = (v >> 24) & 0xff;
@@ -194,6 +204,9 @@ png_write_band(fz_context *ctx, fz_band_writer *writer_, int stride, int band_st
 		writer->csize = compressBound(writer->usize);
 		writer->udata = fz_malloc(ctx, writer->usize);
 		writer->cdata = fz_malloc(ctx, writer->csize);
+		writer->stream.opaque = ctx;
+		writer->stream.zalloc = zalloc_outpng;
+		writer->stream.zfree = zfree_outpng;
 		err = deflateInit(&writer->stream, Z_DEFAULT_COMPRESSION);
 		if (err != Z_OK)
 			fz_throw(ctx, FZ_ERROR_GENERIC, "compression error %d", err);
diff --git a/source/fitz/output-ps.c b/source/fitz/output-ps.c
index 2d9ccf1f..b270fcd4 100644
--- a/source/fitz/output-ps.c
+++ b/source/fitz/output-ps.c
@@ -13,6 +13,16 @@ typedef struct ps_band_writer_s
 	unsigned char *output;
 } ps_band_writer;
 
+static void *zalloc_ps(void *opaque, unsigned int items, unsigned int size)
+{
+	return fz_malloc_array_no_throw(opaque, items, size);
+}
+
+static void zfree_ps(void *opaque, void *address)
+{
+	fz_free(opaque, address);
+}
+
 void
 fz_write_ps_file_header(fz_context *ctx, fz_output *out)
 {
@@ -70,6 +80,10 @@ ps_write_header(fz_context *ctx, fz_band_writer *writer_, const fz_colorspace *c
 	writer->super.h = h;
 	writer->super.n = n;
 
+	writer->stream.zalloc = zalloc_ps;
+	writer->stream.zfree = zfree_ps;
+	writer->stream.opaque = ctx;
+
 	err = deflateInit(&writer->stream, Z_DEFAULT_COMPRESSION);
 	if (err != Z_OK)
 		fz_throw(ctx, FZ_ERROR_GENERIC, "compression error %d", err);
diff --git a/source/fitz/unzip.c b/source/fitz/unzip.c
index c5798fa8..4eb90dda 100644
--- a/source/fitz/unzip.c
+++ b/source/fitz/unzip.c
@@ -37,6 +37,16 @@ struct fz_zip_archive_s
 	zip_entry *entries;
 };
 
+static void *zalloc_zip(void *opaque, unsigned int items, unsigned int size)
+{
+	return fz_malloc_array_no_throw(opaque, items, size);
+}
+
+static void zfree_zip(void *opaque, void *address)
+{
+	fz_free(opaque, address);
+}
+
 static void drop_zip_archive(fz_context *ctx, fz_archive *arch)
 {
 	fz_zip_archive *zip = (fz_zip_archive *) arch;
@@ -323,8 +333,8 @@ static fz_buffer *read_zip_entry(fz_context *ctx, fz_archive *arch, const char *
 		{
 			fz_read(ctx, file, cbuf, ent->csize);
 
-			z.zalloc = (alloc_func) fz_malloc_array;
-			z.zfree = (free_func) fz_free;
+			z.zalloc = zalloc_zip;
+			z.zfree = zfree_zip;
 			z.opaque = ctx;
 			z.next_in = cbuf;
 			z.avail_in = ent->csize;