vboot_kernel: Validate GPT header before using

In AllocAndReadGptData() the code was changed to use the GPT header
to determine the LBA of the GPT entries.  This change did not account
for devices that have an invalid header and it can attempt to read
from invalid block addresses on a device.

This commit happened here:
a2d72f7 vboot: cgpt: Refer to partition entries by entries_lba.
https://chromium-review.googlesource.com/213861

The subsequent steps in vboot, LoadKernel->GptInit->GptRepair will
fix a missing header and entries, so it is only necessary for one of
the headers to be valid.

This is commonly the case with a new USB stick that has an image
written to it as only the primary header will be valid in this case.
However it is also true if the primary header has been corrupted and
the secondary header is still valid.

The code has been changed to call CheckHeader() on the primary and
secondary headers before attempting to use the 'entries_lba' field
to read the entries from the device.  AllocAndReadGptData() now only
fails if both headers are invalid.

A number of new unit tests are created to check for these failure
conditions.  In order to support this I had to extend the vboot_kernel
test infrastructure to have a buffer for the mocked disk data instead
of just ignoring reads and writes.  This is because many of the existing
tests assumed they could have an invalid GPT header and still pass.
Now that the header is checked it is necessary for a valid header to
be created before the tests can pass.

BUG=chrome-os-partner:32386
BRANCH=samus,auron
TEST=All unit tests pass when running 'make runtests'
In addition real-world testing was done by corrupting the primary
and/or secondary headers of USB stick to ensure that it will
successfully boot if one of the headers is valid.

Change-Id: I7f840a44742fa3ba9a124df29ab5749e4c5a40c1
Signed-off-by: Duncan Laurie <dlaurie@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/220757
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Reviewed-by: Nam Nguyen <namnguyen@chromium.org>
(cherry picked from commit 162f7885964b39419419f07a576fc390068678a9)
Reviewed-on: https://chromium-review.googlesource.com/267454
Tested-by: Zhuo-hao Lee <zhuo-hao.lee@intel.com>
Commit-Queue: Zhuo-hao Lee <zhuo-hao.lee@intel.com>
diff --git a/firmware/lib/vboot_kernel.c b/firmware/lib/vboot_kernel.c
index df32fc0..8c985a9 100644
--- a/firmware/lib/vboot_kernel.c
+++ b/firmware/lib/vboot_kernel.c
@@ -39,6 +39,7 @@
 int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 {
 	uint64_t entries_sectors = TOTAL_ENTRIES_SIZE / gptdata->sector_bytes;
+	int primary_valid = 0, secondary_valid = 0;
 
 	/* No data to be written yet */
 	gptdata->modified = 0;
@@ -56,23 +57,43 @@
 	    gptdata->secondary_entries == NULL)
 		return 1;
 
-	/* Read data from the drive, skipping the protective MBR */
+	/* Read primary header from the drive, skipping the protective MBR */
 	if (0 != VbExDiskRead(disk_handle, 1, 1, gptdata->primary_header))
 		return 1;
-	GptHeader* primary_header = (GptHeader*)gptdata->primary_header;
-	if (0 != VbExDiskRead(disk_handle, primary_header->entries_lba,
-			      entries_sectors, gptdata->primary_entries))
-		return 1;
 
+	/* Only read primary GPT if the primary header is valid */
+	GptHeader* primary_header = (GptHeader*)gptdata->primary_header;
+	if (0 == CheckHeader(primary_header, 0, gptdata->drive_sectors)) {
+		primary_valid = 1;
+		if (0 != VbExDiskRead(disk_handle,
+				      primary_header->entries_lba,
+				      entries_sectors,
+				      gptdata->primary_entries))
+			return 1;
+	} else {
+		VBDEBUG(("Primary GPT header invalid!\n"));
+	}
+
+	/* Read secondary header from the end of the drive */
 	if (0 != VbExDiskRead(disk_handle, gptdata->drive_sectors - 1, 1,
 			      gptdata->secondary_header))
 		return 1;
-	GptHeader* secondary_header = (GptHeader*)gptdata->secondary_header;
-	if (0 != VbExDiskRead(disk_handle, secondary_header->entries_lba,
-			      entries_sectors, gptdata->secondary_entries))
-		return 1;
 
