vboot: cgpt: Treat drive_path as the GPT storage

Previously, "cgpt" called out to "flashrom" directly to read and write
NOR area. This CL removes that dependency and always treats "drive_path"
as the storage of GPT structs. This makes it consistent that whatever
device that cgpt reads from or writes to is always the device that
stores GPT structs. We only need to pass in the size of the drive that
contains the partitions, but we do not need to access to that drive.

More information is in the bug.

BUG=chromium:432611
BRANCH=none
TEST=unittest
CQ-DEPEND=CL:228942

Change-Id: Id0139adf70463cec4f2924de8b9a4725dbec822b
Reviewed-on: https://chromium-review.googlesource.com/229736
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Commit-Queue: Nam Nguyen <namnguyen@chromium.org>
Tested-by: Nam Nguyen <namnguyen@chromium.org>

24 files changed, 309 insertions, 552 deletions

diff --git a/Makefile b/Makefile
index bfdbf85b..07aabaf1 100644
--- a/Makefile
+++ b/Makefile
@@ -361,7 +361,6 @@ UTILLIB_SRCS = \
 	cgpt/cgpt_repair.c \
 	cgpt/cgpt_prioritize.c \
 	cgpt/cgpt_common.c \
-	cgpt/drive.c \
 	cgpt/flash_ts.c \
 	cgpt/flash_ts_drv.c \
 	firmware/lib/cgptlib/mtdlib.c \
@@ -391,7 +390,6 @@ HOSTLIB_SRCS = \
 	cgpt/cgpt_common.c \
 	cgpt/cgpt_create.c \
 	cgpt/cgpt_prioritize.c \
-	cgpt/drive.c \
 	cgpt/flash_ts.c \
 	cgpt/flash_ts_drv.c \
 	firmware/lib/cgptlib/cgptlib_internal.c \
@@ -424,7 +422,6 @@ TINYHOSTLIB_SRCS = \
 	cgpt/cgpt_common.c \
 	cgpt/cgpt_create.c \
 	cgpt/cgpt_prioritize.c \
-	cgpt/drive.c \
 	cgpt/flash_ts.c \
 	cgpt/flash_ts_drv.c \
 	firmware/lib/cgptlib/cgptlib_internal.c \
@@ -432,8 +429,7 @@ TINYHOSTLIB_SRCS = \
 	firmware/lib/cgptlib/mtdlib.c \
 	firmware/lib/utility_string.c \
 	firmware/stub/utility_stub.c \
-	futility/dump_kernel_config_lib.c \
-	host/lib/fmap.c
+	futility/dump_kernel_config_lib.c
 
 TINYHOSTLIB_OBJS = ${TINYHOSTLIB_SRCS:%.c=${BUILD}/%.o}
 
@@ -1134,7 +1130,7 @@ runcgpttests: test_setup
 runtestscripts: test_setup genfuzztestcases
 	tests/load_kernel_tests.sh
 	tests/run_cgpt_tests.sh ${BUILD_RUN}/cgpt/cgpt
-	tests/run_cgpt_tests.sh ${BUILD_RUN}/cgpt/cgpt -N=512,32,1,3
+	tests/run_cgpt_tests.sh ${BUILD_RUN}/cgpt/cgpt -D 358400
 	tests/run_preamble_tests.sh
 	tests/run_rsa_tests.sh
 	tests/run_vbutil_kernel_arg_tests.sh
diff --git a/cgpt/cgpt.h b/cgpt/cgpt.h
index ea459cb1..bb3438a0 100644
--- a/cgpt/cgpt.h
+++ b/cgpt/cgpt.h
@@ -12,7 +12,6 @@
 #include <stdlib.h>
 #include "cgpt_endian.h"
 #include "cgptlib.h"
-#include "drive.h"
 #include "gpt.h"
 #include "mtdlib.h"
 
