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>

3 files changed, 18 insertions, 41 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");