Fixing the bug of CGPT when primary entry table is invalid.

http://code.google.com/p/chromium-os/issues/detail?id=9279
This issue disclosed a bug of cgpt. The bug comes from the 'show' command always
reads the primary entry table when '-i partition' is specified. I added an
ANY_VALID constant for GetEntry to automatically select valid entry table.

Also fixed the bugs in cmd_boot.c and cmd_find.c. In cmd_add.c, stop user to
continue if any header/entry table is invalid.

Also fixed the bug that untrusted header size could cause segmentation failure.

Hungte, this is FYI. But welcome to do review.

BUG=chromium-os:9279
TEST=RUNTESTS=1 emerge-x86-generic vboot_reference
     Manually tested:
       cgpt show /tmp/test -i 1 -b
       cgpt show /tmp/test
       cgpt add /tmp/test -i 1 -l TEST
       cgpt find /tmp/test -l STATE
       cgpt boot /tmp/test -i 1

Change-Id: Iaba9c635754096a82b3ec74634af184362d4e264
Change-Id: I6f3e87e3998457676e3388d2a6ed36c0564796d8

Review URL: http://codereview.chromium.org/5115002

5 files changed, 27 insertions, 13 deletions

diff --git a/cgpt/cgpt_common.c b/cgpt/cgpt_common.c
index 4b9a5a14..0e466fdc 100644
--- a/cgpt/cgpt_common.c
+++ b/cgpt/cgpt_common.c
@@ -480,8 +480,16 @@ GptEntry *GetEntry(GptData *gpt, int secondary, uint32_t entry_index) {
 
   if (secondary == PRIMARY) {
     entries = gpt->primary_entries;
-  } else {
+  } else if (secondary == SECONDARY) {
     entries = gpt->secondary_entries;
+  } else {  /* ANY_VALID */
+    require(secondary == ANY_VALID);
+    if (gpt->valid_entries & MASK_PRIMARY) {
+      entries = gpt->primary_entries;
+    } else {
+      require(gpt->valid_entries & MASK_SECONDARY);
+      entries = gpt->secondary_entries;
+    }
   }
 
   return (GptEntry*)(&entries[stride * entry_index]);
@@ -570,12 +578,12 @@ void UpdateCrc(GptData *gpt) {
   if (gpt->modified & GPT_MODIFIED_HEADER1) {
     primary_header->header_crc32 = 0;
     primary_header->header_crc32 = Crc32(
-        (const uint8_t *)primary_header, primary_header->size);
+        (const uint8_t *)primary_header, sizeof(GptHeader));
   }
   if (gpt->modified & GPT_MODIFIED_HEADER2) {
     secondary_header->header_crc32 = 0;
     secondary_header->header_crc32 = Crc32(
-        (const uint8_t *)secondary_header, secondary_header->size);
+        (const uint8_t *)secondary_header, sizeof(GptHeader));
   }
 }
 /* Two headers are NOT bitwise identical. For example, my_lba pointers to header
@@ -662,14 +670,14 @@ uint8_t RepairHeader(GptData *gpt, const uint32_t valid_headers) {
       return GPT_MODIFIED_HEADER2;
     }
   } else if (valid_headers == MASK_PRIMARY) {
-    memcpy(secondary_header, primary_header, primary_header->size);
+    memcpy(secondary_header, primary_header, sizeof(GptHeader));
     secondary_header->my_lba = gpt->drive_sectors - 1;  /* the last sector */
     secondary_header->alternate_lba = primary_header->my_lba;
     secondary_header->entries_lba = secondary_header->my_lba -
         GPT_ENTRIES_SECTORS;
     return GPT_MODIFIED_HEADER2;
   } else if (valid_headers == MASK_SECONDARY) {
-    memcpy(primary_header, secondary_header, secondary_header->size);
+    memcpy(primary_header, secondary_header, sizeof(GptHeader));
     primary_header->my_lba = GPT_PMBR_SECTOR;  /* the second sector on drive */
     primary_header->alternate_lba = secondary_header->my_lba;
     primary_header->entries_lba = primary_header->my_lba + GPT_HEADER_SECTOR;
@@ -693,4 +701,3 @@ void PMBRToStr(struct pmbr *pmbr, char *str, unsigned int buflen) {
     require(snprintf(str, buflen, "PMBR (Boot GUID: %s)", buf) < buflen);
   }
 }
