1 files changed, 55 insertions, 27 deletions

diff --git a/cgpt/cgpt_create.c b/cgpt/cgpt_create.c
index 8c899d19..b567432f 100644
--- a/cgpt/cgpt_create.c
+++ b/cgpt/cgpt_create.c
@@ -9,42 +9,33 @@
 #include "cgptlib_internal.h"
 #include "cgpt_params.h"
 
-int cgpt_create(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 cgpt_create(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 cgpt_create(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