MTD modification commands: add & prioritize

Additionally renames the new functions in add to be CamelCase style

BUG=chromium:221745
TEST=MTD version of run_cgpt_tests.sh passes
BRANCH=cros/embedded
CQ-DEPEND=CL:47172

Change-Id: Ic173f99d7ca4af025403ab11f37061c33c9d59ea
Reviewed-on: https://gerrit.chromium.org/gerrit/47173
Tested-by: Albert Chaulk <achaulk@chromium.org>
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Commit-Queue: Albert Chaulk <achaulk@chromium.org>

2 files changed, 108 insertions, 49 deletions

diff --git a/cgpt/cgpt_add.c b/cgpt/cgpt_add.c
index 3b44e9aa..58bdce1c 100644
--- a/cgpt/cgpt_add.c
+++ b/cgpt/cgpt_add.c
@@ -64,9 +64,9 @@ static const char* DumpCgptAddParams(const CgptAddParams *params) {
 }
 
 // This is the implementation-specific helper function.
-static int gpt_set_entry_attributes(struct drive *drive,
-                                    uint32_t index,
-                                    CgptAddParams *params) {
+static int GptSetEntryAttributes(struct drive *drive,
+                                 uint32_t index,
+                                 CgptAddParams *params) {
   GptEntry *entry;
 
   entry = GetEntry(&drive->gpt, PRIMARY, index);
@@ -95,11 +95,27 @@ static int gpt_set_entry_attributes(struct drive *drive,
   return 0;
 }
 
+static int MtdSetEntryAttributes(struct drive *drive,
+                                 uint32_t index,
+                                 CgptAddParams *params) {
+  MtdDiskPartition *entry;
+
+  entry = MtdGetEntry(&drive->mtd, PRIMARY, index);
+  if (params->set_begin)
+    entry->starting_lba = params->begin;
+  if (params->set_size)
+    entry->ending_lba = entry->starting_lba + params->size - 1;
+  if (params->set_type)
+    MtdSetEntryType(entry, LookupMtdTypeForGuid(&params->type_guid));
+
+  return 0;
+}
+
 // This is an internal helper function which assumes no NULL args are passed.
 // It sets the given attribute values for a single entry at the given index.
-static int set_entry_attributes(struct drive *drive,
-                                uint32_t index,
-                                CgptAddParams *params) {
+static int SetEntryAttributes(struct drive *drive,
+                              uint32_t index,
+                              CgptAddParams *params) {
   if (params->set_raw) {
     SetRaw(drive, PRIMARY, index, params->raw_value);
   } else {
@@ -126,25 +142,32 @@ static int set_entry_attributes(struct drive *drive,
   return 0;
 }
 
-static int cgpt_check_add_validity(struct drive *drive) {
-  int gpt_retval;
-  if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive->gpt))) {
-    Error("GptSanityCheck() returned %d: %s\n",
-          gpt_retval, GptError(gpt_retval));
-    return -1;
-  }
+static int CgptCheckAddValidity(struct drive *drive) {
+  if (drive->is_mtd) {
+    if (drive->mtd.primary.crc32 != MtdHeaderCrc(&drive->mtd.primary)) {
+      Error("MTD header CRC is invalid\n");
+      return -1;
+    }
+  } else {
+    int gpt_retval;
+    if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive->gpt))) {
+      Error("GptSanityCheck() returned %d: %s\n",
+            gpt_retval, GptError(gpt_retval));
+      return -1;
+    }
 
-  if (((drive->gpt.valid_headers & MASK_BOTH) != MASK_BOTH) ||
-      ((drive->gpt.valid_entries & MASK_BOTH) != MASK_BOTH)) {
-    Error("one of the GPT header/entries is invalid.\n"
-          "please run 'cgpt repair' before adding anything.\n");
-    return -1;
+    if (((drive->gpt.valid_headers & MASK_BOTH) != MASK_BOTH) ||
+        ((drive->gpt.valid_entries & MASK_BOTH) != MASK_BOTH)) {
+      Error("one of the GPT header/entries is invalid.\n"
+            "please run 'cgpt repair' before adding anything.\n");
+      return -1;
+    }
   }
   return 0;
 }
 
-static int cgpt_get_unused_partition(struct drive *drive, uint32_t *index,
-                                     CgptAddParams *params) {
+static int CgptGetUnusedPartition(struct drive *drive, uint32_t *index,
+                                  CgptAddParams *params) {
   uint32_t i;
   uint32_t max_part = GetNumberOfEntries(drive);
   if (params->partition) {
@@ -168,49 +191,78 @@ static int cgpt_get_unused_partition(struct drive *drive, uint32_t *index,
   }
 }
 
-int cgpt_add(CgptAddParams *params) {
-  struct drive drive;
-
+int GptAdd(struct drive *drive, CgptAddParams *params, uint32_t index) {
   GptEntry *entry, backup;
-  uint32_t index;
   int rv;
 
-  if (params == NULL)
-    return CGPT_FAILED;
-
-  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
-    return CGPT_FAILED;
-
-  if (cgpt_check_add_validity(&drive)) {
-    goto bad;
-  }
-
-  if (cgpt_get_unused_partition(&drive, &index, params)) {
-    goto bad;
-  }
-
-  entry = GetEntry(&drive.gpt, PRIMARY, index);
+  entry = GetEntry(&drive->gpt, PRIMARY, index);
   memcpy(&backup, entry, sizeof(backup));
 
-  if (set_entry_attributes(&drive, index, params) ||
-      gpt_set_entry_attributes(&drive, index, params)) {
+  if (SetEntryAttributes(drive, index, params) ||
+      GptSetEntryAttributes(drive, index, params)) {
     memcpy(entry, &backup, sizeof(*entry));
-    goto bad;
+    return -1;
   }
 
-  UpdateAllEntries(&drive);
+  UpdateAllEntries(drive);
 
-  rv = CheckEntries((GptEntry*)drive.gpt.primary_entries,
-                    (GptHeader*)drive.gpt.primary_header);
+  rv = CheckEntries((GptEntry*)drive->gpt.primary_entries,
+                    (GptHeader*)drive->gpt.primary_header);
 
   if (0 != rv) {
     // If the modified entry is illegal, recover it and return error.
     memcpy(entry, &backup, sizeof(*entry));
     Error("%s\n", GptErrorText(rv));
     Error(DumpCgptAddParams(params));
+    return -1;
+  }
+
+  return 0;
+}
+
+int MtdAdd(struct drive *drive, CgptAddParams *params, uint32_t index) {
+  MtdDiskPartition *entry, backup;
+  entry = MtdGetEntry(&drive->mtd, PRIMARY, index);
+  memcpy(&backup, entry, sizeof(backup));
+
+  if (SetEntryAttributes(drive, index, params) ||
+      MtdSetEntryAttributes(drive, index, params)) {
+    memcpy(entry, &backup, sizeof(*entry));
+    return -1;
+  }
+
+  UpdateAllEntries(drive);
+
+  return 0;
+}
+
+
+int cgpt_add(CgptAddParams *params) {
+  struct drive drive;
+  uint32_t index;
+
+  if (params == NULL)
+    return CGPT_FAILED;
+
+  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
+    return CGPT_FAILED;
+
+  if (CgptCheckAddValidity(&drive)) {
     goto bad;
   }
 
+  if (CgptGetUnusedPartition(&drive, &index, params)) {
+    goto bad;
+  }
+
+  if (drive.is_mtd) {
+    if (MtdAdd(&drive, params, index))
+      goto bad;
+  } else {
+    if (GptAdd(&drive, params, index))
+      goto bad;
+  }
+
   // Write it all out.
   return DriveClose(&drive, 1);
 
diff --git a/cgpt/cgpt_prioritize.c b/cgpt/cgpt_prioritize.c
index b62c45b3..eb04ae9c 100644
--- a/cgpt/cgpt_prioritize.c
+++ b/cgpt/cgpt_prioritize.c
@@ -110,10 +110,17 @@ int cgpt_prioritize(CgptPrioritizeParams *params) {
   if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR))
     return CGPT_FAILED;
 
-  if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive.gpt))) {
-    Error("GptSanityCheck() returned %d: %s\n",
-          gpt_retval, GptError(gpt_retval));
-    return CGPT_FAILED;
+  if (drive.is_mtd) {
+    if (drive.mtd.primary.crc32 != MtdHeaderCrc(&drive.mtd.primary)) {
+      Error("MTD header crc failure\n");
+      return CGPT_FAILED;
+    }
+  } else {
+    if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive.gpt))) {
+      Error("GptSanityCheck() returned %d: %s\n",
+            gpt_retval, GptError(gpt_retval));
+      return CGPT_FAILED;
+    }
   }
 
   max_part = GetNumberOfEntries(&drive);