-
diff --git a/cgpt/cmd_add.c b/cgpt/cmd_add.c
index 21085e71..dafcc50f 100644
--- a/cgpt/cmd_add.c
+++ b/cgpt/cmd_add.c
@@ -198,6 +198,13 @@ int cmd_add(int argc, char *argv[]) {
     return CGPT_FAILED;
   }
 
+  if (((drive.gpt.valid_headers & MASK_BOTH) != MASK_BOTH) ||
+      ((drive.gpt.valid_entries & MASK_BOTH) != MASK_BOTH)) {
+    Error("one of the GPT header/entries is invalid.\n"
+          "please run 'cgpt repair' before adding anything.\n");
+    return CGPT_FAILED;
+  }
+
   uint32_t max_part = GetNumberOfEntries(&drive.gpt);
   if (partition) {
     if (partition > max_part) {
diff --git a/cgpt/cmd_boot.c b/cgpt/cmd_boot.c
index 56ab1d58..669e9ae7 100644
--- a/cgpt/cmd_boot.c
+++ b/cgpt/cmd_boot.c
@@ -132,7 +132,7 @@ int cmd_boot(int argc, char *argv[]) {
     }
 
     uint32_t index = partition - 1;
-    GptEntry *entry = GetEntry(&drive.gpt, PRIMARY, index);
+    GptEntry *entry = GetEntry(&drive.gpt, ANY_VALID, index);
     memcpy(&drive.pmbr.boot_guid, &entry->unique, sizeof(Guid));
   }
 
diff --git a/cgpt/cmd_find.c b/cgpt/cmd_find.c
index e8fe9b31..40f10ba0 100644
--- a/cgpt/cmd_find.c
+++ b/cgpt/cmd_find.c
@@ -171,7 +171,7 @@ static int do_search(char *filename) {
   }
 
   for (i = 0; i < GetNumberOfEntries(&drive.gpt); ++i) {
-    entry = GetEntry(&drive.gpt, PRIMARY, i);
+    entry = GetEntry(&drive.gpt, ANY_VALID, i);
 
     if (IsZero(&entry->type))
       continue;
diff --git a/cgpt/cmd_show.c b/cgpt/cmd_show.c
index a9963b5b..920a30fe 100644
--- a/cgpt/cmd_show.c
+++ b/cgpt/cmd_show.c
@@ -274,7 +274,7 @@ int cmd_show(int argc, char *argv[]) {
     }
 
     uint32_t index = partition - 1;
-    GptEntry *entry = GetEntry(&drive.gpt, PRIMARY, index);
+    GptEntry *entry = GetEntry(&drive.gpt, ANY_VALID, index);
     char buf[256];                      // scratch buffer for string conversion
 
     if (single_item) {
@@ -299,13 +299,13 @@ int cmd_show(int argc, char *argv[]) {
         printf("%s\n", buf);
         break;
       case 'S':
-        printf("%d\n", GetSuccessful(&drive.gpt, PRIMARY, index));
+        printf("%d\n", GetSuccessful(&drive.gpt, ANY_VALID, index));
         break;
       case 'T':
-        printf("%d\n", GetTries(&drive.gpt, PRIMARY, index));
+        printf("%d\n", GetTries(&drive.gpt, ANY_VALID, index));
         break;
       case 'P':
-        printf("%d\n", GetPriority(&drive.gpt, PRIMARY, index));
+        printf("%d\n", GetPriority(&drive.gpt, ANY_VALID, index));
         break;
       case 'A':
         printf("0x%x\n", entry->attrs.fields.gpt_att);
@@ -322,7 +322,7 @@ int cmd_show(int argc, char *argv[]) {
     char type[GUID_STRLEN];
 
     for (i = 0; i < GetNumberOfEntries(&drive.gpt); ++i) {
-      entry = GetEntry(&drive.gpt, PRIMARY, i);
+      entry = GetEntry(&drive.gpt, ANY_VALID, i);
 
       if (IsZero(&entry->type))
         continue;