-	return 0;
+	/* Only read secondary GPT if the secondary header is valid */
+	GptHeader* secondary_header = (GptHeader*)gptdata->secondary_header;
+	if (0 == CheckHeader(secondary_header, 1, gptdata->drive_sectors)) {
+		secondary_valid = 1;
+		if (0 != VbExDiskRead(disk_handle,
+				      secondary_header->entries_lba,
+				      entries_sectors,
+				      gptdata->secondary_entries))
+			return 1;
+	} else {
+		VBDEBUG(("Secondary GPT header invalid!\n"));
+	}
+
+	/* Return 0 if least one GPT header was valid */
+	return (primary_valid || secondary_valid) ? 0 : 1;
 }
 
 /**
diff --git a/tests/vboot_kernel_tests.c b/tests/vboot_kernel_tests.c
index 935dd86..3712d0d 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vboot_kernel_tests.c
@@ -11,6 +11,7 @@
 #include <string.h>
 
 #include "cgptlib.h"
+#include "cgptlib_internal.h"
 #include "gbb_header.h"
 #include "gpt.h"
 #include "host_common.h"
@@ -24,6 +25,9 @@
 #define LOGCALL(fmt, args...) sprintf(call_log + strlen(call_log), fmt, ##args)
 #define TEST_CALLS(expect_log) TEST_STR_EQ(call_log, expect_log, "  calls")
 
+#define MOCK_SECTOR_SIZE  512
+#define MOCK_SECTOR_COUNT 1024
+
 /* Mock kernel partition */
 struct mock_part {
 	uint32_t start;
@@ -57,6 +61,43 @@
 static VbKeyBlockHeader kbh;
 static VbKernelPreambleHeader kph;
 static VbCommonParams cparams;
+static uint8_t mock_disk[MOCK_SECTOR_SIZE * MOCK_SECTOR_COUNT];
+static GptHeader *mock_gpt_primary =
+	(GptHeader*)&mock_disk[MOCK_SECTOR_SIZE * 1];
+static GptHeader *mock_gpt_secondary =
+	(GptHeader*)&mock_disk[MOCK_SECTOR_SIZE * (MOCK_SECTOR_COUNT - 1)];
+
+
+/**
+ * Prepare a valid GPT header that will pass CheckHeader() tests
+ */
+static void SetupGptHeader(GptHeader *h, int is_secondary)
+{
+	Memset(h, '\0', MOCK_SECTOR_SIZE);
+
+	/* "EFI PART" */
+	memcpy(h->signature, GPT_HEADER_SIGNATURE, GPT_HEADER_SIGNATURE_SIZE);
+	h->revision = GPT_HEADER_REVISION;
+	h->size = MIN_SIZE_OF_HEADER;
+
+	/* 16KB: 128 entries of 128 bytes */
+	h->size_of_entry = sizeof(GptEntry);
+	h->number_of_entries = TOTAL_ENTRIES_SIZE / h->size_of_entry;
+
+	/* Set LBA pointers for primary or secondary header */
+	if (is_secondary) {
+		h->my_lba = MOCK_SECTOR_COUNT - GPT_HEADER_SECTORS;
+		h->entries_lba = h->my_lba - GPT_ENTRIES_SECTORS;
+	} else {
+		h->my_lba = GPT_PMBR_SECTORS;
+		h->entries_lba = h->my_lba + 1;
+	}
+
+	h->first_usable_lba = 2 + GPT_ENTRIES_SECTORS;
+	h->last_usable_lba = MOCK_SECTOR_COUNT - 2 - GPT_ENTRIES_SECTORS;
+
+	h->header_crc32 = HeaderCrc(h);
+}
 
 static void ResetCallLog(void)
 {
@@ -70,6 +111,10 @@
 {
 	ResetCallLog();
 
+	memset(&mock_disk, 0, sizeof(mock_disk));
+	SetupGptHeader(mock_gpt_primary, 0);
+	SetupGptHeader(mock_gpt_secondary, 1);
+
 	disk_read_to_fail = -1;
 	disk_write_to_fail = -1;
 
@@ -137,16 +182,9 @@
 	if ((int)lba_start == disk_read_to_fail)
 		return VBERROR_SIMULATED;
 
-	/* Keep valgrind happy */
-	Memset(buffer, '\0', lba_count);
-	/* Fix up entries_lba in GPT header. */
-	if (lba_start == 1 || lba_start == 1024 - 1) {
-		GptHeader* h = (GptHeader*)buffer;
-		if (lba_start == 1)
-			h->entries_lba = 1 + 1;
-		else
-			h->entries_lba = (1024 - 1 - 32);
-	}
+	memcpy(buffer, &mock_disk[lba_start * MOCK_SECTOR_SIZE],
+	       lba_count * MOCK_SECTOR_SIZE);
+
 	return VBERROR_SUCCESS;
 }
 
@@ -158,6 +196,9 @@
 	if ((int)lba_start == disk_write_to_fail)
 		return VBERROR_SIMULATED;
 
+	memcpy(&mock_disk[lba_start * MOCK_SECTOR_SIZE], buffer,
+	       lba_count * MOCK_SECTOR_SIZE);
+
 	return VBERROR_SUCCESS;
 }
 
@@ -246,8 +287,9 @@
 	GptData g;
 	GptHeader *h;
 
-	g.sector_bytes = 512;
-	g.drive_sectors = 1024;
+	g.sector_bytes = MOCK_SECTOR_SIZE;
+	g.drive_sectors = MOCK_SECTOR_COUNT;
+	g.valid_headers = g.valid_entries = MASK_BOTH;
 
 	ResetMocks();
 	TEST_EQ(AllocAndReadGptData(handle, &g), 0, "AllocAndRead");
@@ -264,6 +306,105 @@
 	TEST_EQ(WriteAndFreeGptData(handle, &g), 0, "WriteAndFree");
 	TEST_CALLS("");
 
+	/*
+	 * Invalidate primary GPT header,
+	 * check that AllocAndReadGptData still succeeds
+	 */
+	ResetMocks();
+	Memset(mock_gpt_primary, '\0', sizeof(*mock_gpt_primary));
+	TEST_EQ(AllocAndReadGptData(handle, &g), 0,
+		"AllocAndRead primary invalid");
+	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.drive_sectors), 1,
+		"Primary header is invalid");
+	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.drive_sectors), 0,
+		"Secondary header is valid");
+	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
+		   "VbExDiskRead(h, 1023, 1)\n"
+		   "VbExDiskRead(h, 991, 32)\n");
+	WriteAndFreeGptData(handle, &g);
+
+	/*
+	 * Invalidate secondary GPT header,
+	 * check that AllocAndReadGptData still succeeds
+	 */
+	ResetMocks();
+	Memset(mock_gpt_secondary, '\0', sizeof(*mock_gpt_secondary));
+	TEST_EQ(AllocAndReadGptData(handle, &g), 0,
+		"AllocAndRead secondary invalid");
+	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.drive_sectors), 0,
+		"Primary header is valid");
+	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.drive_sectors), 1,
+		"Secondary header is invalid");
+	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
+		   "VbExDiskRead(h, 2, 32)\n"
+		   "VbExDiskRead(h, 1023, 1)\n");
+	WriteAndFreeGptData(handle, &g);
+
+	/*
+	 * Invalidate primary AND secondary GPT header,
+	 * check that AllocAndReadGptData fails.
+	 */
+	ResetMocks();
+	Memset(mock_gpt_primary, '\0', sizeof(*mock_gpt_primary));
+	Memset(mock_gpt_secondary, '\0', sizeof(*mock_gpt_secondary));
+	TEST_EQ(AllocAndReadGptData(handle, &g), 1,
+		"AllocAndRead primary and secondary invalid");
+	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.drive_sectors), 1,
+		"Primary header is invalid");
+	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.drive_sectors), 1,
+		"Secondary header is invalid");
+	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
+		   "VbExDiskRead(h, 1023, 1)\n");
+	WriteAndFreeGptData(handle, &g);
+
+	/*
+	 * Invalidate primary GPT header and check that it is
+	 * repaired by GptRepair().
+	 *
+	 * This would normally be called by LoadKernel()->GptInit()
+	 * but this callback is mocked in these tests.
+	 */
+	ResetMocks();
+	Memset(mock_gpt_primary, '\0', sizeof(*mock_gpt_primary));
+	TEST_EQ(AllocAndReadGptData(handle, &g), 0,
+		"Fix Primary GPT: AllocAndRead");
+	/* Call GptRepair() with input indicating secondary GPT is valid */
+	g.valid_headers = g.valid_entries = MASK_SECONDARY;
+	GptRepair(&g);
+	TEST_EQ(WriteAndFreeGptData(handle, &g), 0,
+		"Fix Primary GPT: WriteAndFreeGptData");
+	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
+		   "VbExDiskRead(h, 1023, 1)\n"
+		   "VbExDiskRead(h, 991, 32)\n"
+		   "VbExDiskWrite(h, 1, 1)\n"
+		   "VbExDiskWrite(h, 2, 32)\n");
+	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.drive_sectors), 0,
+		"Fix Primary GPT: Primary header is valid");
+
+	/*
+	 * Invalidate secondary GPT header and check that it can be
+	 * repaired by GptRepair().
+	 *
+	 * This would normally be called by LoadKernel()->GptInit()
+	 * but this callback is mocked in these tests.
+	 */
+	ResetMocks();
+	Memset(mock_gpt_secondary, '\0', sizeof(*mock_gpt_secondary));
+	TEST_EQ(AllocAndReadGptData(handle, &g), 0,
+		"Fix Secondary GPT: AllocAndRead");
+	/* Call GptRepair() with input indicating primary GPT is valid */
+	g.valid_headers = g.valid_entries = MASK_PRIMARY;
+	GptRepair(&g);
+	TEST_EQ(WriteAndFreeGptData(handle, &g), 0,
+		"Fix Secondary GPT: WriteAndFreeGptData");
+	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
+		   "VbExDiskRead(h, 2, 32)\n"
+		   "VbExDiskRead(h, 1023, 1)\n"
+		   "VbExDiskWrite(h, 1023, 1)\n"
+		   "VbExDiskWrite(h, 991, 32)\n");
+	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.drive_sectors), 0,
+		"Fix Secondary GPT: Secondary header is valid");
+
 	/* Data which is changed is written */
 	ResetMocks();
 	AllocAndReadGptData(handle, &g);