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>

9 files changed, 165 insertions, 104 deletions

diff --git a/cgpt/cgpt_common.c b/cgpt/cgpt_common.c
index f6108c1e..a5550b4c 100644
--- a/cgpt/cgpt_common.c
+++ b/cgpt/cgpt_common.c
@@ -191,11 +191,13 @@ 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.streaming_drive_sectors,
                   drive->gpt.gpt_drive_sectors,
-                  drive->gpt.flags) == 0) {
+                  drive->gpt.flags,
+                  drive->gpt.sector_bytes) == 0) {
     if (CGPT_OK != Load(drive, &drive->gpt.primary_entries,
                         primary_header->entries_lba,
                         drive->gpt.sector_bytes,
-                        CalculateEntriesSectors(primary_header))) {
+                        CalculateEntriesSectors(primary_header,
+                          drive->gpt.sector_bytes))) {
       Error("Cannot read primary partition entry array\n");
       return -1;
     }
@@ -209,11 +211,13 @@ 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.streaming_drive_sectors,
                   drive->gpt.gpt_drive_sectors,
-                  drive->gpt.flags) == 0) {
+                  drive->gpt.flags,
+                  drive->gpt.sector_bytes) == 0) {
     if (CGPT_OK != Load(drive, &drive->gpt.secondary_entries,
                         secondary_header->entries_lba,
                         drive->gpt.sector_bytes,
-                        CalculateEntriesSectors(secondary_header))) {
+                        CalculateEntriesSectors(secondary_header,
+                          drive->gpt.sector_bytes))) {
       Error("Cannot read secondary partition entry array\n");
       return -1;
     }
@@ -244,7 +248,8 @@ static int GptSave(struct drive *drive) {
       if (CGPT_OK != Save(drive, drive->gpt.primary_entries,
                           primary_header->entries_lba,
                           drive->gpt.sector_bytes,
-                          CalculateEntriesSectors(primary_header))) {
+                          CalculateEntriesSectors(primary_header,
+                            drive->gpt.sector_bytes))) {
         errors++;
         Error("Cannot write primary entries: %s\n", strerror(errno));
       }
@@ -273,7 +278,8 @@ static int GptSave(struct drive *drive) {
       if (CGPT_OK != Save(drive, drive->gpt.secondary_entries,
                           secondary_header->entries_lba,
                           drive->gpt.sector_bytes,
-                          CalculateEntriesSectors(secondary_header))) {
+                          CalculateEntriesSectors(secondary_header,
+                            drive->gpt.sector_bytes))) {
         errors++;
         Error("Cannot write secondary entries: %s\n", strerror(errno));
       }
@@ -343,7 +349,6 @@ int DriveOpen(const char *drive_path, struct drive *drive, int mode,
     return CGPT_FAILED;
   }
 
