diff --git a/image/main.c b/image/main.c
index ec8dacd6c1c5..13cd621eb92c 100644
--- a/image/main.c
+++ b/image/main.c
@@ -1959,6 +1959,126 @@ static int read_chunk_block(struct mdrestore_struct *mdres, u8 *buffer,
 	return ret;
 }
 
+static int search_chunk_uncompressed(struct mdrestore_struct *mdres,
+				     struct meta_cluster_item *item,
+				     u64 current_cluster)
+{
+	u32 item_size = le32_to_cpu(item->size);
+	u64 item_bytenr = le64_to_cpu(item->bytenr);
+	int bufsize = SZ_512K;
+	int read_size;
+	u32 offset = 0;
+	u8 *buffer;
+	int ret = 0;
+
+	ASSERT(mdres->compress_method == COMPRESS_NONE);
+	buffer = malloc(bufsize);
+	if (!buffer)
+		return -ENOMEM;
+
+	while (offset < item_size) {
+		read_size = min_t(u32, bufsize, item_size - offset);
+		ret = fread(buffer, read_size, 1, mdres->in);
+		if (ret != 1) {
+			error("read error: %m");
+			ret = -EIO;
+			goto out;
+		}
+		ret = read_chunk_block(mdres, buffer, item_bytenr, read_size,
+					current_cluster);
+		if (ret < 0) {
+			error(
+	"failed to search tree blocks in item bytenr %llu size %u",
+				item_bytenr, item_size);
+			goto out;
+		}
+		offset += read_size;
+	}
+out:
+	free(buffer);
+	return ret;
+}
+
+static int search_chunk_compressed(struct mdrestore_struct *mdres,
+				   struct meta_cluster_item *item,
+				   u64 current_cluster)
+{
+	z_stream strm;
+	u32 item_size = le32_to_cpu(item->size);
+	u64 item_bytenr = le64_to_cpu(item->bytenr);
+	int bufsize = SZ_512K;
+	int read_size;
+	u8 *out_buf = NULL;	/* uncompressed data */
+	u8 *in_buf = NULL;	/* compressed data */
+	bool end = false;
+	int ret;
+
+	ASSERT(mdres->compress_method != COMPRESS_NONE);
+	strm.zalloc = Z_NULL;
+	strm.zfree = Z_NULL;
+	strm.opaque = Z_NULL;
+	strm.avail_in = 0;
+	strm.next_in = Z_NULL;
+	strm.avail_out = 0;
+	strm.next_out = Z_NULL;
+	ret = inflateInit(&strm);
+	if (ret != Z_OK) {
+		error("failed to initialize decompress parameters: %d", ret);
+		return ret;
+	}
+
+	out_buf = malloc(bufsize);
+	in_buf = malloc(item_size);
+	if (!in_buf || !out_buf) {
+		ret = -ENOMEM;
+		goto out;
+	}
+
+	ret = fread(in_buf, item_size, 1, mdres->in);
+	if (ret != 1) {
+		error("read error: %m");
+		ret = -EIO;
+		goto out;
+	}
+	strm.avail_in = item_size;
+	strm.next_in = in_buf;
+	while (!end) {
+		if (strm.avail_out == 0) {
+			strm.avail_out = bufsize;
+			strm.next_out = out_buf;
+		}
+		ret = inflate(&strm, Z_NO_FLUSH);
+		switch (ret) {
+		case Z_NEED_DICT:
+			ret = Z_DATA_ERROR; /* fallthrough */
+			__attribute__ ((fallthrough));
+		case Z_DATA_ERROR:
+		case Z_MEM_ERROR:
+			goto out;
+		}
+		if (ret == Z_STREAM_END) {
+			ret = 0;
+			end = true;
+		}
+		read_size = bufsize - strm.avail_out;
+
+		ret = read_chunk_block(mdres, out_buf, item_bytenr, read_size,
+					current_cluster);
+		if (ret < 0) {
+			error(
+	"failed to search tree blocks in item bytenr %llu size %u",
+				item_bytenr, item_size);
+			goto out;
+		}
+	}
+
+out:
+	free(in_buf);
+	free(out_buf);
+	inflateEnd(&strm);
+	return ret;
+}
+
 /*
  * This function will try to find all chunk items in the dump image.
  *
@@ -2044,8 +2164,6 @@ static int search_for_chunk_blocks(struct mdrestore_struct *mdres)
 
 		/* Search items for tree blocks in sys chunks */
 		for (i = 0; i < nritems; i++) {
-			size_t size;
-
 			item = &cluster->items[i];
 			bufsize = le32_to_cpu(item->size);
 			item_bytenr = le64_to_cpu(item->bytenr);
@@ -2070,41 +2188,16 @@ static int search_for_chunk_blocks(struct mdrestore_struct *mdres)
 			}
 
 			if (mdres->compress_method == COMPRESS_ZLIB) {
-				ret = fread(tmp, bufsize, 1, mdres->in);
-				if (ret != 1) {
-					error("read error: %m");
-					ret = -EIO;
-					goto out;
-				}
-
-				size = max_size;
-				ret = uncompress(buffer,
-						 (unsigned long *)&size, tmp,
-						 bufsize);
-				if (ret != Z_OK) {
-					error("decompression failed with %d",
-							ret);
-					ret = -EIO;
-					goto out;
-				}
+				ret = search_chunk_compressed(mdres, item,
+						current_cluster);
 			} else {
-				ret = fread(buffer, bufsize, 1, mdres->in);
-				if (ret != 1) {
-					error("read error: %m");
-					ret = -EIO;
-					goto out;
-				}
-				size = bufsize;
+				ret = search_chunk_uncompressed(mdres, item,
+						current_cluster);
 			}
-			ret = 0;
-
-			ret = read_chunk_block(mdres, buffer,
-					       item_bytenr, size,
-					       current_cluster);
 			if (ret < 0) {
 				error(
-	"failed to search tree blocks in item bytenr %llu size %lu",
-					item_bytenr, size);
+	"failed to search tree blocks in item bytenr %llu size %u",
+					item_bytenr, bufsize);
 				goto out;
 			}
 			bytenr += bufsize;
