cgpt repair: fix segfault which occurs when one of the headers is bad

When one of GPT headers is invalid the corresponding partition table
is not loaded and corresponding pointers in GptData are NULL.
GptRepair will try to memcpy one entries table to another which
results in SIGSEGV.

This change fixes it by freeing and then reallocating bad copy of
partition table. This potentially fixes problems which would occur
if two tables have different size.

Change that initially introduced this problem by not always allocating
secondary_entries:
https://chromium-review.googlesource.com/223800

TEST="cgpt repair" works where it previously didn't
TEST=make runtests
BUG=brillo:1203
BRANCH=none

Change-Id: Ibb2fcf33faa5ba157b0865d04c90ee3f26eee113
Reviewed-on: https://chromium-review.googlesource.com/276766
Reviewed-by: Randall Spangler <rspangler@chromium.org>
Commit-Queue: Andrey Ulanov <andreyu@google.com>
Tested-by: Andrey Ulanov <andreyu@google.com>
diff --git a/cgpt/cgpt_repair.c b/cgpt/cgpt_repair.c
index 1880ea5..ebe034b 100644
--- a/cgpt/cgpt_repair.c
+++ b/cgpt/cgpt_repair.c
@@ -24,6 +24,24 @@
     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/include/gpt_misc.h b/firmware/include/gpt_misc.h
index a94224f..e5a366b 100644
--- a/firmware/include/gpt_misc.h
+++ b/firmware/include/gpt_misc.h
@@ -84,9 +84,9 @@
 	uint8_t *primary_header;
 	/* GPT secondary header, from last sector of disk (size: 512 bytes) */
 	uint8_t *secondary_header;
-	/* Primary GPT table, follows primary header (size: 16 KB) */
+	/* Primary GPT table, follows primary header */
 	uint8_t *primary_entries;
-	/* Secondary GPT table, precedes secondary header (size: 16 KB) */
+	/* Secondary GPT table, precedes secondary header */
 	uint8_t *secondary_entries;
 	/* Size of a LBA sector, in bytes */
 	uint32_t sector_bytes;
diff --git a/firmware/lib/cgptlib/include/cgptlib_internal.h b/firmware/lib/cgptlib/include/cgptlib_internal.h
index 0be2bc5..50477c4 100644
--- a/firmware/lib/cgptlib/include/cgptlib_internal.h
+++ b/firmware/lib/cgptlib/include/cgptlib_internal.h
@@ -125,6 +125,10 @@
  * Repair GPT data by copying from one set of valid headers/entries to the
  * other.  Assumes GptSanityCheck() has been run to determine which headers
  * and/or entries are already valid.
+ *
+ * The caller must make sure that even if one of the entries table is invalid
+ * then corresponding buffer is allocated and big enough to accommodate entries
+ * from the other (good) table.
  */
 void GptRepair(GptData *gpt);
 
diff --git a/tests/run_cgpt_tests.sh b/tests/run_cgpt_tests.sh
index 5176017..8a46d03 100755
--- a/tests/run_cgpt_tests.sh
+++ b/tests/run_cgpt_tests.sh
@@ -269,6 +269,22 @@
 $CGPT prioritize $MTD -i 1 -f ${DEV}
 assert_pri 15 15 13 12 14 11 10 10  9  9  8  8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 1 1 1 1 0
 
+echo "Test cgpt repair command"
+$CGPT repair $MTD ${DEV}
+($CGPT show $MTD ${DEV} | grep -q INVALID) && error
+
+# Zero primary header and partition table and try to repair it.
+dd if=/dev/zero of=${DEV} conv=notrunc bs=512 count=33 2>/dev/null
+$CGPT show $MTD ${DEV} | grep -q INVALID
+$CGPT repair $MTD ${DEV}
+($CGPT show $MTD ${DEV} | grep -q INVALID) && error
+
+# Zero secondary header and partition table and try to repair it.
+dd if=/dev/zero of=${DEV} seek=$(($NUM_SECTORS - 33)) conv=notrunc bs=512 count=33 2>/dev/null
+$CGPT show $MTD ${DEV} | grep -q INVALID
+$CGPT repair $MTD ${DEV}
+($CGPT show $MTD ${DEV} | grep -q INVALID) && error
+
 # Now make sure that we don't need write access if we're just looking.
 echo "Test read vs read-write access..."
 chmod 0444 ${DEV}