@@ -48,28 +47,7 @@ struct drive {
   GptData gpt;
   MtdData mtd;
   struct pmbr pmbr;
-  /*
-   * For use with regular file or block device.
-   * GPT structures will occupy the first and last few blocks.
-   */
-  struct {
-    int fd;       /* file descriptor */
-  };
-  /*
-   * For use with flash.
-   * GPT structures will be stored in flash, while the partitions are in
-   * /dev/mtd*.
-   */
-  struct {
-    off_t current_position;  /* for used by DriveSeekFunc */
-    uint32_t flash_start;    /* offset where we can write to flash, in FMAP */
-    uint32_t flash_size;     /* size of that area, in bytes, in FMAP */
-  };  /* for use with flashrom */
-  DriveSeekFunc seek;
-  DriveReadFunc read;
-  DriveWriteFunc write;
-  DriveSyncFunc sync;
-  DriveCloseFunc close;
+  int fd;       /* file descriptor */
 };
 
 struct nand_layout {
@@ -83,11 +61,16 @@ void EnableNandImage(int bytes_per_page, int pages_per_block,
                      int fts_block_offset, int fts_block_size);
 
 // Opens a block device or file, loads raw GPT data from it.
-// mode should be O_RDONLY or O_RDWR
+// 'mode' should be O_RDONLY or O_RDWR.
+// If 'drive_size' is 0, both the partitions and GPT structs reside on the same
+// 'drive_path'.
+// Otherwise, 'drive_size' is taken as the size of the device that all
+// partitions will reside on, and 'drive_path' is where we store GPT structs.
 //
 // Returns CGPT_FAILED if any error happens.
-// Returns CGPT_OK if success and information are stored in 'drive'.
-int DriveOpen(const char *drive_path, struct drive *drive, int mode);
+// Returns CGPT_OK if success and information are stored in 'drive'. */
+int DriveOpen(const char *drive_path, struct drive *drive, int mode,
+              uint64_t drive_size);
 int DriveClose(struct drive *drive, int update_as_needed);
 int CheckValid(const struct drive *drive);
 
diff --git a/cgpt/cgpt_add.c b/cgpt/cgpt_add.c
index 1ff9e560..9bbc3fb5 100644
--- a/cgpt/cgpt_add.c
+++ b/cgpt/cgpt_add.c
@@ -207,7 +207,8 @@ int CgptSetAttributes(CgptAddParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (CgptCheckAddValidity(&drive)) {
@@ -244,7 +245,8 @@ int CgptGetPartitionDetails(CgptAddParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (CgptCheckAddValidity(&drive)) {
@@ -359,7 +361,8 @@ int CgptAdd(CgptAddParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (CgptCheckAddValidity(&drive)) {
diff --git a/cgpt/cgpt_boot.c b/cgpt/cgpt_boot.c
index d23c303a..928210b5 100644
--- a/cgpt/cgpt_boot.c
+++ b/cgpt/cgpt_boot.c
@@ -20,7 +20,8 @@ int CgptGetBootPartitionNumber(CgptBootParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDONLY))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDONLY,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive.gpt))) {
@@ -73,7 +74,8 @@ int CgptBoot(CgptBootParams *params) {
   if (params->create_pmbr || params->partition || params->bootfile)
     mode = O_RDWR;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, mode)) {
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, mode,
+                           params->drive_size)) {
     return CGPT_FAILED;
   }
 
diff --git a/cgpt/cgpt_common.c b/cgpt/cgpt_common.c
index caee434a..1e7b5ab4 100644
--- a/cgpt/cgpt_common.c
+++ b/cgpt/cgpt_common.c
@@ -94,12 +94,12 @@ int Load(struct drive *drive, uint8_t **buf,
   *buf = malloc(count);
   require(*buf);
 
-  if (-1 == drive->seek(drive, sector * sector_bytes, SEEK_SET)) {
+  if (-1 == lseek(drive->fd, sector * sector_bytes, SEEK_SET)) {
     Error("Can't seek: %s\n", strerror(errno));
     goto error_free;
   }
 
-  nread = drive->read(drive, *buf, count);
+  nread = read(drive->fd, *buf, count);
   if (nread < count) {
     Error("Can't read enough: %d, not %d\n", nread, count);
     goto error_free;
@@ -115,10 +115,10 @@ error_free:
 
 
 int ReadPMBR(struct drive *drive) {
-  if (-1 == drive->seek(drive, 0, SEEK_SET))
+  if (-1 == lseek(drive->fd, 0, SEEK_SET))
     return CGPT_FAILED;
 
-  int nread = drive->read(drive, &drive->pmbr, sizeof(struct pmbr));
+  int nread = read(drive->fd, &drive->pmbr, sizeof(struct pmbr));
   if (nread != sizeof(struct pmbr))
     return CGPT_FAILED;
 
@@ -126,10 +126,10 @@ int ReadPMBR(struct drive *drive) {
 }
 
 int WritePMBR(struct drive *drive) {
-  if (-1 == drive->seek(drive, 0, SEEK_SET))
+  if (-1 == lseek(drive->fd, 0, SEEK_SET))
     return CGPT_FAILED;
 
-  int nwrote = drive->write(drive, &drive->pmbr, sizeof(struct pmbr));
+  int nwrote = write(drive->fd, &drive->pmbr, sizeof(struct pmbr));
   if (nwrote != sizeof(struct pmbr))
     return CGPT_FAILED;
 
@@ -146,10 +146,10 @@ int Save(struct drive *drive, const uint8_t *buf,
   require(buf);
   count = sector_bytes * sector_count;
 
-  if (-1 == drive->seek(drive, sector * sector_bytes, SEEK_SET))
+  if (-1 == lseek(drive->fd, sector * sector_bytes, SEEK_SET))
     return CGPT_FAILED;
 
-  nwrote = drive->write(drive, buf, count);
+  nwrote = write(drive->fd, buf, count);
   if (nwrote < count)
     return CGPT_FAILED;
 
@@ -420,14 +420,31 @@ static int GptSave(struct drive *drive) {
   return errors ? -1 : 0;
 }
 
-
-// Opens a block device or file, loads raw GPT data from it.
-// mode should be O_RDONLY or O_RDWR
-//
-// Returns CGPT_FAILED if any error happens.
-// Returns CGPT_OK if success and information are stored in 'drive'. */
-int DriveOpen(const char *drive_path, struct drive *drive, int mode) {
+/*
+ * Query drive size and bytes per sector. Return zero on success. On error,
+ * -1 is returned and errno is set appropriately.
+ */
+static int ObtainDriveSize(int fd, uint64_t* size, uint32_t* sector_bytes) {
   struct stat stat;
+  if (fstat(fd, &stat) == -1) {
+    return -1;
+  }
+  if ((stat.st_mode & S_IFMT) != S_IFREG) {
+    if (ioctl(fd, BLKGETSIZE64, size) < 0) {
+      return -1;
+    }
+    if (ioctl(fd, BLKSSZGET, sector_bytes) < 0) {
+      return -1;
+    }
+  } else {
+    *sector_bytes = 512;  /* bytes */
+    *size = stat.st_size;
+  }
+  return 0;
+}
+
+int DriveOpen(const char *drive_path, struct drive *drive, int mode,
+              uint64_t drive_size) {
   uint32_t sector_bytes;
   int is_mtd = nand.enabled;
 
@@ -448,50 +465,21 @@ int DriveOpen(const char *drive_path, struct drive *drive, int mode) {
       return CGPT_FAILED;
     }
 
-    drive->seek = FileSeek;
-    drive->read = FileRead;
-    drive->write = FileWrite;
-    drive->sync = FileSync;
-    drive->close = FileClose;
-    if (fstat(drive->fd, &stat) == -1) {
-      Error("Can't fstat %s: %s\n", drive_path, strerror(errno));
+    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",
+            drive_path, strerror(errno));
       goto error_close;
     }
-    if (major(stat.st_rdev) == MTD_CHAR_MAJOR) {
-      mtd_info_t mtd_info;
-      if (ioctl(drive->fd, MEMGETINFO, &mtd_info) != 0) {
-        Error("Can't get the size of the MTD device\n");
-        goto error_close;
-      }
-      drive->size = mtd_info.size;
-
-      if (FlashInit(drive) != 0) {
-        Error("Can't obtain NOR flash info with flashrom\n");
-        goto error_close;
-      }
 
-      sector_bytes = 512;
-      drive->gpt.stored_on_device = GPT_STORED_OFF_DEVICE;
-      drive->gpt.gpt_drive_sectors = drive->flash_size / sector_bytes;
-      drive->seek = FlashSeek;
-      drive->read = FlashRead;
-      drive->write = FlashWrite;
-      drive->sync = FlashSync;
-      drive->close = FlashClose;
-    } else if ((stat.st_mode & S_IFMT) != S_IFREG) {
-      if (ioctl(drive->fd, BLKGETSIZE64, &drive->size) < 0) {
-        Error("Can't read drive size from %s: %s\n", drive_path,
-              strerror(errno));
-        goto error_close;
-      }
-      if (ioctl(drive->fd, BLKSSZGET, &sector_bytes) < 0) {
-        Error("Can't read sector size from %s: %s\n",
-              drive_path, strerror(errno));
-        goto error_close;
-      }
+    drive->gpt.gpt_drive_sectors = gpt_drive_size / sector_bytes;
+    if (drive_size == 0) {
+      drive->size = gpt_drive_size;
+      drive->gpt.stored_on_device = GPT_STORED_ON_DEVICE;
     } else {
-      sector_bytes = 512;  /* bytes */
-      drive->size = stat.st_size;
+      drive->size = drive_size;
+      drive->gpt.stored_on_device = GPT_STORED_OFF_DEVICE;
     }
   }
   drive->is_mtd = is_mtd;
@@ -537,9 +525,9 @@ int DriveClose(struct drive *drive, int update_as_needed) {
   // Sync early! Only sync file descriptor here, and leave the whole system sync
   // outside cgpt because whole system sync would trigger tons of disk accesses
   // and timeout tests.
-  drive->sync(drive);
+  fsync(drive->fd);
 
-  drive->close(drive);
+  close(drive->fd);
 
   return errors ? CGPT_FAILED : CGPT_OK;
 }
@@ -1234,7 +1222,8 @@ int CgptGetNumNonEmptyPartitions(CgptShowParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDONLY))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDONLY,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive.gpt))) {
diff --git a/cgpt/cgpt_create.c b/cgpt/cgpt_create.c
index 6a21a8b0..18e2f51b 100644
--- a/cgpt/cgpt_create.c
+++ b/cgpt/cgpt_create.c
@@ -52,7 +52,7 @@ static int GptCreate(struct drive *drive, CgptCreateParams *params) {
       h->last_usable_lba = (drive->gpt.drive_sectors - GPT_HEADER_SECTORS -
                             GPT_ENTRIES_SECTORS - 1);
     } else {
-      h->first_usable_lba = 0;
+      h->first_usable_lba = params->padding;
       h->last_usable_lba = (drive->gpt.drive_sectors - 1);
     }
     if (CGPT_OK != GenerateGuid(&h->disk_uuid)) {
@@ -63,18 +63,18 @@ static int GptCreate(struct drive *drive, CgptCreateParams *params) {
     h->number_of_entries = TOTAL_ENTRIES_SIZE / h->size_of_entry;
     if (drive->gpt.stored_on_device != GPT_STORED_ON_DEVICE) {
       // We might have smaller space for the GPT table. Scale accordingly.
-      size_t half_size = drive->flash_size / 2;
-      size_t header_block_size = GPT_HEADER_SECTORS * drive->gpt.sector_bytes;
-      if (half_size < header_block_size) {
+      size_t half_size_sectors = drive->gpt.gpt_drive_sectors / 2;
+      if (half_size_sectors < GPT_HEADER_SECTORS) {
         Error("Not enough space for a GPT header.\n");
         return -1;
       }
-      half_size -= header_block_size;
-      if (half_size < MIN_NUMBER_OF_ENTRIES * h->size_of_entry) {
+      half_size_sectors -= GPT_HEADER_SECTORS;
+      size_t half_size = half_size_sectors * drive->gpt.sector_bytes;
+      if (half_size < (MIN_NUMBER_OF_ENTRIES * h->size_of_entry)) {
         Error("Not enough space for minimum number of entries.\n");
         return -1;
       }
-      if (128 > half_size / h->size_of_entry) {
+      if (128 > (half_size / h->size_of_entry)) {
         h->number_of_entries = half_size / h->size_of_entry;
       }
     }
@@ -119,7 +119,8 @@ int CgptCreate(CgptCreateParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (drive.is_mtd) {
diff --git a/cgpt/cgpt_find.c b/cgpt/cgpt_find.c
index fc223a02..7d9ec8df 100644
--- a/cgpt/cgpt_find.c
+++ b/cgpt/cgpt_find.c
@@ -211,7 +211,7 @@ static int do_search(CgptFindParams *params, char *fileName) {
   int retval;
   struct drive drive;
 
-  if (CGPT_OK != DriveOpen(fileName, &drive, O_RDONLY))
+  if (CGPT_OK != DriveOpen(fileName, &drive, O_RDONLY, params->drive_size))
     return 0;
 
   if (drive.is_mtd) {
diff --git a/cgpt/cgpt_legacy.c b/cgpt/cgpt_legacy.c
index e736bd4c..07607690 100644
--- a/cgpt/cgpt_legacy.c
+++ b/cgpt/cgpt_legacy.c
@@ -15,7 +15,8 @@ int CgptLegacy(CgptLegacyParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (drive.is_mtd) {
diff --git a/cgpt/cgpt_prioritize.c b/cgpt/cgpt_prioritize.c
index 8fc1843e..019638d2 100644
--- a/cgpt/cgpt_prioritize.c
+++ b/cgpt/cgpt_prioritize.c
@@ -106,7 +106,8 @@ int CgptPrioritize(CgptPrioritizeParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (drive.is_mtd) {
diff --git a/cgpt/cgpt_repair.c b/cgpt/cgpt_repair.c
index 9b4d2dfb..98899026 100644
--- a/cgpt/cgpt_repair.c
+++ b/cgpt/cgpt_repair.c
@@ -15,7 +15,8 @@ int CgptRepair(CgptRepairParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (drive.is_mtd) {
diff --git a/cgpt/cgpt_show.c b/cgpt/cgpt_show.c
index 5b7974f7..32d62d38 100644
--- a/cgpt/cgpt_show.c
+++ b/cgpt/cgpt_show.c
@@ -482,7 +482,8 @@ int CgptShow(CgptShowParams *params) {
   if (params == NULL)
     return CGPT_FAILED;
 
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDONLY))
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDONLY,
+                           params->drive_size))
     return CGPT_FAILED;
 
   if (drive.is_mtd) {
diff --git a/cgpt/cmd_add.c b/cgpt/cmd_add.c
index 2bb5df43..0d89065b 100644
--- a/cgpt/cmd_add.c
+++ b/cgpt/cmd_add.c
@@ -15,6 +15,9 @@ static void Usage(void)
   printf("\nUsage: %s add [OPTIONS] DRIVE\n\n"
          "Add, edit, or remove a partition entry.\n\n"
          "Options:\n"
+         "  -D NUM       Size (in bytes) of the disk where partitions reside\n"
+         "                 default 0, meaning partitions and GPT structs are\n"
+         "                 both on DRIVE\n"
          "  -i NUM       Specify partition (default is next available)\n"
          "  -b NUM       Beginning sector\n"
          "  -s NUM       Size in sectors\n"
@@ -42,10 +45,18 @@ int cmd_add(int argc, char *argv[]) {
   char *e = 0;
 
   opterr = 0;                     // quiet, you
-  while ((c=getopt(argc, argv, ":hi:b:s:t:u:l:S:T:P:A:")) != -1)
+  while ((c=getopt(argc, argv, ":hi:b:s:t:u:l:S:T:P:A:D:")) != -1)
   {
     switch (c)
     {
+    case 'D':
+      params.drive_size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
     case 'i':
       params.partition = (uint32_t)strtoul(optarg, &e, 0);
       if (!*optarg || (e && *e))
diff --git a/cgpt/cmd_boot.c b/cgpt/cmd_boot.c
index 2a65f202..2f815663 100644
--- a/cgpt/cmd_boot.c
+++ b/cgpt/cmd_boot.c
@@ -15,6 +15,9 @@ static void Usage(void)
   printf("\nUsage: %s boot [OPTIONS] DRIVE\n\n"
          "Edit the PMBR sector for legacy BIOSes\n\n"
          "Options:\n"
+         "  -D NUM       Size (in bytes) of the disk where partitions reside\n"
+         "                 default 0, meaning partitions and GPT structs are\n"
+         "                 both on DRIVE\n"
          "  -i NUM       Set bootable partition\n"
          "  -b FILE      Install bootloader code in the PMBR\n"
          "  -p           Create legacy PMBR partition table\n"
@@ -34,10 +37,18 @@ int cmd_boot(int argc, char *argv[]) {
   char *e = 0;
 
   opterr = 0;                     // quiet, you
-  while ((c=getopt(argc, argv, ":hi:b:p")) != -1)
+  while ((c=getopt(argc, argv, ":hi:b:pD:")) != -1)
   {
     switch (c)
     {
+    case 'D':
+      params.drive_size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
     case 'i':
       params.partition = (uint32_t)strtoul(optarg, &e, 0);
       if (!*optarg || (e && *e))
diff --git a/cgpt/cmd_create.c b/cgpt/cmd_create.c
index ca1b815d..47137bdc 100644
--- a/cgpt/cmd_create.c
+++ b/cgpt/cmd_create.c
@@ -15,10 +15,13 @@ static void Usage(void)
   printf("\nUsage: %s create [OPTIONS] DRIVE\n\n"
          "Create or reset an empty GPT.\n\n"
          "Options:\n"
+         "  -D NUM       Size (in bytes) of the disk where partitions reside\n"
+         "                 default 0, meaning partitions and GPT structs are\n"
+         "                 both on DRIVE\n"
          "  -z           Zero the sectors of the GPT table and entries\n"
-         "  -s           Size (in byes) of the disk (MTD only)\n"
-         "  -p           Size (in blocks) of the disk to pad between the\n"
-         "               primary GPT header and its entries, default 0\n"
+         "  -s NUM       Size (in bytes) of the disk (MTD only)\n"
+         "  -p NUM       Size (in blocks) of the disk to pad between the\n"
+         "                 primary GPT header and its entries, default 0\n"
          "\n", progname);
 }
 
@@ -31,10 +34,18 @@ int cmd_create(int argc, char *argv[]) {
   char *e = 0;
 
   opterr = 0;                     // quiet, you
-  while ((c=getopt(argc, argv, ":hzsp:")) != -1)
+  while ((c=getopt(argc, argv, ":hzs:p:D:")) != -1)
   {
     switch (c)
     {
+    case 'D':
+      params.drive_size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
     case 'z':
       params.zap = 1;
       break;
@@ -43,8 +54,12 @@ int cmd_create(int argc, char *argv[]) {
       break;
     case 'p':
       params.padding = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
       break;
-
     case 'h':
       Usage();
       return CGPT_OK;
diff --git a/cgpt/cmd_find.c b/cgpt/cmd_find.c
index bdd745c2..d0f048f2 100644
--- a/cgpt/cmd_find.c
+++ b/cgpt/cmd_find.c
@@ -16,6 +16,9 @@ static void Usage(void)
          "Find a partition by its UUID or label. With no specified DRIVE\n"
          "it scans all physical drives.\n\n"
          "Options:\n"
+         "  -D NUM       Size (in bytes) of the disk where partitions reside\n"
+         "                 default 0, meaning partitions and GPT structs are\n"
+         "                 both on DRIVE\n"
          "  -t GUID      Search for Partition Type GUID\n"
          "  -u GUID      Search for Partition Unique ID\n"
          "  -l LABEL     Search for Label\n"
@@ -71,10 +74,18 @@ int cmd_find(int argc, char *argv[]) {
   int c;
 
   opterr = 0;                     // quiet, you
-  while ((c=getopt(argc, argv, ":hv1nt:u:l:M:O:")) != -1)
+  while ((c=getopt(argc, argv, ":hv1nt:u:l:M:O:D:")) != -1)
   {
     switch (c)
     {
+    case 'D':
+      params.drive_size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
     case 'v':
       params.verbose++;
       break;
diff --git a/cgpt/cmd_legacy.c b/cgpt/cmd_legacy.c
index 955ce3f4..5fd742cc 100644
--- a/cgpt/cmd_legacy.c
+++ b/cgpt/cmd_legacy.c
@@ -15,6 +15,9 @@ static void Usage(void)
   printf("\nUsage: %s legacy [OPTIONS] DRIVE\n\n"
          "Switch GPT header signature to \"CHROMEOS\".\n\n"
          "Options:\n"
+         "  -D NUM       Size (in bytes) of the disk where partitions reside\n"
+         "                 default 0, meaning partitions and GPT structs are\n"
+         "                 both on DRIVE\n"
          "  -e           Switch GPT header signature back to \"EFI PART\"\n"
          "\n", progname);
 }
@@ -24,13 +27,22 @@ int cmd_legacy(int argc, char *argv[]) {
   memset(&params, 0, sizeof(params));
 
   int c;
+  char* e = 0;
   int errorcnt = 0;
 
   opterr = 0;                     // quiet, you
-  while ((c=getopt(argc, argv, ":he")) != -1)
+  while ((c=getopt(argc, argv, ":heD:")) != -1)
   {
     switch (c)
     {
+    case 'D':
+      params.drive_size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
     case 'e':
       params.efipart = 1;
       break;
diff --git a/cgpt/cmd_prioritize.c b/cgpt/cmd_prioritize.c
index 9bf292e4..ab4e2e36 100644
--- a/cgpt/cmd_prioritize.c
+++ b/cgpt/cmd_prioritize.c
@@ -18,6 +18,9 @@ static void Usage(void)
   printf("\nUsage: %s prioritize [OPTIONS] DRIVE\n\n"
          "Reorder the priority of all active ChromeOS Kernel partitions.\n\n"
          "Options:\n"
+         "  -D NUM       Size (in bytes) of the disk where partitions reside\n"
+         "                 default 0, meaning partitions and GPT structs are\n"
+         "                 both on DRIVE\n"
          "  -P NUM       Highest priority to use in the new ordering. The\n"
          "                 other partitions will be ranked in decreasing\n"
          "                 priority while preserving their original order.\n"
@@ -43,10 +46,18 @@ int cmd_prioritize(int argc, char *argv[]) {
   char *e = 0;
 
   opterr = 0;                     // quiet, you
-  while ((c=getopt(argc, argv, ":hi:fP:")) != -1)
+  while ((c=getopt(argc, argv, ":hi:fP:D:")) != -1)
   {
     switch (c)
     {
+    case 'D':
+      params.drive_size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
     case 'i':
       params.set_partition = (uint32_t)strtoul(optarg, &e, 0);
       if (!*optarg || (e && *e))
diff --git a/cgpt/cmd_repair.c b/cgpt/cmd_repair.c
index b5f7d31e..528c650a 100644
--- a/cgpt/cmd_repair.c
+++ b/cgpt/cmd_repair.c
@@ -15,6 +15,9 @@ static void Usage(void)
   printf("\nUsage: %s repair [OPTIONS] DRIVE\n\n"
          "Repair damaged GPT headers and tables.\n\n"
          "Options:\n"
+         "  -D NUM       Size (in bytes) of the disk where partitions reside\n"
+         "                 default 0, meaning partitions and GPT structs are\n"
+         "                 both on DRIVE\n"
          "  -v           Verbose\n"
          "\n", progname);
 }
@@ -24,13 +27,22 @@ int cmd_repair(int argc, char *argv[]) {
   memset(&params, 0, sizeof(params));
 
   int c;
+  char* e = 0;
   int errorcnt = 0;
 
   opterr = 0;                     // quiet, you
-  while ((c=getopt(argc, argv, ":hv")) != -1)
+  while ((c=getopt(argc, argv, ":hvD:")) != -1)
   {
     switch (c)
     {
+    case 'D':
+      params.drive_size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
     case 'v':
       params.verbose++;
       break;
diff --git a/cgpt/cmd_show.c b/cgpt/cmd_show.c
index df4cba1f..dcfefdd3 100644
--- a/cgpt/cmd_show.c
+++ b/cgpt/cmd_show.c
@@ -17,6 +17,9 @@ static void Usage(void)
   printf("\nUsage: %s show [OPTIONS] DRIVE\n\n"
          "Display the GPT table\n\n"
          "Options:\n"
+         "  -D NUM       Size (in bytes) of the disk where partitions reside\n"
+         "                 default 0, meaning partitions and GPT structs are\n"
+         "                 both on DRIVE\n"
          "  -n           Numeric output only\n"
          "  -v           Verbose output\n"
          "  -q           Quick output\n"
@@ -43,10 +46,18 @@ int cmd_show(int argc, char *argv[]) {
   char *e = 0;
 
   opterr = 0;                     // quiet, you
-  while ((c=getopt(argc, argv, ":hnvqi:bstulSTPAd")) != -1)
+  while ((c=getopt(argc, argv, ":hnvqi:bstulSTPAdD:")) != -1)
   {
     switch (c)
     {
+    case 'D':
+      params.drive_size = strtoull(optarg, &e, 0);
+      if (!*optarg || (e && *e))
+      {
+        Error("invalid argument to -%c: \"%s\"\n", c, optarg);
+        errorcnt++;
+      }
+      break;
     case 'n':
       params.numeric = 1;
       break;
diff --git a/cgpt/drive.c b/cgpt/drive.c
deleted file mode 100644
index db6601d1..00000000
--- a/cgpt/drive.c
+++ /dev/null
@@ -1,305 +0,0 @@
-/* Copyright 2014 The Chromium OS Authors. All rights reserved.
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- */
-
-#include <errno.h>
-#include <stdarg.h>
-#include <string.h>
-#include <unistd.h>
-
-#include "cgpt.h"
-#include "fmap.h"
-
-// TODO(namnguyen): Remove RW_UNUSED
-#ifdef DEBUG
-static const char FMAP_GPT_SECTION[] = "RW_UNUSED";
-#else
-static const char FMAP_GPT_SECTION[] = "RW_GPT";
-#endif
-
-off_t FileSeek(struct drive* drive, off_t offset, int whence) {
-  return lseek(drive->fd, offset, whence);
-}
-
-ssize_t FileRead(struct drive* drive, void* buf, size_t count) {
-  return read(drive->fd, buf, count);
-}
-
-ssize_t FileWrite(struct drive* drive, const void* buf, size_t count) {
-  return write(drive->fd, buf, count);
-}
-
-int FileSync(struct drive* drive) {
-  return fsync(drive->fd);
-}
-
-int FileClose(struct drive* drive) {
-  return close(drive->fd);
-}
-
-// Always terminate the buffer after snprintf.
-static int tsnprintf(char *buf, size_t size, const char* fmt, ...) {
-  if (size == 0) {
-    // No space for the null char.
-    errno = ENOSPC;
-    return -1;
-  }
-  va_list ap;
-  va_start(ap, fmt);
-  int ret = vsnprintf(buf, size, fmt, ap);
-  va_end(ap);
-  if (ret >= 0) {
-    buf[size - 1] = '\x00';
-  }
-  return ret;
-}
-
-int FlashInit(struct drive* drive) {
-  int return_code = 1;
-  char tempdir[] = "/tmp/cgptXXXXXX";
-  if (mkdtemp(tempdir) == NULL) {
-    Error("Cannot create temp directory for flashrom work.\n");
-    return return_code;
-  }
-
-  char cmd[256];
-  char fmap_name[28];
-  tsnprintf(fmap_name, sizeof(fmap_name), "%s/fmap", tempdir);
-  tsnprintf(cmd, sizeof(cmd), "/usr/sbin/flashrom -p host -i FMAP:%s -r "
-            "> /dev/null 2>&1", fmap_name);
-  return_code++;
-  if (system(cmd) != 0) {
-    Error("Cannot dump FMAP section from flash.\n");
-    goto cleanup;
-  };
-
-  return_code++;
-  int fmap_fd = open(fmap_name, O_RDONLY);
-  if (fmap_fd < 0) {
-    Error("Cannot open %s.\n", fmap_name);
-    goto cleanup;
-  }
-  // Allocate 4096 bytes. ChromeOS FMAP is usually 2048 bytes.
-  return_code++;
-  const size_t fmap_alloc_size = 4096;
-  uint8_t* fmap = malloc(fmap_alloc_size);
-  if (!fmap) {
-    Error("Cannot read fmap.\n");
-    goto cleanup2;
-  }
-  return_code++;
-  int fmap_size = read(fmap_fd, fmap, fmap_alloc_size);
-  if (fmap_size < 0) {
-    Error("Cannot read from %s.\n", fmap_name);
-    goto cleanup3;
-  }
-
-  return_code++;
-  FmapAreaHeader* gpt_area;
-  if (fmap_find_by_name(fmap, fmap_size, NULL,
-                        FMAP_GPT_SECTION, &gpt_area) == NULL) {
-    Error("Cannot find GPT section in the FMAP.\n");
-    goto cleanup3;
-  }
-
-  drive->flash_start = gpt_area->area_offset;
-  drive->flash_size = gpt_area->area_size;
-  drive->current_position = 0;
-
-  return_code = 0;
-
-cleanup3:
-  free(fmap);
-cleanup2:
-  close(fmap_fd);
-cleanup:
-  tsnprintf(cmd, sizeof(cmd), "/bin/rm -rf %s", tempdir);
-  if (system(cmd)) {
-    Warning("Cannot remove temp directory", tempdir);
-  }
-  return return_code;
-}
-
-off_t FlashSeek(struct drive* drive, off_t offset, int whence) {
-  off_t new_position;
-  switch (whence) {
-    case SEEK_SET:
-      new_position = offset;
-      break;
-    case SEEK_CUR:
-      new_position = drive->current_position + offset;
-      break;
-    case SEEK_END:
-      new_position = drive->size + offset;
-      break;
-    default:
-      errno = EINVAL;
-      return -1;
-  }
-  if (new_position < 0 || new_position > drive->size) {
-    errno = EINVAL;
-    return -1;
-  }
-  drive->current_position = new_position;
-  return new_position;
-}
-
-// Translate |position| to an address in flash.
-// We only use a small area in flash to store the GPT structures. This area is
-// identified in FMAP. So the idea is to map |position| from 0 to flash_size to
-// the physical position in flash linearly.
-// This function returns 0 for success.
-static int TranslateToFlash(struct drive* drive, off_t position, size_t count,
-                            off_t* translated) {
-  if (position < 0 || position + count > drive->flash_size) {
-    return -1;
-  }
-  *translated = position + drive->flash_start;
-  return 0;
-}
-
-static int CreateLayout(char* file_name, off_t position, size_t count) {
-  int fd = mkstemp(file_name);
-  if (fd < 0) {
-    Error("Cannot create layout file.\n");
-    return -1;
-  }
-  char buf[128];
-  tsnprintf(buf, sizeof(buf), "%08X:%08X landmark\n", (unsigned int) position,
-            (unsigned int) (position + count - 1));
-  int layout_len = strlen(buf);
-  int nr_written = write(fd, buf, layout_len);
-  close(fd);
-  if (nr_written != layout_len) {
-    Error("Cannot write out layout for flashrom.\n");
-    return -1;
-  }
-
-  return 0;
-}
-
-ssize_t FlashRead(struct drive* drive, void* buf, size_t count) {
-  off_t offset;
-  if (TranslateToFlash(drive, drive->current_position, count, &offset) != 0) {
-    Error("Cannot translate disk address %08X to SPI address.\n",
-          drive->current_position);
-    errno = EINVAL;
-    return -1;
-  }
-
-  char tempdir[] = "/tmp/cgptXXXXXX";
-  if (mkdtemp(tempdir) == NULL) {
-    Error("Cannot create temp directory for flashrom work.\n");
-    errno = EIO;
-    return -1;
-  }
-
-  int return_value = -1;
-  char layout_file[40];
-  tsnprintf(layout_file, sizeof(layout_file), "%s/layoutXXXXXX", tempdir);
-  if (CreateLayout(layout_file, offset, count) != 0) {
-    Error("Cannot create layout file for flashrom.\n");
-    goto cleanup;
-  }
-
-  char content_file[40];
-  tsnprintf(content_file, sizeof(content_file), "%s/contentXXXXXX", tempdir);
-  int fd = mkstemp(content_file);
-  if (fd < 0) {
-    goto cleanup;
-  }
-
-  char cmd[256];
-  tsnprintf(cmd, sizeof(cmd), "/usr/sbin/flashrom -p host -l %s -i landmark:%s "
-            "-r > /dev/null 2>&1", layout_file, content_file);
-  if (system(cmd) != 0) {
-    Error("Cannot read from SPI flash.\n");
-    goto cleanup2;
-  }
-
-  return_value = read(fd, buf, count);
-  if (return_value != count) {
-    Error("Cannot read from retrieved content file.\n");
-    return_value = -1;
-  } else {
-    drive->current_position += return_value;
-  }
-
-cleanup2:
-  close(fd);
-cleanup:
-  tsnprintf(cmd, sizeof(cmd), "/bin/rm -rf %s", tempdir);
-  if (system(cmd)) {
-    Warning("Cannot remove temp directory", tempdir);
-  }
-  errno = EIO;
-  return return_value;
-}
-
-ssize_t FlashWrite(struct drive* drive, const void* buf, size_t count) {
-  off_t offset;
-  if (TranslateToFlash(drive, drive->current_position, count, &offset) != 0) {
-    Error("Cannot translate disk address %08X to SPI address.\n",
-          drive->current_position);
-    errno = EINVAL;
-    return -1;
-  }
-
-  char tempdir[] = "/tmp/cgptXXXXXX";
-  if (mkdtemp(tempdir) == NULL) {
-    Error("Cannot create temp directory for flashrom work.\n");
-    errno = EIO;
-    return -1;
-  }
-
-  int return_value = -1;
-  char layout_file[40];
-  tsnprintf(layout_file, sizeof(layout_file), "%s/layoutXXXXXX", tempdir);
-  if (CreateLayout(layout_file, offset, count) != 0) {
-    Error("Cannot create layout file for flashrom.\n");
-    goto cleanup;
-  }
-
-  char content_file[40];
-  tsnprintf(content_file, sizeof(content_file), "%s/contentXXXXXX", tempdir);
-  int fd = mkstemp(content_file);
-  if (fd < 0) {
-    goto cleanup;
-  }
-
-  return_value = write(fd, buf, count);
-  close(fd);
-  if (return_value != count) {
-    Error("Cannot prepare content file for flashrom.\n");
-    return_value = -1;
-    goto cleanup;
-  }
-
-  char cmd[256];
-  // TODO(namnguyen): Add --fast-verify after 428475 is fixed
-  tsnprintf(cmd, sizeof(cmd), "/usr/sbin/flashrom -p host -l %s -i landmark:%s "
-            "-w > /dev/null 2>&1", layout_file, content_file);
-  if (system(cmd) != 0) {
-    Error("Cannot write to SPI flash.\n");
-    return_value = -1;
-  } else {
-    drive->current_position += return_value;
-  }
-
-cleanup:
-  tsnprintf(cmd, sizeof(cmd), "/bin/rm -rf %s", tempdir);
-  if (system(cmd)) {
-    Warning("Cannot remove temp directory", tempdir);
-  }
-  errno = EIO;
-  return return_value;
-}
-
-int FlashSync(struct drive* drive) {
-  return 0;
-}
-
-int FlashClose(struct drive* drive) {
-  return 0;
-}
diff --git a/cgpt/drive.h b/cgpt/drive.h
deleted file mode 100644
index d03bfe74..00000000
--- a/cgpt/drive.h
+++ /dev/null
@@ -1,29 +0,0 @@
-/* Copyright 2014 The Chromium OS Authors. All rights reserved.
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- */
-
-#ifndef VBOOT_REFERENCE_UTILITY_CGPT_DRIVE_H_
-#define VBOOT_REFERENCE_UTILITY_CGPT_DRIVE_H_
-
-struct drive;
-typedef off_t (*DriveSeekFunc)(struct drive*, off_t offset, int whence);
-typedef ssize_t (*DriveReadFunc)(struct drive*, void* buf, size_t count);
-typedef ssize_t (*DriveWriteFunc)(struct drive*, const void* buf, size_t count);
-typedef int (*DriveCloseFunc)(struct drive*);
-typedef int (*DriveSyncFunc)(struct drive*);
-
-off_t FileSeek(struct drive* drive, off_t offset, int whence);
-ssize_t FileRead(struct drive* drive, void* buf, size_t count);
-ssize_t FileWrite(struct drive* drive, const void* buf, size_t count);
-int FileSync(struct drive* drive);
-int FileClose(struct drive* drive);
-
-int FlashInit(struct drive* drive);
-off_t FlashSeek(struct drive* drive, off_t offset, int whence);
-ssize_t FlashRead(struct drive* drive, void* buf, size_t count);
-ssize_t FlashWrite(struct drive* drive, const void* buf, size_t count);
-int FlashSync(struct drive* drive);
-int FlashClose(struct drive* drive);
-
-#endif  // VBOOT_REFERENCE_UTILITY_CGPT_DRIVE_H_
diff --git a/firmware/lib/cgptlib/cgptlib_internal.c b/firmware/lib/cgptlib/cgptlib_internal.c
index 48b527ac..9bcfbd81 100644
--- a/firmware/lib/cgptlib/cgptlib_internal.c
+++ b/firmware/lib/cgptlib/cgptlib_internal.c
@@ -34,7 +34,7 @@ int CheckParameters(GptData *gpt)
 	 * too small to contain basic GPT structure (PMBR + Headers + Entries),
 	 * the value is wrong.
 	 */
-	if (gpt->drive_sectors < (1 + 2 * (1 + GPT_ENTRIES_SECTORS)))
+	if (gpt->gpt_drive_sectors < (1 + 2 * (1 + GPT_ENTRIES_SECTORS)))
 		return GPT_ERROR_INVALID_SECTOR_NUMBER;
 
 	return GPT_SUCCESS;
diff --git a/host/include/cgpt_params.h b/host/include/cgpt_params.h
index ac437744..e10d15e7 100644
--- a/host/include/cgpt_params.h
+++ b/host/include/cgpt_params.h
@@ -15,6 +15,7 @@ enum {
 
 typedef struct CgptCreateParams {
   char *drive_name;
+  uint64_t drive_size;
   int zap;
   uint64_t size;
   uint64_t padding;
@@ -22,6 +23,7 @@ typedef struct CgptCreateParams {
 
 typedef struct CgptAddParams {
   char *drive_name;
+  uint64_t drive_size;
   uint32_t partition;
   uint64_t begin;
   uint64_t size;
@@ -44,6 +46,7 @@ typedef struct CgptAddParams {
 
 typedef struct CgptShowParams {
   char *drive_name;
+  uint64_t drive_size;
   int numeric;
   int verbose;
   int quick;
@@ -55,11 +58,13 @@ typedef struct CgptShowParams {
 
 typedef struct CgptRepairParams {
   char *drive_name;
+  uint64_t drive_size;
   int verbose;
 } CgptRepairParams;
 
 typedef struct CgptBootParams {
   char *drive_name;
+  uint64_t drive_size;
   uint32_t partition;
   char *bootfile;
   int create_pmbr;
@@ -67,6 +72,7 @@ typedef struct CgptBootParams {
 
 typedef struct CgptPrioritizeParams {
   char *drive_name;
+  uint64_t drive_size;
   uint32_t set_partition;
   int set_friends;
   int max_priority;
@@ -75,6 +81,7 @@ typedef struct CgptPrioritizeParams {
 
 typedef struct CgptFindParams {
   char *drive_name;
+  uint64_t drive_size;
   int verbose;
   int set_unique;
   int set_type;
@@ -94,6 +101,7 @@ typedef struct CgptFindParams {
 
 typedef struct CgptLegacyParams {
   char *drive_name;
+  uint64_t drive_size;
   int efipart;
 } CgptLegacyParams;
 
diff --git a/tests/run_cgpt_tests.sh b/tests/run_cgpt_tests.sh
index 95840fbb..6619a6b8 100755
--- a/tests/run_cgpt_tests.sh
+++ b/tests/run_cgpt_tests.sh
@@ -12,8 +12,7 @@
 CGPT=$(readlink -f "$1")
 [ -x "$CGPT" ] || error "Can't execute $CGPT"
 
-MTD="${2:-}"
-CGPT="${CGPT} ${MTD}"
+MTD="${@:2}"
 
 # Run tests in a dedicated directory for easy cleanup or debugging.
 DIR="${TEST_DIR}/cgpt_test_dir"
@@ -21,11 +20,17 @@ DIR="${TEST_DIR}/cgpt_test_dir"
 warning "testing $CGPT in $DIR"
 cd "$DIR"
 
+assert_fail() {
+  set +e
+  "$@" 2>/dev/null
+  if [ $? == 0 ]; then
+    error "$*" " should have failed but did not"
+  fi
+  set -e
+}
+
 # Test failure on non existing file.
-set +e
-${CGPT} show blah_404_haha
-[ $? != 0 ] || error "CGPT should fail on non existing file."
-set -e
+assert_fail ${CGPT} show $MTD blah_404_haha
 
 echo "Create an empty file to use as the device..."
 NUM_SECTORS=1000
@@ -71,85 +76,82 @@ RANDOM_LABEL="random stuff"
 RANDOM_GUID='2364a860-bf63-42fb-a83d-9ad3e057fcf5'
 RANDOM_NUM=6
 
-$CGPT create ${DEV}
+$CGPT create $MTD ${DEV}
 
-$CGPT add -b ${DATA_START} -s ${DATA_SIZE} -t ${DATA_GUID} \
+$CGPT add $MTD -b ${DATA_START} -s ${DATA_SIZE} -t ${DATA_GUID} \
   -l "${DATA_LABEL}" ${DEV}
-$CGPT add -b ${KERN_START} -s ${KERN_SIZE} -t ${KERN_GUID} \
+$CGPT add $MTD -b ${KERN_START} -s ${KERN_SIZE} -t ${KERN_GUID} \
   -l "${KERN_LABEL}" ${DEV}
-$CGPT add -b ${ROOTFS_START} -s ${ROOTFS_SIZE} -t ${ROOTFS_GUID} \
+$CGPT add $MTD -b ${ROOTFS_START} -s ${ROOTFS_SIZE} -t ${ROOTFS_GUID} \
   -l "${ROOTFS_LABEL}" ${DEV}
-$CGPT add -b ${ESP_START} -s ${ESP_SIZE} -t ${ESP_GUID} \
+$CGPT add $MTD -b ${ESP_START} -s ${ESP_SIZE} -t ${ESP_GUID} \
   -l "${ESP_LABEL}" ${DEV}
-$CGPT add -b ${FUTURE_START} -s ${FUTURE_SIZE} -t ${FUTURE_GUID} \
+$CGPT add $MTD -b ${FUTURE_START} -s ${FUTURE_SIZE} -t ${FUTURE_GUID} \
   -l "${FUTURE_LABEL}" ${DEV}
-$CGPT add -b ${RANDOM_START} -s ${RANDOM_SIZE} -t ${RANDOM_GUID} \
+$CGPT add $MTD -b ${RANDOM_START} -s ${RANDOM_SIZE} -t ${RANDOM_GUID} \
   -l "${RANDOM_LABEL}" ${DEV}
 
 
 echo "Extract the start and size of given partitions..."
 
-X=$($CGPT show -b -i $DATA_NUM ${DEV})
-Y=$($CGPT show -s -i $DATA_NUM ${DEV})
+X=$($CGPT show $MTD -b -i $DATA_NUM ${DEV})
+Y=$($CGPT show $MTD -s -i $DATA_NUM ${DEV})
 [ "$X $Y" = "$DATA_START $DATA_SIZE" ] || error
 
-X=$($CGPT show -b -i $KERN_NUM ${DEV})
-Y=$($CGPT show -s -i $KERN_NUM ${DEV})
+X=$($CGPT show $MTD -b -i $KERN_NUM ${DEV})
+Y=$($CGPT show $MTD -s -i $KERN_NUM ${DEV})
 [ "$X $Y" = "$KERN_START $KERN_SIZE" ] || error
 
-X=$($CGPT show -b -i $ROOTFS_NUM ${DEV})
-Y=$($CGPT show -s -i $ROOTFS_NUM ${DEV})
+X=$($CGPT show $MTD -b -i $ROOTFS_NUM ${DEV})
+Y=$($CGPT show $MTD -s -i $ROOTFS_NUM ${DEV})
 [ "$X $Y" = "$ROOTFS_START $ROOTFS_SIZE" ] || error
 
-X=$($CGPT show -b -i $ESP_NUM ${DEV})
-Y=$($CGPT show -s -i $ESP_NUM ${DEV})
+X=$($CGPT show $MTD -b -i $ESP_NUM ${DEV})
+Y=$($CGPT show $MTD -s -i $ESP_NUM ${DEV})
 [ "$X $Y" = "$ESP_START $ESP_SIZE" ] || error
 
-X=$($CGPT show -b -i $FUTURE_NUM ${DEV})
-Y=$($CGPT show -s -i $FUTURE_NUM ${DEV})
+X=$($CGPT show $MTD -b -i $FUTURE_NUM ${DEV})
+Y=$($CGPT show $MTD -s -i $FUTURE_NUM ${DEV})
 [ "$X $Y" = "$FUTURE_START $FUTURE_SIZE" ] || error
 
-X=$($CGPT show -b -i $RANDOM_NUM ${DEV})
-Y=$($CGPT show -s -i $RANDOM_NUM ${DEV})
+X=$($CGPT show $MTD -b -i $RANDOM_NUM ${DEV})
+Y=$($CGPT show $MTD -s -i $RANDOM_NUM ${DEV})
 [ "$X $Y" = "$RANDOM_START $RANDOM_SIZE" ] || error
 
 
 echo "Change the beginning..."
 DATA_START=$((DATA_START + 10))
-$CGPT add -i 1 -b ${DATA_START} ${DEV} || error
-X=$($CGPT show -b -i 1 ${DEV})
+$CGPT add $MTD -i 1 -b ${DATA_START} ${DEV} || error
+X=$($CGPT show $MTD -b -i 1 ${DEV})
 [ "$X" = "$DATA_START" ] || error
 
 echo "Change the size..."
 DATA_SIZE=$((DATA_SIZE + 10))
-$CGPT add -i 1 -s ${DATA_SIZE} ${DEV} || error
-X=$($CGPT show -s -i 1 ${DEV})
+$CGPT add $MTD -i 1 -s ${DATA_SIZE} ${DEV} || error
+X=$($CGPT show $MTD -s -i 1 ${DEV})
 [ "$X" = "$DATA_SIZE" ] || error
 
 echo "Change the type..."
-$CGPT add -i 1 -t reserved ${DEV} || error
-X=$($CGPT show -t -i 1 ${DEV} | tr 'A-Z' 'a-z')
+$CGPT add $MTD -i 1 -t reserved ${DEV} || error
+X=$($CGPT show $MTD -t -i 1 ${DEV} | tr 'A-Z' 'a-z')
 [ "$X" = "$FUTURE_GUID" ] || error
 # arbitrary value
-if [ -z "$MTD" ]; then
-  $CGPT add -i 1 -t 610a563a-a55c-4ae0-ab07-86e5bb9db67f ${DEV} || error
-  X=$($CGPT show -t -i 1 ${DEV})
-  [ "$X" = "610A563A-A55C-4AE0-AB07-86E5BB9DB67F" ] || error
-fi
-$CGPT add -i 1 -t data ${DEV} || error
-X=$($CGPT show -t -i 1 ${DEV} | tr 'A-Z' 'a-z')
+$CGPT add $MTD -i 1 -t 610a563a-a55c-4ae0-ab07-86e5bb9db67f ${DEV} || error
+X=$($CGPT show $MTD -t -i 1 ${DEV})
+[ "$X" = "610A563A-A55C-4AE0-AB07-86E5BB9DB67F" ] || error
+
+$CGPT add $MTD -i 1 -t data ${DEV} || error
+X=$($CGPT show $MTD -t -i 1 ${DEV} | tr 'A-Z' 'a-z')
 [ "$X" = "$DATA_GUID" ] || error
 
 
-if [ -z "$MTD" ]; then
-  echo "Set the boot partition.."
-  $CGPT boot -i ${KERN_NUM} ${DEV} >/dev/null
+echo "Set the boot partition.."
+$CGPT boot $MTD -i ${KERN_NUM} ${DEV} >/dev/null
 
-  echo "Check the PMBR's idea of the boot partition..."
-  X=$($CGPT boot ${DEV})
-  Y=$($CGPT show -u -i $KERN_NUM $DEV)
-  [ "$X" = "$Y" ] || error
-fi
+echo "Check the PMBR's idea of the boot partition..."
+X=$($CGPT boot $MTD ${DEV})
+Y=$($CGPT show $MTD -u -i $KERN_NUM $DEV)
+[ "$X" = "$Y" ] || error
 
 echo "Test the cgpt prioritize command..."
 
@@ -157,18 +159,18 @@ echo "Test the cgpt prioritize command..."
 # Output: ${DEV} has kernel partitions with the given priorities
 make_pri() {
   local idx=0
-  $CGPT create ${DEV}
+  $CGPT create $MTD ${DEV}
   for pri in "$@"; do
     idx=$((idx+1))
-    $CGPT add -t kernel -l "kern$idx" -b $((100 + 2 * $idx)) -s 1 -P $pri ${DEV}
+    $CGPT add $MTD -t kernel -l "kern$idx" -b $((100 + 2 * $idx)) -s 1 -P $pri ${DEV}
   done
 }
 
 # Output: returns string containing priorities of all kernels
 get_pri() {
   echo $(
-  for idx in $($CGPT find -t kernel ${DEV} | sed -e s@${DEV}@@); do
-    $CGPT show -i $idx -P ${DEV}
+  for idx in $($CGPT find $MTD -t kernel ${DEV} | sed -e s@${DEV}@@); do
+    $CGPT show $MTD -i $idx -P ${DEV}
   done
   )
 }
@@ -183,117 +185,126 @@ assert_pri() {
 }
 
 # no kernels at all. This should do nothing.
-$CGPT create ${DEV}
-$CGPT add -t rootfs -b 100 -s 1 ${DEV}
-$CGPT prioritize ${DEV}
+$CGPT create $MTD ${DEV}
+$CGPT add $MTD -t rootfs -b 100 -s 1 ${DEV}
+$CGPT prioritize $MTD ${DEV}
 assert_pri ""
 
 # common install/upgrade sequence
 make_pri   2 0 0
-$CGPT prioritize -i 1 ${DEV}
+$CGPT prioritize $MTD -i 1 ${DEV}
 assert_pri 1 0 0
-$CGPT prioritize -i 2 ${DEV}
+$CGPT prioritize $MTD -i 2 ${DEV}
 assert_pri 1 2 0
-$CGPT prioritize -i 1 ${DEV}
+$CGPT prioritize $MTD -i 1 ${DEV}
 assert_pri 2 1 0
-$CGPT prioritize -i 2 ${DEV}
+$CGPT prioritize $MTD -i 2 ${DEV}
 assert_pri 1 2 0
-
 # lots of kernels, all same starting priority, should go to priority 1
 make_pri   8 8 8 8 8 8 8 8 8 8 8 0 0 8
-$CGPT prioritize ${DEV}
+$CGPT prioritize $MTD ${DEV}
 assert_pri 1 1 1 1 1 1 1 1 1 1 1 0 0 1
 
 # now raise them all up again
-$CGPT prioritize -P 4 ${DEV}
+$CGPT prioritize $MTD -P 4 ${DEV}
 assert_pri 4 4 4 4 4 4 4 4 4 4 4 0 0 4
 
 # set one of them higher, should leave the rest alone
-$CGPT prioritize -P 5 -i 3 ${DEV}
+$CGPT prioritize $MTD -P 5 -i 3 ${DEV}
 assert_pri 4 4 5 4 4 4 4 4 4 4 4 0 0 4
 
 # set one of them lower, should bring the rest down
-$CGPT prioritize -P 3 -i 4 ${DEV}
+$CGPT prioritize $MTD -P 3 -i 4 ${DEV}
 assert_pri 1 1 2 3 1 1 1 1 1 1 1 0 0 1
 
 # raise a group by including the friends of one partition
-$CGPT prioritize -P 6 -i 1 -f ${DEV}
+$CGPT prioritize $MTD -P 6 -i 1 -f ${DEV}
 assert_pri 6 6 4 5 6 6 6 6 6 6 6 0 0 6
 
 # resurrect one, should not affect the others
 make_pri   0 0 0 0 0 0 0 0 0 0 0 0 0 0
-$CGPT prioritize -i 2 ${DEV}
+$CGPT prioritize $MTD -i 2 ${DEV}
 assert_pri 0 1 0 0 0 0 0 0 0 0 0 0 0 0
 
 # resurrect one and all its friends
 make_pri   0 0 0 0 0 0 0 0 1 2 0 0 0 0
-$CGPT prioritize -P 5 -i 2 -f ${DEV}
+$CGPT prioritize $MTD -P 5 -i 2 -f ${DEV}
 assert_pri 5 5 5 5 5 5 5 5 3 4 5 5 5 5
 
 # no options should maintain the same order
-$CGPT prioritize ${DEV}
+$CGPT prioritize $MTD ${DEV}
 assert_pri 3 3 3 3 3 3 3 3 1 2 3 3 3 3
 
 # squish all the ranks
 make_pri   1 1 2 2 3 3 4 4 5 5 0 6 7 7
-$CGPT prioritize -P 6 ${DEV}
+$CGPT prioritize $MTD -P 6 ${DEV}
 assert_pri 1 1 1 1 2 2 3 3 4 4 0 5 6 6
 
 # squish the ranks by not leaving room
 make_pri   1 1 2 2 3 3 4 4 5 5 0 6 7 7
-$CGPT prioritize -P 7 -i 3 ${DEV}
+$CGPT prioritize $MTD -P 7 -i 3 ${DEV}
 assert_pri 1 1 7 1 2 2 3 3 4 4 0 5 6 6
 
 # squish the ranks while bringing the friends along
 make_pri   1 1 2 2 3 3 4 4 5 5 0 6 7 7
-$CGPT prioritize -P 6 -i 3 -f ${DEV}
+$CGPT prioritize $MTD -P 6 -i 3 -f ${DEV}
 assert_pri 1 1 6 6 1 1 2 2 3 3 0 4 5 5
 
 # squish them pretty hard
 make_pri   1 1 2 2 3 3 4 4 5 5 0 6 7 7
-$CGPT prioritize -P 2 ${DEV}
+$CGPT prioritize $MTD -P 2 ${DEV}
 assert_pri 1 1 1 1 1 1 1 1 1 1 0 1 2 2
 
 # squish them really really hard (nobody gets reduced to zero, though)
 make_pri   1 1 2 2 3 3 4 4 5 5 0 6 7 7
-$CGPT prioritize -P 1 -i 3 ${DEV}
+$CGPT prioritize $MTD -P 1 -i 3 ${DEV}
 assert_pri 1 1 1 1 1 1 1 1 1 1 0 1 1 1
 
-if [ -z "$MTD" ]; then # MTD doesn't support this many partitions
-  # squish if we try to go too high
-  make_pri   15 15 14 14 13 13 12 12 11 11 10 10 9 9 8 8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 0
-  $CGPT prioritize -i 3 ${DEV}
-  assert_pri 14 14 15 13 12 12 11 11 10 10  9  9 8 8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 1 1 0
-  $CGPT prioritize -i 5 ${DEV}
-  assert_pri 13 13 14 12 15 11 10 10  9  9  8  8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 1 1 1 1 0
-  # but if I bring friends I don't have to squish
-  $CGPT prioritize -i 1 -f ${DEV}
-  assert_pri 15 15 13 12 14 11 10 10  9  9  8  8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 1 1 1 1 0
-fi
-
+make_pri   15 15 14 14 13 13 12 12 11 11 10 10 9 9 8 8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 0
+$CGPT prioritize $MTD -i 3 ${DEV}
+assert_pri 14 14 15 13 12 12 11 11 10 10  9  9 8 8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 1 1 0
+$CGPT prioritize $MTD -i 5 ${DEV}
+assert_pri 13 13 14 12 15 11 10 10  9  9  8  8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 1 1 1 1 0
+# but if I bring friends I don't have to squish
+$CGPT prioritize $MTD -i 1 -f ${DEV}
+assert_pri 15 15 13 12 14 11 10 10  9  9  8  8 7 7 6 6 5 5 4 4 3 3 2 2 1 1 1 1 1 1 0
 
 # Now make sure that we don't need write access if we're just looking.
 echo "Test read vs read-write access..."
 chmod 0444 ${DEV}
 
 # These should fail
-$CGPT create -z ${DEV} 2>/dev/null && error
-$CGPT add -i 2 -P 3 ${DEV} 2>/dev/null && error
-$CGPT repair ${DEV} 2>/dev/null && error
-$CGPT prioritize -i 3 ${DEV} 2>/dev/null && error
+$CGPT create $MTD -z ${DEV} 2>/dev/null && error
+$CGPT add $MTD -i 2 -P 3 ${DEV} 2>/dev/null && error
+$CGPT repair $MTD ${DEV} 2>/dev/null && error
+$CGPT prioritize $MTD -i 3 ${DEV} 2>/dev/null && error
 
 # Most 'boot' usage should fail too.
-$CGPT boot -p ${DEV} 2>/dev/null && error
+$CGPT boot $MTD -p ${DEV} 2>/dev/null && error
 dd if=/dev/zero of=fake_mbr.bin bs=100 count=1 2>/dev/null
-$CGPT boot -b fake_mbr.bin ${DEV} 2>/dev/null && error
-$CGPT boot -i 2 ${DEV} 2>/dev/null && error
-
-# These should pass
-if [ -z "$MTD" ]; then
-  $CGPT boot ${DEV} >/dev/null
-fi
-$CGPT show ${DEV} >/dev/null
-$CGPT find -t kernel ${DEV} >/dev/null
+$CGPT boot $MTD -b fake_mbr.bin ${DEV} 2>/dev/null && error
+$CGPT boot $MTD -i 2 ${DEV} 2>/dev/null && error
+
+$CGPT boot $MTD ${DEV} >/dev/null
+$CGPT show $MTD ${DEV} >/dev/null
+$CGPT find $MTD -t kernel ${DEV} >/dev/null
+
+# Enable write access again to test boundary in off device storage
+chmod 600 ${DEV}
+# Drive size is not multiple of 512
+assert_fail $CGPT create -D 511 ${DEV}
+assert_fail $CGPT create -D 513 ${DEV}
+MTD="-D 1024"
+# Create a GPT table for a device of 1024 bytes (2 sectors)
+$CGPT create $MTD ${DEV}
+# This fails because header verification is off due to different drive size
+assert_fail $CGPT show ${DEV}
+# But this passes because we pass in correct drive size
+$CGPT show $MTD ${DEV}
+# This fails because beginning sector is over the size of the device
+assert_fail $CGPT add $MTD -b 2 -s 1 -t data ${DEV}
+# This fails because partition size is over the size of the device
+assert_fail $CGPT add $MTD -b 0 -s 3 -t data ${DEV}
 
 echo "Done."