cgptlib: Minor edge case fixes

This patch fixes a sanitizer issue in cgpt where a GPT entries array may
have been passed even though it was not loaded from disk (parsing an
uninitialized buffer). The GPT library seems to have been written with
the assumption that both headers and entries would always be loaded and
it could recover even if only the primary header and the secondary
entries were valid. In practice, this doesn't really work because the
caller doesn't know how to read entries for an invalid header.
Therefore, change the code so that entries are only assumed to be loaded
for valid headers.

Also fix some minor problems with loading GPTs by aligning sizes up (not
down) to the next sector boundary and making sure we always allocate the
maximum amount of space for entry arrays, even if the current header may
not need that much (in case a repair wants to overwrite it).

This practically reverts CL:276766 which becomes obsolete (and was
really just a dirty hack to hide an underlying problem).

BRANCH=none
BUG=chromium:1017797
TEST=make runtests

Change-Id: I86c601dc074261d53f013b98ae214efdc44f3563
Signed-off-by: Julius Werner <jwerner@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/1885098
Reviewed-by: Mattias Nissler <mnissler@chromium.org>
Reviewed-by: Randall Spangler <rspangler@chromium.org>

7 files changed, 38 insertions, 93 deletions

diff --git a/cgpt/cgpt.h b/cgpt/cgpt.h
index 2cf9ef5a..23f2f90c 100644
--- a/cgpt/cgpt.h
+++ b/cgpt/cgpt.h
@@ -65,18 +65,16 @@ int DriveClose(struct drive *drive, int update_as_needed);
 int CheckValid(const struct drive *drive);
 
 /* Loads sectors from 'drive'.
- * *buf is pointed to an allocated memory when returned, and should be
- * freed.
  *
  *   drive -- open drive.
- *   buf -- pointer to buffer pointer
+ *   buf -- pointer to buffer of at least (sector_bytes * sector_count) size
  *   sector -- offset of starting sector (in sectors)
  *   sector_bytes -- bytes per sector
  *   sector_count -- number of sectors to load
  *
  * Returns CGPT_OK for successful. Aborts if any error occurs.
  */
