From c20f01c4ea0e1b30072bd53f47482df1378b36dc Mon Sep 17 00:00:00 2001
From: Marcel Moolenaar <marcel@FreeBSD.org>
Date: Fri, 7 Aug 2015 04:27:51 +0000
Subject: [PATCH] Fix the dynamic VHD format to work with qemu. The size of the
 disk is taken to match the geometry and only when the geometry is max'd out,
 is the actual recorded size taken.

Note that qemu has the same logic for the fixed VHD format. However
that is known to conflict with Microsoft Azure, where the recorded
size of the image is what counts.

Pointed out by: gjb@
---
 usr.bin/mkimg/vhd.c | 92 +++++++++++++++++++++++++++------------------
 1 file changed, 55 insertions(+), 37 deletions(-)

diff --git a/usr.bin/mkimg/vhd.c b/usr.bin/mkimg/vhd.c
index 054f82facd81..c4c1d1d45b8b 100644
--- a/usr.bin/mkimg/vhd.c
+++ b/usr.bin/mkimg/vhd.c
@@ -159,6 +159,34 @@ vhd_geometry(uint64_t image_size, struct vhd_geom *geom)
 	geom->cylinders = cth / geom->heads;
 }
 
+static uint64_t
+vhd_resize(uint64_t origsz)
+{
+	struct vhd_geom geom;
+	uint64_t newsz;
+
+	/*
+	 * Round the image size to the pre-determined geometry that
+	 * matches the image size. This circular dependency implies
+	 * that we need to loop to handle boundary conditions.
+	 * The first time, newsz equals origsz and the geometry will
+	 * typically yield a new size that's smaller. We keep adding
+	 * cylinder's worth of sectors to the new size until its
+	 * larger or equal or origsz. But during those iterations,
+	 * the geometry can change, so we need to account for that.
+	 */
+	newsz = origsz;
+	while (1) {
+		vhd_geometry(newsz, &geom);
+		newsz = (int64_t)geom.cylinders * geom.heads *
+		    geom.sectors * VHD_SECTOR_SIZE;
+		if (newsz >= origsz)
+			break;
+		newsz += geom.heads * geom.sectors * VHD_SECTOR_SIZE;
+	}
+	return (newsz);
+}
+
 static uint32_t
 vhd_timestamp(void)
 {
@@ -256,8 +284,7 @@ vhd_dyn_resize(lba_t imgsz)
 {
 	uint64_t imagesz;
 
-	imagesz = imgsz * secsz;
-	imagesz = (imagesz + VHD_BLOCK_SIZE - 1) & ~(VHD_BLOCK_SIZE - 1);
+	imagesz = vhd_resize(imgsz * secsz);
 	return (image_set_size(imagesz / secsz));
 }
 
@@ -266,7 +293,7 @@ vhd_dyn_write(int fd)
 {
 	struct vhd_footer footer;
 	struct vhd_dyn_header header;
-	uint64_t imgsz;
+	uint64_t imgsz, rawsz;
 	lba_t blk, blkcnt, nblks;
 	uint32_t *bat;
 	void *bitmap;
@@ -274,13 +301,14 @@ vhd_dyn_write(int fd)
 	uint32_t sector;
 	int bat_entries, error, entry;
 
-	imgsz = image_get_size() * secsz;
-	bat_entries = imgsz / VHD_BLOCK_SIZE;
+	rawsz = image_get_size() * secsz;
+	imgsz = (rawsz + VHD_BLOCK_SIZE - 1) & ~(VHD_BLOCK_SIZE - 1);
 
-	vhd_make_footer(&footer, imgsz, VHD_DISK_TYPE_DYNAMIC, sizeof(footer));
+	vhd_make_footer(&footer, rawsz, VHD_DISK_TYPE_DYNAMIC, sizeof(footer));
 	if (sparse_write(fd, &footer, sizeof(footer)) < 0)
 		return (errno);
 
+	bat_entries = imgsz / VHD_BLOCK_SIZE;
 	memset(&header, 0, sizeof(header));
 	be64enc(&header.cookie, VHD_HEADER_COOKIE);
 	be64enc(&header.data_offset, ~0ULL);
@@ -321,7 +349,7 @@ vhd_dyn_write(int fd)
 	blk = 0;
 	blkcnt = VHD_BLOCK_SIZE / secsz;
 	error = 0;
-	nblks = image_get_size();
+	nblks = rawsz / secsz;
 	while (blk < nblks) {
 		if (!image_data(blk, blkcnt)) {
 			blk += blkcnt;
@@ -331,15 +359,20 @@ vhd_dyn_write(int fd)
 			error = errno;
 			break;
 		}
+		/* Handle partial last block */
+		if (blk + blkcnt > nblks)
+			blkcnt = nblks - blk;
 		error = image_copyout_region(fd, blk, blkcnt);
 		if (error)
 			break;
 		blk += blkcnt;
 	}
 	free(bitmap);
-	if (blk != nblks)
+	if (error)
+		return (error);
+	error = image_copyout_zeroes(fd, imgsz - rawsz);
+	if (error)
 		return (error);
-
 	if (sparse_write(fd, &footer, sizeof(footer)) < 0)
 		return (errno);
 
@@ -362,24 +395,9 @@ FORMAT_DEFINE(vhd_dyn_format);
 static int
 vhd_fix_resize(lba_t imgsz)
 {
-	struct vhd_geom geom;
-	int64_t imagesz;
+	uint64_t imagesz;
 
-	/*
-	 * Round the image size to the pre-determined geometry that
-	 * matches the image size. This circular dependency implies
-	 * that we need to loop to handle boundary conditions.
-	 */
-	imgsz *= secsz;
-	imagesz = imgsz;
-	while (1) {
-		vhd_geometry(imagesz, &geom);
-		imagesz = (int64_t)geom.cylinders * geom.heads *
-		    geom.sectors * VHD_SECTOR_SIZE;
-		if (imagesz >= imgsz)
-			break;
-		imagesz += geom.heads * geom.sectors * VHD_SECTOR_SIZE;
-	}
+	imagesz = vhd_resize(imgsz * secsz);
 	/*
 	 * Azure demands that images are a whole number of megabytes.
 	 */
@@ -391,24 +409,24 @@ static int
 vhd_fix_write(int fd)
 {
 	struct vhd_footer footer;
-	uint64_t imgsz;
+	uint64_t imagesz;
 	int error;
 
 	error = image_copyout(fd);
-	if (!error) {
-		imgsz = image_get_size() * secsz;
-		vhd_make_footer(&footer, imgsz, VHD_DISK_TYPE_FIXED, ~0ULL);
-		if (sparse_write(fd, &footer, sizeof(footer)) < 0)
-			error = errno;
-	}
+	if (error)
+		return (error);
+
+	imagesz = image_get_size() * secsz;
+	vhd_make_footer(&footer, imagesz, VHD_DISK_TYPE_FIXED, ~0ULL);
+	error = (sparse_write(fd, &footer, sizeof(footer)) < 0) ? errno : 0;
 	return (error);
 }
 
 static struct mkimg_format vhd_fix_format = {
-        .name = "vhdf",
-        .description = "Fixed Virtual Hard Disk",
-        .resize = vhd_fix_resize,
-        .write = vhd_fix_write,
+	.name = "vhdf",
+	.description = "Fixed Virtual Hard Disk",
+	.resize = vhd_fix_resize,
+	.write = vhd_fix_write,
 };
 
 FORMAT_DEFINE(vhd_fix_format);