1 files changed, 8 insertions, 0 deletions
diff --git a/cgpt/cgpt_create.c b/cgpt/cgpt_create.c
index e7cbadf8..90b746c6 100644
--- a/cgpt/cgpt_create.c
+++ b/cgpt/cgpt_create.c
@@ -64,6 +64,14 @@ int MtdCreate(struct drive *drive, CgptCreateParams *params) {
     h->last_offset = (drive->mtd.drive_sectors * drive->mtd.sector_bytes) - 1;
     h->crc32 = MtdHeaderCrc(h);
   }
+  if (params->size) {
+    h->last_offset = params->size - 1;
+    drive->size = params->size;
+    drive->mtd.drive_sectors = drive->size / drive->mtd.sector_bytes;
+  } else if (!drive->mtd.drive_sectors) {
+    Error("MTD create with params->size == 0 && drive->mtd.drive_sectors == 0");
+    return -1;
+  }
 
   return 0;
 }