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("/EFI/BOOT/mmx64.efi");
 316             fl_remove("/ENROLL_THIS_KEY_IN_MOKMANAGER.cer");
 317
 318                         file = fl_fopen("/EFI/BOOT/BOOTX64.EFI", "wb");
 319                         printf("Open bootx64 efi file %p\n", file);
 320                         if (file)
 321                         {
 322                                 if (filebuf)
 323                                 {
 324                                         fl_fwrite(filebuf, 1, size, file);
 325                                 }
 326
 327                                 fl_fflush(file);
 328                                 fl_fclose(file);
 329                         }
 330
 331                         if (filebuf)
 332                         {
 333                                 free(filebuf);
 334                         }
 335                 }
 336
 337         file = fl_fopen("/EFI/BOOT/grubia32_real.efi", "rb");
 338         printf("Open ventoy ia32 efi file %p\n", file);
 339         if (file)
 340         {
 341             fl_fseek(file, 0, SEEK_END);
 342             size = (int)fl_ftell(file);
 343             fl_fseek(file, 0, SEEK_SET);
 344
 345             printf("ventoy efi file size %d ...\n", size);
 346
 347             filebuf = (char *)malloc(size);
 348             if (filebuf)
 349             {
 350                 fl_fread(filebuf, 1, size, file);
 351             }
 352
 353             fl_fclose(file);
 354
 355             fl_remove("/EFI/BOOT/BOOTIA32.EFI");
 356             fl_remove("/EFI/BOOT/grubia32.efi");
 357             fl_remove("/EFI/BOOT/grubia32_real.efi");
 358             fl_remove("/EFI/BOOT/mmia32.efi");
 359
 360             file = fl_fopen("/EFI/BOOT/BOOTIA32.EFI", "wb");
 361             printf("Open bootia32 efi file %p\n", file);
 362             if (file)
 363             {
 364                 if (filebuf)
 365                 {
 366                     fl_fwrite(filebuf, 1, size, file);
 367                 }
 368
 369                 fl_fflush(file);
 370                 fl_fclose(file);
 371             }
 372
 373             if (filebuf)
 374             {
 375                 free(filebuf);
 376             }
 377         }
 378
 379         }
 380         else
 381         {
 382                 rc = 1;
 383         }
 384
 385         fl_shutdown();
 386     fsync(fd);
 387
 388         return rc;
 389 }
 390
 391 static int VentoyFillMBRLocation(UINT64 DiskSizeInBytes, UINT32 StartSectorId, UINT32 SectorCount, PART_TABLE *Table)
 392 {
 393     UINT8 Head;
 394     UINT8 Sector;
 395     UINT8 nSector = 63;
 396     UINT8 nHead = 8;
 397     UINT32 Cylinder;
 398     UINT32 EndSectorId;
 399
 400     while (nHead != 0 && (DiskSizeInBytes / 512 / nSector / nHead) > 1024)
 401     {
 402         nHead = (UINT8)nHead * 2;
 403     }
 404
 405     if (nHead == 0)
 406     {
 407         nHead = 255;
 408     }
 409
 410     Cylinder = StartSectorId / nSector / nHead;
 411     Head = StartSectorId / nSector % nHead;
 412     Sector = StartSectorId % nSector + 1;
 413
 414     Table->StartHead = Head;
 415     Table->StartSector = Sector;
 416     Table->StartCylinder = Cylinder;
 417
 418     EndSectorId = StartSectorId + SectorCount - 1;
 419     Cylinder = EndSectorId / nSector / nHead;
 420     Head = EndSectorId / nSector % nHead;
 421     Sector = EndSectorId % nSector + 1;
 422
 423     Table->EndHead = Head;
 424     Table->EndSector = Sector;
 425     Table->EndCylinder = Cylinder;
 426
 427     Table->StartSectorId = StartSectorId;
 428     Table->SectorCount = SectorCount;
 429
 430     return 0;
 431 }
 432
 433
 434 static int WriteDataToPhyDisk(int fd, UINT64 offset, void *buffer, int len)
 435 {
 436     ssize_t wrlen;
 437     off_t newseek;
 438
 439     newseek = lseek(fd, offset, SEEK_SET);
 440     if (newseek != offset)
 441     {
 442         printf("Failed to lseek %llu %lld %d\n", offset, (long long)newseek, errno);
 443         return 0;
 444     }
 445
 446     wrlen = write(fd, buffer, len);
 447     if ((int)wrlen != len)
 448     {
 449         printf("Failed to write %d %d %d\n", len, (int)wrlen, errno);
 450         return 0;
 451     }
 452
 453     return 1;
 454 }
 455
 456 static int VentoyFillBackupGptHead(VTOY_GPT_INFO *pInfo, VTOY_GPT_HDR *pHead)
 457 {
 458     UINT64 LBA;
 459     UINT64 BackupLBA;
 460
 461     memcpy(pHead, &pInfo->Head, sizeof(VTOY_GPT_HDR));
 462
 463     LBA = pHead->EfiStartLBA;
 464     BackupLBA = pHead->EfiBackupLBA;
 465
 466     pHead->EfiStartLBA = BackupLBA;
 467     pHead->EfiBackupLBA = LBA;
 468     pHead->PartTblStartLBA = BackupLBA + 1 - 33;
 469
 470     pHead->Crc = 0;
 471     pHead->Crc = VtoyCrc32(pHead, pHead->Length);
 472
 473     return 0;
 474 }
 475
 476 static int update_part_table(char *disk, UINT64 part2start)
 477 {
 478     int i;
 479     int j;
 480     int fd = -1;
 481     int rc = 1;
 482     int PartStyle = 0;
 483     ssize_t len = 0;
 484     UINT64 DiskSizeInBytes;
 485     VTOY_GPT_INFO *pGPT = NULL;
 486     VTOY_GPT_HDR *pBack = NULL;
 487
 488     DiskSizeInBytes = get_disk_size_in_byte(disk);
 489     if (DiskSizeInBytes == 0)
 490     {
 491         printf("Failed to get disk size of %s\n", disk);
 492         goto out;
 493     }
 494
 495     fd = open(disk, O_RDWR);
 496     if (fd < 0)
 497     {
 498         printf("Failed to open %s\n", disk);
 499         goto out;
 500     }
 501
 502     pGPT = malloc(sizeof(VTOY_GPT_INFO) + sizeof(VTOY_GPT_HDR));
 503     if (NULL == pGPT)
 504     {
 505         goto out;
 506     }
 507     memset(pGPT, 0, sizeof(VTOY_GPT_INFO) + sizeof(VTOY_GPT_HDR));
 508
 509     pBack = (VTOY_GPT_HDR *)(pGPT + 1);
 510
 511     len = read(fd, pGPT, sizeof(VTOY_GPT_INFO));
 512     if (len != (ssize_t)sizeof(VTOY_GPT_INFO))
 513     {
 514         printf("Failed to read partition table %d err:%d\n", (int)len, errno);
 515         goto out;
 516     }
 517
 518     if (pGPT->MBR.PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
 519         {
 520         PartStyle = 1;
 521         }
 522         else
 523         {
 524         PartStyle = 0;
 525         }
 526
 527     if (PartStyle == 0)
 528     {
 529                 PART_TABLE *PartTbl = pGPT->MBR.PartTbl;
 530
 531         for (i = 1; i < 4; i++)
 532         {
 533             if (PartTbl[i].SectorCount == 0)
 534             {
 535                                 break;
 536             }
 537         }
 538
 539                 if (i >= 4)
 540                 {
 541                         printf("###[FAIL] Can not find a free MBR partition table.\n");
 542                         goto out;
 543                 }
 544
 545         for (j = i - 1; j > 0; j--)
 546                 {
 547                         printf("Move MBR partition table %d --> %d\n", j + 1, j + 2);
 548                         memcpy(PartTbl + (j + 1), PartTbl + j, sizeof(PART_TABLE));
 549                 }
 550
 551         memset(PartTbl + 1, 0, sizeof(PART_TABLE));
 552         VentoyFillMBRLocation(DiskSizeInBytes, (UINT32)part2start, VENTOY_EFI_PART_SIZE / 512, PartTbl + 1);
 553                 PartTbl[1].Active = 0x00;
 554                 PartTbl[1].FsFlag = 0xEF; // EFI System Partition
 555
 556         PartTbl[0].Active = 0x80; // bootable
 557         PartTbl[0].SectorCount = (UINT32)part2start - 2048;
 558
 559         if (!WriteDataToPhyDisk(fd, 0, &(pGPT->MBR), 512))
 560                 {
 561                         printf("MBR write MBR failed\n");
 562                         goto out;
 563                 }
 564
 565         fsync(fd);
 566         printf("MBR update partition table success.\n");
 567         rc = 0;
 568     }
 569     else
 570     {
 571                 VTOY_GPT_PART_TBL *PartTbl = pGPT->PartTbl;
 572
 573         for (i = 1; i < 128; i++)
 574         {
 575             if (memcmp(&(PartTbl[i].PartGuid), &g_ZeroGuid, sizeof(GUID)) == 0)
 576             {
 577                                 break;
 578             }
 579         }
 580
 581                 if (i >= 128)
 582                 {
 583                         printf("###[FAIL] Can not find a free GPT partition table.\n");
 584                         goto out;
 585                 }
 586
 587                 for (j = i - 1; j > 0; j--)
 588                 {
 589                         printf("Move GPT partition table %d --> %d\n", j + 1, j + 2);
 590                         memcpy(PartTbl + (j + 1), PartTbl + j, sizeof(VTOY_GPT_PART_TBL));
 591                 }
 592
 593         // to fix windows issue
 594         memset(PartTbl + 1, 0, sizeof(VTOY_GPT_PART_TBL));
 595                 memcpy(&(PartTbl[1].PartType), &g_WindowsDataPartGuid, sizeof(GUID));
 596                 ventoy_gen_preudo_uuid(&(PartTbl[1].PartGuid));
 597
 598         PartTbl[0].LastLBA = part2start - 1;
 599
 600         PartTbl[1].StartLBA = PartTbl[0].LastLBA + 1;
 601                 PartTbl[1].LastLBA = PartTbl[1].StartLBA + VENTOY_EFI_PART_SIZE / 512 - 1;
 602                 PartTbl[1].Attr = 0xC000000000000001ULL;
 603         PartTbl[1].Name[0] = 'V';
 604         PartTbl[1].Name[1] = 'T';
 605         PartTbl[1].Name[2] = 'O';
 606         PartTbl[1].Name[3] = 'Y';
 607         PartTbl[1].Name[4] = 'E';
 608         PartTbl[1].Name[5] = 'F';
 609         PartTbl[1].Name[6] = 'I';
 610         PartTbl[1].Name[7] = 0;
 611
 612                 //Update CRC
 613                 pGPT->Head.PartTblCrc = VtoyCrc32(pGPT->PartTbl, sizeof(pGPT->PartTbl));
 614                 pGPT->Head.Crc = 0;
 615                 pGPT->Head.Crc = VtoyCrc32(&(pGPT->Head), pGPT->Head.Length);
 616
 617                 printf("pGPT->Head.EfiStartLBA=%llu\n", pGPT->Head.EfiStartLBA);
 618                 printf("pGPT->Head.EfiBackupLBA=%llu\n", pGPT->Head.EfiBackupLBA);
 619
 620                 VentoyFillBackupGptHead(pGPT, pBack);
 621                 if (!WriteDataToPhyDisk(fd, pGPT->Head.EfiBackupLBA * 512, pBack, 512))
 622                 {
 623                         printf("GPT write backup head failed\n");
 624                         goto out;
 625                 }
 626
 627                 if (!WriteDataToPhyDisk(fd, (pGPT->Head.EfiBackupLBA - 32) * 512, pGPT->PartTbl, 512 * 32))
 628                 {
 629                         printf("GPT write backup partition table failed\n");
 630                         goto out;
 631                 }
 632
 633                 if (!WriteDataToPhyDisk(fd, 0, pGPT, 512 * 34))
 634                 {
 635                         printf("GPT write MBR & Main partition table failed\n");
 636                         goto out;
 637                 }
 638
 639         fsync(fd);
 640         printf("GPT update partition table success.\n");
 641         rc = 0;
 642     }
 643
 644 out:
 645     check_close(fd);
 646     check_free(pGPT);
 647     return rc;
 648 }
 649
 650 int partresize_main(int argc, char **argv)
 651 {
 652     UINT64 sector;
 653
 654     if (argc != 3 && argc != 4)
 655     {
 656         printf("usage: partresize -c/-f /dev/sdb\n");
 657         return 1;
 658     }
 659
 660     if (strcmp(argv[1], "-c") == 0)
 661     {
 662         return part_check(argv[2]);
 663     }
 664     else if (strcmp(argv[1], "-s") == 0)
 665     {
 666         sector = strtoull(argv[3], NULL, 10);
 667         return secureboot_proc(argv[2], sector);
 668     }
 669     else if (strcmp(argv[1], "-p") == 0)
 670     {
 671         sector = strtoull(argv[3], NULL, 10);
 672         return update_part_table(argv[2], sector);
 673     }
 674     else if (strcmp(argv[1], "-t") == 0)
 675     {
 676         return gpt_check(argv[2]);
 677     }
 678     else
 679     {
 680         return 1;
 681     }
 682 }
 683