-int Load(struct drive *drive, uint8_t **buf,
+int Load(struct drive *drive, uint8_t *buf,
                 const uint64_t sector,
                 const uint64_t sector_bytes,
                 const uint64_t sector_count);
diff --git a/cgpt/cgpt_common.c b/cgpt/cgpt_common.c
index 04eab333..6d3dcdb9 100644
--- a/cgpt/cgpt_common.c
+++ b/cgpt/cgpt_common.c
@@ -76,7 +76,7 @@ int CheckValid(const struct drive *drive) {
   return CGPT_OK;
 }
 
-int Load(struct drive *drive, uint8_t **buf,
+int Load(struct drive *drive, uint8_t *buf,
                 const uint64_t sector,
                 const uint64_t sector_bytes,
                 const uint64_t sector_count) {
@@ -96,26 +96,19 @@ int Load(struct drive *drive, uint8_t **buf,
     return CGPT_FAILED;
   }
   count = sector_bytes * sector_count;
-  *buf = malloc(count);
-  require(*buf);
 
   if (-1 == lseek(drive->fd, sector * sector_bytes, SEEK_SET)) {
     Error("Can't seek: %s\n", strerror(errno));
-    goto error_free;
+    return CGPT_FAILED;
   }
 
-  nread = read(drive->fd, *buf, count);
+  nread = read(drive->fd, buf, count);
   if (nread < count) {
     Error("Can't read enough: %d, not %d\n", nread, count);
-    goto error_free;
+    return CGPT_FAILED;
   }
 
   return CGPT_OK;
-
-error_free:
-  free(*buf);
-  *buf = 0;
-  return CGPT_FAILED;
 }
 
 
@@ -170,19 +163,27 @@ static int GptLoad(struct drive *drive, uint32_t sector_bytes) {
   }
   drive->gpt.streaming_drive_sectors = drive->size / drive->gpt.sector_bytes;
 
+  drive->gpt.primary_header = malloc(drive->gpt.sector_bytes);
+  drive->gpt.secondary_header = malloc(drive->gpt.sector_bytes);
+  drive->gpt.primary_entries = malloc(GPT_ENTRIES_ALLOC_SIZE);
+  drive->gpt.secondary_entries = malloc(GPT_ENTRIES_ALLOC_SIZE);
+  if (!drive->gpt.primary_header || !drive->gpt.secondary_header ||
+      !drive->gpt.primary_entries || !drive->gpt.secondary_entries)
+    return -1;
+
   /* TODO(namnguyen): Remove this and totally trust gpt_drive_sectors. */
   if (!(drive->gpt.flags & GPT_FLAG_EXTERNAL)) {
     drive->gpt.gpt_drive_sectors = drive->gpt.streaming_drive_sectors;
   } /* Else, we trust gpt.gpt_drive_sectors. */
 
   // Read the data.
-  if (CGPT_OK != Load(drive, &drive->gpt.primary_header,
+  if (CGPT_OK != Load(drive, drive->gpt.primary_header,
                       GPT_PMBR_SECTORS,
                       drive->gpt.sector_bytes, GPT_HEADER_SECTORS)) {
     Error("Cannot read primary GPT header\n");
     return -1;
   }
-  if (CGPT_OK != Load(drive, &drive->gpt.secondary_header,
+  if (CGPT_OK != Load(drive, drive->gpt.secondary_header,
                       drive->gpt.gpt_drive_sectors - GPT_PMBR_SECTORS,
                       drive->gpt.sector_bytes, GPT_HEADER_SECTORS)) {
     Error("Cannot read secondary GPT header\n");
@@ -193,7 +194,7 @@ static int GptLoad(struct drive *drive, uint32_t sector_bytes) {
                   drive->gpt.gpt_drive_sectors,
                   drive->gpt.flags,
                   drive->gpt.sector_bytes) == 0) {
-    if (CGPT_OK != Load(drive, &drive->gpt.primary_entries,
+    if (CGPT_OK != Load(drive, drive->gpt.primary_entries,
                         primary_header->entries_lba,
                         drive->gpt.sector_bytes,
                         CalculateEntriesSectors(primary_header,
@@ -205,15 +206,13 @@ static int GptLoad(struct drive *drive, uint32_t sector_bytes) {
     Warning("Primary GPT header is %s\n",
       memcmp(primary_header->signature, GPT_HEADER_SIGNATURE_IGNORED,
              GPT_HEADER_SIGNATURE_SIZE) ? "invalid" : "being ignored");
-    drive->gpt.primary_entries = calloc(MAX_NUMBER_OF_ENTRIES,
-                                        sizeof(GptEntry));
   }
   GptHeader* secondary_header = (GptHeader*)drive->gpt.secondary_header;
   if (CheckHeader(secondary_header, 1, drive->gpt.streaming_drive_sectors,
                   drive->gpt.gpt_drive_sectors,
                   drive->gpt.flags,
                   drive->gpt.sector_bytes) == 0) {
-    if (CGPT_OK != Load(drive, &drive->gpt.secondary_entries,
+    if (CGPT_OK != Load(drive, drive->gpt.secondary_entries,
                         secondary_header->entries_lba,
                         drive->gpt.sector_bytes,
                         CalculateEntriesSectors(secondary_header,
@@ -225,8 +224,6 @@ static int GptLoad(struct drive *drive, uint32_t sector_bytes) {
     Warning("Secondary GPT header is %s\n",
       memcmp(primary_header->signature, GPT_HEADER_SIGNATURE_IGNORED,
              GPT_HEADER_SIGNATURE_SIZE) ? "invalid" : "being ignored");
-    drive->gpt.secondary_entries = calloc(MAX_NUMBER_OF_ENTRIES,
-                                          sizeof(GptEntry));
   }
   return 0;
 }
diff --git a/cgpt/cgpt_repair.c b/cgpt/cgpt_repair.c
index fc650883..f06d118e 100644
--- a/cgpt/cgpt_repair.c
+++ b/cgpt/cgpt_repair.c
@@ -24,24 +24,6 @@ int CgptRepair(CgptRepairParams *params) {
     printf("GptSanityCheck() returned %d: %s\n",
            gpt_retval, GptError(gpt_retval));
 
-  GptHeader *header;
-  if (MASK_PRIMARY == drive.gpt.valid_headers ||
-      MASK_BOTH == drive.gpt.valid_headers) {
-    header = (GptHeader *)(drive.gpt.primary_header);
-  } else {
-    header = (GptHeader *)(drive.gpt.secondary_header);
-  }
-
-  if (MASK_PRIMARY == drive.gpt.valid_entries) {
-    free(drive.gpt.secondary_entries);
-    drive.gpt.secondary_entries =
-        malloc(header->size_of_entry * header->number_of_entries);
-  } else if (MASK_SECONDARY == drive.gpt.valid_entries) {
-    free(drive.gpt.primary_entries);
-    drive.gpt.primary_entries =
-        malloc(header->size_of_entry * header->number_of_entries);
-  }
-
   GptRepair(&drive.gpt);
   if (drive.gpt.modified & GPT_MODIFIED_HEADER1)
     printf("Primary Header is updated.\n");
diff --git a/firmware/lib/cgptlib/cgptlib_internal.c b/firmware/lib/cgptlib/cgptlib_internal.c
index 50ad936a..f4d1c1ec 100644
--- a/firmware/lib/cgptlib/cgptlib_internal.c
+++ b/firmware/lib/cgptlib/cgptlib_internal.c
@@ -261,6 +261,8 @@ int GptSanityCheck(GptData *gpt)
 			     gpt->sector_bytes)) {
 		gpt->valid_headers |= MASK_PRIMARY;
 		goodhdr = header1;
+		if (0 == CheckEntries(entries1, goodhdr))
+			gpt->valid_entries |= MASK_PRIMARY;
 	} else if (header1 && !memcmp(header1->signature,
 		   GPT_HEADER_SIGNATURE_IGNORED, GPT_HEADER_SIGNATURE_SIZE)) {
 		gpt->ignored |= MASK_PRIMARY;
@@ -271,6 +273,9 @@ int GptSanityCheck(GptData *gpt)
 		gpt->valid_headers |= MASK_SECONDARY;
 		if (!goodhdr)
 			goodhdr = header2;
+		/* Check header1+entries2 if it was good, to catch mismatch. */
+		if (0 == CheckEntries(entries2, goodhdr))
+			gpt->valid_entries |= MASK_SECONDARY;
 	} else if (header2 && !memcmp(header2->signature,
 		   GPT_HEADER_SIGNATURE_IGNORED, GPT_HEADER_SIGNATURE_SIZE)) {
 		gpt->ignored |= MASK_SECONDARY;
@@ -280,18 +285,6 @@ int GptSanityCheck(GptData *gpt)
 		return GPT_ERROR_INVALID_HEADERS;
 
 	/*
-	 * Check if entries are valid.
-	 *
-	 * Note that we use the same header in both checks.  This way we'll
-	 * catch the case where (header1,entries1) and (header2,entries2) are
-	 * both valid, but (entries1 != entries2).
-	 */
-	if (0 == CheckEntries(entries1, goodhdr))
-		gpt->valid_entries |= MASK_PRIMARY;
-	if (0 == CheckEntries(entries2, goodhdr))
-		gpt->valid_entries |= MASK_SECONDARY;
-
-	/*
 	 * If both headers are good but neither entries were good, check the
 	 * entries with the secondary header.
 	 */
diff --git a/firmware/lib/cgptlib/include/cgptlib_internal.h b/firmware/lib/cgptlib/include/cgptlib_internal.h
index d58a35f6..cc01f4ca 100644
--- a/firmware/lib/cgptlib/include/cgptlib_internal.h
+++ b/firmware/lib/cgptlib/include/cgptlib_internal.h
@@ -60,6 +60,9 @@
 #define MIN_NUMBER_OF_ENTRIES 16
 #define MAX_NUMBER_OF_ENTRIES 128
 
+/* All GptData.(primary|secondary)_entries must be allocated to this size! */
+#define GPT_ENTRIES_ALLOC_SIZE (MAX_NUMBER_OF_ENTRIES * sizeof(GptEntry))
+
 /* Defines GPT sizes */
 #define GPT_PMBR_SECTORS 1  /* size (in sectors) of PMBR */
 #define GPT_HEADER_SECTORS 1
diff --git a/firmware/lib/gpt_misc.c b/firmware/lib/gpt_misc.c
index dfe9084b..6d91ec15 100644
--- a/firmware/lib/gpt_misc.c
+++ b/firmware/lib/gpt_misc.c
@@ -22,7 +22,6 @@
  */
 int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 {
-	uint64_t max_entries_bytes = MAX_NUMBER_OF_ENTRIES * sizeof(GptEntry);
 	int primary_valid = 0, secondary_valid = 0;
 
 	/* No data to be written yet */
@@ -34,8 +33,8 @@ int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 	gptdata->primary_header = (uint8_t *)malloc(gptdata->sector_bytes);
 	gptdata->secondary_header =
 		(uint8_t *)malloc(gptdata->sector_bytes);
-	gptdata->primary_entries = (uint8_t *)malloc(max_entries_bytes);
-	gptdata->secondary_entries = (uint8_t *)malloc(max_entries_bytes);
+	gptdata->primary_entries = (uint8_t *)malloc(GPT_ENTRIES_ALLOC_SIZE);
+	gptdata->secondary_entries = (uint8_t *)malloc(GPT_ENTRIES_ALLOC_SIZE);
 
 	if (gptdata->primary_header == NULL ||
 	    gptdata->secondary_header == NULL ||
@@ -60,8 +59,9 @@ int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 		uint64_t entries_bytes =
 				(uint64_t)primary_header->number_of_entries
 				* primary_header->size_of_entry;
-		uint64_t entries_sectors = entries_bytes
-					/ gptdata->sector_bytes;
+		uint64_t entries_sectors =
+				(entries_bytes + gptdata->sector_bytes - 1)
+				/ gptdata->sector_bytes;
 		if (0 != VbExDiskRead(disk_handle,
 				      primary_header->entries_lba,
 				      entries_sectors,
@@ -95,7 +95,8 @@ int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 		uint64_t entries_bytes =
 				(uint64_t)secondary_header->number_of_entries
 				* secondary_header->size_of_entry;
-		uint64_t entries_sectors = entries_bytes
+		uint64_t entries_sectors =
+				(entries_bytes + gptdata->sector_bytes - 1)
 				/ gptdata->sector_bytes;
 		if (0 != VbExDiskRead(disk_handle,
 				      secondary_header->entries_lba,
diff --git a/tests/cgptlib_test.c b/tests/cgptlib_test.c
index 669f7969..821ffb13 100644
--- a/tests/cgptlib_test.c
+++ b/tests/cgptlib_test.c
@@ -41,7 +41,7 @@
 #define DEFAULT_SECTOR_SIZE 512
 #define MAX_SECTOR_SIZE 4096
 #define DEFAULT_DRIVE_SECTORS 467
-#define TOTAL_ENTRIES_SIZE (MAX_NUMBER_OF_ENTRIES * sizeof(GptEntry)) /* 16384 */
+#define TOTAL_ENTRIES_SIZE GPT_ENTRIES_ALLOC_SIZE /* 16384 */
 #define PARTITION_ENTRIES_SIZE TOTAL_ENTRIES_SIZE /* 16384 */
 
 static const Guid guid_zero = {{{0, 0, 0, 0, 0, {0, 0, 0, 0, 0, 0}}}};
@@ -911,27 +911,27 @@ static int SanityCheckTest(void)
 	gpt->primary_header[0]++;
 	EXPECT(GPT_SUCCESS == GptSanityCheck(gpt));
 	EXPECT(MASK_SECONDARY == gpt->valid_headers);
-	EXPECT(MASK_BOTH == gpt->valid_entries);
+	EXPECT(MASK_SECONDARY == gpt->valid_entries);
 	EXPECT(0 == gpt->ignored);
 	GptRepair(gpt);
 	EXPECT(GPT_SUCCESS == GptSanityCheck(gpt));
 	EXPECT(MASK_BOTH == gpt->valid_headers);
 	EXPECT(MASK_BOTH == gpt->valid_entries);
 	EXPECT(0 == gpt->ignored);
-	EXPECT(GPT_MODIFIED_HEADER1 == gpt->modified);
+	EXPECT((GPT_MODIFIED_HEADER1 | GPT_MODIFIED_ENTRIES1) == gpt->modified);
 
 	BuildTestGptData(gpt);
 	gpt->secondary_header[0]++;
 	EXPECT(GPT_SUCCESS == GptSanityCheck(gpt));
 	EXPECT(MASK_PRIMARY == gpt->valid_headers);
-	EXPECT(MASK_BOTH == gpt->valid_entries);
+	EXPECT(MASK_PRIMARY == gpt->valid_entries);
 	EXPECT(0 == gpt->ignored);
 	GptRepair(gpt);
 	EXPECT(GPT_SUCCESS == GptSanityCheck(gpt));
 	EXPECT(MASK_BOTH == gpt->valid_headers);
 	EXPECT(MASK_BOTH == gpt->valid_entries);
 	EXPECT(0 == gpt->ignored);
-	EXPECT(GPT_MODIFIED_HEADER2 == gpt->modified);
+	EXPECT((GPT_MODIFIED_HEADER2 | GPT_MODIFIED_ENTRIES2) == gpt->modified);
 
 	/*
 	 * Modify header1 and update its CRC.  Since header2 is now different
@@ -1039,35 +1039,6 @@ static int SanityCheckTest(void)
 	EXPECT(0 == gpt->ignored);
 	EXPECT((GPT_MODIFIED_HEADER2 | GPT_MODIFIED_ENTRIES2) == gpt->modified);
 
-	/* Test cross-correction (h1+e2, h2+e1) */
-	BuildTestGptData(gpt);
-	gpt->primary_header[0]++;
-	gpt->secondary_entries[0]++;
-	EXPECT(GPT_SUCCESS == GptSanityCheck(gpt));
-	EXPECT(MASK_SECONDARY == gpt->valid_headers);
-	EXPECT(MASK_PRIMARY == gpt->valid_entries);
-	EXPECT(0 == gpt->ignored);
-	GptRepair(gpt);
-	EXPECT(GPT_SUCCESS == GptSanityCheck(gpt));
-	EXPECT(MASK_BOTH == gpt->valid_headers);
-	EXPECT(MASK_BOTH == gpt->valid_entries);
-	EXPECT(0 == gpt->ignored);
-	EXPECT((GPT_MODIFIED_HEADER1 | GPT_MODIFIED_ENTRIES2) == gpt->modified);
-
-	BuildTestGptData(gpt);
-	gpt->secondary_header[0]++;
-	gpt->primary_entries[0]++;
-	EXPECT(GPT_SUCCESS == GptSanityCheck(gpt));
-	EXPECT(MASK_PRIMARY == gpt->valid_headers);
-	EXPECT(MASK_SECONDARY == gpt->valid_entries);
-	EXPECT(0 == gpt->ignored);
-	GptRepair(gpt);
-	EXPECT(GPT_SUCCESS == GptSanityCheck(gpt));
-	EXPECT(MASK_BOTH == gpt->valid_headers);
-	EXPECT(MASK_BOTH == gpt->valid_entries);
-	EXPECT(0 == gpt->ignored);
-	EXPECT((GPT_MODIFIED_HEADER2 | GPT_MODIFIED_ENTRIES1) == gpt->modified);
-
 	/*
 	 * Test mismatched pairs (h1+e1 valid, h2+e2 valid but different.  This
 	 * simulates a partial update of the drive.