1 files changed, 83 insertions, 55 deletions

diff --git a/c_src/qcow2.c b/c_src/qcow2.c
index 30a6e056..53959a00 100644
--- a/c_src/qcow2.c
+++ b/c_src/qcow2.c
@@ -31,24 +31,21 @@ struct qcow2_hdr {
 	u64			snapshots_offset;
 };
 
-struct qcow2_image {
-	int			fd;
-	u32			block_size;
-	u64			*l1_table;
-	u64			l1_offset;
-	u32			l1_index;
-	u64			*l2_table;
-	u64			offset;
-};
+static void __qcow2_write_buf(struct qcow2_image *img, void *buf, unsigned len)
+{
+	assert(!(len % img->block_size));
+
+	xpwrite(img->outfd, buf, len, img->offset, "qcow2 data");
+	img->offset += len;
+}
 
 static void flush_l2(struct qcow2_image *img)
 {
 	if (img->l1_index != -1) {
 		img->l1_table[img->l1_index] =
 			cpu_to_be64(img->offset|QCOW_OFLAG_COPIED);
-		xpwrite(img->fd, img->l2_table, img->block_size, img->offset,
-			"qcow2 l2 table");
-		img->offset += img->block_size;
+
+		__qcow2_write_buf(img, img->l2_table, img->block_size);
 
 		memset(img->l2_table, 0, img->block_size);
 		img->l1_index = -1;
@@ -69,66 +66,97 @@ static void add_l2(struct qcow2_image *img, u64 src_blk, u64 dst_offset)
 	img->l2_table[l2_index] = cpu_to_be64(dst_offset|QCOW_OFLAG_COPIED);
 }
 
-void qcow2_write_image(int infd, int outfd, ranges *data,
-		       unsigned block_size)
+void qcow2_write_buf(struct qcow2_image *img, void *buf, unsigned len, u64 src_offset)
 {
-	u64 image_size = get_size(infd);
-	unsigned l2_size = block_size / sizeof(u64);
-	unsigned l1_size = DIV_ROUND_UP(image_size, (u64) block_size * l2_size);
-	struct qcow2_hdr hdr = { 0 };
-	struct qcow2_image img = {
-		.fd		= outfd,
-		.block_size	= block_size,
-		.l2_table	= xcalloc(l2_size, sizeof(u64)),
-		.l1_table	= xcalloc(l1_size, sizeof(u64)),
-		.l1_index	= -1,
-		.offset		= round_up(sizeof(hdr), block_size),
-	};
-	char *buf = xmalloc(block_size);
-	u64 src_offset, dst_offset;
-
-	assert(is_power_of_2(block_size));
+	u64 dst_offset = img->offset;
+	__qcow2_write_buf(img, buf, len);
+
+	while (len) {
+		add_l2(img, src_offset / img->block_size, dst_offset);
+		dst_offset += img->block_size;
+		src_offset += img->block_size;
+		len -= img->block_size;
+	}
+}
 
-	ranges_roundup(data, block_size);
+void qcow2_write_ranges(struct qcow2_image *img, ranges *data)
+{
+	ranges_roundup(data, img->block_size);
 	ranges_sort_merge(data);
 
+	char *buf = xmalloc(img->block_size);
+
 	/* Write data: */
 	darray_for_each(*data, r)
-		for (src_offset = r->start;
+		for (u64 src_offset = r->start;
 		     src_offset < r->end;
-		     src_offset += block_size) {
-			dst_offset = img.offset;
-			img.offset += img.block_size;
+		     src_offset += img->block_size) {
+			xpread(img->infd, buf, img->block_size, src_offset);
+			qcow2_write_buf(img, buf, img->block_size, src_offset);
+		}
+
+	free(buf);
+}
 
-			xpread(infd, buf, block_size, src_offset);
-			xpwrite(outfd, buf, block_size, dst_offset,
-				"qcow2 data");
+void qcow2_image_init(struct qcow2_image *img, int infd, int outfd, unsigned block_size)
+{
+	assert(is_power_of_2(block_size));
 
-			add_l2(&img, src_offset / block_size, dst_offset);
-		}
+	u64 image_size = get_size(infd);
+	unsigned l2_size = block_size / sizeof(u64);
+	unsigned l1_size = DIV_ROUND_UP(image_size, (u64) block_size * l2_size);
+
+	*img = (struct qcow2_image) {
+		.infd		= infd,
+		.outfd		= outfd,
+		.image_size	= image_size,
+		.block_size	= block_size,
+		.l1_size	= l1_size,
+		.l1_table	= xcalloc(l1_size, sizeof(u64)),
+		.l1_index	= -1,
+		.l2_table	= xcalloc(l2_size, sizeof(u64)),
+		.offset		= round_up(sizeof(struct qcow2_hdr), block_size),
+	};
+}
+
+void qcow2_image_finish(struct qcow2_image *img)
+{
+	char *buf = xmalloc(img->block_size);
 
-	flush_l2(&img);
+	flush_l2(img);
 
 	/* Write L1 table: */
-	dst_offset		= img.offset;
-	img.offset		+= round_up(l1_size * sizeof(u64), block_size);
-	xpwrite(img.fd, img.l1_table, l1_size * sizeof(u64), dst_offset,
+	u64 dst_offset		= img->offset;
+	img->offset		+= round_up(img->l1_size * sizeof(u64), img->block_size);
+	xpwrite(img->outfd, img->l1_table, img->l1_size * sizeof(u64), dst_offset,
 		"qcow2 l1 table");
 
 	/* Write header: */
-	hdr.magic		= cpu_to_be32(QCOW_MAGIC);
-	hdr.version		= cpu_to_be32(QCOW_VERSION);
-	hdr.block_bits		= cpu_to_be32(ilog2(block_size));
-	hdr.size		= cpu_to_be64(image_size);
-	hdr.l1_size		= cpu_to_be32(l1_size);
-	hdr.l1_table_offset	= cpu_to_be64(dst_offset);
-
-	memset(buf, 0, block_size);
+	struct qcow2_hdr hdr = {
+		.magic			= cpu_to_be32(QCOW_MAGIC),
+		.version		= cpu_to_be32(QCOW_VERSION),
+		.block_bits		= cpu_to_be32(ilog2(img->block_size)),
+		.size			= cpu_to_be64(img->image_size),
+		.l1_size		= cpu_to_be32(img->l1_size),
+		.l1_table_offset	= cpu_to_be64(dst_offset),
+	};
+
+	memset(buf, 0, img->block_size);
 	memcpy(buf, &hdr, sizeof(hdr));
-	xpwrite(img.fd, buf, block_size, 0,
+	xpwrite(img->outfd, buf, img->block_size, 0,
 		"qcow2 header");
 
-	free(img.l2_table);
-	free(img.l1_table);
+	free(img->l2_table);
+	free(img->l1_table);
 	free(buf);
 }
+
+void qcow2_write_image(int infd, int outfd, ranges *data,
+		       unsigned block_size)
+{
+	struct qcow2_image img;
+
+	qcow2_image_init(&img, infd, outfd, block_size);
+	qcow2_write_ranges(&img, data);
+	qcow2_image_finish(&img);
+}