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>

20 files changed, 190 insertions, 448 deletions

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_