cgpt: Remove hard coded 512 block size.

Remove 512 sector block size restriction so that UFS, with sector block size 4096
or greater, can be used. The sector block size is queried from the kernel with
ioctl(BLKSSZGET) or queried from depthcharge with VbExDiskGetInfo().

BUG=b:77540192
BRANCH=none
TEST=manual
make runtests passed.
Tested firmware on Kevin and boot to kernel from disk.
Executed cgpt show /dev/mmcblk0 on eve device and verified output was correct.
Should be tested on device with sector block size greater than 512.

Change-Id: I8165c8ee4da68180eecc8d12b3fb501cc5c60a5d
Reviewed-on: https://chromium-review.googlesource.com/1007498
Commit-Ready: Sam Hurst <shurst@google.com>
Tested-by: Sam Hurst <shurst@google.com>
Reviewed-by: Julius Werner <jwerner@chromium.org>

2 files changed, 110 insertions, 66 deletions

diff --git a/tests/cgptlib_test.c b/tests/cgptlib_test.c
index f37a2725..0336c8e7 100644
--- a/tests/cgptlib_test.c
+++ b/tests/cgptlib_test.c
@@ -256,11 +256,16 @@ static int ParameterTests(void)
 	} cases[] = {
 		{512, DEFAULT_DRIVE_SECTORS, GPT_SUCCESS},
 		{520, DEFAULT_DRIVE_SECTORS, GPT_ERROR_INVALID_SECTOR_SIZE},
+		{123, DEFAULT_DRIVE_SECTORS, GPT_ERROR_INVALID_SECTOR_SIZE},
+		{4097, DEFAULT_DRIVE_SECTORS, GPT_ERROR_INVALID_SECTOR_SIZE},
+		{256, DEFAULT_DRIVE_SECTORS, GPT_ERROR_INVALID_SECTOR_SIZE},
 		{512, 0, GPT_ERROR_INVALID_SECTOR_NUMBER},
 		{512, 10, GPT_ERROR_INVALID_SECTOR_NUMBER},
 		{512, GPT_PMBR_SECTORS + GPT_HEADER_SECTORS * 2 +
 		 TOTAL_ENTRIES_SIZE / DEFAULT_SECTOR_SIZE * 2, GPT_SUCCESS},
-		{4096, DEFAULT_DRIVE_SECTORS, GPT_ERROR_INVALID_SECTOR_SIZE},
+		{4096, DEFAULT_DRIVE_SECTORS, GPT_SUCCESS},
+		{2048, DEFAULT_DRIVE_SECTORS, GPT_SUCCESS},
+		{8192, DEFAULT_DRIVE_SECTORS, GPT_SUCCESS},
 	};
 	int i;
 
@@ -364,15 +369,18 @@ static int SignatureTest(void)
 	GptHeader *h2 = (GptHeader *)gpt->secondary_header;
 	int i;
 
-	EXPECT(1 == CheckHeader(NULL, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(NULL, 0, gpt->streaming_drive_sectors,
+				gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	for (i = 0; i < 8; ++i) {
 		BuildTestGptData(gpt);
 		h1->signature[i] ^= 0xff;
 		h2->signature[i] ^= 0xff;
 		RefreshCrc32(gpt);
-		EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-		EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+		EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+				gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+		EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+				gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 	}
 
 	return TEST_OK;
@@ -406,10 +414,12 @@ static int RevisionTest(void)
 		h2->revision = cases[i].value_to_test;
 		RefreshCrc32(gpt);
 
-		EXPECT(CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0) ==
-		       cases[i].expect_rv);
-		EXPECT(CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0) ==
-		       cases[i].expect_rv);
+		EXPECT(CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+			gpt->gpt_drive_sectors, 0, gpt->sector_bytes) ==
+			cases[i].expect_rv);
+		EXPECT(CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+			gpt->gpt_drive_sectors, 0,
+			gpt->sector_bytes) == cases[i].expect_rv);
 	}
 	return TEST_OK;
 }
