Bugfixes & MTD create command

Fix some bugs in the cgpt implementation of the flash I/O functions & load
logic, it was validating too much at load time.

Implement the create command for MTD

BUG=chromium:221745
TEST=MTD version of run_cgpt_tests.sh passes
BRANCH=none

Original-Change-Id: I2f52637d82962f4d805aa827c5c37685f10e76ea
Reviewed-on: https://gerrit.chromium.org/gerrit/47172
Tested-by: Albert Chaulk <achaulk@chromium.org>
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Reviewed-by: Randall Spangler <rspangler@chromium.org>
Commit-Queue: Albert Chaulk <achaulk@chromium.org>
(cherry picked from commit 931544744ba410dad267064b87d504b0b4c24772)

Change-Id: If9364155fb2c030645adc6ee6f3fbe5373bcc153
Reviewed-on: https://gerrit.chromium.org/gerrit/49793
Commit-Queue: Albert Chaulk <achaulk@chromium.org>
Reviewed-by: Albert Chaulk <achaulk@chromium.org>
Tested-by: Albert Chaulk <achaulk@chromium.org>

4 files changed, 79 insertions, 54 deletions

diff --git a/cgpt/cgpt.h b/cgpt/cgpt.h
index f6580868..7595c0e7 100644
--- a/cgpt/cgpt.h
+++ b/cgpt/cgpt.h
@@ -145,6 +145,7 @@ int ResolveType(const Guid *type, char *buf);
 int SupportedType(const char *name, Guid *type);
 void PrintTypes(void);
 void EntryDetails(GptEntry *entry, uint32_t index, int raw);
+void MtdEntryDetails(MtdDiskPartition *entry, uint32_t index, int raw);
 
 uint32_t GetNumberOfEntries(const struct drive *drive);
 GptEntry *GetEntry(GptData *gpt, int secondary, uint32_t entry_index);
@@ -173,6 +174,7 @@ int IsSynonymous(const GptHeader* a, const GptHeader* b);
 int IsUnused(struct drive *drive, int secondary, uint32_t index);
 int IsKernel(struct drive *drive, int secondary, uint32_t index);
 int LookupMtdTypeForGuid(const Guid *type);
+const Guid *LookupGuidForMtdType(int type);
 
 // For usage and error messages.
 extern const char* progname;
