ptgen: allow image generation for a specified disk size

Sometimes we know an exact size of the disk and want to create a proper
disk image. This patch allows such operation.

A special case of zero disk size is supported. In this case the disk
size will be automatically calculated on the base of provided partitions
list.

Signed-off-by: Mikhail Kshevetskiy <mikhail.kshevetskiy@iopsys.eu>

1 files changed, 57 insertions, 27 deletions

diff --git a/src/ptgen.c b/src/ptgen.c
index ed9efa3..d675abd 100644
--- a/src/ptgen.c
+++ b/src/ptgen.c
@@ -172,7 +172,9 @@ bool use_guid_partition_table = false;
 struct partinfo parts[GPT_ENTRY_MAX];
 char *filename = NULL;
 
+int gpt_alternate = false;
 uint64_t gpt_first_entry_sector = GPT_FIRST_ENTRY_SECTOR;
+uint64_t gpt_last_usable_sector = 0;
 
 /*
  * parse the size argument, which is either
@@ -441,6 +443,12 @@ static int gen_gptable(uint32_t signature, guid_t guid, unsigned nr)
 		} else if (kb_align != 0) {
 			start = round_to_kb(start);
 		}
+		if ((gpt_last_usable_sector > 0) &&
+		    (start + parts[i].size * 2 > gpt_last_usable_sector + 1)) {
+				fprintf(stderr, "Partition %d ends after last usable sector %ld\n",
+					i, gpt_last_usable_sector);
+				return ret;
+		}
 		parts[i].actual_start = start;
 		gpte[i].start = cpu_to_le64(start);
 
@@ -488,7 +496,10 @@ static int gen_gptable(uint32_t signature, guid_t guid, unsigned nr)
 		gpte[GPT_ENTRY_MAX - 1].guid.b[sizeof(guid_t) -1] += GPT_ENTRY_MAX;
 	}
 
-	end = sect + GPT_SIZE;
+	if (gpt_last_usable_sector == 0)
+		gpt_last_usable_sector = sect - 1;
+
+	end = gpt_last_usable_sector + GPT_SIZE + 1;
 
 	pte[0].type = 0xEE;
 	pte[0].start = cpu_to_le32(GPT_HEADER_SECTOR);
@@ -496,14 +507,14 @@ static int gen_gptable(uint32_t signature, guid_t guid, unsigned nr)
 	to_chs(GPT_HEADER_SECTOR, pte[0].chs_start);
 	to_chs(end, pte[0].chs_end);
 
-	gpth.last_usable = cpu_to_le64(end - GPT_SIZE - 1);
+	gpth.last_usable = cpu_to_le64(gpt_last_usable_sector);
 	gpth.alternate = cpu_to_le64(end);
 	gpth.entry_crc32 = cpu_to_le32(gpt_crc32(gpte, GPT_ENTRY_SIZE * GPT_ENTRY_MAX));
 	gpth.crc32 = cpu_to_le32(gpt_crc32((char *)&gpth, GPT_HEADER_SIZE));
 
 	if (verbose)
 		fprintf(stderr, "PartitionEntryLBA=%" PRIu64 ", FirstUsableLBA=%" PRIu64 ", LastUsableLBA=%" PRIu64 "\n",
-			gpt_first_entry_sector, gpt_first_entry_sector + GPT_SIZE, end - GPT_SIZE - 1);
+			gpt_first_entry_sector, gpt_first_entry_sector + GPT_SIZE, gpt_last_usable_sector);
 
 	if ((fd = open(filename, O_WRONLY|O_CREAT|O_TRUNC, 0644)) < 0) {
 		fprintf(stderr, "Can't open output file '%s'\n",filename);
@@ -539,30 +550,30 @@ static int gen_gptable(uint32_t signature, guid_t guid, unsigned nr)
 		goto fail;
 	}
 
-#ifdef WANT_ALTERNATE_PTABLE
-	/* The alternate partition table (We omit it by default) */
-	swap(gpth.self, gpth.alternate);
-	gpth.first_entry = cpu_to_le64(end - GPT_ENTRY_SIZE * GPT_ENTRY_MAX / DISK_SECTOR_SIZE),
-	gpth.crc32 = 0;
-	gpth.crc32 = cpu_to_le32(gpt_crc32(&gpth, GPT_HEADER_SIZE));
-
-	lseek(fd, end * DISK_SECTOR_SIZE - GPT_ENTRY_SIZE * GPT_ENTRY_MAX, SEEK_SET);
-	if (write(fd, &gpte, GPT_ENTRY_SIZE * GPT_ENTRY_MAX) != GPT_ENTRY_SIZE * GPT_ENTRY_MAX) {
-		fputs("write failed.\n", stderr);
-		goto fail;
-	}
+	if (gpt_alternate) {
+		/* The alternate partition table (We omit it by default) */
+		swap(gpth.self, gpth.alternate);
+		gpth.first_entry = cpu_to_le64(end - GPT_ENTRY_SIZE * GPT_ENTRY_MAX / DISK_SECTOR_SIZE),
+		gpth.crc32 = 0;
+		gpth.crc32 = cpu_to_le32(gpt_crc32(&gpth, GPT_HEADER_SIZE));
+
+		lseek(fd, end * DISK_SECTOR_SIZE - GPT_ENTRY_SIZE * GPT_ENTRY_MAX, SEEK_SET);
+		if (write(fd, &gpte, GPT_ENTRY_SIZE * GPT_ENTRY_MAX) != GPT_ENTRY_SIZE * GPT_ENTRY_MAX) {
+			fputs("write failed.\n", stderr);
+			goto fail;
+		}
 
-	lseek(fd, end * DISK_SECTOR_SIZE, SEEK_SET);
-	if (write(fd, &gpth, GPT_HEADER_SIZE) != GPT_HEADER_SIZE) {
-		fputs("write failed.\n", stderr);
-		goto fail;
-	}
-	lseek(fd, (end + 1) * DISK_SECTOR_SIZE -1, SEEK_SET);
-	if (write(fd, "\x00", 1) != 1) {
-		fputs("write failed.\n", stderr);
-		goto fail;
+		lseek(fd, end * DISK_SECTOR_SIZE, SEEK_SET);
+		if (write(fd, &gpth, GPT_HEADER_SIZE) != GPT_HEADER_SIZE) {
+			fputs("write failed.\n", stderr);
+			goto fail;
+		}
+		lseek(fd, (end + 1) * DISK_SECTOR_SIZE -1, SEEK_SET);
+		if (write(fd, "\x00", 1) != 1) {
+			fputs("write failed.\n", stderr);
+			goto fail;
+		}
 	}
-#endif
 
 	ret = 0;
 fail:
@@ -574,7 +585,7 @@ static void usage(char *prog)
 {
 	fprintf(stderr, "Usage: %s [-v] [-n] [-g] -h <heads> -s <sectors> -o <outputfile>\n"
 			"          [-a <part number>] [-l <align kB>] [-G <guid>]\n"
-			"          [-e <gpt_entry_offset>]\n"
+			"          [-e <gpt_entry_offset>] [-d <gpt_disk_size>]\n"
 			"          [[-t <type> | -T <GPT part type>] [-r] [-N <name>] -p <size>[@<start>]...] \n", prog);
 
 	exit(EXIT_FAILURE);
@@ -609,11 +620,12 @@ int main (int argc, char **argv)
 	int part = 0;
 	char *name = NULL;
 	unsigned short int hybrid = 0, required = 0;
+	uint64_t total_sectors;
 	uint32_t signature = 0x5452574F; /* 'OWRT' */
 	guid_t guid = GUID_INIT( signature, 0x2211, 0x4433, \
 			0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0x00);
 
-	while ((ch = getopt(argc, argv, "h:s:p:a:t:T:o:vnHN:gl:rS:G:e:")) != -1) {
+	while ((ch = getopt(argc, argv, "h:s:p:a:t:T:o:vnHN:gl:rS:G:e:d:")) != -1) {
 		switch (ch) {
 		case 'o':
 			filename = optarg;
@@ -639,6 +651,24 @@ int main (int argc, char **argv)
 				exit(EXIT_FAILURE);
 			}
 			break;
+		case 'd':
+			/*
+			 * Zero disk_size is specially allowed. It means: find a disk size
+			 * on the base of provided partitions list.
+			 *
+			 * based on DISK_SECTOR_SIZE = 512
+			 */
+			gpt_alternate = true;
+			total_sectors = 2 * to_kbytes(optarg);
+			if (total_sectors != 0) {
+				if (total_sectors <= 2 * GPT_SIZE + 3) {
+					fprintf(stderr, "GPT disk size must be larger than %d KBytes\n",
+						(2 * GPT_SIZE + 3) * DISK_SECTOR_SIZE / 1024);
+					exit(EXIT_FAILURE);
+				}
+				gpt_last_usable_sector = total_sectors - GPT_SIZE - 2;
+			}
+			break;
 		case 'h':
 			heads = (int)strtoul(optarg, NULL, 0);
 			break;