@@ -439,10 +449,12 @@ static int SizeTest(void)
 		h2->size = cases[i].value_to_test;
 		RefreshCrc32(gpt);
 
-		EXPECT(CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0) ==
-		       cases[i].expect_rv);
-		EXPECT(CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0) ==
-		       cases[i].expect_rv);
+		EXPECT(CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+			gpt->gpt_drive_sectors, 0, gpt->sector_bytes) ==
+			cases[i].expect_rv);
+		EXPECT(CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+			gpt->gpt_drive_sectors, 0, gpt->sector_bytes) ==
+			cases[i].expect_rv);
 	}
 	return TEST_OK;
 }
@@ -458,12 +470,16 @@ 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->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 	/* Refresh the CRC; should pass now */
 	RefreshCrc32(gpt);
-	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(0 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(0 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	return TEST_OK;
 }
@@ -479,8 +495,10 @@ static int ReservedFieldsTest(void)
 	h1->reserved_zero ^= 0x12345678;  /* whatever random */
 	h2->reserved_zero ^= 0x12345678;  /* whatever random */
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 #ifdef PADDING_CHECKED
 	/* TODO: padding check is currently disabled */
@@ -488,8 +506,10 @@ static int ReservedFieldsTest(void)
 	h1->padding[12] ^= 0x34;  /* whatever random */
 	h2->padding[56] ^= 0x78;  /* whatever random */
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 #endif
 
 	return TEST_OK;
@@ -527,9 +547,11 @@ static int SizeOfPartitionEntryTest(void) {
 			cases[i].value_to_test;
 		RefreshCrc32(gpt);
 
-		EXPECT(CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0) ==
+		EXPECT(CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+			gpt->gpt_drive_sectors, 0, gpt->sector_bytes) ==
 		       cases[i].expect_rv);
-		EXPECT(CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0) ==
+		EXPECT(CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+			gpt->gpt_drive_sectors, 0, gpt->sector_bytes) ==
 		       cases[i].expect_rv);
 	}
 
@@ -550,13 +572,17 @@ static int NumberOfPartitionEntriesTest(void)
 	h1->number_of_entries--;
 	h2->number_of_entries /= 2;
 	/* Because we halved h2 entries, its entries_lba is going to change. */
-	h2->entries_lba = h2->my_lba - CalculateEntriesSectors(h2);
+	h2->entries_lba = h2->my_lba - CalculateEntriesSectors(h2, gpt->sector_bytes);
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 	/* But it's okay to have less if the GPT structs are stored elsewhere. */
-	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL));
-	EXPECT(0 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL, gpt->sector_bytes));
+	EXPECT(0 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL, gpt->sector_bytes));
 
 	return TEST_OK;
 }
@@ -571,37 +597,47 @@ static int MyLbaTest(void)
 
 	/* myLBA depends on primary vs secondary flag */
 	BuildTestGptData(gpt);
