vboot: cgpt: fix my_lba of the secondary GPT

Previously, my_lba of the secondary GPT was recorded as if that
GPT was written at the end of the device. This patch tweaks
my_lba to report where it is in the random-access GPT address space,
namely at the end of that space.

TEST=Compiled it into the firmware and observed the firmware to update
the my_lba field of the secondary GPT.
BRANCH=none
BUG=chromium:425677

Signed-off-by: Dan Ehrenberg <dehrenberg@chromium.org>

Change-Id: I88791fb2cff1086351ca8a3adeef675c4a88cc9a
Reviewed-on: https://chromium-review.googlesource.com/228942
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Commit-Queue: Daniel Ehrenberg <dehrenberg@chromium.org>
Tested-by: Daniel Ehrenberg <dehrenberg@chromium.org>

6 files changed, 77 insertions, 63 deletions

diff --git a/cgpt/cgpt_common.c b/cgpt/cgpt_common.c
index 43312fe6..caee434a 100644
--- a/cgpt/cgpt_common.c
+++ b/cgpt/cgpt_common.c
@@ -340,6 +340,7 @@ static int GptLoad(struct drive *drive, uint32_t sector_bytes) {
   }
   GptHeader* primary_header = (GptHeader*)drive->gpt.primary_header;
   if (CheckHeader(primary_header, 0, drive->gpt.drive_sectors,
+                  drive->gpt.gpt_drive_sectors,
                   drive->gpt.stored_on_device) == 0) {
     if (CGPT_OK != Load(drive, &drive->gpt.primary_entries,
                         primary_header->entries_lba,
@@ -352,6 +353,7 @@ static int GptLoad(struct drive *drive, uint32_t sector_bytes) {
   }
   GptHeader* secondary_header = (GptHeader*)drive->gpt.secondary_header;
   if (CheckHeader(secondary_header, 1, drive->gpt.drive_sectors,
+                  drive->gpt.gpt_drive_sectors,
                   drive->gpt.stored_on_device) == 0) {
     if (CGPT_OK != Load(drive, &drive->gpt.secondary_entries,
                         secondary_header->entries_lba,
diff --git a/firmware/lib/cgptlib/cgptlib_internal.c b/firmware/lib/cgptlib/cgptlib_internal.c
index e9f27c96..48b527ac 100644
--- a/firmware/lib/cgptlib/cgptlib_internal.c
+++ b/firmware/lib/cgptlib/cgptlib_internal.c
@@ -54,7 +54,7 @@ uint32_t HeaderCrc(GptHeader *h)
 }
 
 int CheckHeader(GptHeader *h, int is_secondary, uint64_t drive_sectors,
-	uint8_t stored_on_device)
+		uint64_t gpt_drive_sectors, uint8_t stored_on_device)
 {
 	if (!h)
 		return 1;
@@ -101,7 +101,7 @@ int CheckHeader(GptHeader *h, int is_secondary, uint64_t drive_sectors,
 	 * secondary is at the end of the drive, preceded by its entries.
 	 */
 	if (is_secondary) {
-		if (h->my_lba != drive_sectors - GPT_HEADER_SECTORS)
+		if (h->my_lba != gpt_drive_sectors - GPT_HEADER_SECTORS)
 			return 1;
 		if (h->entries_lba != h->my_lba - GPT_ENTRIES_SECTORS)
 			return 1;
@@ -245,11 +245,13 @@ int GptSanityCheck(GptData *gpt)
 		return retval;
 
 	/* Check both headers; we need at least one valid header. */
-	if (0 == CheckHeader(header1, 0, gpt->drive_sectors, gpt->stored_on_device)) {
+	if (0 == CheckHeader(header1, 0, gpt->drive_sectors,
+			     gpt->gpt_drive_sectors, gpt->stored_on_device)) {
 		gpt->valid_headers |= MASK_PRIMARY;
 		goodhdr = header1;
 	}
-	if (0 == CheckHeader(header2, 1, gpt->drive_sectors, gpt->stored_on_device)) {
+	if (0 == CheckHeader(header2, 1, gpt->drive_sectors,
+			     gpt->gpt_drive_sectors, gpt->stored_on_device)) {
 		gpt->valid_headers |= MASK_SECONDARY;
 		if (!goodhdr)
 			goodhdr = header2;
@@ -320,7 +322,7 @@ void GptRepair(GptData *gpt)
 	if (MASK_PRIMARY == gpt->valid_headers) {
 		/* Primary is good, secondary is bad */
 		Memcpy(header2, header1, sizeof(GptHeader));
-		header2->my_lba = gpt->drive_sectors - GPT_HEADER_SECTORS;
+		header2->my_lba = gpt->gpt_drive_sectors - GPT_HEADER_SECTORS;
 		header2->alternate_lba = GPT_PMBR_SECTORS;  /* Second sector. */
 		header2->entries_lba = header2->my_lba - GPT_ENTRIES_SECTORS;
 		header2->header_crc32 = HeaderCrc(header2);
diff --git a/firmware/lib/cgptlib/include/cgptlib_internal.h b/firmware/lib/cgptlib/include/cgptlib_internal.h
index 4a52420c..825bbbb2 100644
--- a/firmware/lib/cgptlib/include/cgptlib_internal.h
+++ b/firmware/lib/cgptlib/include/cgptlib_internal.h
@@ -91,7 +91,7 @@ int CheckParameters(GptData* gpt);
  * Returns 0 if header is valid, 1 if invalid.
  */
 int CheckHeader(GptHeader *h, int is_secondary, uint64_t drive_sectors,
-                uint8_t stored_on_device);
+                uint64_t gpt_drive_sectors, uint8_t stored_on_device);
 
 /**
  * Calculate and return the header CRC.
diff --git a/firmware/lib/gpt_misc.c b/firmware/lib/gpt_misc.c
index 499cc9e7..dc15d915 100644
--- a/firmware/lib/gpt_misc.c
+++ b/firmware/lib/gpt_misc.c
@@ -48,7 +48,8 @@ int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 
 	/* Only read primary GPT if the primary header is valid */
 	GptHeader* primary_header = (GptHeader*)gptdata->primary_header;
-	if (0 == CheckHeader(primary_header, 0, gptdata->gpt_drive_sectors,
+	if (0 == CheckHeader(primary_header, 0, gptdata->drive_sectors,
+			gptdata->gpt_drive_sectors,
 			gptdata->stored_on_device)) {
 		primary_valid = 1;
 		if (0 != VbExDiskRead(disk_handle,
@@ -67,7 +68,8 @@ int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 
 	/* Only read secondary GPT if the secondary header is valid */
 	GptHeader* secondary_header = (GptHeader*)gptdata->secondary_header;
-	if (0 == CheckHeader(secondary_header, 1, gptdata->gpt_drive_sectors,
+	if (0 == CheckHeader(secondary_header, 1, gptdata->drive_sectors,
+			gptdata->gpt_drive_sectors,
 			gptdata->stored_on_device)) {
 		secondary_valid = 1;
 		if (0 != VbExDiskRead(disk_handle,
diff --git a/tests/cgptlib_test.c b/tests/cgptlib_test.c
index 12b366ad..fc166fda 100644
--- a/tests/cgptlib_test.c
+++ b/tests/cgptlib_test.c
@@ -461,15 +461,15 @@ static int SignatureTest(void)
 	GptHeader *h2 = (GptHeader *)gpt->secondary_header;
 	int i;
 
-	EXPECT(1 == CheckHeader(NULL, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(NULL, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	for (i = 0; i < 8; ++i) {
 		BuildTestGptData(gpt);
 		h1->signature[i] ^= 0xff;
 		h2->signature[i] ^= 0xff;
 		RefreshCrc32(gpt);
-		EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-		EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+		EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+		EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 	}
 
 	return TEST_OK;
@@ -503,9 +503,9 @@ static int RevisionTest(void)
 		h2->revision = cases[i].value_to_test;
 		RefreshCrc32(gpt);
 
-		EXPECT(CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE) ==
+		EXPECT(CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE) ==
 		       cases[i].expect_rv);
-		EXPECT(CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE) ==
+		EXPECT(CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE) ==
 		       cases[i].expect_rv);
 	}
 	return TEST_OK;
@@ -536,9 +536,9 @@ static int SizeTest(void)
 		h2->size = cases[i].value_to_test;
 		RefreshCrc32(gpt);
 
-		EXPECT(CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE) ==
+		EXPECT(CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE) ==
 		       cases[i].expect_rv);
-		EXPECT(CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE) ==
+		EXPECT(CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE) ==
 		       cases[i].expect_rv);
 	}
 	return TEST_OK;
@@ -555,12 +555,12 @@ static int CrcFieldTest(void)
 	/* Modify a field that the header verification doesn't care about */
 	h1->entries_crc32++;
 	h2->entries_crc32++;
-	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 	/* Refresh the CRC; should pass now */
 	RefreshCrc32(gpt);
-	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(0 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(0 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	return TEST_OK;
 }
@@ -576,8 +576,8 @@ static int ReservedFieldsTest(void)
 	h1->reserved_zero ^= 0x12345678;  /* whatever random */
 	h2->reserved_zero ^= 0x12345678;  /* whatever random */
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 #ifdef PADDING_CHECKED
 	/* TODO: padding check is currently disabled */
@@ -585,8 +585,8 @@ static int ReservedFieldsTest(void)
 	h1->padding[12] ^= 0x34;  /* whatever random */
 	h2->padding[56] ^= 0x78;  /* whatever random */
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 #endif
 
 	return TEST_OK;
@@ -624,9 +624,9 @@ static int SizeOfPartitionEntryTest(void) {
 			cases[i].value_to_test;
 		RefreshCrc32(gpt);
 
-		EXPECT(CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE) ==
+		EXPECT(CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE) ==
 		       cases[i].expect_rv);
-		EXPECT(CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE) ==
+		EXPECT(CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE) ==
 		       cases[i].expect_rv);
 	}
 
@@ -647,11 +647,11 @@ static int NumberOfPartitionEntriesTest(void)
 	h1->number_of_entries--;
 	h2->number_of_entries /= 2;
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 	/* But it's okay to have less if the GPT structs are stored elsewhere. */
-	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_OFF_DEVICE));
-	EXPECT(0 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_OFF_DEVICE));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_OFF_DEVICE));
+	EXPECT(0 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_OFF_DEVICE));
 
 	return TEST_OK;
 }
@@ -666,37 +666,37 @@ static int MyLbaTest(void)
 
 	/* myLBA depends on primary vs secondary flag */
 	BuildTestGptData(gpt);
-	EXPECT(1 == CheckHeader(h1, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(1 == CheckHeader(h2, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h1, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	BuildTestGptData(gpt);
 	h1->my_lba--;
 	h2->my_lba--;
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	BuildTestGptData(gpt);
 	h1->my_lba = 2;
 	h2->my_lba--;
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	/* We should ignore the alternate_lba field entirely */
 	BuildTestGptData(gpt);
 	h1->alternate_lba++;
 	h2->alternate_lba++;
 	RefreshCrc32(gpt);
-	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(0 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(0 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	BuildTestGptData(gpt);
 	h1->alternate_lba--;
 	h2->alternate_lba--;
 	RefreshCrc32(gpt);
-	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(0 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(0 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	BuildTestGptData(gpt);
 	h1->entries_lba++;
@@ -706,19 +706,19 @@ static int MyLbaTest(void)
 	 * We support a padding between primary GPT header and its entries. So
 	 * this still passes.
 	 */
-	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 	/*
 	 * But the secondary table should fail because it would overlap the
 	 * header, which is now lying after its entry array.
 	 */
-	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	BuildTestGptData(gpt);
 	h1->entries_lba--;
 	h2->entries_lba--;
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 
 	return TEST_OK;
 }
@@ -767,9 +767,9 @@ static int FirstUsableLbaAndLastUsableLbaTest(void)
 		h2->last_usable_lba = cases[i].secondary_last_usable_lba;
 		RefreshCrc32(gpt);
 
-		EXPECT(CheckHeader(h1, 0, gpt->drive_sectors, GPT_STORED_ON_DEVICE) ==
+		EXPECT(CheckHeader(h1, 0, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE) ==
 		       cases[i].primary_rv);
-		EXPECT(CheckHeader(h2, 1, gpt->drive_sectors, GPT_STORED_ON_DEVICE) ==
+		EXPECT(CheckHeader(h2, 1, gpt->drive_sectors, gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE) ==
 		       cases[i].secondary_rv);
 	}
 
@@ -1981,26 +1981,26 @@ static int CheckHeaderOffDevice()
 	// GPT is stored on the same device so first usable lba should not
 	// start at 0.
 	EXPECT(1 == CheckHeader(primary_header, 0, gpt->drive_sectors,
-		GPT_STORED_ON_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 	// But off device, it is okay to accept this GPT header.
 	EXPECT(0 == CheckHeader(primary_header, 0, gpt->drive_sectors,
-		GPT_STORED_OFF_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_OFF_DEVICE));
 
 	BuildTestGptData(gpt);
 	primary_header->number_of_entries = 100;
 	RefreshCrc32(gpt);
 	// Normally, number of entries is 128. So this should fail.
 	EXPECT(1 == CheckHeader(primary_header, 0, gpt->drive_sectors,
-		GPT_STORED_ON_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 	// But off device, it is okay.
 	EXPECT(0 == CheckHeader(primary_header, 0, gpt->drive_sectors,
-		GPT_STORED_OFF_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_OFF_DEVICE));
 
 	primary_header->number_of_entries = MIN_NUMBER_OF_ENTRIES - 1;
 	RefreshCrc32(gpt);
 	// However, too few entries is not good.
 	EXPECT(1 == CheckHeader(primary_header, 0, gpt->drive_sectors,
-		GPT_STORED_OFF_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_OFF_DEVICE));
 
 	// Repeat for secondary header.
 	BuildTestGptData(gpt);
@@ -2008,22 +2008,22 @@ static int CheckHeaderOffDevice()
 	secondary_header->first_usable_lba = 0;
 	RefreshCrc32(gpt);
 	EXPECT(1 == CheckHeader(secondary_header, 1, gpt->drive_sectors,
-		GPT_STORED_ON_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 	EXPECT(0 == CheckHeader(secondary_header, 1, gpt->drive_sectors,
-		GPT_STORED_OFF_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_OFF_DEVICE));
 
 	BuildTestGptData(gpt);
 	secondary_header->number_of_entries = 100;
 	RefreshCrc32(gpt);
 	EXPECT(1 == CheckHeader(secondary_header, 1, gpt->drive_sectors,
-		GPT_STORED_ON_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_ON_DEVICE));
 	EXPECT(0 == CheckHeader(secondary_header, 1, gpt->drive_sectors,
-		GPT_STORED_OFF_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_OFF_DEVICE));
 
 	secondary_header->number_of_entries = MIN_NUMBER_OF_ENTRIES - 1;
 	RefreshCrc32(gpt);
 	EXPECT(1 == CheckHeader(secondary_header, 1, gpt->drive_sectors,
-		GPT_STORED_OFF_DEVICE));
+		gpt->gpt_drive_sectors, GPT_STORED_OFF_DEVICE));
 
 	return TEST_OK;
 }
diff --git a/tests/vboot_kernel_tests.c b/tests/vboot_kernel_tests.c
index 05035f28..c546ed12 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vboot_kernel_tests.c
@@ -316,9 +316,11 @@ static void ReadWriteGptTest(void)
 	TEST_EQ(AllocAndReadGptData(handle, &g), 0,
 		"AllocAndRead primary invalid");
 	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.drive_sectors,
-		GPT_STORED_ON_DEVICE), 1, "Primary header is invalid");
+                g.gpt_drive_sectors, GPT_STORED_ON_DEVICE),
+                1, "Primary header is invalid");
 	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.drive_sectors,
-		GPT_STORED_ON_DEVICE), 0, "Secondary header is valid");
+		g.gpt_drive_sectors, GPT_STORED_ON_DEVICE),
+                0, "Secondary header is valid");
 	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
 		   "VbExDiskRead(h, 1023, 1)\n"
 		   "VbExDiskRead(h, 991, 32)\n");
@@ -333,9 +335,11 @@ static void ReadWriteGptTest(void)
 	TEST_EQ(AllocAndReadGptData(handle, &g), 0,
 		"AllocAndRead secondary invalid");
 	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.drive_sectors,
-		GPT_STORED_ON_DEVICE), 0, "Primary header is valid");
+		g.gpt_drive_sectors, GPT_STORED_ON_DEVICE),
+                0, "Primary header is valid");
 	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.drive_sectors,
-		GPT_STORED_ON_DEVICE), 1, "Secondary header is invalid");
+		g.gpt_drive_sectors, GPT_STORED_ON_DEVICE),
+                1, "Secondary header is invalid");
 	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
 		   "VbExDiskRead(h, 2, 32)\n"
 		   "VbExDiskRead(h, 1023, 1)\n");
@@ -351,9 +355,11 @@ static void ReadWriteGptTest(void)
 	TEST_EQ(AllocAndReadGptData(handle, &g), 1,
 		"AllocAndRead primary and secondary invalid");
 	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.drive_sectors,
-		GPT_STORED_ON_DEVICE), 1, "Primary header is invalid");
+		g.gpt_drive_sectors, GPT_STORED_ON_DEVICE),
+                1, "Primary header is invalid");
 	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.drive_sectors,
