From fb64d9328870737502c6ff94b38a084c28bfd008 Mon Sep 17 00:00:00 2001
From: Samuel Dionne-Riel <samuel@dionne-riel.com>
Date: Sat, 17 Apr 2021 02:17:53 -0400
Subject: [PATCH] libfdisk: Include table-length in first-lba checks

Otherwise, sfdisk would error out on its own dump, as first-lba would be
earlier than "plausible" assuming the default table size.

This is because it did not look for `table-length` when using an sfdisk
script. This means that using `--dump` on a valid image, could not be
unserialized back into a valid image.

This is helpful in exotic situations, where a partition has to be
created earlier than the conventional "first-lba" (2048). For example,
on Allwinner A64 boards, where a firmware is stored at LBA16. With this
it is possible to ask for `first-lba` at or earlier than 16, by reducing
the number of entries to e.g. 48.
---
 libfdisk/src/gpt.c | 36 +++++++++++++++++++++++-------------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/libfdisk/src/gpt.c b/libfdisk/src/gpt.c
index 6b07d9841..1a086a61d 100644
--- a/libfdisk/src/gpt.c
+++ b/libfdisk/src/gpt.c
@@ -747,7 +747,7 @@ static int count_first_last_lba(struct fdisk_context *cxt,
 {
 	int rc = 0;
 	uint64_t flba = 0, llba = 0;
-	uint32_t nents;
+	uint64_t nents = GPT_NPARTITIONS;
 
 	assert(cxt);
 	assert(first);
@@ -755,18 +755,28 @@ static int count_first_last_lba(struct fdisk_context *cxt,
 
 	*first = *last = 0;
 
-	/* The default is GPT_NPARTITIONS, if the device is not large enough
-	 * than reduce this number of partitions and try to recalculate it
-	 * again, until we get something useful or return error.
-	 */
-	for (nents = GPT_NPARTITIONS; nents > 0; nents--) {
-		rc = gpt_calculate_last_lba(NULL, nents, &llba, cxt);
-		if (rc == 0)
-			rc = gpt_calculate_first_lba(NULL, nents, &flba, cxt);
-		if (llba < flba)
-			rc = -ENOSPC;
-		else if (rc == 0)
-			break;
+	/* Get the table length from the script, if given */
+	if (cxt->script) {
+		rc = get_script_u64(cxt, &nents, "table-length");
+		if (rc < 0)
+			return rc;
+	}
+
+	/* The table length was not changed by the script, compute it. */
+	if (flba == 0) {
+		/* If the device is not large enough reduce this number of
+		 * partitions and try to recalculate it again, until we get
+		 * something useful or return error.
+		 */
+		for (; nents > 0; nents--) {
+			rc = gpt_calculate_last_lba(NULL, nents, &llba, cxt);
+			if (rc == 0)
+				rc = gpt_calculate_first_lba(NULL, nents, &flba, cxt);
+			if (llba < flba)
+				rc = -ENOSPC;
+			else if (rc == 0)
+				break;
+		}
 	}
 
 	if (rc)