-	EXPECT(1 == CheckHeader(h1, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(1 == CheckHeader(h2, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h1, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(1 == CheckHeader(h2, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	BuildTestGptData(gpt);
 	h1->my_lba--;
 	h2->my_lba--;
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	BuildTestGptData(gpt);
 	h1->my_lba = 2;
 	h2->my_lba--;
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	/* We should ignore the alternate_lba field entirely */
 	BuildTestGptData(gpt);
 	h1->alternate_lba++;
 	h2->alternate_lba++;
 	RefreshCrc32(gpt);
-	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(0 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(0 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	BuildTestGptData(gpt);
 	h1->alternate_lba--;
 	h2->alternate_lba--;
 	RefreshCrc32(gpt);
-	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(0 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(0 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	BuildTestGptData(gpt);
 	h1->entries_lba++;
@@ -611,19 +647,23 @@ 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->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(0 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 	/*
 	 * 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->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	BuildTestGptData(gpt);
 	h1->entries_lba--;
 	h2->entries_lba--;
 	RefreshCrc32(gpt);
-	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
-	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0));
+	EXPECT(1 == CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
+	EXPECT(1 == CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 
 	return TEST_OK;
 }
@@ -672,9 +712,11 @@ static int FirstUsableLbaAndLastUsableLbaTest(void)
 		h2->last_usable_lba = cases[i].secondary_last_usable_lba;
 		RefreshCrc32(gpt);
 
-		EXPECT(CheckHeader(h1, 0, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0) ==
+		EXPECT(CheckHeader(h1, 0, gpt->streaming_drive_sectors,
+			gpt->gpt_drive_sectors, 0, gpt->sector_bytes) ==
 		       cases[i].primary_rv);
-		EXPECT(CheckHeader(h2, 1, gpt->streaming_drive_sectors, gpt->gpt_drive_sectors, 0) ==
+		EXPECT(CheckHeader(h2, 1, gpt->streaming_drive_sectors,
+			gpt->gpt_drive_sectors, 0, gpt->sector_bytes) ==
 		       cases[i].secondary_rv);
 	}
 
@@ -846,7 +888,7 @@ static int SanityCheckTest(void)
 
 	/* Invalid sector size should fail */
 	BuildTestGptData(gpt);
-	gpt->sector_bytes = 1024;
+	gpt->sector_bytes = 1023;
 	EXPECT(GPT_ERROR_INVALID_SECTOR_SIZE == GptSanityCheck(gpt));
 
 	/* Modify headers */
@@ -1614,26 +1656,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->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, 0));
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 	// But off device, it is okay to accept this GPT header.
 	EXPECT(0 == CheckHeader(primary_header, 0, gpt->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL));
+		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL, gpt->sector_bytes));
 
 	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->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, 0));
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 	// But off device, it is okay.
 	EXPECT(0 == CheckHeader(primary_header, 0, gpt->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL));
+		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL, gpt->sector_bytes));
 
 	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->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL));
+		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL, gpt->sector_bytes));
 
 	// Repeat for secondary header.
 	BuildTestGptData(gpt);
@@ -1641,25 +1683,25 @@ static int CheckHeaderOffDevice()
 	secondary_header->first_usable_lba = 0;
 	RefreshCrc32(gpt);
 	EXPECT(1 == CheckHeader(secondary_header, 1, gpt->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, 0));
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 	EXPECT(0 == CheckHeader(secondary_header, 1, gpt->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL));
+		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL, gpt->sector_bytes));
 
 	BuildTestGptData(gpt);
 	secondary_header->number_of_entries = 100;
 	/* Because we change number of entries, we need to also update entrie_lba. */
 	secondary_header->entries_lba = secondary_header->my_lba -
-		CalculateEntriesSectors(secondary_header);
+		CalculateEntriesSectors(secondary_header, gpt->sector_bytes);
 	RefreshCrc32(gpt);
 	EXPECT(1 == CheckHeader(secondary_header, 1, gpt->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, 0));
+		gpt->gpt_drive_sectors, 0, gpt->sector_bytes));
 	EXPECT(0 == CheckHeader(secondary_header, 1, gpt->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL));
+		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL, gpt->sector_bytes));
 
 	secondary_header->number_of_entries = MIN_NUMBER_OF_ENTRIES - 1;
 	RefreshCrc32(gpt);
 	EXPECT(1 == CheckHeader(secondary_header, 1, gpt->streaming_drive_sectors,
-		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL));
+		gpt->gpt_drive_sectors, GPT_FLAG_EXTERNAL, gpt->sector_bytes));
 
 	return TEST_OK;
 }
