More cleanup

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

9 files changed, 144 insertions, 192 deletions

diff --git a/vboot_firmware/lib/cgptlib/cgptlib.c b/vboot_firmware/lib/cgptlib/cgptlib.c
index 3c92384a..4856311a 100644
--- a/vboot_firmware/lib/cgptlib/cgptlib.c
+++ b/vboot_firmware/lib/cgptlib/cgptlib.c
@@ -95,7 +95,6 @@ int GptUpdateKernelEntry(GptData* gpt, uint32_t update_type) {
   GptEntry* e = entries + gpt->current_kernel;
   uint64_t previous_attr = e->attributes;
 
-  /* TODO: need a better return code for these errors? */
   if (gpt->current_kernel == CGPT_KERNEL_ENTRY_NOT_FOUND)
     return GPT_ERROR_INVALID_UPDATE_TYPE;
   if (!IsKernelEntry(e))
diff --git a/vboot_firmware/lib/cgptlib/cgptlib_internal.c b/vboot_firmware/lib/cgptlib/cgptlib_internal.c
index 0a4f46d5..4f022936 100644
--- a/vboot_firmware/lib/cgptlib/cgptlib_internal.c
+++ b/vboot_firmware/lib/cgptlib/cgptlib_internal.c
@@ -60,7 +60,7 @@ int CheckHeader(GptHeader *h, int is_secondary, uint64_t drive_sectors) {
   if (h->reserved_zero)
     return 1;
 
-  /* TODO: Padding must be set to zero. */
+  /* Could check that padding is zero, but that doesn't matter to us. */
 
   /* If entry size is different than our struct, we won't be able to
    * parse it.  Technically, any size 2^N where N>=7 is valid. */
diff --git a/vboot_firmware/lib/cgptlib/include/cgptlib.h b/vboot_firmware/lib/cgptlib/include/cgptlib.h
index acf91dc0..320b5ed8 100644
--- a/vboot_firmware/lib/cgptlib/include/cgptlib.h
+++ b/vboot_firmware/lib/cgptlib/include/cgptlib.h
@@ -24,6 +24,10 @@ enum {
 #define GPT_MODIFIED_ENTRIES1 0x04
 #define GPT_MODIFIED_ENTRIES2 0x08
 
+#define TOTAL_ENTRIES_SIZE 16384  /* Size of GptData.primary_entries
+                                   * and secondary_entries: 128
+                                   * bytes/entry * 128 entries. */
+
 /* The 'update_type' of GptUpdateKernelEntry()
  * We expose TRY and BAD only because those are what verified boot needs.
  * For more precise control on GPT attribute bits, please refer to
diff --git a/vboot_firmware/lib/cgptlib/include/cgptlib_internal.h b/vboot_firmware/lib/cgptlib/include/cgptlib_internal.h
index 0b993c36..b8d85aa9 100644
--- a/vboot_firmware/lib/cgptlib/include/cgptlib_internal.h
+++ b/vboot_firmware/lib/cgptlib/include/cgptlib_internal.h
@@ -44,7 +44,6 @@
                                       CGPT_ATTRIBUTE_PRIORITY_OFFSET)
 
 /* Defines ChromeOS-specific limitation on GPT */
-/* TODO: Move these to cgptlib_internal.h */
 #define MIN_SIZE_OF_HEADER 92
 #define MAX_SIZE_OF_HEADER 512
 #define MIN_SIZE_OF_ENTRY 128
@@ -52,7 +51,6 @@
 #define SIZE_OF_ENTRY_MULTIPLE 8
 #define MIN_NUMBER_OF_ENTRIES 32
 #define MAX_NUMBER_OF_ENTRIES 512
-#define TOTAL_ENTRIES_SIZE 16384  /* usual case is 128 bytes * 128 entries */
 
 /* Defines GPT sizes */
 #define GPT_PMBR_SECTOR 1  /* size (in sectors) of PMBR */
diff --git a/vboot_firmware/lib/include/vboot_common.h b/vboot_firmware/lib/include/vboot_common.h
index 41c6500b..407746b9 100644
--- a/vboot_firmware/lib/include/vboot_common.h
+++ b/vboot_firmware/lib/include/vboot_common.h
@@ -13,16 +13,20 @@
 #include "cryptolib.h"
 #include "vboot_struct.h"
 
-/* Error Codes for VerifyFirmware. */
-#define VBOOT_SUCCESS 0
-#define VBOOT_INVALID_IMAGE 1
-#define VBOOT_KEY_SIGNATURE_FAILED 2
-#define VBOOT_INVALID_ALGORITHM 3
-#define VBOOT_PREAMBLE_SIGNATURE_FAILED 4
-#define VBOOT_SIGNATURE_FAILED 5
-#define VBOOT_WRONG_MAGIC 6
-#define VBOOT_ERROR_MAX 7  /* Generic catch-all. */
-
+/* Error Codes for all common functions. */
+enum {
+  VBOOT_SUCCESS = 0,
+  VBOOT_KEY_BLOCK_INVALID,    /* Key block internal structure is
+                               * invalid, or not a key block */
+  VBOOT_KEY_BLOCK_SIGNATURE,  /* Key block signature check failed */
+  VBOOT_KEY_BLOCK_HASH,      /* Key block hash check failed */
+  VBOOT_PUBLIC_KEY_INVALID,  /* Invalid public key passed to a
+                              * signature verficiation function. */
+  VBOOT_PREAMBLE_INVALID,    /* Preamble internal structure is
+                              * invalid */
+  VBOOT_PREAMBLE_SIGNATURE,  /* Preamble signature check failed */
+  VBOOT_ERROR_MAX,
+};
 extern char* kVbootErrors[VBOOT_ERROR_MAX];
 
 
diff --git a/vboot_firmware/lib/include/vboot_kernel.h b/vboot_firmware/lib/include/vboot_kernel.h
index f55fd2c5..ffe417c3 100644
--- a/vboot_firmware/lib/include/vboot_kernel.h
+++ b/vboot_firmware/lib/include/vboot_kernel.h
@@ -12,12 +12,7 @@
 #include <stdint.h>
 
 #include "cgptlib.h"
-#include "cryptolib.h"
 #include "load_kernel_fw.h"
-#include "vboot_common.h"
-
-/* TODO: temporary hack */
-void FakePartitionAttributes(GptData* gpt);
 
 /* Allocates and reads GPT data from the drive.  The sector_bytes and
  * drive_sectors fields should be filled on input.  The primary and
@@ -28,7 +23,7 @@ int AllocAndReadGptData(GptData* gptdata);
 
 /* Writes any changes for the GPT data back to the drive, then frees the
  * buffers. */
-void WriteAndFreeGptData(GptData* gptdata);
+int WriteAndFreeGptData(GptData* gptdata);
 
 /* Alternate LoadKernel() implementation; see load_kernel_fw.h */
 int LoadKernel2(LoadKernelParams* params);
diff --git a/vboot_firmware/lib/load_kernel_fw.c b/vboot_firmware/lib/load_kernel_fw.c
index 2a17c6cd..88870679 100644
--- a/vboot_firmware/lib/load_kernel_fw.c
+++ b/vboot_firmware/lib/load_kernel_fw.c
@@ -13,6 +13,7 @@
 #include "kernel_image_fw.h"
 #include "rollback_index.h"
 #include "utility.h"
+#include "vboot_kernel.h"
 
 #define GPT_ENTRIES_SIZE 16384 /* Bytes to read for GPT entries */
 
@@ -20,134 +21,9 @@
 // TODO: for testing
 #include <stdio.h>
 #include <inttypes.h>  /* For PRIu64 macro */
-#endif
-
-/* TODO: Remove this terrible hack which fakes partition attributes
- * for the kernel partitions so that GptNextKernelEntry() won't
- * choke. */
 #include "cgptlib_internal.h"
-void FakePartitionAttributes(GptData* gpt) {
-  GptHeader* h = (GptHeader*)gpt->primary_header;
-  GptEntry* entries = (GptEntry*)gpt->primary_entries;
-  GptEntry* e;
-  int i;
-
-  for (i = 0, e = entries; i < h->number_of_entries; i++, e++) {
-    if (!IsKernelEntry(e))
-      continue;
-
-#ifdef PRINT_DEBUG_INFO
-
-    printf("%2d %08x %04x %04x %02x %02x %02x %02x %02x %02x %02x %02x",
-           i,
-           e->type.u.Uuid.time_low,
-           e->type.u.Uuid.time_mid,
-           e->type.u.Uuid.time_high_and_version,
-           e->type.u.Uuid.clock_seq_high_and_reserved,
-           e->type.u.Uuid.clock_seq_low,
-           e->type.u.Uuid.node[0],
-           e->type.u.Uuid.node[1],
-           e->type.u.Uuid.node[2],
-           e->type.u.Uuid.node[3],
-           e->type.u.Uuid.node[4],
-           e->type.u.Uuid.node[5]
-           );
-    printf(" %8" PRIu64 " %8" PRIu64"\n", e->starting_lba,
-           e->ending_lba - e->starting_lba + 1);
-    printf("Hacking attributes for kernel partition %d\n", i);
 #endif
 
-    SetEntryPriority(e, 2);
-    SetEntrySuccessful(e, 1);
-  }
-}
-
-
-/* Allocates and reads GPT data from the drive.  The sector_bytes and
- * drive_sectors fields should be filled on input.  The primary and
- * secondary header and entries are filled on output.
- *
- * Returns 0 if successful, 1 if error. */
-int AllocAndReadGptData(GptData* gptdata) {
-
-  uint64_t entries_sectors = GPT_ENTRIES_SIZE / gptdata->sector_bytes;
-
-  /* No data to be written yet */
-  gptdata->modified = 0;
-
-  /* Allocate all buffers */
-  gptdata->primary_header = (uint8_t*)Malloc(gptdata->sector_bytes);
-  gptdata->secondary_header = (uint8_t*)Malloc(gptdata->sector_bytes);
-  gptdata->primary_entries = (uint8_t*)Malloc(GPT_ENTRIES_SIZE);
-  gptdata->secondary_entries = (uint8_t*)Malloc(GPT_ENTRIES_SIZE);
-
-  if (gptdata->primary_header == NULL || gptdata->secondary_header == NULL ||
-      gptdata->primary_entries == NULL || gptdata->secondary_entries == NULL)
-    return 1;
-
-  /* Read data from the drive, skipping the protective MBR */
-  if (0 != BootDeviceReadLBA(1, 1, gptdata->primary_header))
-    return 1;
-  if (0 != BootDeviceReadLBA(2, entries_sectors, gptdata->primary_entries))
-    return 1;
-  if (0 != BootDeviceReadLBA(gptdata->drive_sectors - entries_sectors - 1,
-                             entries_sectors, gptdata->secondary_entries))
-    return 1;
-  if (0 != BootDeviceReadLBA(gptdata->drive_sectors - 1,
-                             1, gptdata->secondary_header))
-    return 1;
-
-  return 0;
-}
-
-
-/* Writes any changes for the GPT data back to the drive, then frees
- * the buffers.
- *
- * Returns 0 if successful, 1 if error. */
-int WriteAndFreeGptData(GptData* gptdata) {
-
-  uint64_t entries_sectors = GPT_ENTRIES_SIZE / gptdata->sector_bytes;
-
-  if (gptdata->primary_header) {
-    if (gptdata->modified & GPT_MODIFIED_HEADER1) {
-      if (0 != BootDeviceWriteLBA(1, 1, gptdata->primary_header))
-        return 1;
-    }
-    Free(gptdata->primary_header);
-  }
-
-  if (gptdata->primary_entries) {
-    if (gptdata->modified & GPT_MODIFIED_ENTRIES1) {
-      if (0 != BootDeviceWriteLBA(2, entries_sectors,
-                                  gptdata->primary_entries))
-        return 1;
-    }
-    Free(gptdata->primary_entries);
-  }
-
-  if (gptdata->secondary_entries) {
-    if (gptdata->modified & GPT_MODIFIED_ENTRIES2) {
-      if (0 != BootDeviceWriteLBA(gptdata->drive_sectors - entries_sectors - 1,
-                                  entries_sectors, gptdata->secondary_entries))
-        return 1;
-    }
-    Free(gptdata->secondary_entries);
-  }
-
-  if (gptdata->secondary_header) {
-    if (gptdata->modified & GPT_MODIFIED_HEADER2) {
-      if (0 != BootDeviceWriteLBA(gptdata->drive_sectors - 1, 1,
-                                  gptdata->secondary_header))
-        return 1;
-    }
-    Free(gptdata->secondary_header);
-  }
-
-  /* Success */
-  return 0;
-}
-
 
 #define KBUF_SIZE 65536  /* Bytes to read at start of kernel partition */
 
@@ -194,9 +70,6 @@ int LoadKernel(LoadKernelParams* params) {
     if (GPT_SUCCESS != GptInit(&gpt))
       break;
 
-    /* TODO: TERRIBLE KLUDGE - fake partition attributes */
-    FakePartitionAttributes(&gpt);
-
     /* Allocate kernel header and image work buffers */
     kbuf = (uint8_t*)Malloc(KBUF_SIZE);
     if (!kbuf)
diff --git a/vboot_firmware/lib/vboot_common.c b/vboot_firmware/lib/vboot_common.c
index fb643bdc..3eac471e 100644
--- a/vboot_firmware/lib/vboot_common.c
+++ b/vboot_firmware/lib/vboot_common.c
@@ -6,21 +6,19 @@
  * (Firmware portion)
  */
 
-/* TODO: change all 'return 0', 'return 1' into meaningful return codes */
 
 #include "vboot_common.h"
 #include "utility.h"
 
-#include <stdio.h> /* TODO: FOR TESTING */
 
 char* kVbootErrors[VBOOT_ERROR_MAX] = {
   "Success.",
-  "Invalid Image.",
-  "Kernel Key Signature Failed.",
-  "Invalid Kernel Verification Algorithm.",
-  "Preamble Signature Failed.",
-  "Kernel Signature Failed.",
-  "Wrong Kernel Magic.",
+  "Key block invalid.",
+  "Key block signature failed.",
+  "Key block hash failed.",
+  "Public key invalid.",
+  "Preamble invalid.",
+  "Preamble signature check failed.",
 };
 
 
@@ -132,15 +130,15 @@ int KeyBlockVerify(const VbKeyBlockHeader* block, uint64_t size,
   /* Sanity checks before attempting signature of data */
   if (SafeMemcmp(block->magic, KEY_BLOCK_MAGIC, KEY_BLOCK_MAGIC_SIZE)) {
     debug("Not a valid verified boot key block.\n");
-    return 1;
+    return VBOOT_KEY_BLOCK_INVALID;
   }
   if (block->header_version_major != KEY_BLOCK_HEADER_VERSION_MAJOR) {
     debug("Incompatible key block header version.\n");
-    return 1;
+    return VBOOT_KEY_BLOCK_INVALID;
   }
   if (size < block->key_block_size) {
     debug("Not enough data for key block.\n");
-    return 1;
+    return VBOOT_KEY_BLOCK_INVALID;
   }
 
   /* Check signature or hash, depending on whether we have a key. */
@@ -153,18 +151,17 @@ int KeyBlockVerify(const VbKeyBlockHeader* block, uint64_t size,
 
     if (VerifySignatureInside(block, block->key_block_size, sig)) {
       debug("Key block signature off end of block\n");
-      return 1;
+      return VBOOT_KEY_BLOCK_INVALID;
     }
 
     if (!((rsa = PublicKeyToRSA(key)))) {
       debug("Invalid public key\n");
-      return 1;
+      return VBOOT_PUBLIC_KEY_INVALID;
     }
     rv = VerifyData((const uint8_t*)block, sig, rsa);
     RSAPublicKeyFree(rsa);
-
     if (rv)
-      return rv;
+      return VBOOT_KEY_BLOCK_SIGNATURE;
 
   } else {
     /* Check hash */
@@ -175,11 +172,11 @@ int KeyBlockVerify(const VbKeyBlockHeader* block, uint64_t size,
 
     if (VerifySignatureInside(block, block->key_block_size, sig)) {
       debug("Key block hash off end of block\n");
-      return 1;
+      return VBOOT_KEY_BLOCK_INVALID;
     }
     if (sig->sig_size != SHA512_DIGEST_SIZE) {
       debug("Wrong hash size for key block.\n");
-      return 1;
+      return VBOOT_KEY_BLOCK_INVALID;
     }
 
     header_checksum = DigestBuf((const uint8_t*)block, sig->data_size,
@@ -189,28 +186,28 @@ int KeyBlockVerify(const VbKeyBlockHeader* block, uint64_t size,
     Free(header_checksum);
     if (rv) {
       debug("Invalid key block hash.\n");
-      return 1;
+      return VBOOT_KEY_BLOCK_HASH;
     }
   }
 
   /* Verify we signed enough data */
   if (sig->data_size < sizeof(VbKeyBlockHeader)) {
     debug("Didn't sign enough data\n");
-    return 1;
+    return VBOOT_KEY_BLOCK_INVALID;
   }
 
   /* Verify data key is inside the block and inside signed data */
   if (VerifyPublicKeyInside(block, block->key_block_size, &block->data_key)) {
     debug("Data key off end of key block\n");
-    return 1;
+    return VBOOT_KEY_BLOCK_INVALID;
   }
   if (VerifyPublicKeyInside(block, sig->data_size, &block->data_key)) {
     debug("Data key off end of signed data\n");
-    return 1;
+    return VBOOT_KEY_BLOCK_INVALID;
   }
 
   /* Success */
-  return 0;
+  return VBOOT_SUCCESS;
 }
 
 
@@ -219,51 +216,49 @@ int VerifyFirmwarePreamble2(const VbFirmwarePreambleHeader* preamble,
 
   const VbSignature* sig = &preamble->preamble_signature;
 
-  /* TODO: caller needs to make sure key version is valid */
-
   /* Sanity checks before attempting signature of data */
   if (preamble->header_version_major !=
       FIRMWARE_PREAMBLE_HEADER_VERSION_MAJOR) {
     debug("Incompatible firmware preamble header version.\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
   if (size < preamble->preamble_size) {
     debug("Not enough data for preamble.\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
 
   /* Check signature */
   if (VerifySignatureInside(preamble, preamble->preamble_size, sig)) {
     debug("Preamble signature off end of preamble\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
   if (VerifyData((const uint8_t*)preamble, sig, key)) {
     debug("Preamble signature validation failed\n");
-    return 1;
+    return VBOOT_PREAMBLE_SIGNATURE;
   }
 
   /* Verify we signed enough data */
   if (sig->data_size < sizeof(VbFirmwarePreambleHeader)) {
     debug("Didn't sign enough data\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
 
   /* Verify body signature is inside the block */
   if (VerifySignatureInside(preamble, preamble->preamble_size,
                             &preamble->body_signature)) {
     debug("Firmware body signature off end of preamble\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
 
   /* Verify kernel subkey is inside the block */
   if (VerifyPublicKeyInside(preamble, preamble->preamble_size,
                             &preamble->kernel_subkey)) {
     debug("Kernel subkey off end of preamble\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
 
   /* Success */
-  return 0;
+  return VBOOT_SUCCESS;
 }
 
 
@@ -272,41 +267,39 @@ int VerifyKernelPreamble2(const VbKernelPreambleHeader* preamble,
 
   const VbSignature* sig = &preamble->preamble_signature;
 
-  /* TODO: caller needs to make sure key version is valid */
-
   /* Sanity checks before attempting signature of data */
   if (preamble->header_version_major != KERNEL_PREAMBLE_HEADER_VERSION_MAJOR) {
     debug("Incompatible kernel preamble header version.\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
   if (size < preamble->preamble_size) {
     debug("Not enough data for preamble.\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
 
   /* Check signature */
   if (VerifySignatureInside(preamble, preamble->preamble_size, sig)) {
     debug("Preamble signature off end of preamble\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
   if (VerifyData((const uint8_t*)preamble, sig, key)) {
     debug("Preamble signature validation failed\n");
-    return 1;
+    return VBOOT_PREAMBLE_SIGNATURE;
   }
 
   /* Verify we signed enough data */
   if (sig->data_size < sizeof(VbKernelPreambleHeader)) {
     debug("Didn't sign enough data\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
 
   /* Verify body signature is inside the block */
   if (VerifySignatureInside(preamble, preamble->preamble_size,
                             &preamble->body_signature)) {
     debug("Kernel body signature off end of preamble\n");
-    return 1;
+    return VBOOT_PREAMBLE_INVALID;
   }
 
   /* Success */
-  return 0;
+  return VBOOT_SUCCESS;
 }
diff --git a/vboot_firmware/lib/vboot_kernel.c b/vboot_firmware/lib/vboot_kernel.c
index ca8ba99c..9e5dc52f 100644
--- a/vboot_firmware/lib/vboot_kernel.c
+++ b/vboot_firmware/lib/vboot_kernel.c
@@ -13,9 +13,98 @@
 #include "load_kernel_fw.h"
 #include "rollback_index.h"
 #include "utility.h"
+#include "vboot_common.h"
+
 
 #define KBUF_SIZE 65536  /* Bytes to read at start of kernel partition */
 
+
+/* Allocates and reads GPT data from the drive.  The sector_bytes and
+ * drive_sectors fields should be filled on input.  The primary and
+ * secondary header and entries are filled on output.
+ *
+ * Returns 0 if successful, 1 if error. */
+int AllocAndReadGptData(GptData* gptdata) {
+
+  uint64_t entries_sectors = TOTAL_ENTRIES_SIZE / gptdata->sector_bytes;
+
+  /* No data to be written yet */
+  gptdata->modified = 0;
+
+  /* Allocate all buffers */
+  gptdata->primary_header = (uint8_t*)Malloc(gptdata->sector_bytes);
+  gptdata->secondary_header = (uint8_t*)Malloc(gptdata->sector_bytes);
+  gptdata->primary_entries = (uint8_t*)Malloc(TOTAL_ENTRIES_SIZE);
+  gptdata->secondary_entries = (uint8_t*)Malloc(TOTAL_ENTRIES_SIZE);
+
+  if (gptdata->primary_header == NULL || gptdata->secondary_header == NULL ||
+      gptdata->primary_entries == NULL || gptdata->secondary_entries == NULL)
+    return 1;
+
+  /* Read data from the drive, skipping the protective MBR */
+  if (0 != BootDeviceReadLBA(1, 1, gptdata->primary_header))
+    return 1;
+  if (0 != BootDeviceReadLBA(2, entries_sectors, gptdata->primary_entries))
+    return 1;
+  if (0 != BootDeviceReadLBA(gptdata->drive_sectors - entries_sectors - 1,
+                             entries_sectors, gptdata->secondary_entries))
+    return 1;
+  if (0 != BootDeviceReadLBA(gptdata->drive_sectors - 1,
+                             1, gptdata->secondary_header))
+    return 1;
+
+  return 0;
+}
+
+
+/* Writes any changes for the GPT data back to the drive, then frees
+ * the buffers.
+ *
+ * Returns 0 if successful, 1 if error. */
+int WriteAndFreeGptData(GptData* gptdata) {
+
+  uint64_t entries_sectors = TOTAL_ENTRIES_SIZE / gptdata->sector_bytes;
+
+  if (gptdata->primary_header) {
+    if (gptdata->modified & GPT_MODIFIED_HEADER1) {
+      if (0 != BootDeviceWriteLBA(1, 1, gptdata->primary_header))
+        return 1;
+    }
+    Free(gptdata->primary_header);
+  }
+
+  if (gptdata->primary_entries) {
+    if (gptdata->modified & GPT_MODIFIED_ENTRIES1) {
+      if (0 != BootDeviceWriteLBA(2, entries_sectors,
+                                  gptdata->primary_entries))
+        return 1;
+    }
+    Free(gptdata->primary_entries);
+  }
+
+  if (gptdata->secondary_entries) {
+    if (gptdata->modified & GPT_MODIFIED_ENTRIES2) {
+      if (0 != BootDeviceWriteLBA(gptdata->drive_sectors - entries_sectors - 1,
+                                  entries_sectors, gptdata->secondary_entries))
+        return 1;
+    }
+    Free(gptdata->secondary_entries);
+  }
+
+  if (gptdata->secondary_header) {
+    if (gptdata->modified & GPT_MODIFIED_HEADER2) {
+      if (0 != BootDeviceWriteLBA(gptdata->drive_sectors - 1, 1,
+                                  gptdata->secondary_header))
+        return 1;
+    }
+    Free(gptdata->secondary_header);
+  }
+
+  /* Success */
+  return 0;
+}
+
+
 int LoadKernel2(LoadKernelParams* params) {
 
   VbPublicKey* kernel_subkey = (VbPublicKey*)params->header_sign_key_blob;
@@ -65,9 +154,6 @@ int LoadKernel2(LoadKernelParams* params) {
     if (GPT_SUCCESS != GptInit(&gpt))
       break;
 
-    /* TODO: TERRIBLE KLUDGE - fake partition attributes */
-    FakePartitionAttributes(&gpt);
-
     /* Allocate kernel header buffers */
     kbuf = (uint8_t*)Malloc(KBUF_SIZE);
     if (!kbuf)