-		GPT_STORED_ON_DEVICE), 1, "Secondary header is invalid");
+		g.gpt_drive_sectors, GPT_STORED_ON_DEVICE),
+                1, "Secondary header is invalid");
 	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
 		   "VbExDiskRead(h, 1023, 1)\n");
 	WriteAndFreeGptData(handle, &g);
@@ -380,7 +386,8 @@ static void ReadWriteGptTest(void)
 		   "VbExDiskWrite(h, 1, 1)\n"
 		   "VbExDiskWrite(h, 2, 32)\n");
 	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.drive_sectors,
-		GPT_STORED_ON_DEVICE), 0, "Fix Primary GPT: Primary header is valid");
+		g.gpt_drive_sectors, GPT_STORED_ON_DEVICE),
+                0, "Fix Primary GPT: Primary header is valid");
 
 	/*
 	 * Invalidate secondary GPT header and check that it can be
@@ -404,7 +411,8 @@ static void ReadWriteGptTest(void)
 		   "VbExDiskWrite(h, 1023, 1)\n"
 		   "VbExDiskWrite(h, 991, 32)\n");
 	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.drive_sectors,
-		GPT_STORED_ON_DEVICE), 0, "Fix Secondary GPT: Secondary header is valid");
+		g.gpt_drive_sectors, GPT_STORED_ON_DEVICE),
+                0, "Fix Secondary GPT: Secondary header is valid");
 
 	/* Data which is changed is written */
 	ResetMocks();