-  sector_bytes = 512;
   uint64_t gpt_drive_size;
   if (ObtainDriveSize(drive->fd, &gpt_drive_size, &sector_bytes) != 0) {
     Error("Can't get drive size and bytes per sector for %s: %s\n",
@@ -1015,7 +1020,7 @@ uint8_t RepairHeader(GptData *gpt, const uint32_t valid_headers) {
     secondary_header->my_lba = gpt->gpt_drive_sectors - 1;  /* the last sector */
     secondary_header->alternate_lba = primary_header->my_lba;
     secondary_header->entries_lba = secondary_header->my_lba -
-        CalculateEntriesSectors(primary_header);
+        CalculateEntriesSectors(primary_header, gpt->sector_bytes);
     return GPT_MODIFIED_HEADER2;
   } else if (valid_headers == MASK_SECONDARY) {
     memcpy(primary_header, secondary_header, sizeof(GptHeader));
diff --git a/cgpt/cgpt_create.c b/cgpt/cgpt_create.c
index c34cfe31..a56a9762 100644
--- a/cgpt/cgpt_create.c
+++ b/cgpt/cgpt_create.c
@@ -81,9 +81,11 @@ static int GptCreate(struct drive *drive, CgptCreateParams *params) {
     h->entries_lba = h->my_lba + GPT_HEADER_SECTORS;
     if (!(drive->gpt.flags & GPT_FLAG_EXTERNAL)) {
       h->entries_lba += params->padding;
-      h->first_usable_lba = h->entries_lba + CalculateEntriesSectors(h);
-      h->last_usable_lba = (drive->gpt.streaming_drive_sectors - GPT_HEADER_SECTORS -
-                            CalculateEntriesSectors(h) - 1);
+      h->first_usable_lba = h->entries_lba + CalculateEntriesSectors(h,
+                                               drive->gpt.sector_bytes);
+      h->last_usable_lba =
+        (drive->gpt.streaming_drive_sectors - GPT_HEADER_SECTORS -
+          CalculateEntriesSectors(h, drive->gpt.sector_bytes) - 1);
     } else {
       h->first_usable_lba = params->padding;
       h->last_usable_lba = (drive->gpt.streaming_drive_sectors - 1);
diff --git a/cgpt/cgpt_find.c b/cgpt/cgpt_find.c
index a02ad7ae..e207fb3a 100644
--- a/cgpt/cgpt_find.c
+++ b/cgpt/cgpt_find.c
@@ -14,9 +14,6 @@
 #include "vboot_host.h"
 
 #define BUFSIZE 1024
-// FIXME: currently we only support 512-byte sectors.
-#define LBA_SIZE 512
-
 
 // fill comparebuf with the data to be examined, returning true on success.
 static int FillBuffer(CgptFindParams *params, int fd, uint64_t pos,
@@ -48,15 +45,15 @@ static int match_content(CgptFindParams *params, struct drive *drive,
     return 1;
 
   // Ensure that the region we want to match against is inside the partition.
-  part_size = LBA_SIZE * (entry->ending_lba - entry->starting_lba + 1);
+  part_size = drive->gpt.sector_bytes *
+    (entry->ending_lba - entry->starting_lba + 1);
   if (params->matchoffset + params->matchlen > part_size) {
     return 0;
   }
 
   // Read the partition data.
-  if (!FillBuffer(params,
-                  drive->fd,
-                  (LBA_SIZE * entry->starting_lba) + params->matchoffset,
+  if (!FillBuffer(params, drive->fd,
+    (drive->gpt.sector_bytes * entry->starting_lba) + params->matchoffset,
                   params->matchlen)) {
     Error("unable to read partition data\n");
     return 0;
diff --git a/cgpt/cgpt_show.c b/cgpt/cgpt_show.c
index 60d73ec7..72217419 100644
--- a/cgpt/cgpt_show.c
+++ b/cgpt/cgpt_show.c
@@ -302,7 +302,8 @@ static int GptShow(struct drive *drive, CgptShowParams *params) {
 
       GptHeader* primary_header = (GptHeader*)drive->gpt.primary_header;
       printf(GPT_FMT, (uint64_t)primary_header->entries_lba,
-             (uint64_t)CalculateEntriesSectors(primary_header),
+             (uint64_t)CalculateEntriesSectors(primary_header,
+                         drive->gpt.sector_bytes),
              drive->gpt.valid_entries & MASK_PRIMARY ? "" : "INVALID",
              "Pri GPT table");
 
@@ -319,7 +320,8 @@ static int GptShow(struct drive *drive, CgptShowParams *params) {
     } else {
       GptHeader* secondary_header = (GptHeader*)drive->gpt.secondary_header;
       printf(GPT_FMT, (uint64_t)secondary_header->entries_lba,
-             (uint64_t)CalculateEntriesSectors(secondary_header),
+             (uint64_t)CalculateEntriesSectors(secondary_header,
+                         drive->gpt.sector_bytes),
              drive->gpt.valid_entries & MASK_SECONDARY ? "" : "INVALID",
              "Sec GPT table");
       /* We show secondary table details if any of following is true.
diff --git a/firmware/lib/cgptlib/cgptlib_internal.c b/firmware/lib/cgptlib/cgptlib_internal.c
index 9d7bd0c9..77a0de2b 100644
--- a/firmware/lib/cgptlib/cgptlib_internal.c
+++ b/firmware/lib/cgptlib/cgptlib_internal.c
@@ -12,19 +12,20 @@
 #include "gpt_misc.h"
 #include "utility.h"
 
-const static int SECTOR_SIZE = 512;
+const static int MIN_SECTOR_SIZE = 512;
 
-size_t CalculateEntriesSectors(GptHeader* h)
+size_t CalculateEntriesSectors(GptHeader* h, uint32_t sector_bytes)
 {
 	size_t bytes = h->number_of_entries * h->size_of_entry;
-	size_t ret = (bytes + SECTOR_SIZE - 1) / SECTOR_SIZE;
+	size_t ret = (bytes + sector_bytes - 1) / sector_bytes;
 	return ret;
 }
 
 int CheckParameters(GptData *gpt)
 {
-	/* Currently, we only support 512-byte sectors. */
-	if (gpt->sector_bytes != SECTOR_SIZE)
+	/* Only support 512-byte or larger sectors that are a power of 2 */
+	if (gpt->sector_bytes < MIN_SECTOR_SIZE ||
+			(gpt->sector_bytes  & (gpt->sector_bytes  - 1)) != 0)
 		return GPT_ERROR_INVALID_SECTOR_SIZE;
 
 	/*
@@ -45,7 +46,7 @@ int CheckParameters(GptData *gpt)
 	 */
 	if (gpt->gpt_drive_sectors <
 		(1 + 2 * (1 + MIN_NUMBER_OF_ENTRIES /
-				(SECTOR_SIZE / sizeof(GptEntry)))))
+				(gpt->sector_bytes / sizeof(GptEntry)))))
 		return GPT_ERROR_INVALID_SECTOR_NUMBER;
 
 	return GPT_SUCCESS;
@@ -66,7 +67,8 @@ uint32_t HeaderCrc(GptHeader *h)
 
 int CheckHeader(GptHeader *h, int is_secondary,
 		uint64_t streaming_drive_sectors,
-		uint64_t gpt_drive_sectors, uint32_t flags)
+		uint64_t gpt_drive_sectors, uint32_t flags,
+		uint32_t sector_bytes)
 {
 	if (!h)
 		return 1;
@@ -115,7 +117,8 @@ int CheckHeader(GptHeader *h, int is_secondary,
 	if (is_secondary) {
 		if (h->my_lba != gpt_drive_sectors - GPT_HEADER_SECTORS)
 			return 1;
-		if (h->entries_lba != h->my_lba - CalculateEntriesSectors(h))
+		if (h->entries_lba != h->my_lba - CalculateEntriesSectors(h,
+								sector_bytes))
 			return 1;
 	} else {
 		if (h->my_lba != GPT_PMBR_SECTORS)
@@ -141,10 +144,11 @@ int CheckHeader(GptHeader *h, int is_secondary,
 	 * array.
 	 */
 	/* TODO(namnguyen): Also check for padding between header & entries. */
-	if (h->first_usable_lba < 2 + CalculateEntriesSectors(h))
+	if (h->first_usable_lba < 2 + CalculateEntriesSectors(h, sector_bytes))
 		return 1;
 	if (h->last_usable_lba >=
-			streaming_drive_sectors - 1 - CalculateEntriesSectors(h))
+			streaming_drive_sectors - 1 - CalculateEntriesSectors(h,
+								sector_bytes))
 		return 1;
 
 	/* Success */
@@ -254,7 +258,8 @@ int GptSanityCheck(GptData *gpt)
 
 	/* Check both headers; we need at least one valid header. */
 	if (0 == CheckHeader(header1, 0, gpt->streaming_drive_sectors,
-			     gpt->gpt_drive_sectors, gpt->flags)) {
+			     gpt->gpt_drive_sectors, gpt->flags,
+			     gpt->sector_bytes)) {
 		gpt->valid_headers |= MASK_PRIMARY;
 		goodhdr = header1;
 	} else if (header1 && !memcmp(header1->signature,
@@ -262,7 +267,8 @@ int GptSanityCheck(GptData *gpt)
 		gpt->ignored |= MASK_PRIMARY;
 	}
 	if (0 == CheckHeader(header2, 1, gpt->streaming_drive_sectors,
-			     gpt->gpt_drive_sectors, gpt->flags)) {
+			     gpt->gpt_drive_sectors, gpt->flags,
+			     gpt->sector_bytes)) {
 		gpt->valid_headers |= MASK_SECONDARY;
 		if (!goodhdr)
 			goodhdr = header2;
@@ -347,7 +353,8 @@ void GptRepair(GptData *gpt)
 		memcpy(header2, header1, sizeof(GptHeader));
 		header2->my_lba = gpt->gpt_drive_sectors - GPT_HEADER_SECTORS;
 		header2->alternate_lba = GPT_PMBR_SECTORS;  /* Second sector. */
-		header2->entries_lba = header2->my_lba - CalculateEntriesSectors(header1);
+		header2->entries_lba = header2->my_lba -
+			CalculateEntriesSectors(header1, gpt->sector_bytes);
 		header2->header_crc32 = HeaderCrc(header2);
 		gpt->modified |= GPT_MODIFIED_HEADER2;
 	}
diff --git a/firmware/lib/cgptlib/include/cgptlib_internal.h b/firmware/lib/cgptlib/include/cgptlib_internal.h
index eaa24312..f19ef819 100644
--- a/firmware/lib/cgptlib/include/cgptlib_internal.h
+++ b/firmware/lib/cgptlib/include/cgptlib_internal.h
@@ -93,7 +93,8 @@ int CheckParameters(GptData* gpt);
  */
 int CheckHeader(GptHeader *h, int is_secondary,
 		uint64_t streaming_drive_sectors,
-		uint64_t gpt_drive_sectors, uint32_t flags);
+		uint64_t gpt_drive_sectors, uint32_t flags,
+		uint32_t sector_bytes);
 
 /**
  * Calculate and return the header CRC.
@@ -160,8 +161,9 @@ void GetCurrentKernelUniqueGuid(GptData *gpt, void *dest);
 const char *GptErrorText(int error_code);
 
 /**
- * Return number of 512-byte sectors required to store the entries table.
+ * Return number of sectors required to store the entries table. Where
+ * a sector has size sector_bytes.
  */
-size_t CalculateEntriesSectors(GptHeader* h);
+size_t CalculateEntriesSectors(GptHeader* h, uint32_t sector_bytes);
 
 #endif /* VBOOT_REFERENCE_CGPTLIB_INTERNAL_H_ */
diff --git a/firmware/lib/gpt_misc.c b/firmware/lib/gpt_misc.c
index ca16f220..45885834 100644
--- a/firmware/lib/gpt_misc.c
+++ b/firmware/lib/gpt_misc.c
@@ -57,7 +57,8 @@ int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 	if (0 == CheckHeader(primary_header, 0,
 			gptdata->streaming_drive_sectors,
 			gptdata->gpt_drive_sectors,
-			gptdata->flags)) {
+			gptdata->flags,
+			gptdata->sector_bytes)) {
 		primary_valid = 1;
 		uint64_t entries_bytes =
 				(uint64_t)primary_header->number_of_entries
@@ -91,7 +92,8 @@ int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
 	if (0 == CheckHeader(secondary_header, 1,
 			gptdata->streaming_drive_sectors,
 			gptdata->gpt_drive_sectors,
-			gptdata->flags)) {
+			gptdata->flags,
+			gptdata->sector_bytes)) {
 		secondary_valid = 1;
 		uint64_t entries_bytes =
 				(uint64_t)secondary_header->number_of_entries
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 */