diff --git a/cgpt/cgpt_common.c b/cgpt/cgpt_common.c
index dd8e510d..7b206d68 100644
--- a/cgpt/cgpt_common.c
+++ b/cgpt/cgpt_common.c
@@ -193,11 +193,11 @@ int FlashSet(const char *key, const uint8_t *data, uint32_t bufsz) {
 
 int MtdLoad(struct drive *drive, int sector_bytes) {
   int ret;
-  uint32_t old_crc, new_crc;
   uint32_t sz;
   MtdData *mtd = &drive->mtd;
 
   mtd->sector_bytes = sector_bytes;
+  mtd->drive_sectors = drive->size / mtd->sector_bytes;
 
   ret = flash_ts_init(mtd->fts_block_offset,
                       mtd->fts_block_size,
@@ -216,21 +216,7 @@ int MtdLoad(struct drive *drive, int sector_bytes) {
 
   /* Read less than expected */
   if (sz < MTD_DRIVE_V1_SIZE)
-    return -1;
-
-  if (memcmp(mtd->primary.signature, MTD_DRIVE_SIGNATURE,
-             sizeof(mtd->primary.signature))) {
-    return -1;
-  }
-
-  old_crc = mtd->primary.crc32;
-  mtd->primary.crc32 = 0;
-  new_crc = Crc32(&mtd->primary, MTD_DRIVE_V1_SIZE);
-  mtd->primary.crc32 = old_crc;
-
-  if (old_crc != new_crc) {
-    return -1;
-  }
+    memset(&mtd->primary, 0, sizeof(mtd->primary));
 
   mtd->current_kernel = -1;
   mtd->current_priority = 0;
@@ -241,6 +227,9 @@ int MtdLoad(struct drive *drive, int sector_bytes) {
 int MtdSave(struct drive *drive) {
   MtdData *mtd = &drive->mtd;
 
+  if (!mtd->modified)
+    return 0;
+
   mtd->primary.crc32 = 0;
   mtd->primary.crc32 = Crc32(&mtd->primary, MTD_DRIVE_V1_SIZE);
 
@@ -737,6 +726,16 @@ int LookupMtdTypeForGuid(const Guid *type) {
   return MTD_PARTITION_TYPE_OTHER;
 }
 
+const Guid *LookupGuidForMtdType(int type) {
+  int i;
+  for (i = 0; i < ARRAY_COUNT(supported_types); ++i) {
+    if (supported_types[i].mtd_type == type) {
+      return supported_types[i].type;
+    }
+  }
+  return NULL;
+}
+
 /* Resolves human-readable GPT type.
  * Returns CGPT_OK if found.
  * Returns CGPT_FAILED if no known type found. */
diff --git a/cgpt/cgpt_create.c b/cgpt/cgpt_create.c
index b0256c51..0b116b2b 100644
--- a/cgpt/cgpt_create.c
+++ b/cgpt/cgpt_create.c
@@ -9,42 +9,33 @@
 #include "cgptlib_internal.h"
 #include "vboot_host.h"
 
-int CgptCreate(CgptCreateParams *params) {
-  struct drive drive;
-
-  if (params == NULL)
-    return CGPT_FAILED;
-
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
-    return CGPT_FAILED;
-
+int GptCreate(struct drive *drive, CgptCreateParams *params) {
   // Erase the data
-  memset(drive.gpt.primary_header, 0,
-         drive.gpt.sector_bytes * GPT_HEADER_SECTOR);
-  memset(drive.gpt.secondary_header, 0,
-         drive.gpt.sector_bytes * GPT_HEADER_SECTOR);
-  memset(drive.gpt.primary_entries, 0,
-         drive.gpt.sector_bytes * GPT_ENTRIES_SECTORS);
-  memset(drive.gpt.secondary_entries, 0,
-         drive.gpt.sector_bytes * GPT_ENTRIES_SECTORS);
-
-  drive.gpt.modified |= (GPT_MODIFIED_HEADER1 | GPT_MODIFIED_ENTRIES1 |
+  memset(drive->gpt.primary_header, 0,
+         drive->gpt.sector_bytes * GPT_HEADER_SECTOR);
+  memset(drive->gpt.secondary_header, 0,
+         drive->gpt.sector_bytes * GPT_HEADER_SECTOR);
+  memset(drive->gpt.primary_entries, 0,
+         drive->gpt.sector_bytes * GPT_ENTRIES_SECTORS);
+  memset(drive->gpt.secondary_entries, 0,
+         drive->gpt.sector_bytes * GPT_ENTRIES_SECTORS);
+
+  drive->gpt.modified |= (GPT_MODIFIED_HEADER1 | GPT_MODIFIED_ENTRIES1 |
                          GPT_MODIFIED_HEADER2 | GPT_MODIFIED_ENTRIES2);
 
   // Initialize a blank set
-  if (!params->zap)
-  {
-    GptHeader *h = (GptHeader *)drive.gpt.primary_header;
+  if (!params->zap) {
+    GptHeader *h = (GptHeader *)drive->gpt.primary_header;
     memcpy(h->signature, GPT_HEADER_SIGNATURE, GPT_HEADER_SIGNATURE_SIZE);
     h->revision = GPT_HEADER_REVISION;
     h->size = sizeof(GptHeader);
     h->my_lba = 1;
-    h->alternate_lba = drive.gpt.drive_sectors - 1;
+    h->alternate_lba = drive->gpt.drive_sectors - 1;
     h->first_usable_lba = 1 + 1 + GPT_ENTRIES_SECTORS;
-    h->last_usable_lba = drive.gpt.drive_sectors - 1 - GPT_ENTRIES_SECTORS - 1;
+    h->last_usable_lba = drive->gpt.drive_sectors - 1 - GPT_ENTRIES_SECTORS - 1;
     if (!uuid_generator) {
       Error("Unable to generate new GUID. uuid_generator not set.\n");
-      goto bad;
+      return -1;
     }
     (*uuid_generator)((uint8_t *)&h->disk_uuid);
     h->entries_lba = 2;
@@ -52,9 +43,46 @@ int CgptCreate(CgptCreateParams *params) {
     h->size_of_entry = sizeof(GptEntry);
 
     // Copy to secondary
-    RepairHeader(&drive.gpt, MASK_PRIMARY);
+    RepairHeader(&drive->gpt, MASK_PRIMARY);
+
+    UpdateCrc(&drive->gpt);
+  }
+
+  return 0;
+}
+
+int MtdCreate(struct drive *drive, CgptCreateParams *params) {
+  MtdDiskLayout *h = &drive->mtd.primary;
+  memset(h, 0, sizeof(*h));
+  drive->mtd.modified = 1;
 
-    UpdateCrc(&drive.gpt);
+  if (!params->zap) {
+    // Prep basic parameters
+    memcpy(h->signature, MTD_DRIVE_SIGNATURE, sizeof(h->signature));
+    h->size = sizeof(*h);
+    h->first_lba = 0;
+    h->last_lba = drive->mtd.drive_sectors - 1;
+    h->crc32 = MtdHeaderCrc(h);
+  }
+
+  return 0;
+}
+
+int CgptCreate(CgptCreateParams *params) {
+  struct drive drive;
+
+  if (params == NULL)
+    return CGPT_FAILED;
+
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+    return CGPT_FAILED;
+
+  if (drive.is_mtd) {
+    if (MtdCreate(&drive, params))
+      goto bad;
+  } else {
+    if (GptCreate(&drive, params))
+      goto bad;
   }
 
   // Write it all out
diff --git a/cgpt/flash_ts_drv.c b/cgpt/flash_ts_drv.c
index 8c918866..01a68951 100644
--- a/cgpt/flash_ts_drv.c
+++ b/cgpt/flash_ts_drv.c
@@ -20,9 +20,6 @@ inline int page_to_sector(const nand_geom *nand, int page) {
 int nand_read_page(const nand_geom *nand, int page, void *buf, int size) {
   uint8_t *page_buff;
 
-  if (size > nand->szofpg) {
-    return -1;
-  }
   if (Load((struct drive *)nand->user, &page_buff,
            page_to_sector(nand, page), nand->szofsector,
            (size + nand->szofsector - 1) / nand->szofsector)) {
@@ -37,22 +34,21 @@ int nand_read_page(const nand_geom *nand, int page, void *buf, int size) {
 }
 
 int nand_write_page(const nand_geom *nand,
-                    int page, const void *buf, int size) {
+                    int page, const void *buf, int buf_size) {
   void *page_buff;
   int ret;
+  int sectors = (buf_size + nand->szofsector - 1) / nand->szofsector;
+  int size = nand->szofsector * sectors;
 
-  if (size > nand->szofpg) {
-    return -1;
-  }
-  page_buff  = malloc(nand->szofpg);
+  page_buff  = malloc(size);
   if (!page_buff)
     return -1;
 
-  memset(page_buff, 0xff, nand->szofpg);
-  memcpy(page_buff, buf, size < nand->szofpg ? size : nand->szofpg);
+  memset(page_buff, 0xff, size);
+  memcpy(page_buff, buf, buf_size);
 
   ret = Save((struct drive *)nand->user, page_buff, page_to_sector(nand, page),
-             nand->szofsector, nand->szofpg / nand->szofsector);
+             nand->szofsector, sectors);
   free(page_buff);
   return ret;
 }