vtoycli/partresize.c

   1 /******************************************************************************
   2  * partresize.c  ---- ventoy part resize util
   3  *
   4  * Copyright (c) 2021, longpanda <admin@ventoy.net>
   5  *
   6  * This program is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU General Public License as
   8  * published by the Free Software Foundation; either version 3 of the
   9  * License, or (at your option) any later version.
  10  *
  11  * This program is distributed in the hope that it will be useful, but
  12  * WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, see <http://www.gnu.org/licenses/>.
  18  *
  19  */
  20 #include <stdio.h>
  21 #include <stdio.h>
  22 #include <stdlib.h>
  23 #include <string.h>
  24 #include <errno.h>
  25 #include <unistd.h>
  26 #include <fcntl.h>
  27 #include <sys/types.h>
  28 #include <sys/mman.h>
  29 #include <sys/ioctl.h>
  30 #include <sys/stat.h>
  31 #include <sys/types.h>
  32 #include <dirent.h>
  33 #include <fat_filelib.h>
  34 #include "vtoycli.h"
  35
  36 static int g_disk_fd = 0;
  37 static UINT64 g_disk_offset = 0;
  38 static GUID g_ZeroGuid = {0};
  39 static GUID g_WindowsDataPartGuid = { 0xebd0a0a2, 0xb9e5, 0x4433, { 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7 } };
  40
  41 static int vtoy_disk_read(uint32 sector, uint8 *buffer, uint32 sector_count)
  42 {
  43     UINT64 offset = sector * 512ULL;
  44
  45     lseek(g_disk_fd, g_disk_offset + offset, SEEK_SET);
  46     read(g_disk_fd, buffer, sector_count * 512);
  47
  48     return 1;
  49 }
  50
  51 static int vtoy_disk_write(uint32 sector, uint8 *buffer, uint32 sector_count)
  52 {
  53     UINT64 offset = sector * 512ULL;
  54
  55     lseek(g_disk_fd, g_disk_offset + offset, SEEK_SET);
  56     write(g_disk_fd, buffer, sector_count * 512);
  57
  58     return 1;
  59 }
  60
  61
  62 static int gpt_check(const char *disk)
  63 {
  64     int fd = -1;
  65     int rc = 1;
  66     VTOY_GPT_INFO *pGPT = NULL;
  67
  68     fd = open(disk, O_RDONLY);
  69     if (fd < 0)
  70     {
  71         printf("Failed to open %s\n", disk);
  72         goto out;
  73     }
  74
  75     pGPT = malloc(sizeof(VTOY_GPT_INFO));
  76     if (NULL == pGPT)
  77     {
  78         goto out;
  79     }
  80     memset(pGPT, 0, sizeof(VTOY_GPT_INFO));
  81
  82     read(fd, pGPT, sizeof(VTOY_GPT_INFO));
  83
  84     if (pGPT->MBR.PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
  85         {
  86         rc = 0;
  87         }
  88
  89 out:
  90     check_close(fd);
  91     check_free(pGPT);
  92     return rc;
  93 }
  94
  95 static int part_check(const char *disk)
  96 {
  97     int i;
  98     int fd = -1;
  99     int rc = 0;
 100     int Index = 0;
 101     int Count = 0;
 102     int PartStyle = 0;
 103     UINT64 Part1Start;
 104     UINT64 Part1End;
 105     UINT64 NextPartStart;
 106     UINT64 DiskSizeInBytes;
 107     VTOY_GPT_INFO *pGPT = NULL;
 108
 109     DiskSizeInBytes = get_disk_size_in_byte(disk);
 110     if (DiskSizeInBytes == 0)
 111     {
 112         printf("Failed to get disk size of %s\n", disk);
 113         goto out;
 114     }
 115
 116     fd = open(disk, O_RDONLY);
 117     if (fd < 0)
 118     {
 119         printf("Failed to open %s\n", disk);
 120         goto out;
 121     }
 122
 123     pGPT = malloc(sizeof(VTOY_GPT_INFO));
 124     if (NULL == pGPT)
 125     {
 126         goto out;
 127     }
 128     memset(pGPT, 0, sizeof(VTOY_GPT_INFO));
 129
 130     read(fd, pGPT, sizeof(VTOY_GPT_INFO));
 131
 132     if (pGPT->MBR.PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
 133         {
 134         PartStyle = 1;
 135         }
 136         else
 137         {
 138         PartStyle = 0;
 139         }
 140
 141     if (PartStyle == 0)
 142     {
 143                 PART_TABLE *PartTbl = pGPT->MBR.PartTbl;
 144
 145         for (Count = 0, i = 0; i < 4; i++)
 146         {
 147             if (PartTbl[i].SectorCount > 0)
 148             {
 149                                 printf("MBR Part%d SectorStart:%u SectorCount:%u\n", i + 1, PartTbl[i].StartSectorId, PartTbl[i].SectorCount);
 150                 Count++;
 151             }
 152         }
 153
 154                 //We must have a free partition table for VTOYEFI partition
 155                 if (Count >= 4)
 156                 {
 157                         printf("###[FAIL] 4 MBR partition tables are all used.\n");
 158                         goto out;
 159                 }
 160
 161                 if (PartTbl[0].SectorCount > 0)
 162                 {
 163                         Part1Start = PartTbl[0].StartSectorId;
 164                         Part1End = PartTbl[0].SectorCount + Part1Start;
 165                 }
 166                 else
 167                 {
 168                         printf("###[FAIL] MBR Partition 1 is invalid\n");
 169                         goto out;
 170                 }
 171
 172                 Index = -1;
 173                 NextPartStart = DiskSizeInBytes / 512ULL;
 174                 for (i = 1; i < 4; i++)
 175                 {
 176                         if (PartTbl[i].SectorCount > 0 && NextPartStart > PartTbl[i].StartSectorId)
 177                         {
 178                                 Index = i;
 179                                 NextPartStart = PartTbl[i].StartSectorId;
 180                         }
 181                 }
 182
 183         NextPartStart *= 512ULL;
 184                 printf("DiskSize:%llu NextPartStart:%llu(LBA:%llu) Index:%d\n",
 185             DiskSizeInBytes, NextPartStart, NextPartStart / 512ULL, Index);
 186     }
 187     else
 188     {
 189                 VTOY_GPT_PART_TBL *PartTbl = pGPT->PartTbl;
 190
 191         for (Count = 0, i = 0; i < 128; i++)
 192         {
 193             if (memcmp(&(PartTbl[i].PartGuid), &g_ZeroGuid, sizeof(GUID)))
 194             {
 195                                 printf("GPT Part%d StartLBA:%llu LastLBA:%llu\n", i + 1, PartTbl[i].StartLBA, PartTbl[i].LastLBA);
 196                 Count++;
 197             }
 198         }
 199
 200                 if (Count >= 128)
 201                 {
 202                         printf("###[FAIL] 128 GPT partition tables are all used.\n");
 203                         goto out;
 204                 }
 205
 206                 if (memcmp(&(PartTbl[0].PartGuid), &g_ZeroGuid, sizeof(GUID)))
 207                 {
 208                         Part1Start = PartTbl[0].StartLBA;
 209                         Part1End = PartTbl[0].LastLBA + 1;
 210                 }
 211                 else
 212                 {
 213                         printf("###[FAIL] GPT Partition 1 is invalid\n");
 214                         goto out;
 215                 }
 216
 217                 Index = -1;
 218                 NextPartStart = (pGPT->Head.PartAreaEndLBA + 1);
 219                 for (i = 1; i < 128; i++)
 220                 {
 221                         if (memcmp(&(PartTbl[i].PartGuid), &g_ZeroGuid, sizeof(GUID)) && NextPartStart > PartTbl[i].StartLBA)
 222                         {
 223                                 Index = i;
 224                                 NextPartStart = PartTbl[i].StartLBA;
 225                         }
 226                 }
 227
 228                 NextPartStart *= 512ULL;
 229                 printf("DiskSize:%llu NextPartStart:%llu(LBA:%llu) Index:%d\n",
 230             DiskSizeInBytes, NextPartStart, NextPartStart / 512ULL, Index);
 231     }
 232
 233         printf("Valid partition table (%s): Valid partition count:%d\n", (PartStyle == 0) ? "MBR" : "GPT", Count);
 234
 235         //Partition 1 MUST start at 1MB
 236         Part1Start *= 512ULL;
 237         Part1End *= 512ULL;
 238
 239         printf("Partition 1 start at: %llu %lluKB, end:%llu, NextPartStart:%llu\n",
 240                 Part1Start, Part1Start / 1024, Part1End, NextPartStart);
 241     if (Part1Start != SIZE_1MB)
 242     {
 243         printf("###[FAIL] Partition 1 is not start at 1MB\n");
 244         goto out;
 245     }
 246
 247
 248         //If we have free space after partition 1
 249         if (NextPartStart - Part1End >= VENTOY_EFI_PART_SIZE)
 250         {
 251                 printf("Free space after partition 1 (%llu) is enough for VTOYEFI part\n", NextPartStart - Part1End);
 252                 rc = 1;
 253         }
 254         else if (NextPartStart == Part1End)
 255         {
 256                 printf("There is no free space after partition 1\n");
 257         rc = 2;
 258         }
 259         else
 260         {
 261                 printf("The free space after partition 1 is not enough\n");
 262         rc = 2;
 263         }
 264
 265 out:
 266     check_close(fd);
 267     check_free(pGPT);
 268     return rc;
 269 }
 270
 271 static int secureboot_proc(char *disk, UINT64 part2start)
 272 {
 273         int rc = 0;
 274         int size;
 275     int fd = -1;
 276         char *filebuf = NULL;
 277         void *file = NULL;
 278
 279     fd = open(disk, O_RDWR);
 280     if (fd < 0)
 281     {
 282         printf("Failed to open %s\n", disk);
 283         return 1;
 284     }
 285
 286     g_disk_fd = fd;
 287     g_disk_offset = part2start * 512ULL;
 288
 289         fl_init();
 290
 291         if (0 == fl_attach_media(vtoy_disk_read, vtoy_disk_write))
 292         {
 293                 file = fl_fopen("/EFI/BOOT/grubx64_real.efi", "rb");
 294                 printf("Open ventoy efi file %p\n", file);
 295                 if (file)
 296                 {
 297                         fl_fseek(file, 0, SEEK_END);
 298                         size = (int)fl_ftell(file);
 299                         fl_fseek(file, 0, SEEK_SET);
 300
 301                         printf("ventoy x64 efi file size %d ...\n", size);
 302
 303                         filebuf = (char *)malloc(size);
 304                         if (filebuf)
 305                         {
 306                                 fl_fread(filebuf, 1, size, file);
 307                         }
 308
 309                         fl_fclose(file);
 310
 311                         fl_remove("/EFI/BOOT/BOOTX64.EFI");
 312                         fl_remove("/EFI/BOOT/grubx64.efi");
 313                         fl_remove("/EFI/BOOT/grubx64_real.efi");
 314                         fl_remove("/EFI/BOOT/MokManager.efi");
 315             fl_remove("/ENROLL_THIS_KEY_IN_MOKMANAGER.cer");
 316
 317                         file = fl_fopen("/EFI/BOOT/BOOTX64.EFI", "wb");
 318                         printf("Open bootx64 efi file %p\n", file);
 319                         if (file)
 320                         {
 321                                 if (filebuf)
 322                                 {
 323                                         fl_fwrite(filebuf, 1, size, file);
 324                                 }
 325
 326                                 fl_fflush(file);
 327                                 fl_fclose(file);
 328                         }
 329
 330                         if (filebuf)
 331                         {
 332                                 free(filebuf);
 333                         }
 334                 }
 335
 336         file = fl_fopen("/EFI/BOOT/grubia32_real.efi", "rb");
 337         printf("Open ventoy ia32 efi file %p\n", file);
 338         if (file)
 339         {
 340             fl_fseek(file, 0, SEEK_END);
 341             size = (int)fl_ftell(file);
 342             fl_fseek(file, 0, SEEK_SET);
 343
 344             printf("ventoy efi file size %d ...\n", size);
 345
 346             filebuf = (char *)malloc(size);
 347             if (filebuf)
 348             {
 349                 fl_fread(filebuf, 1, size, file);
 350             }
 351
 352             fl_fclose(file);
 353
 354             fl_remove("/EFI/BOOT/BOOTIA32.EFI");
 355             fl_remove("/EFI/BOOT/grubia32.efi");
 356             fl_remove("/EFI/BOOT/grubia32_real.efi");
 357             fl_remove("/EFI/BOOT/mmia32.efi");
 358
 359             file = fl_fopen("/EFI/BOOT/BOOTIA32.EFI", "wb");
 360             printf("Open bootia32 efi file %p\n", file);
 361             if (file)
 362             {
 363                 if (filebuf)
 364                 {
 365                     fl_fwrite(filebuf, 1, size, file);
 366                 }
 367
 368                 fl_fflush(file);
 369                 fl_fclose(file);
 370             }
 371
 372             if (filebuf)
 373             {
 374                 free(filebuf);
 375             }
 376         }
 377
 378         }
 379         else
 380         {
 381                 rc = 1;
 382         }
 383
 384         fl_shutdown();
 385     fsync(fd);
 386
 387         return rc;
 388 }
 389
 390 static int VentoyFillMBRLocation(UINT64 DiskSizeInBytes, UINT32 StartSectorId, UINT32 SectorCount, PART_TABLE *Table)
 391 {
 392     UINT8 Head;
 393     UINT8 Sector;
 394     UINT8 nSector = 63;
 395     UINT8 nHead = 8;
 396     UINT32 Cylinder;
 397     UINT32 EndSectorId;
 398
 399     while (nHead != 0 && (DiskSizeInBytes / 512 / nSector / nHead) > 1024)
 400     {
 401         nHead = (UINT8)nHead * 2;
 402     }
 403
 404     if (nHead == 0)
 405     {
 406         nHead = 255;
 407     }
 408
 409     Cylinder = StartSectorId / nSector / nHead;
 410     Head = StartSectorId / nSector % nHead;
 411     Sector = StartSectorId % nSector + 1;
 412
 413     Table->StartHead = Head;
 414     Table->StartSector = Sector;
 415     Table->StartCylinder = Cylinder;
 416
 417     EndSectorId = StartSectorId + SectorCount - 1;
 418     Cylinder = EndSectorId / nSector / nHead;
 419     Head = EndSectorId / nSector % nHead;
 420     Sector = EndSectorId % nSector + 1;
 421
 422     Table->EndHead = Head;
 423     Table->EndSector = Sector;
 424     Table->EndCylinder = Cylinder;
 425
 426     Table->StartSectorId = StartSectorId;
 427     Table->SectorCount = SectorCount;
 428
 429     return 0;
 430 }
 431
 432
 433 static int WriteDataToPhyDisk(int fd, UINT64 offset, void *buffer, int len)
 434 {
 435     ssize_t wrlen;
 436     off_t newseek;
 437
 438     newseek = lseek(fd, offset, SEEK_SET);
 439     if (newseek != offset)
 440     {
 441         printf("Failed to lseek %llu %lld %d\n", offset, (long long)newseek, errno);
 442         return 0;
 443     }
 444
 445     wrlen = write(fd, buffer, len);
 446     if ((int)wrlen != len)
 447     {
 448         printf("Failed to write %d %d %d\n", len, (int)wrlen, errno);
 449         return 0;
 450     }
 451
 452     return 1;
 453 }
 454
 455 static int VentoyFillBackupGptHead(VTOY_GPT_INFO *pInfo, VTOY_GPT_HDR *pHead)
 456 {
 457     UINT64 LBA;
 458     UINT64 BackupLBA;
 459
 460     memcpy(pHead, &pInfo->Head, sizeof(VTOY_GPT_HDR));
 461
 462     LBA = pHead->EfiStartLBA;
 463     BackupLBA = pHead->EfiBackupLBA;
 464
 465     pHead->EfiStartLBA = BackupLBA;
 466     pHead->EfiBackupLBA = LBA;
 467     pHead->PartTblStartLBA = BackupLBA + 1 - 33;
 468
 469     pHead->Crc = 0;
 470     pHead->Crc = VtoyCrc32(pHead, pHead->Length);
 471
 472     return 0;
 473 }
 474
 475 static int update_part_table(char *disk, UINT64 part2start)
 476 {
 477     int i;
 478     int j;
 479     int fd = -1;
 480     int rc = 1;
 481     int PartStyle = 0;
 482     ssize_t len = 0;
 483     UINT64 DiskSizeInBytes;
 484     VTOY_GPT_INFO *pGPT = NULL;
 485     VTOY_GPT_HDR *pBack = NULL;
 486
 487     DiskSizeInBytes = get_disk_size_in_byte(disk);
 488     if (DiskSizeInBytes == 0)
 489     {
 490         printf("Failed to get disk size of %s\n", disk);
 491         goto out;
 492     }
 493
 494     fd = open(disk, O_RDWR);
 495     if (fd < 0)
 496     {
 497         printf("Failed to open %s\n", disk);
 498         goto out;
 499     }
 500
 501     pGPT = malloc(sizeof(VTOY_GPT_INFO) + sizeof(VTOY_GPT_HDR));
 502     if (NULL == pGPT)
 503     {
 504         goto out;
 505     }
 506     memset(pGPT, 0, sizeof(VTOY_GPT_INFO) + sizeof(VTOY_GPT_HDR));
 507
 508     pBack = (VTOY_GPT_HDR *)(pGPT + 1);
 509
 510     len = read(fd, pGPT, sizeof(VTOY_GPT_INFO));
 511     if (len != (ssize_t)sizeof(VTOY_GPT_INFO))
 512     {
 513         printf("Failed to read partition table %d err:%d\n", (int)len, errno);
 514         goto out;
 515     }
 516
 517     if (pGPT->MBR.PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
 518         {
 519         PartStyle = 1;
 520         }
 521         else
 522         {
 523         PartStyle = 0;
 524         }
 525
 526     if (PartStyle == 0)
 527     {
 528                 PART_TABLE *PartTbl = pGPT->MBR.PartTbl;
 529
 530         for (i = 1; i < 4; i++)
 531         {
 532             if (PartTbl[i].SectorCount == 0)
 533             {
 534                                 break;
 535             }
 536         }
 537
 538                 if (i >= 4)
 539                 {
 540                         printf("###[FAIL] Can not find a free MBR partition table.\n");
 541                         goto out;
 542                 }
 543
 544         for (j = i - 1; j > 0; j--)
 545                 {
 546                         printf("Move MBR partition table %d --> %d\n", j + 1, j + 2);
 547                         memcpy(PartTbl + (j + 1), PartTbl + j, sizeof(PART_TABLE));
 548                 }
 549
 550         memset(PartTbl + 1, 0, sizeof(PART_TABLE));
 551         VentoyFillMBRLocation(DiskSizeInBytes, (UINT32)part2start, VENTOY_EFI_PART_SIZE / 512, PartTbl + 1);
 552                 PartTbl[1].Active = 0x00;
 553                 PartTbl[1].FsFlag = 0xEF; // EFI System Partition
 554
 555         PartTbl[0].Active = 0x80; // bootable
 556         PartTbl[0].SectorCount = (UINT32)part2start - 2048;
 557
 558         if (!WriteDataToPhyDisk(fd, 0, &(pGPT->MBR), 512))
 559                 {
 560                         printf("MBR write MBR failed\n");
 561                         goto out;
 562                 }
 563
 564         fsync(fd);
 565         printf("MBR update partition table success.\n");
 566         rc = 0;
 567     }
 568     else
 569     {
 570                 VTOY_GPT_PART_TBL *PartTbl = pGPT->PartTbl;
 571
 572         for (i = 1; i < 128; i++)
 573         {
 574             if (memcmp(&(PartTbl[i].PartGuid), &g_ZeroGuid, sizeof(GUID)) == 0)
 575             {
 576                                 break;
 577             }
 578         }
 579
 580                 if (i >= 128)
 581                 {
 582                         printf("###[FAIL] Can not find a free GPT partition table.\n");
 583                         goto out;
 584                 }
 585
 586                 for (j = i - 1; j > 0; j--)
 587                 {
 588                         printf("Move GPT partition table %d --> %d\n", j + 1, j + 2);
 589                         memcpy(PartTbl + (j + 1), PartTbl + j, sizeof(VTOY_GPT_PART_TBL));
 590                 }
 591
 592         // to fix windows issue
 593         memset(PartTbl + 1, 0, sizeof(VTOY_GPT_PART_TBL));
 594                 memcpy(&(PartTbl[1].PartType), &g_WindowsDataPartGuid, sizeof(GUID));
 595                 ventoy_gen_preudo_uuid(&(PartTbl[1].PartGuid));
 596
 597         PartTbl[0].LastLBA = part2start - 1;
 598
 599         PartTbl[1].StartLBA = PartTbl[0].LastLBA + 1;
 600                 PartTbl[1].LastLBA = PartTbl[1].StartLBA + VENTOY_EFI_PART_SIZE / 512 - 1;
 601                 PartTbl[1].Attr = 0xC000000000000001ULL;
 602         PartTbl[1].Name[0] = 'V';
 603         PartTbl[1].Name[1] = 'T';
 604         PartTbl[1].Name[2] = 'O';
 605         PartTbl[1].Name[3] = 'Y';
 606         PartTbl[1].Name[4] = 'E';
 607         PartTbl[1].Name[5] = 'F';
 608         PartTbl[1].Name[6] = 'I';
 609         PartTbl[1].Name[7] = 0;
 610
 611                 //Update CRC
 612                 pGPT->Head.PartTblCrc = VtoyCrc32(pGPT->PartTbl, sizeof(pGPT->PartTbl));
 613                 pGPT->Head.Crc = 0;
 614                 pGPT->Head.Crc = VtoyCrc32(&(pGPT->Head), pGPT->Head.Length);
 615
 616                 printf("pGPT->Head.EfiStartLBA=%llu\n", pGPT->Head.EfiStartLBA);
 617                 printf("pGPT->Head.EfiBackupLBA=%llu\n", pGPT->Head.EfiBackupLBA);
 618
 619                 VentoyFillBackupGptHead(pGPT, pBack);
 620                 if (!WriteDataToPhyDisk(fd, pGPT->Head.EfiBackupLBA * 512, pBack, 512))
 621                 {
 622                         printf("GPT write backup head failed\n");
 623                         goto out;
 624                 }
 625
 626                 if (!WriteDataToPhyDisk(fd, (pGPT->Head.EfiBackupLBA - 32) * 512, pGPT->PartTbl, 512 * 32))
 627                 {
 628                         printf("GPT write backup partition table failed\n");
 629                         goto out;
 630                 }
 631
 632                 if (!WriteDataToPhyDisk(fd, 0, pGPT, 512 * 34))
 633                 {
 634                         printf("GPT write MBR & Main partition table failed\n");
 635                         goto out;
 636                 }
 637
 638         fsync(fd);
 639         printf("GPT update partition table success.\n");
 640         rc = 0;
 641     }
 642
 643 out:
 644     check_close(fd);
 645     check_free(pGPT);
 646     return rc;
 647 }
 648
 649 int partresize_main(int argc, char **argv)
 650 {
 651     UINT64 sector;
 652
 653     if (argc != 3 && argc != 4)
 654     {
 655         printf("usage: partresize -c/-f /dev/sdb\n");
 656         return 1;
 657     }
 658
 659     if (strcmp(argv[1], "-c") == 0)
 660     {
 661         return part_check(argv[2]);
 662     }
 663     else if (strcmp(argv[1], "-s") == 0)
 664     {
 665         sector = strtoull(argv[3], NULL, 10);
 666         return secureboot_proc(argv[2], sector);
 667     }
 668     else if (strcmp(argv[1], "-p") == 0)
 669     {
 670         sector = strtoull(argv[3], NULL, 10);
 671         return update_part_table(argv[2], sector);
 672     }
 673     else if (strcmp(argv[1], "-t") == 0)
 674     {
 675         return gpt_check(argv[2]);
 676     }
 677     else
 678     {
 679         return 1;
 680     }
 681 }
 682