diff --git a/tests/vboot_kernel_tests.c b/tests/vboot_kernel_tests.c
index 14bffc91..c9907597 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vboot_kernel_tests.c
@@ -96,14 +96,16 @@ static void SetupGptHeader(GptHeader *h, int is_secondary)
 	/* 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 - CalculateEntriesSectors(h);
+		h->entries_lba = h->my_lba - CalculateEntriesSectors(h,
+							MOCK_SECTOR_SIZE);
 	} else {
 		h->my_lba = GPT_PMBR_SECTORS;
 		h->entries_lba = h->my_lba + 1;
 	}
 
-	h->first_usable_lba = 2 + CalculateEntriesSectors(h);
-	h->last_usable_lba = MOCK_SECTOR_COUNT - 2 - CalculateEntriesSectors(h);
+	h->first_usable_lba = 2 + CalculateEntriesSectors(h, MOCK_SECTOR_SIZE);
+	h->last_usable_lba = MOCK_SECTOR_COUNT - 2 - CalculateEntriesSectors(h,
+								MOCK_SECTOR_SIZE);
 
 	h->header_crc32 = HeaderCrc(h);
 }
@@ -350,10 +352,10 @@ static void ReadWriteGptTest(void)
 	TEST_EQ(AllocAndReadGptData(handle, &g), 0,
 		"AllocAndRead primary invalid");
 	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.streaming_drive_sectors,
-                g.gpt_drive_sectors, 0),
+                g.gpt_drive_sectors, 0, g.sector_bytes),
                 1, "Primary header is invalid");
 	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.streaming_drive_sectors,
-		g.gpt_drive_sectors, 0),
+		g.gpt_drive_sectors, 0, g.sector_bytes),
                 0, "Secondary header is valid");
 	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
 		   "VbExDiskRead(h, 1023, 1)\n"
@@ -369,10 +371,10 @@ static void ReadWriteGptTest(void)
 	TEST_EQ(AllocAndReadGptData(handle, &g), 0,
 		"AllocAndRead secondary invalid");
 	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.streaming_drive_sectors,
-		g.gpt_drive_sectors, 0),
+		g.gpt_drive_sectors, 0, g.sector_bytes),
                 0, "Primary header is valid");
 	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.streaming_drive_sectors,
-		g.gpt_drive_sectors, 0),
+		g.gpt_drive_sectors, 0, g.sector_bytes),
                 1, "Secondary header is invalid");
 	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
 		   "VbExDiskRead(h, 2, 32)\n"
@@ -389,10 +391,10 @@ static void ReadWriteGptTest(void)
 	TEST_EQ(AllocAndReadGptData(handle, &g), 1,
 		"AllocAndRead primary and secondary invalid");
 	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.streaming_drive_sectors,
-		g.gpt_drive_sectors, 0),
+		g.gpt_drive_sectors, 0, g.sector_bytes),
                 1, "Primary header is invalid");
 	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.streaming_drive_sectors,
-		g.gpt_drive_sectors, 0),
+		g.gpt_drive_sectors, 0, g.sector_bytes),
                 1, "Secondary header is invalid");
 	TEST_CALLS("VbExDiskRead(h, 1, 1)\n"
 		   "VbExDiskRead(h, 1023, 1)\n");
@@ -420,7 +422,7 @@ static void ReadWriteGptTest(void)
 		   "VbExDiskWrite(h, 1, 1)\n"
 		   "VbExDiskWrite(h, 2, 32)\n");
 	TEST_EQ(CheckHeader(mock_gpt_primary, 0, g.streaming_drive_sectors,
-		g.gpt_drive_sectors, 0),
+		g.gpt_drive_sectors, 0, g.sector_bytes),
                 0, "Fix Primary GPT: Primary header is valid");
 
 	/*
@@ -445,7 +447,7 @@ static void ReadWriteGptTest(void)
 		   "VbExDiskWrite(h, 1023, 1)\n"
 		   "VbExDiskWrite(h, 991, 32)\n");
 	TEST_EQ(CheckHeader(mock_gpt_secondary, 1, g.streaming_drive_sectors,
-		g.gpt_drive_sectors, 0),
+		g.gpt_drive_sectors, 0, g.sector_bytes),
                 0, "Fix Secondary GPT: Secondary header is valid");
 
 	/* Data which is changed is written */