cgpt: Validate GPT headers before loading them

This CL validates the GPT headers before continue loading its fields.

BRANCH=none
BUG=chromium:422469
TEST=unittest
TEST=cpgt show on a random file. There should be some warnings.
TEST=boot from SD/USB on a device. cgpt show that boot device. It should
     not fail.

Change-Id: I1e5e986cc46620643ec8ec6914fa696a3d04d23a
Reviewed-on: https://chromium-review.googlesource.com/223800
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Commit-Queue: Nam Nguyen <namnguyen@chromium.org>
Tested-by: Nam Nguyen <namnguyen@chromium.org>
diff --git a/cgpt/cgpt_create.c b/cgpt/cgpt_create.c
index cec6077..7d3c059 100644
--- a/cgpt/cgpt_create.c
+++ b/cgpt/cgpt_create.c
@@ -9,16 +9,30 @@
 #include "cgptlib_internal.h"
 #include "vboot_host.h"
 
+static void AllocAndClear(uint8_t **buf, uint64_t size) {
+  if (*buf) {
+    memset(*buf, 0, size);
+  } else {
+    *buf = calloc(1, size);
+    if (!*buf) {
+      Error("Cannot allocate %u bytes.\n", size);
+      abort();
+    }
+  }
+}
+
 static int GptCreate(struct drive *drive, CgptCreateParams *params) {
-  // Erase the data
-  memset(drive->gpt.primary_header, 0,
-         drive->gpt.sector_bytes * GPT_HEADER_SECTORS);
-  memset(drive->gpt.secondary_header, 0,
-         drive->gpt.sector_bytes * GPT_HEADER_SECTORS);
-  memset(drive->gpt.primary_entries, 0,
-         drive->gpt.sector_bytes * GPT_ENTRIES_SECTORS);
-  memset(drive->gpt.secondary_entries, 0,
-         drive->gpt.sector_bytes * GPT_ENTRIES_SECTORS);
+  // Allocate and/or erase the data.
+  // We cannot assume the GPT headers or entry arrays have been allocated
+  // by GptLoad() because those fields might have failed validation checks.
+  AllocAndClear(&drive->gpt.primary_header,
+                drive->gpt.sector_bytes * GPT_HEADER_SECTORS);
+  AllocAndClear(&drive->gpt.secondary_header,
+                drive->gpt.sector_bytes * GPT_HEADER_SECTORS);
+  AllocAndClear(&drive->gpt.primary_entries,
+                drive->gpt.sector_bytes * GPT_ENTRIES_SECTORS);
+  AllocAndClear(&drive->gpt.secondary_entries,
+                drive->gpt.sector_bytes * GPT_ENTRIES_SECTORS);
 
   drive->gpt.modified |= (GPT_MODIFIED_HEADER1 | GPT_MODIFIED_ENTRIES1 |
                          GPT_MODIFIED_HEADER2 | GPT_MODIFIED_ENTRIES2);