Ventoy2Disk/Ventoy2Disk/PhyDrive.c

   1 /******************************************************************************
   2  * PhyDrive.c
   3  *
   4  * Copyright (c) 2020, longpanda <admin@ventoy.net>
   5  * Copyright (c) 2011-2020, Pete Batard <pete@akeo.ie>
   6  *
   7  * This program is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU General Public License as
   9  * published by the Free Software Foundation; either version 3 of the
  10  * License, or (at your option) any later version.
  11  *
  12  * This program is distributed in the hope that it will be useful, but
  13  * WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15  * General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU General Public License
  18  * along with this program; if not, see <http://www.gnu.org/licenses/>.
  19  *
  20  */
  21
  22 #include <Windows.h>
  23 #include <time.h>
  24 #include <winternl.h>
  25 #include <commctrl.h>
  26 #include <initguid.h>
  27 #include "resource.h"
  28 #include "Language.h"
  29 #include "Ventoy2Disk.h"
  30 #include "fat_filelib.h"
  31 #include "ff.h"
  32 #include "DiskService.h"
  33
  34 static int g_backup_bin_index = 0;
  35
  36 static DWORD GetVentoyVolumeName(int PhyDrive, UINT64 StartSectorId, CHAR *NameBuf, UINT32 BufLen, BOOL DelSlash)
  37 {
  38     size_t len;
  39     BOOL bRet;
  40     DWORD dwSize;
  41     HANDLE hDrive;
  42     HANDLE hVolume;
  43     UINT64 PartOffset;
  44     DWORD Status = ERROR_NOT_FOUND;
  45     DISK_EXTENT *pExtents = NULL;
  46     CHAR VolumeName[MAX_PATH] = { 0 };
  47     VOLUME_DISK_EXTENTS DiskExtents;
  48
  49     PartOffset = 512ULL * StartSectorId;
  50
  51         Log("GetVentoyVolumeName PhyDrive %d SectorStart:%llu PartOffset:%llu", PhyDrive, (ULONGLONG)StartSectorId, (ULONGLONG)PartOffset);
  52
  53     hVolume = FindFirstVolumeA(VolumeName, sizeof(VolumeName));
  54     if (hVolume == INVALID_HANDLE_VALUE)
  55     {
  56         return 1;
  57     }
  58
  59     do {
  60
  61         len = strlen(VolumeName);
  62         Log("Find volume:%s", VolumeName);
  63
  64         VolumeName[len - 1] = 0;
  65
  66         hDrive = CreateFileA(VolumeName, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
  67         if (hDrive == INVALID_HANDLE_VALUE)
  68         {
  69             continue;
  70         }
  71
  72         bRet = DeviceIoControl(hDrive,
  73             IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
  74             NULL,
  75             0,
  76             &DiskExtents,
  77             (DWORD)(sizeof(DiskExtents)),
  78             (LPDWORD)&dwSize,
  79             NULL);
  80
  81         Log("IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS bRet:%u code:%u", bRet, LASTERR);
  82         Log("NumberOfDiskExtents:%u DiskNumber:%u", DiskExtents.NumberOfDiskExtents, DiskExtents.Extents[0].DiskNumber);
  83
  84         if (bRet && DiskExtents.NumberOfDiskExtents == 1)
  85         {
  86             pExtents = DiskExtents.Extents;
  87
  88             Log("This volume DiskNumber:%u offset:%llu", pExtents->DiskNumber, (ULONGLONG)pExtents->StartingOffset.QuadPart);
  89             if ((int)pExtents->DiskNumber == PhyDrive && pExtents->StartingOffset.QuadPart == PartOffset)
  90             {
  91                 Log("This volume match");
  92
  93                 if (!DelSlash)
  94                 {
  95                     VolumeName[len - 1] = '\\';
  96                 }
  97
  98                 sprintf_s(NameBuf, BufLen, "%s", VolumeName);
  99                 Status = ERROR_SUCCESS;
 100                 CloseHandle(hDrive);
 101                 break;
 102             }
 103         }
 104
 105         CloseHandle(hDrive);
 106     } while (FindNextVolumeA(hVolume, VolumeName, sizeof(VolumeName)));
 107
 108     FindVolumeClose(hVolume);
 109
 110     Log("GetVentoyVolumeName return %u", Status);
 111     return Status;
 112 }
 113
 114 static int GetLettersBelongPhyDrive(int PhyDrive, char *DriveLetters, size_t Length)
 115 {
 116     int n = 0;
 117     DWORD DataSize = 0;
 118     CHAR *Pos = NULL;
 119     CHAR *StringBuf = NULL;
 120
 121     DataSize = GetLogicalDriveStringsA(0, NULL);
 122     StringBuf = (CHAR *)malloc(DataSize + 1);
 123     if (StringBuf == NULL)
 124     {
 125         return 1;
 126     }
 127
 128     GetLogicalDriveStringsA(DataSize, StringBuf);
 129
 130     for (Pos = StringBuf; *Pos; Pos += strlen(Pos) + 1)
 131     {
 132         if (n < (int)Length && PhyDrive == GetPhyDriveByLogicalDrive(Pos[0]))
 133         {
 134             Log("%C: is belong to phydrive%d", Pos[0], PhyDrive);
 135             DriveLetters[n++] = Pos[0];
 136         }
 137     }
 138
 139     free(StringBuf);
 140     return 0;
 141 }
 142
 143 static HANDLE GetPhysicalHandle(int Drive, BOOLEAN bLockDrive, BOOLEAN bWriteAccess, BOOLEAN bWriteShare)
 144 {
 145     int i;
 146     DWORD dwSize;
 147     DWORD LastError;
 148     UINT64 EndTime;
 149     HANDLE hDrive = INVALID_HANDLE_VALUE;
 150     CHAR PhyDrive[128];
 151     CHAR DevPath[MAX_PATH] = { 0 };
 152
 153     safe_sprintf(PhyDrive, "\\\\.\\PhysicalDrive%d", Drive);
 154
 155     if (0 == QueryDosDeviceA(PhyDrive + 4, DevPath, sizeof(DevPath)))
 156     {
 157         Log("QueryDosDeviceA failed error:%u", GetLastError());
 158         strcpy_s(DevPath, sizeof(DevPath), "???");
 159     }
 160     else
 161     {
 162         Log("QueryDosDeviceA success %s", DevPath);
 163     }
 164
 165     for (i = 0; i < DRIVE_ACCESS_RETRIES; i++)
 166     {
 167         // Try without FILE_SHARE_WRITE (unless specifically requested) so that
 168         // we won't be bothered by the OS or other apps when we set up our data.
 169         // However this means we might have to wait for an access gap...
 170         // We keep FILE_SHARE_READ though, as this shouldn't hurt us any, and is
 171         // required for enumeration.
 172         hDrive = CreateFileA(PhyDrive,
 173             GENERIC_READ | (bWriteAccess ? GENERIC_WRITE : 0),
 174             FILE_SHARE_READ | (bWriteShare ? FILE_SHARE_WRITE : 0),
 175             NULL,
 176             OPEN_EXISTING,
 177             FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
 178             NULL);
 179
 180         LastError = GetLastError();
 181         Log("[%d] CreateFileA %s code:%u %p", i, PhyDrive, LastError, hDrive);
 182
 183         if (hDrive != INVALID_HANDLE_VALUE)
 184         {
 185             break;
 186         }
 187
 188         if ((LastError != ERROR_SHARING_VIOLATION) && (LastError != ERROR_ACCESS_DENIED))
 189         {
 190             break;
 191         }
 192
 193         if (i == 0)
 194         {
 195             Log("Waiting for access on %s [%s]...", PhyDrive, DevPath);
 196         }
 197         else if (!bWriteShare && (i > DRIVE_ACCESS_RETRIES / 3))
 198         {
 199             // If we can't seem to get a hold of the drive for some time, try to enable FILE_SHARE_WRITE...
 200             Log("Warning: Could not obtain exclusive rights. Retrying with write sharing enabled...");
 201             bWriteShare = TRUE;
 202
 203             // Try to report the process that is locking the drive
 204             // We also use bit 6 as a flag to indicate that SearchProcess was called.
 205             //access_mask = SearchProcess(DevPath, SEARCH_PROCESS_TIMEOUT, TRUE, TRUE, FALSE) | 0x40;
 206
 207         }
 208         Sleep(DRIVE_ACCESS_TIMEOUT / DRIVE_ACCESS_RETRIES);
 209     }
 210
 211     if (hDrive == INVALID_HANDLE_VALUE)
 212     {
 213         Log("Could not open %s %u", PhyDrive, LASTERR);
 214         goto End;
 215     }
 216
 217     if (bWriteAccess)
 218     {
 219         Log("Opened %s for %s write access", PhyDrive, bWriteShare ? "shared" : "exclusive");
 220     }
 221
 222     if (bLockDrive)
 223     {
 224         if (DeviceIoControl(hDrive, FSCTL_ALLOW_EXTENDED_DASD_IO, NULL, 0, NULL, 0, &dwSize, NULL))
 225         {
 226             Log("I/O boundary checks disabled");
 227         }
 228
 229         EndTime = GetTickCount64() + DRIVE_ACCESS_TIMEOUT;
 230
 231         do {
 232             if (DeviceIoControl(hDrive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL))
 233             {
 234                 Log("FSCTL_LOCK_VOLUME success");
 235                 goto End;
 236             }
 237             Sleep(DRIVE_ACCESS_TIMEOUT / DRIVE_ACCESS_RETRIES);
 238         } while (GetTickCount64() < EndTime);
 239
 240         // If we reached this section, either we didn't manage to get a lock or the user cancelled
 241         Log("Could not lock access to %s %u", PhyDrive, LASTERR);
 242
 243         // See if we can report the processes are accessing the drive
 244         //if (!IS_ERROR(FormatStatus) && (access_mask == 0))
 245         //    access_mask = SearchProcess(DevPath, SEARCH_PROCESS_TIMEOUT, TRUE, TRUE, FALSE);
 246         // Try to continue if the only access rights we saw were for read-only
 247         //if ((access_mask & 0x07) != 0x01)
 248         //    safe_closehandle(hDrive);
 249
 250         CHECK_CLOSE_HANDLE(hDrive);
 251     }
 252
 253 End:
 254
 255     if (hDrive == INVALID_HANDLE_VALUE)
 256     {
 257         Log("Can get handle of %s, maybe some process control it.", DevPath);
 258     }
 259
 260     return hDrive;
 261 }
 262
 263 int GetPhyDriveByLogicalDrive(int DriveLetter)
 264 {
 265     BOOL Ret;
 266     DWORD dwSize;
 267     HANDLE Handle;
 268     VOLUME_DISK_EXTENTS DiskExtents;
 269     CHAR PhyPath[128];
 270
 271     safe_sprintf(PhyPath, "\\\\.\\%C:", (CHAR)DriveLetter);
 272
 273     Handle = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 274     if (Handle == INVALID_HANDLE_VALUE)
 275     {
 276         Log("Could not open the disk<%s>, error:%u", PhyPath, LASTERR);
 277         return -1;
 278     }
 279
 280     Ret = DeviceIoControl(Handle,
 281         IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 282         NULL,
 283         0,
 284         &DiskExtents,
 285         (DWORD)(sizeof(DiskExtents)),
 286         (LPDWORD)&dwSize,
 287         NULL);
 288
 289     if (!Ret || DiskExtents.NumberOfDiskExtents == 0)
 290     {
 291         Log("DeviceIoControl IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS failed %s, error:%u", PhyPath, LASTERR);
 292         CHECK_CLOSE_HANDLE(Handle);
 293         return -1;
 294     }
 295     CHECK_CLOSE_HANDLE(Handle);
 296
 297     Log("LogicalDrive:%s PhyDrive:%d Offset:%llu ExtentLength:%llu",
 298         PhyPath,
 299         DiskExtents.Extents[0].DiskNumber,
 300         DiskExtents.Extents[0].StartingOffset.QuadPart,
 301         DiskExtents.Extents[0].ExtentLength.QuadPart
 302         );
 303
 304     return (int)DiskExtents.Extents[0].DiskNumber;
 305 }
 306
 307 int GetAllPhysicalDriveInfo(PHY_DRIVE_INFO *pDriveList, DWORD *pDriveCount)
 308 {
 309     int i;
 310     int Count;
 311     int id;
 312     int Letter = 'A';
 313     BOOL  bRet;
 314     DWORD dwBytes;
 315     DWORD DriveCount = 0;
 316     HANDLE Handle = INVALID_HANDLE_VALUE;
 317     CHAR PhyDrive[128];
 318     PHY_DRIVE_INFO *CurDrive = pDriveList;
 319     GET_LENGTH_INFORMATION LengthInfo;
 320     STORAGE_PROPERTY_QUERY Query;
 321     STORAGE_DESCRIPTOR_HEADER DevDescHeader;
 322     STORAGE_DEVICE_DESCRIPTOR *pDevDesc;
 323     int PhyDriveId[VENTOY_MAX_PHY_DRIVE];
 324
 325     Count = GetPhysicalDriveCount();
 326
 327     for (i = 0; i < Count && i < VENTOY_MAX_PHY_DRIVE; i++)
 328     {
 329         PhyDriveId[i] = i;
 330     }
 331
 332     dwBytes = GetLogicalDrives();
 333     Log("Logical Drives: 0x%x", dwBytes);
 334     while (dwBytes)
 335     {
 336         if (dwBytes & 0x01)
 337         {
 338             id = GetPhyDriveByLogicalDrive(Letter);
 339             Log("%C --> %d", Letter, id);
 340             if (id >= 0)
 341             {
 342                 for (i = 0; i < Count; i++)
 343                 {
 344                     if (PhyDriveId[i] == id)
 345                     {
 346                         break;
 347                     }
 348                 }
 349
 350                 if (i >= Count)
 351                 {
 352                     Log("Add phy%d to list", i);
 353                     PhyDriveId[Count] = id;
 354                     Count++;
 355                 }
 356             }
 357         }
 358
 359         Letter++;
 360         dwBytes >>= 1;
 361     }
 362
 363     for (i = 0; i < Count && DriveCount < VENTOY_MAX_PHY_DRIVE; i++)
 364     {
 365         CHECK_CLOSE_HANDLE(Handle);
 366
 367         safe_sprintf(PhyDrive, "\\\\.\\PhysicalDrive%d", PhyDriveId[i]);
 368         Handle = CreateFileA(PhyDrive, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
 369         Log("Create file Handle:%p %s status:%u", Handle, PhyDrive, LASTERR);
 370
 371         if (Handle == INVALID_HANDLE_VALUE)
 372         {
 373             continue;
 374         }
 375
 376         bRet = DeviceIoControl(Handle,
 377                                IOCTL_DISK_GET_LENGTH_INFO, NULL,
 378                                0,
 379                                &LengthInfo,
 380                                sizeof(LengthInfo),
 381                                &dwBytes,
 382                                NULL);
 383         if (!bRet)
 384         {
 385             Log("DeviceIoControl IOCTL_DISK_GET_LENGTH_INFO failed error:%u", LASTERR);
 386             continue;
 387         }
 388
 389         Log("PHYSICALDRIVE%d size %llu bytes", i, (ULONGLONG)LengthInfo.Length.QuadPart);
 390
 391         Query.PropertyId = StorageDeviceProperty;
 392         Query.QueryType = PropertyStandardQuery;
 393
 394         bRet = DeviceIoControl(Handle,
 395                                IOCTL_STORAGE_QUERY_PROPERTY,
 396                                &Query,
 397                                sizeof(Query),
 398                                &DevDescHeader,
 399                                sizeof(STORAGE_DESCRIPTOR_HEADER),
 400                                &dwBytes,
 401                                NULL);
 402         if (!bRet)
 403         {
 404             Log("DeviceIoControl1 error:%u dwBytes:%u", LASTERR, dwBytes);
 405             continue;
 406         }
 407
 408         if (DevDescHeader.Size < sizeof(STORAGE_DEVICE_DESCRIPTOR))
 409         {
 410             Log("Invalid DevDescHeader.Size:%u", DevDescHeader.Size);
 411             continue;
 412         }
 413
 414         pDevDesc = (STORAGE_DEVICE_DESCRIPTOR *)malloc(DevDescHeader.Size);
 415         if (!pDevDesc)
 416         {
 417             Log("failed to malloc error:%u len:%u", LASTERR, DevDescHeader.Size);
 418             continue;
 419         }
 420
 421         bRet = DeviceIoControl(Handle,
 422                                IOCTL_STORAGE_QUERY_PROPERTY,
 423                                &Query,
 424                                sizeof(Query),
 425                                pDevDesc,
 426                                DevDescHeader.Size,
 427                                &dwBytes,
 428                                NULL);
 429         if (!bRet)
 430         {
 431             Log("DeviceIoControl2 error:%u dwBytes:%u", LASTERR, dwBytes);
 432             free(pDevDesc);
 433             continue;
 434         }
 435
 436         CurDrive->PhyDrive = i;
 437         CurDrive->SizeInBytes = LengthInfo.Length.QuadPart;
 438         CurDrive->DeviceType = pDevDesc->DeviceType;
 439         CurDrive->RemovableMedia = pDevDesc->RemovableMedia;
 440         CurDrive->BusType = pDevDesc->BusType;
 441
 442         if (pDevDesc->VendorIdOffset)
 443         {
 444             safe_strcpy(CurDrive->VendorId, (char *)pDevDesc + pDevDesc->VendorIdOffset);
 445             TrimString(CurDrive->VendorId);
 446         }
 447
 448         if (pDevDesc->ProductIdOffset)
 449         {
 450             safe_strcpy(CurDrive->ProductId, (char *)pDevDesc + pDevDesc->ProductIdOffset);
 451             TrimString(CurDrive->ProductId);
 452         }
 453
 454         if (pDevDesc->ProductRevisionOffset)
 455         {
 456             safe_strcpy(CurDrive->ProductRev, (char *)pDevDesc + pDevDesc->ProductRevisionOffset);
 457             TrimString(CurDrive->ProductRev);
 458         }
 459
 460         if (pDevDesc->SerialNumberOffset)
 461         {
 462             safe_strcpy(CurDrive->SerialNumber, (char *)pDevDesc + pDevDesc->SerialNumberOffset);
 463             TrimString(CurDrive->SerialNumber);
 464         }
 465
 466         CurDrive++;
 467         DriveCount++;
 468
 469         free(pDevDesc);
 470
 471         CHECK_CLOSE_HANDLE(Handle);
 472     }
 473
 474     for (i = 0, CurDrive = pDriveList; i < (int)DriveCount; i++, CurDrive++)
 475     {
 476         Log("PhyDrv:%d BusType:%-4s Removable:%u Size:%dGB(%llu) Name:%s %s",
 477             CurDrive->PhyDrive, GetBusTypeString(CurDrive->BusType), CurDrive->RemovableMedia,
 478             GetHumanReadableGBSize(CurDrive->SizeInBytes), CurDrive->SizeInBytes,
 479             CurDrive->VendorId, CurDrive->ProductId);
 480     }
 481
 482     *pDriveCount = DriveCount;
 483
 484     return 0;
 485 }
 486
 487
 488 static HANDLE g_FatPhyDrive;
 489 static UINT64 g_Part2StartSec;
 490 static int GetVentoyVersionFromFatFile(CHAR *VerBuf, size_t BufLen)
 491 {
 492     int rc = 1;
 493     int size = 0;
 494     char *buf = NULL;
 495     void *flfile = NULL;
 496
 497     flfile = fl_fopen("/grub/grub.cfg", "rb");
 498     if (flfile)
 499     {
 500         fl_fseek(flfile, 0, SEEK_END);
 501         size = (int)fl_ftell(flfile);
 502
 503         fl_fseek(flfile, 0, SEEK_SET);
 504
 505         buf = (char *)malloc(size + 1);
 506         if (buf)
 507         {
 508             fl_fread(buf, 1, size, flfile);
 509             buf[size] = 0;
 510
 511             rc = 0;
 512             sprintf_s(VerBuf, BufLen, "%s", ParseVentoyVersionFromString(buf));
 513             free(buf);
 514         }
 515
 516         fl_fclose(flfile);
 517     }
 518
 519     return rc;
 520 }
 521
 522 static int VentoyFatDiskRead(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 523 {
 524     DWORD dwSize;
 525     BOOL bRet;
 526     DWORD ReadSize;
 527     LARGE_INTEGER liCurrentPosition;
 528
 529     liCurrentPosition.QuadPart = Sector + g_Part2StartSec;
 530     liCurrentPosition.QuadPart *= 512;
 531     SetFilePointerEx(g_FatPhyDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 532
 533     ReadSize = (DWORD)(SectorCount * 512);
 534
 535     bRet = ReadFile(g_FatPhyDrive, Buffer, ReadSize, &dwSize, NULL);
 536     if (bRet == FALSE || dwSize != ReadSize)
 537     {
 538         Log("ReadFile error bRet:%u WriteSize:%u dwSize:%u ErrCode:%u\n", bRet, ReadSize, dwSize, LASTERR);
 539     }
 540
 541     return 1;
 542 }
 543
 544
 545 int GetVentoyVerInPhyDrive(const PHY_DRIVE_INFO *pDriveInfo, UINT64 Part2StartSector, CHAR *VerBuf, size_t BufLen, BOOL *pSecureBoot)
 546 {
 547     int rc = 0;
 548     HANDLE hDrive;
 549     void *flfile;
 550
 551     hDrive = GetPhysicalHandle(pDriveInfo->PhyDrive, FALSE, FALSE, FALSE);
 552     if (hDrive == INVALID_HANDLE_VALUE)
 553     {
 554         return 1;
 555     }
 556
 557     g_FatPhyDrive = hDrive;
 558         g_Part2StartSec = Part2StartSector;
 559
 560     Log("Parse FAT fs...");
 561
 562     fl_init();
 563
 564     if (0 == fl_attach_media(VentoyFatDiskRead, NULL))
 565     {
 566         Log("attach media success...");
 567         rc = GetVentoyVersionFromFatFile(VerBuf, BufLen);
 568     }
 569     else
 570     {
 571         Log("attach media failed...");
 572         rc = 1;
 573     }
 574
 575     Log("GetVentoyVerInPhyDrive rc=%d...", rc);
 576     if (rc == 0)
 577     {
 578         Log("VentoyVerInPhyDrive %d is <%s>...", pDriveInfo->PhyDrive, VerBuf);
 579
 580         flfile = fl_fopen("/EFI/BOOT/grubx64_real.efi", "rb");
 581         if (flfile)
 582         {
 583             *pSecureBoot = TRUE;
 584             fl_fclose(flfile);
 585         }
 586     }
 587
 588     fl_shutdown();
 589
 590     CHECK_CLOSE_HANDLE(hDrive);
 591
 592     return rc;
 593 }
 594
 595
 596
 597
 598
 599 static unsigned int g_disk_unxz_len = 0;
 600 static BYTE *g_part_img_pos = NULL;
 601 static BYTE *g_part_img_buf[VENTOY_EFI_PART_SIZE / SIZE_1MB];
 602
 603
 604 static int VentoyFatMemRead(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 605 {
 606         uint32 i;
 607         uint32 offset;
 608         BYTE *MbBuf = NULL;
 609
 610         for (i = 0; i < SectorCount; i++)
 611         {
 612                 offset = (Sector + i) * 512;
 613
 614                 if (g_part_img_buf[1] == NULL)
 615                 {
 616                         MbBuf = g_part_img_buf[0] + offset;
 617                         memcpy(Buffer + i * 512, MbBuf, 512);
 618                 }
 619                 else
 620                 {
 621                         MbBuf = g_part_img_buf[offset / SIZE_1MB];
 622                         memcpy(Buffer + i * 512, MbBuf + (offset % SIZE_1MB), 512);
 623                 }
 624         }
 625
 626         return 1;
 627 }
 628
 629
 630 static int VentoyFatMemWrite(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 631 {
 632         uint32 i;
 633         uint32 offset;
 634         BYTE *MbBuf = NULL;
 635
 636         for (i = 0; i < SectorCount; i++)
 637         {
 638                 offset = (Sector + i) * 512;
 639
 640                 if (g_part_img_buf[1] == NULL)
 641                 {
 642                         MbBuf = g_part_img_buf[0] + offset;
 643                         memcpy(MbBuf, Buffer + i * 512, 512);
 644                 }
 645                 else
 646                 {
 647                         MbBuf = g_part_img_buf[offset / SIZE_1MB];
 648                         memcpy(MbBuf + (offset % SIZE_1MB), Buffer + i * 512, 512);
 649                 }
 650         }
 651
 652         return 1;
 653 }
 654
 655 int VentoyProcSecureBoot(BOOL SecureBoot)
 656 {
 657         int rc = 0;
 658         int size;
 659         char *filebuf = NULL;
 660         void *file = NULL;
 661
 662         Log("VentoyProcSecureBoot %d ...", SecureBoot);
 663
 664         if (SecureBoot)
 665         {
 666                 Log("Secure boot is enabled ...");
 667                 return 0;
 668         }
 669
 670         fl_init();
 671
 672         if (0 == fl_attach_media(VentoyFatMemRead, VentoyFatMemWrite))
 673         {
 674                 file = fl_fopen("/EFI/BOOT/grubx64_real.efi", "rb");
 675                 Log("Open ventoy efi file %p ", file);
 676                 if (file)
 677                 {
 678                         fl_fseek(file, 0, SEEK_END);
 679                         size = (int)fl_ftell(file);
 680                         fl_fseek(file, 0, SEEK_SET);
 681
 682                         Log("ventoy efi file size %d ...", size);
 683
 684                         filebuf = (char *)malloc(size);
 685                         if (filebuf)
 686                         {
 687                                 fl_fread(filebuf, 1, size, file);
 688                         }
 689
 690                         fl_fclose(file);
 691
 692                         Log("Now delete all efi files ...");
 693                         fl_remove("/EFI/BOOT/BOOTX64.EFI");
 694                         fl_remove("/EFI/BOOT/grubx64.efi");
 695                         fl_remove("/EFI/BOOT/grubx64_real.efi");
 696                         fl_remove("/EFI/BOOT/MokManager.efi");
 697             fl_remove("/ENROLL_THIS_KEY_IN_MOKMANAGER.cer");
 698
 699                         file = fl_fopen("/EFI/BOOT/BOOTX64.EFI", "wb");
 700                         Log("Open bootx64 efi file %p ", file);
 701                         if (file)
 702                         {
 703                                 if (filebuf)
 704                                 {
 705                                         fl_fwrite(filebuf, 1, size, file);
 706                                 }
 707
 708                                 fl_fflush(file);
 709                                 fl_fclose(file);
 710                         }
 711
 712                         if (filebuf)
 713                         {
 714                                 free(filebuf);
 715                         }
 716                 }
 717
 718         file = fl_fopen("/EFI/BOOT/grubia32_real.efi", "rb");
 719         Log("Open ventoy efi file %p ", file);
 720         if (file)
 721         {
 722             fl_fseek(file, 0, SEEK_END);
 723             size = (int)fl_ftell(file);
 724             fl_fseek(file, 0, SEEK_SET);
 725
 726             Log("ventoy efi file size %d ...", size);
 727
 728             filebuf = (char *)malloc(size);
 729             if (filebuf)
 730             {
 731                 fl_fread(filebuf, 1, size, file);
 732             }
 733
 734             fl_fclose(file);
 735
 736             Log("Now delete all efi files ...");
 737             fl_remove("/EFI/BOOT/BOOTIA32.EFI");
 738             fl_remove("/EFI/BOOT/grubia32.efi");
 739             fl_remove("/EFI/BOOT/grubia32_real.efi");
 740             fl_remove("/EFI/BOOT/mmia32.efi");
 741
 742             file = fl_fopen("/EFI/BOOT/BOOTIA32.EFI", "wb");
 743             Log("Open bootia32 efi file %p ", file);
 744             if (file)
 745             {
 746                 if (filebuf)
 747                 {
 748                     fl_fwrite(filebuf, 1, size, file);
 749                 }
 750
 751                 fl_fflush(file);
 752                 fl_fclose(file);
 753             }
 754
 755             if (filebuf)
 756             {
 757                 free(filebuf);
 758             }
 759         }
 760
 761         }
 762         else
 763         {
 764                 rc = 1;
 765         }
 766
 767         fl_shutdown();
 768
 769         return rc;
 770 }
 771
 772
 773
 774 static int disk_xz_flush(void *src, unsigned int size)
 775 {
 776     unsigned int i;
 777     BYTE *buf = (BYTE *)src;
 778
 779     for (i = 0; i < size; i++)
 780     {
 781         *g_part_img_pos = *buf++;
 782
 783         g_disk_unxz_len++;
 784         if ((g_disk_unxz_len % SIZE_1MB) == 0)
 785         {
 786             g_part_img_pos = g_part_img_buf[g_disk_unxz_len / SIZE_1MB];
 787         }
 788         else
 789         {
 790             g_part_img_pos++;
 791         }
 792     }
 793
 794     return (int)size;
 795 }
 796
 797 static void unxz_error(char *x)
 798 {
 799     Log("%s", x);
 800 }
 801
 802 static BOOL TryWritePart2(HANDLE hDrive, UINT64 StartSectorId)
 803 {
 804     BOOL bRet;
 805     DWORD TrySize = 16 * 1024;
 806     DWORD dwSize;
 807     BYTE *Buffer = NULL;
 808     unsigned char *data = NULL;
 809     LARGE_INTEGER liCurrentPosition;
 810
 811     liCurrentPosition.QuadPart = StartSectorId * 512;
 812     SetFilePointerEx(hDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 813
 814     Buffer = malloc(TrySize);
 815
 816     bRet = WriteFile(hDrive, Buffer, TrySize, &dwSize, NULL);
 817
 818     free(Buffer);
 819
 820     Log("Try write part2 bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 821
 822     if (bRet && dwSize == TrySize)
 823     {
 824         return TRUE;
 825     }
 826
 827     return FALSE;
 828 }
 829
 830 static int FormatPart2Fat(HANDLE hDrive, UINT64 StartSectorId)
 831 {
 832     int i;
 833     int rc = 0;
 834     int len = 0;
 835     int writelen = 0;
 836     int partwrite = 0;
 837     int Pos = PT_WRITE_VENTOY_START;
 838     DWORD dwSize = 0;
 839     BOOL bRet;
 840     unsigned char *data = NULL;
 841     LARGE_INTEGER liCurrentPosition;
 842         LARGE_INTEGER liNewPosition;
 843     BYTE *CheckBuf = NULL;
 844
 845         Log("FormatPart2Fat %llu...", StartSectorId);
 846
 847     CheckBuf = malloc(SIZE_1MB);
 848     if (!CheckBuf)
 849     {
 850         Log("Failed to malloc check buf");
 851         return 1;
 852     }
 853
 854     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&data, &len);
 855     if (rc)
 856     {
 857         Log("Failed to read img file %p %u", data, len);
 858         free(CheckBuf);
 859         return 1;
 860     }
 861
 862     liCurrentPosition.QuadPart = StartSectorId * 512;
 863     SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
 864
 865     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
 866
 867     g_part_img_buf[0] = (BYTE *)malloc(VENTOY_EFI_PART_SIZE);
 868     if (g_part_img_buf[0])
 869     {
 870         Log("Malloc whole img buffer success, now decompress ...");
 871         unxz(data, len, NULL, NULL, g_part_img_buf[0], &writelen, unxz_error);
 872
 873         if (len == writelen)
 874         {
 875             Log("decompress finished success");
 876
 877                         VentoyProcSecureBoot(g_SecureBoot);
 878
 879             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 880             {
 881                 dwSize = 0;
 882                                 bRet = WriteFile(hDrive, g_part_img_buf[0] + i * SIZE_1MB, SIZE_1MB, &dwSize, NULL);
 883                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 884
 885                 if (!bRet)
 886                 {
 887                     rc = 1;
 888                     goto End;
 889                 }
 890
 891                 PROGRESS_BAR_SET_POS(Pos);
 892                 if (i % 2 == 0)
 893                 {
 894                     Pos++;
 895                 }
 896             }
 897
 898             //Read and check the data
 899             liCurrentPosition.QuadPart = StartSectorId * 512;
 900             SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
 901
 902             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 903             {
 904                 bRet = ReadFile(hDrive, CheckBuf, SIZE_1MB, &dwSize, NULL);
 905                 Log("Read part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 906
 907                 if (!bRet || memcmp(CheckBuf, g_part_img_buf[0] + i * SIZE_1MB, SIZE_1MB))
 908                 {
 909                     Log("### [Check Fail] The data write and read does not match");
 910                     rc = 1;
 911                     goto End;
 912                 }
 913
 914                 PROGRESS_BAR_SET_POS(Pos);
 915                 if (i % 2 == 0)
 916                 {
 917                     Pos++;
 918                 }
 919             }
 920         }
 921         else
 922         {
 923             rc = 1;
 924             Log("decompress finished failed");
 925             goto End;
 926         }
 927     }
 928     else
 929     {
 930         Log("Failed to malloc whole img size %u, now split it", VENTOY_EFI_PART_SIZE);
 931
 932         partwrite = 1;
 933         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 934         {
 935             g_part_img_buf[i] = (BYTE *)malloc(SIZE_1MB);
 936             if (g_part_img_buf[i] == NULL)
 937             {
 938                 rc = 1;
 939                 goto End;
 940             }
 941         }
 942
 943         Log("Malloc part img buffer success, now decompress ...");
 944
 945         g_part_img_pos = g_part_img_buf[0];
 946
 947         unxz(data, len, NULL, disk_xz_flush, NULL, NULL, unxz_error);
 948
 949         if (g_disk_unxz_len == VENTOY_EFI_PART_SIZE)
 950         {
 951             Log("decompress finished success");
 952
 953                         VentoyProcSecureBoot(g_SecureBoot);
 954
 955             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 956             {
 957                 dwSize = 0;
 958                 bRet = WriteFile(hDrive, g_part_img_buf[i], SIZE_1MB, &dwSize, NULL);
 959                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 960
 961                 if (!bRet)
 962                 {
 963                     rc = 1;
 964                     goto End;
 965                 }
 966
 967                 PROGRESS_BAR_SET_POS(Pos);
 968                 if (i % 2 == 0)
 969                 {
 970                     Pos++;
 971                 }
 972             }
 973
 974             //Read and check the data
 975             liCurrentPosition.QuadPart = StartSectorId * 512;
 976             SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
 977
 978             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 979             {
 980                 bRet = ReadFile(hDrive, CheckBuf, SIZE_1MB, &dwSize, NULL);
 981                 Log("Read part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 982
 983                 if (!bRet || memcmp(CheckBuf, g_part_img_buf[i], SIZE_1MB))
 984                 {
 985                     Log("### [Check Fail] The data write and read does not match");
 986                     rc = 1;
 987                     goto End;
 988                 }
 989
 990                 PROGRESS_BAR_SET_POS(Pos);
 991                 if (i % 2 == 0)
 992                 {
 993                     Pos++;
 994                 }
 995             }
 996         }
 997         else
 998         {
 999             rc = 1;
1000             Log("decompress finished failed");
1001             goto End;
1002         }
1003     }
1004
1005 End:
1006
1007     if (data) free(data);
1008     if (CheckBuf)free(CheckBuf);
1009
1010     if (partwrite)
1011     {
1012         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1013         {
1014             if (g_part_img_buf[i]) free(g_part_img_buf[i]);
1015         }
1016     }
1017     else
1018     {
1019         if (g_part_img_buf[0]) free(g_part_img_buf[0]);
1020     }
1021
1022     return rc;
1023 }
1024
1025 static int WriteGrubStage1ToPhyDrive(HANDLE hDrive, int PartStyle)
1026 {
1027     int Len = 0;
1028     int readLen = 0;
1029     BOOL bRet;
1030     DWORD dwSize;
1031     BYTE *ImgBuf = NULL;
1032     BYTE *RawBuf = NULL;
1033
1034     Log("WriteGrubStage1ToPhyDrive ...");
1035
1036     RawBuf = (BYTE *)malloc(SIZE_1MB);
1037     if (!RawBuf)
1038     {
1039         return 1;
1040     }
1041
1042     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1043     {
1044         Log("Failed to read stage1 img");
1045         free(RawBuf);
1046         return 1;
1047     }
1048
1049     unxz(ImgBuf, Len, NULL, NULL, RawBuf, &readLen, unxz_error);
1050
1051     if (PartStyle)
1052     {
1053         Log("Write GPT stage1 ...");
1054         RawBuf[500] = 35;//update blocklist
1055         SetFilePointer(hDrive, 512 * 34, NULL, FILE_BEGIN);
1056         bRet = WriteFile(hDrive, RawBuf, SIZE_1MB - 512 * 34, &dwSize, NULL);
1057     }
1058     else
1059     {
1060         Log("Write MBR stage1 ...");
1061         SetFilePointer(hDrive, 512, NULL, FILE_BEGIN);
1062         bRet = WriteFile(hDrive, RawBuf, SIZE_1MB - 512, &dwSize, NULL);
1063     }
1064
1065     Log("WriteFile Ret:%u dwSize:%u ErrCode:%u", bRet, dwSize, GetLastError());
1066
1067     free(RawBuf);
1068     free(ImgBuf);
1069     return 0;
1070 }
1071
1072
1073
1074 static int FormatPart1exFAT(UINT64 DiskSizeBytes)
1075 {
1076     MKFS_PARM Option;
1077     FRESULT Ret;
1078
1079     Option.fmt = FM_EXFAT;
1080     Option.n_fat = 1;
1081     Option.align = 8;
1082     Option.n_root = 1;
1083
1084     // < 32GB select 32KB as cluster size
1085     // > 32GB select 128KB as cluster size
1086     if (DiskSizeBytes / 1024 / 1024 / 1024 <= 32)
1087     {
1088         Option.au_size = 32768;
1089     }
1090     else
1091     {
1092         Option.au_size = 131072;
1093     }
1094
1095     Log("Formatting Part1 exFAT ...");
1096
1097     Ret = f_mkfs(TEXT("0:"), &Option, 0, 8 * 1024 * 1024);
1098     if (FR_OK == Ret)
1099     {
1100         Log("Formatting Part1 exFAT success");
1101         return 0;
1102     }
1103     else
1104     {
1105         Log("Formatting Part1 exFAT failed");
1106         return 1;
1107     }
1108 }
1109
1110
1111
1112 int ClearVentoyFromPhyDrive(HWND hWnd, PHY_DRIVE_INFO *pPhyDrive, char *pDrvLetter)
1113 {
1114     int i;
1115     int rc = 0;
1116     int state = 0;
1117     HANDLE hDrive;
1118     DWORD dwSize;
1119     BOOL bRet;
1120     CHAR MountDrive;
1121     CHAR DriveName[] = "?:\\";
1122     CHAR DriveLetters[MAX_PATH] = { 0 };
1123     LARGE_INTEGER liCurrentPosition;
1124     char *pTmpBuf = NULL;
1125     MBR_HEAD MBR;
1126
1127     *pDrvLetter = 0;
1128
1129     Log("ClearVentoyFromPhyDrive PhyDrive%d <<%s %s %dGB>>",
1130         pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1131         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1132
1133     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1134
1135     Log("Lock disk for clean ............................. ");
1136
1137     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1138     if (hDrive == INVALID_HANDLE_VALUE)
1139     {
1140         Log("Failed to open physical disk");
1141         return 1;
1142     }
1143
1144     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1145
1146     if (DriveLetters[0] == 0)
1147     {
1148         Log("No drive letter was assigned...");
1149         DriveName[0] = GetFirstUnusedDriveLetter();
1150         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1151     }
1152     else
1153     {
1154         // Unmount all mounted volumes that belong to this drive
1155         // Do it in reverse so that we always end on the first volume letter
1156         for (i = (int)strlen(DriveLetters); i > 0; i--)
1157         {
1158             DriveName[0] = DriveLetters[i - 1];
1159             bRet = DeleteVolumeMountPointA(DriveName);
1160             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1161         }
1162     }
1163
1164     MountDrive = DriveName[0];
1165     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1166
1167     // It kind of blows, but we have to relinquish access to the physical drive
1168     // for VDS to be able to delete the partitions that reside on it...
1169     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1170     CHECK_CLOSE_HANDLE(hDrive);
1171
1172     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1173
1174     if (!VDS_DeleteAllPartitions(pPhyDrive->PhyDrive))
1175     {
1176         Log("Notice: Could not delete partitions: %u", GetLastError());
1177     }
1178
1179     Log("Deleting all partitions ......................... OK");
1180
1181     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1182
1183     Log("Lock disk for write ............................. ");
1184     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1185     if (hDrive == INVALID_HANDLE_VALUE)
1186     {
1187         Log("Failed to GetPhysicalHandle for write.");
1188         rc = 1;
1189         goto End;
1190     }
1191
1192     // clear first and last 2MB space
1193     pTmpBuf = malloc(SIZE_2MB);
1194     if (!pTmpBuf)
1195     {
1196         Log("Failed to alloc memory.");
1197         rc = 1;
1198         goto End;
1199     }
1200     memset(pTmpBuf, 0, SIZE_2MB);
1201
1202     SET_FILE_POS(512);
1203     bRet = WriteFile(hDrive, pTmpBuf, SIZE_2MB - 512, &dwSize, NULL);
1204     Log("Write fisrt 1MB ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1205     if (!bRet)
1206     {
1207         rc = 1;
1208         goto End;
1209     }
1210
1211     SET_FILE_POS(pPhyDrive->SizeInBytes - SIZE_2MB);
1212     bRet = WriteFile(hDrive, pTmpBuf, SIZE_2MB, &dwSize, NULL);
1213     Log("Write 2nd 1MB ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1214     if (!bRet)
1215     {
1216         rc = 1;
1217         goto End;
1218     }
1219
1220     SET_FILE_POS(0);
1221
1222     if (pPhyDrive->SizeInBytes > 2199023255552ULL)
1223     {
1224         VTOY_GPT_INFO *pGptInfo;
1225         VTOY_GPT_HDR BackupHead;
1226         LARGE_INTEGER liCurrentPosition;
1227
1228         pGptInfo = (VTOY_GPT_INFO *)pTmpBuf;
1229
1230         VentoyFillWholeGpt(pPhyDrive->SizeInBytes, pGptInfo);
1231
1232         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1233         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1234         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1235         {
1236             rc = 1;
1237             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1238             goto End;
1239         }
1240
1241         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1242         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1243         {
1244             rc = 1;
1245             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1246             goto End;
1247         }
1248
1249         SET_FILE_POS(0);
1250         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1251         {
1252             rc = 1;
1253             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1254             goto End;
1255         }
1256
1257         Log("Write GPT Info OK ...");
1258     }
1259     else
1260     {
1261         bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
1262         Log("Read MBR ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1263         if (!bRet)
1264         {
1265             rc = 1;
1266             goto End;
1267         }
1268
1269         //clear boot code and partition table (reserved disk signature)
1270         memset(MBR.BootCode, 0, 440);
1271         memset(MBR.PartTbl, 0, sizeof(MBR.PartTbl));
1272
1273         VentoyFillMBRLocation(pPhyDrive->SizeInBytes, 2048, (UINT32)(pPhyDrive->SizeInBytes / 512 - 2048), MBR.PartTbl);
1274
1275         MBR.PartTbl[0].Active = 0x00; // bootable
1276         MBR.PartTbl[0].FsFlag = 0x07; // exFAT/NTFS/HPFS
1277
1278         SET_FILE_POS(0);
1279         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
1280         Log("Write MBR ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1281         if (!bRet)
1282         {
1283             rc = 1;
1284             goto End;
1285         }
1286     }
1287
1288     Log("Clear Ventoy successfully finished");
1289
1290         //Refresh Drive Layout
1291         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1292
1293 End:
1294
1295     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1296
1297     if (pTmpBuf)
1298     {
1299         free(pTmpBuf);
1300     }
1301
1302     if (rc == 0)
1303     {
1304         Log("Mounting Ventoy Partition ....................... ");
1305         Sleep(1000);
1306
1307         state = 0;
1308         memset(DriveLetters, 0, sizeof(DriveLetters));
1309         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1310         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1311
1312         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1313         {
1314             DriveName[0] = DriveLetters[i];
1315             Log("%s is ventoy part1, already mounted", DriveName);
1316             state = 1;
1317         }
1318
1319         if (state != 1)
1320         {
1321             Log("need to mount ventoy part1...");
1322             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, 2048, DriveLetters, sizeof(DriveLetters), FALSE))
1323             {
1324                 DriveName[0] = MountDrive;
1325                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1326                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1327
1328                 *pDrvLetter = MountDrive;
1329             }
1330             else
1331             {
1332                 Log("Failed to find ventoy volume");
1333             }
1334         }
1335
1336         Log("OK\n");
1337     }
1338     else
1339     {
1340         FindProcessOccupyDisk(hDrive, pPhyDrive);
1341     }
1342
1343     CHECK_CLOSE_HANDLE(hDrive);
1344     return rc;
1345 }
1346
1347 int InstallVentoy2FileImage(PHY_DRIVE_INFO *pPhyDrive, int PartStyle)
1348 {
1349     int i;
1350     int rc = 1;
1351     int Len = 0;
1352     int dataLen = 0;
1353     UINT size = 0;
1354     UINT segnum = 0;
1355     UINT32 chksum = 0;
1356     UINT64 data_offset = 0;
1357     UINT64 Part2StartSector = 0;
1358     UINT64 Part1StartSector = 0;
1359     UINT64 Part1SectorCount = 0;
1360     UINT8 *pData = NULL;
1361     UINT8 *pBkGptPartTbl = NULL;
1362     BYTE *ImgBuf = NULL;
1363     MBR_HEAD *pMBR = NULL;
1364     VTSI_FOOTER *pImgFooter = NULL;
1365     VTSI_SEGMENT *pSegment = NULL;
1366     VTOY_GPT_INFO *pGptInfo = NULL;
1367     VTOY_GPT_HDR *pBkGptHdr = NULL;
1368     FILE *fp = NULL;
1369
1370     Log("InstallVentoy2FileImage %s PhyDrive%d <<%s %s %dGB>>",
1371         PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1372         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1373
1374     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1375
1376     size = SIZE_1MB + VENTOY_EFI_PART_SIZE + 33 * 512 + VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT) + sizeof(VTSI_FOOTER);
1377
1378     pData = (UINT8 *)malloc(size);
1379     if (!pData)
1380     {
1381         Log("malloc image buffer failed %d.", size);
1382         goto End;
1383     }
1384
1385     pImgFooter = (VTSI_FOOTER *)(pData + size - sizeof(VTSI_FOOTER));
1386     pSegment = (VTSI_SEGMENT *)((UINT8 *)pImgFooter - VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT));
1387     memset(pImgFooter, 0, sizeof(VTSI_FOOTER));
1388     memset(pSegment, 0, VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT));
1389
1390     PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START);
1391
1392     Log("Writing Boot Image ............................. ");
1393     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1394     {
1395         Log("Failed to read stage1 img");
1396         goto End;
1397     }
1398
1399     unxz(ImgBuf, Len, NULL, NULL, pData, &dataLen, unxz_error);
1400     SAFE_FREE(ImgBuf);
1401
1402     Log("decompress %s len:%d", VENTOY_FILE_STG1_IMG, dataLen);
1403
1404     if (PartStyle)
1405     {
1406         pData[500] = 35;//update blocklist
1407         memmove(pData + 34 * 512, pData, SIZE_1MB - 512 * 34);
1408         memset(pData, 0, 34 * 512);
1409
1410         pGptInfo = (VTOY_GPT_INFO *)pData;
1411         memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1412         VentoyFillGpt(pPhyDrive->SizeInBytes, pGptInfo);
1413
1414         pBkGptPartTbl = pData + SIZE_1MB + VENTOY_EFI_PART_SIZE;
1415         memset(pBkGptPartTbl, 0, 33 * 512);
1416
1417         memcpy(pBkGptPartTbl, pGptInfo->PartTbl, 32 * 512);
1418         pBkGptHdr = (VTOY_GPT_HDR *)(pBkGptPartTbl + 32 * 512);
1419         VentoyFillBackupGptHead(pGptInfo, pBkGptHdr);
1420
1421         Part1StartSector = pGptInfo->PartTbl[0].StartLBA;
1422         Part1SectorCount = pGptInfo->PartTbl[0].LastLBA - Part1StartSector + 1;
1423         Part2StartSector = pGptInfo->PartTbl[1].StartLBA;
1424
1425         Log("Write GPT Info OK ...");
1426     }
1427     else
1428     {
1429         memmove(pData + 512, pData, SIZE_1MB - 512);
1430         memset(pData, 0, 512);
1431
1432         pMBR = (MBR_HEAD *)pData;
1433         VentoyFillMBR(pPhyDrive->SizeInBytes, pMBR, PartStyle);
1434         Part1StartSector = pMBR->PartTbl[0].StartSectorId;
1435         Part1SectorCount = pMBR->PartTbl[0].SectorCount;
1436         Part2StartSector = pMBR->PartTbl[1].StartSectorId;
1437
1438         Log("Write MBR OK ...");
1439     }
1440
1441     Log("Writing EFI part Image ............................. ");
1442     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&ImgBuf, &Len);
1443     if (rc)
1444     {
1445         Log("Failed to read img file %p %u", ImgBuf, Len);
1446         goto End;
1447     }
1448
1449     PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START + 28);
1450     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
1451     unxz(ImgBuf, Len, NULL, NULL, pData + SIZE_1MB, &dataLen, unxz_error);
1452     if (dataLen == Len)
1453     {
1454         Log("decompress finished success");
1455         g_part_img_buf[0] = pData + SIZE_1MB;
1456
1457         VentoyProcSecureBoot(g_SecureBoot);
1458     }
1459     else
1460     {
1461         Log("decompress finished failed");
1462         goto End;
1463     }
1464
1465     fopen_s(&fp, "VentoySparseImg.vtsi", "wb+");
1466     if (!fp)
1467     {
1468         Log("Failed to create Ventoy img file");
1469         goto End;
1470     }
1471
1472     Log("Writing stage1 data ............................. ");
1473
1474     fwrite(pData, 1, SIZE_1MB, fp);
1475
1476     pSegment[0].disk_start_sector = 0;
1477     pSegment[0].sector_num = SIZE_1MB / 512;
1478     pSegment[0].data_offset = data_offset;
1479     data_offset += pSegment[0].sector_num * 512;
1480
1481     disk_io_set_param(INVALID_HANDLE_VALUE, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1482     disk_io_set_imghook(fp, pSegment + 1, VTSI_IMG_MAX_SEG - 1, data_offset);
1483
1484     Log("Formatting part1 exFAT ...");
1485     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1486     {
1487         Log("FormatPart1exFAT failed.");
1488         disk_io_reset_imghook(&segnum, &data_offset);
1489         goto End;
1490     }
1491
1492     disk_io_reset_imghook(&segnum, &data_offset);
1493     segnum++;
1494
1495     Log("current segment number:%d dataoff:%ld", segnum, (long)data_offset);
1496
1497     //write data
1498     Log("Writing part2 data ............................. ");
1499     fwrite(pData + SIZE_1MB, 1, VENTOY_EFI_PART_SIZE, fp);
1500     pSegment[segnum].disk_start_sector = Part2StartSector;
1501     pSegment[segnum].sector_num = VENTOY_EFI_PART_SIZE / 512;
1502     pSegment[segnum].data_offset = data_offset;
1503     data_offset += pSegment[segnum].sector_num * 512;
1504     segnum++;
1505
1506     if (PartStyle)
1507     {
1508         Log("Writing backup gpt table ............................. ");
1509         fwrite(pBkGptPartTbl, 1, 33 * 512, fp);
1510         pSegment[segnum].disk_start_sector = pPhyDrive->SizeInBytes / 512 - 33;
1511         pSegment[segnum].sector_num = 33;
1512         pSegment[segnum].data_offset = data_offset;
1513         data_offset += pSegment[segnum].sector_num * 512;
1514         segnum++;
1515     }
1516
1517     Log("Writing segment metadata ............................. ");
1518
1519     for (i = 0; i < (int)segnum; i++)
1520     {
1521         Log("SEG[%d]:  PhySector:%llu SectorNum:%llu DataOffset:%llu(sector:%llu)", i, pSegment[i].disk_start_sector, pSegment[i].sector_num,
1522             pSegment[i].data_offset, pSegment[i].data_offset / 512);
1523     }
1524
1525     dataLen = segnum * sizeof(VTSI_SEGMENT);
1526     fwrite(pSegment, 1, dataLen, fp);
1527
1528     if (dataLen % 512)
1529     {
1530         //pData + SIZE_1MB - 8192 is a temp data buffer with zero
1531         fwrite(pData + SIZE_1MB - 8192, 1, 512 - (dataLen % 512), fp);
1532     }
1533
1534     //Fill footer
1535     pImgFooter->magic = VTSI_IMG_MAGIC;
1536     pImgFooter->version = 1;
1537     pImgFooter->disk_size = pPhyDrive->SizeInBytes;
1538     memcpy(&pImgFooter->disk_signature, pPhyDrive->MBR.BootCode + 0x1b8, 4);
1539     pImgFooter->segment_num = segnum;
1540     pImgFooter->segment_offset = data_offset;
1541
1542     for (i = 0, chksum = 0; i < (int)(segnum * sizeof(VTSI_SEGMENT)); i++)
1543     {
1544         chksum += *((UINT8 *)pSegment + i);
1545     }
1546     pImgFooter->segment_chksum = ~chksum;
1547
1548     for (i = 0, chksum = 0; i < sizeof(VTSI_FOOTER); i++)
1549     {
1550         chksum += *((UINT8 *)pImgFooter + i);
1551     }
1552     pImgFooter->foot_chksum = ~chksum;
1553
1554     Log("Writing footer segnum(%u)  segoffset(%llu) ......................", segnum, data_offset);
1555     Log("disk_size=%llu disk_signature=%lx segment_offset=%llu", pImgFooter->disk_size, pImgFooter->disk_signature, pImgFooter->segment_offset);
1556
1557     fwrite(pImgFooter, 1, sizeof(VTSI_FOOTER), fp);
1558     fclose(fp);
1559
1560     Log("Writing Ventoy image file finished, the file size should be %llu .", data_offset + 512 + ((dataLen + 511) / 512 * 512));
1561
1562     rc = 0;
1563
1564 End:
1565
1566     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1567
1568     Log("retcode:%d\n", rc);
1569
1570     SAFE_FREE(pData);
1571     SAFE_FREE(ImgBuf);
1572
1573     return rc;
1574 }
1575
1576
1577 int InstallVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int PartStyle, int TryId)
1578 {
1579     int i;
1580     int rc = 0;
1581     int state = 0;
1582     HANDLE hDrive;
1583     DWORD dwSize;
1584     BOOL bRet;
1585     CHAR MountDrive;
1586     CHAR DriveName[] = "?:\\";
1587     CHAR DriveLetters[MAX_PATH] = { 0 };
1588     MBR_HEAD MBR;
1589     VTOY_GPT_INFO *pGptInfo = NULL;
1590     UINT64 Part1StartSector = 0;
1591     UINT64 Part1SectorCount = 0;
1592     UINT64 Part2StartSector = 0;
1593
1594     Log("InstallVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
1595         PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1596         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1597
1598     if (PartStyle)
1599     {
1600         pGptInfo = malloc(sizeof(VTOY_GPT_INFO));
1601         memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1602     }
1603
1604     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1605
1606     if (PartStyle)
1607     {
1608         VentoyFillGpt(pPhyDrive->SizeInBytes, pGptInfo);
1609         Part1StartSector = pGptInfo->PartTbl[0].StartLBA;
1610         Part1SectorCount = pGptInfo->PartTbl[0].LastLBA - Part1StartSector + 1;
1611         Part2StartSector = pGptInfo->PartTbl[1].StartLBA;
1612     }
1613     else
1614     {
1615         VentoyFillMBR(pPhyDrive->SizeInBytes, &MBR, PartStyle);
1616         Part1StartSector = MBR.PartTbl[0].StartSectorId;
1617         Part1SectorCount = MBR.PartTbl[0].SectorCount;
1618         Part2StartSector = MBR.PartTbl[1].StartSectorId;
1619     }
1620
1621     Log("Lock disk for clean ............................. ");
1622
1623     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1624     if (hDrive == INVALID_HANDLE_VALUE)
1625     {
1626         Log("Failed to open physical disk");
1627         free(pGptInfo);
1628         return 1;
1629     }
1630
1631     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1632
1633     if (DriveLetters[0] == 0)
1634     {
1635         Log("No drive letter was assigned...");
1636         DriveName[0] = GetFirstUnusedDriveLetter();
1637         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1638     }
1639     else
1640     {
1641         // Unmount all mounted volumes that belong to this drive
1642         // Do it in reverse so that we always end on the first volume letter
1643         for (i = (int)strlen(DriveLetters); i > 0; i--)
1644         {
1645             DriveName[0] = DriveLetters[i - 1];
1646             bRet = DeleteVolumeMountPointA(DriveName);
1647             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1648         }
1649     }
1650
1651     MountDrive = DriveName[0];
1652     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1653
1654     // It kind of blows, but we have to relinquish access to the physical drive
1655     // for VDS to be able to delete the partitions that reside on it...
1656     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1657     CHECK_CLOSE_HANDLE(hDrive);
1658
1659     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1660
1661     if (!VDS_DeleteAllPartitions(pPhyDrive->PhyDrive))
1662     {
1663         Log("Notice: Could not delete partitions: 0x%x", GetLastError());
1664                 rc = 1;
1665                 goto End;
1666     }
1667
1668     Log("Deleting all partitions ......................... OK");
1669
1670     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1671
1672     Log("Lock disk for write ............................. ");
1673     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1674     if (hDrive == INVALID_HANDLE_VALUE)
1675     {
1676         Log("Failed to GetPhysicalHandle for write.");
1677         rc = 1;
1678         goto End;
1679     }
1680
1681     //Refresh Drive Layout
1682     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1683
1684     disk_io_set_param(hDrive, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1685
1686     PROGRESS_BAR_SET_POS(PT_FORMAT_PART1);
1687
1688     if (PartStyle == 1 && pPhyDrive->PartStyle == 0)
1689     {
1690         Log("Wait for format part1 ...");
1691         Sleep(1000 * 5);
1692     }
1693
1694     Log("Formatting part1 exFAT ...");
1695     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1696     {
1697         Log("FormatPart1exFAT failed.");
1698         rc = 1;
1699         goto End;
1700     }
1701
1702     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1703     Log("Writing part2 FAT img ...");
1704
1705     if (0 != FormatPart2Fat(hDrive, Part2StartSector))
1706     {
1707         Log("FormatPart2Fat failed.");
1708         rc = 1;
1709         goto End;
1710     }
1711
1712     PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
1713     Log("Writing Boot Image ............................. ");
1714     if (WriteGrubStage1ToPhyDrive(hDrive, PartStyle) != 0)
1715     {
1716         Log("WriteGrubStage1ToPhyDrive failed.");
1717         rc = 1;
1718         goto End;
1719     }
1720
1721     PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
1722     Log("Writing Partition Table ........................ ");
1723     SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1724
1725     if (PartStyle)
1726     {
1727         VTOY_GPT_HDR BackupHead;
1728         LARGE_INTEGER liCurrentPosition;
1729
1730         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1731         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1732         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1733         {
1734             rc = 1;
1735             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1736             goto End;
1737         }
1738
1739         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1740         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1741         {
1742             rc = 1;
1743             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1744             goto End;
1745         }
1746
1747         SET_FILE_POS(0);
1748         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1749         {
1750             rc = 1;
1751             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1752             goto End;
1753         }
1754
1755         Log("Write GPT Info OK ...");
1756         memcpy(&(pPhyDrive->MBR), &(pGptInfo->MBR), 512);
1757     }
1758     else
1759     {
1760         if (!WriteFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL))
1761         {
1762             rc = 1;
1763             Log("Write MBR Failed, dwSize:%u ErrCode:%u", dwSize, GetLastError());
1764             goto End;
1765         }
1766         Log("Write MBR OK ...");
1767         memcpy(&(pPhyDrive->MBR), &MBR, 512);
1768     }
1769
1770     //Refresh Drive Layout
1771     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1772
1773 End:
1774
1775     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1776
1777     if (rc == 0)
1778     {
1779         Log("Mounting Ventoy Partition ....................... ");
1780         Sleep(1000);
1781
1782         state = 0;
1783         memset(DriveLetters, 0, sizeof(DriveLetters));
1784         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1785         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1786
1787         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1788         {
1789             DriveName[0] = DriveLetters[i];
1790             if (IsVentoyLogicalDrive(DriveName[0]))
1791             {
1792                 Log("%s is ventoy part2, delete mountpoint", DriveName);
1793                 DeleteVolumeMountPointA(DriveName);
1794             }
1795             else
1796             {
1797                 Log("%s is ventoy part1, already mounted", DriveName);
1798                 state = 1;
1799             }
1800         }
1801
1802         if (state != 1)
1803         {
1804             Log("need to mount ventoy part1...");
1805
1806             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, Part1StartSector, DriveLetters, sizeof(DriveLetters), FALSE))
1807             {
1808                 DriveName[0] = MountDrive;
1809                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1810                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1811             }
1812             else
1813             {
1814                 Log("Failed to find ventoy volume");
1815             }
1816         }
1817         Log("OK\n");
1818     }
1819     else
1820     {
1821         FindProcessOccupyDisk(hDrive, pPhyDrive);
1822     }
1823
1824     if (pGptInfo)
1825     {
1826         free(pGptInfo);
1827     }
1828
1829     CHECK_CLOSE_HANDLE(hDrive);
1830     return rc;
1831 }
1832
1833 static BOOL BackupDataBeforeCleanDisk(int PhyDrive, UINT64 DiskSize, BYTE **pBackup)
1834 {
1835         DWORD dwSize;
1836         DWORD dwStatus;
1837         BOOL Return = FALSE;
1838         BOOL ret = FALSE;
1839         BYTE *backup = NULL;
1840         HANDLE hDrive = INVALID_HANDLE_VALUE;
1841         LARGE_INTEGER liCurPosition;
1842         LARGE_INTEGER liNewPosition;
1843
1844         Log("BackupDataBeforeCleanDisk %d", PhyDrive);
1845
1846         backup = malloc(SIZE_1MB * 4);
1847         if (!backup)
1848         {
1849                 goto out;
1850         }
1851
1852         hDrive = GetPhysicalHandle(PhyDrive, FALSE, FALSE, FALSE);
1853         if (hDrive == INVALID_HANDLE_VALUE)
1854         {
1855                 goto out;
1856         }
1857
1858         //read first 2MB
1859         dwStatus = SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1860         if (dwStatus != 0)
1861         {
1862                 goto out;
1863         }
1864
1865         dwSize = 0;
1866         ret = ReadFile(hDrive, backup, SIZE_2MB, &dwSize, NULL);
1867         if ((!ret) || (dwSize != SIZE_2MB))
1868         {
1869                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
1870                 goto out;
1871         }
1872
1873         //read last 2MB
1874         liCurPosition.QuadPart = DiskSize - SIZE_2MB;
1875         liNewPosition.QuadPart = 0;
1876         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
1877                 liNewPosition.QuadPart != liCurPosition.QuadPart)
1878         {
1879                 goto out;
1880         }
1881
1882         dwSize = 0;
1883         ret = ReadFile(hDrive, backup + SIZE_2MB, SIZE_2MB, &dwSize, NULL);
1884         if ((!ret) || (dwSize != SIZE_2MB))
1885         {
1886                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
1887                 goto out;
1888         }
1889
1890         *pBackup = backup;
1891         backup = NULL; //For don't free later
1892         Return = TRUE;
1893
1894 out:
1895         CHECK_CLOSE_HANDLE(hDrive);
1896         if (backup)
1897                 free(backup);
1898
1899         return Return;
1900 }
1901
1902
1903 static BOOL WriteBackupDataToDisk(HANDLE hDrive, UINT64 Offset, BYTE *Data, DWORD Length)
1904 {
1905         DWORD dwSize = 0;
1906         BOOL ret = FALSE;
1907         LARGE_INTEGER liCurPosition;
1908         LARGE_INTEGER liNewPosition;
1909
1910         Log("WriteBackupDataToDisk %llu %p %u", Offset, Data, Length);
1911
1912         liCurPosition.QuadPart = Offset;
1913         liNewPosition.QuadPart = 0;
1914         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
1915                 liNewPosition.QuadPart != liCurPosition.QuadPart)
1916         {
1917                 return FALSE;
1918         }
1919
1920         ret = WriteFile(hDrive, Data, Length, &dwSize, NULL);
1921         if ((!ret) || dwSize != Length)
1922         {
1923                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
1924                 return FALSE;
1925         }
1926
1927         Log("WriteBackupDataToDisk %llu %p %u success", Offset, Data, Length);
1928         return TRUE;
1929 }
1930
1931
1932 int UpdateVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int TryId)
1933 {
1934         int i;
1935         int rc = 0;
1936         int MaxRetry = 4;
1937         BOOL ForceMBR = FALSE;
1938         BOOL Esp2Basic = FALSE;
1939         BOOL ChangeAttr = FALSE;
1940         BOOL CleanDisk = FALSE;
1941         BOOL bWriteBack = TRUE;
1942         HANDLE hVolume;
1943         HANDLE hDrive;
1944         DWORD Status;
1945         DWORD dwSize;
1946         BOOL bRet;
1947         CHAR DriveName[] = "?:\\";
1948         CHAR DriveLetters[MAX_PATH] = { 0 };
1949         CHAR BackBinFile[MAX_PATH];
1950         UINT64 StartSector;
1951         UINT64 ReservedMB = 0;
1952         MBR_HEAD BootImg;
1953         MBR_HEAD MBR;
1954         BYTE *pBackup = NULL;
1955         VTOY_GPT_INFO *pGptInfo = NULL;
1956         UINT8 ReservedData[4096];
1957
1958         Log("UpdateVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
1959                 pPhyDrive->PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1960                 GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1961
1962         PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1963
1964         Log("Lock disk for umount ............................ ");
1965
1966         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1967         if (hDrive == INVALID_HANDLE_VALUE)
1968         {
1969                 Log("Failed to open physical disk");
1970                 return 1;
1971         }
1972
1973         if (pPhyDrive->PartStyle)
1974         {
1975                 pGptInfo = malloc(sizeof(VTOY_GPT_INFO));
1976                 if (!pGptInfo)
1977                 {
1978                         return 1;
1979                 }
1980
1981                 memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1982
1983                 // Read GPT Info
1984                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1985                 ReadFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL);
1986
1987                 //MBR will be used to compare with local boot image
1988                 memcpy(&MBR, &pGptInfo->MBR, sizeof(MBR_HEAD));
1989
1990                 StartSector = pGptInfo->PartTbl[1].StartLBA;
1991                 Log("GPT StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
1992
1993                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512) - 33) / 2048;
1994                 Log("GPT Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
1995         }
1996         else
1997         {
1998                 // Read MBR
1999                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2000                 ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
2001
2002                 StartSector = MBR.PartTbl[1].StartSectorId;
2003                 Log("MBR StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2004
2005                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512)) / 2048;
2006                 Log("MBR Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2007         }
2008
2009         //Read Reserved Data
2010         SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2011         ReadFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2012
2013         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
2014
2015         if (DriveLetters[0] == 0)
2016         {
2017                 Log("No drive letter was assigned...");
2018         }
2019         else
2020         {
2021                 // Unmount all mounted volumes that belong to this drive
2022                 // Do it in reverse so that we always end on the first volume letter
2023                 for (i = (int)strlen(DriveLetters); i > 0; i--)
2024                 {
2025                         DriveName[0] = DriveLetters[i - 1];
2026                         if (IsVentoyLogicalDrive(DriveName[0]))
2027                         {
2028                                 Log("%s is ventoy logical drive", DriveName);
2029                                 bRet = DeleteVolumeMountPointA(DriveName);
2030                                 Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, LASTERR);
2031                                 break;
2032                         }
2033                 }
2034         }
2035
2036         // It kind of blows, but we have to relinquish access to the physical drive
2037         // for VDS to be able to delete the partitions that reside on it...
2038         DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2039         CHECK_CLOSE_HANDLE(hDrive);
2040
2041         if (pPhyDrive->PartStyle == 1)
2042         {
2043                 Log("TryId=%d EFI GPT partition type is 0x%llx", TryId, pPhyDrive->Part2GPTAttr);
2044                 PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
2045
2046                 if (TryId == 1)
2047                 {
2048                         Log("Change GPT partition type to ESP");
2049                         if (VDS_ChangeVtoyEFI2ESP(pPhyDrive->PhyDrive, StartSector * 512))
2050                         {
2051                                 Esp2Basic = TRUE;
2052                                 Sleep(3000);
2053                         }
2054                 }
2055                 else if (TryId == 2)
2056                 {
2057                         Log("Change GPT partition attribute");
2058                         if (VDS_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, 0x8000000000000001))
2059                         {
2060                                 ChangeAttr = TRUE;
2061                                 Sleep(2000);
2062                         }
2063                 }
2064                 else if (TryId == 3)
2065                 {
2066                         Log("Clean disk GPT partition table");
2067                         if (BackupDataBeforeCleanDisk(pPhyDrive->PhyDrive, pPhyDrive->SizeInBytes, &pBackup))
2068                         {
2069                                 sprintf_s(BackBinFile, sizeof(BackBinFile), ".\\ventoy\\phydrive%d_%u_%d.bin",
2070                                         pPhyDrive->PhyDrive, GetCurrentProcessId(), g_backup_bin_index++);
2071                                 SaveBufToFile(BackBinFile, pBackup, 4 * SIZE_1MB);
2072                                 Log("Save backup data to %s", BackBinFile);
2073
2074                                 Log("Success to backup data before clean");
2075                                 CleanDisk = TRUE;
2076                                 if (!VDS_CleanDisk(pPhyDrive->PhyDrive))
2077                                 {
2078                                         Sleep(3000);
2079                                         DSPT_CleanDisk(pPhyDrive->PhyDrive);
2080                                 }
2081                                 Sleep(3000);
2082                         }
2083                         else
2084                         {
2085                                 Log("Failed to backup data before clean");
2086                         }
2087                 }
2088         }
2089
2090     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
2091
2092     Log("Lock disk for update ............................ ");
2093     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2094     if (hDrive == INVALID_HANDLE_VALUE)
2095     {
2096         Log("Failed to GetPhysicalHandle for write.");
2097         rc = 1;
2098         goto End;
2099     }
2100
2101     PROGRESS_BAR_SET_POS(PT_LOCK_VOLUME);
2102
2103     Log("Lock volume for update .......................... ");
2104     hVolume = INVALID_HANDLE_VALUE;
2105
2106         //If we change VTOYEFI to ESP, it can not have s volume name, so don't try to get it.
2107         if (CleanDisk)
2108         {
2109                 //writeback the last 2MB
2110                 if (!WriteBackupDataToDisk(hDrive, pPhyDrive->SizeInBytes - SIZE_2MB, pBackup + SIZE_2MB, SIZE_2MB))
2111                 {
2112                         bWriteBack = FALSE;
2113                 }
2114                 Status = ERROR_NOT_FOUND;
2115         }
2116         else if (Esp2Basic)
2117         {
2118                 Status = ERROR_NOT_FOUND;
2119         }
2120         else
2121         {
2122                 for (i = 0; i < MaxRetry; i++)
2123                 {
2124                         Status = GetVentoyVolumeName(pPhyDrive->PhyDrive, StartSector, DriveLetters, sizeof(DriveLetters), TRUE);
2125                         if (ERROR_SUCCESS == Status)
2126                         {
2127                                 break;
2128                         }
2129                         else
2130                         {
2131                                 Log("==== Volume not found, wait and retry %d... ====", i);
2132                                 Sleep(2);
2133                         }
2134                 }
2135         }
2136
2137     if (ERROR_SUCCESS == Status)
2138     {
2139         Log("Now lock and dismount volume <%s>", DriveLetters);
2140
2141         for (i = 0; i < MaxRetry; i++)
2142         {
2143             hVolume = CreateFileA(DriveLetters,
2144                 GENERIC_READ | GENERIC_WRITE,
2145                 FILE_SHARE_READ,
2146                 NULL,
2147                 OPEN_EXISTING,
2148                 FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
2149                 NULL);
2150
2151             if (hVolume == INVALID_HANDLE_VALUE)
2152             {
2153                 Log("Failed to create file volume, errcode:%u, wait and retry ...", LASTERR);
2154                 Sleep(2000);
2155             }
2156             else
2157             {
2158                 break;
2159             }
2160         }
2161
2162         if (hVolume == INVALID_HANDLE_VALUE)
2163         {
2164             Log("Failed to create file volume, errcode:%u", LASTERR);
2165         }
2166         else
2167         {
2168             bRet = DeviceIoControl(hVolume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2169             Log("FSCTL_LOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2170
2171             bRet = DeviceIoControl(hVolume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2172             Log("FSCTL_DISMOUNT_VOLUME bRet:%u code:%u", bRet, LASTERR);
2173         }
2174     }
2175     else if (ERROR_NOT_FOUND == Status)
2176     {
2177         Log("Volume not found, maybe not supported");
2178     }
2179     else
2180     {
2181         rc = 1;
2182         goto End;
2183     }
2184
2185         bRet = TryWritePart2(hDrive, StartSector);
2186         if (FALSE == bRet && Esp2Basic)
2187         {
2188                 Log("TryWritePart2 agagin ...");
2189                 Sleep(3000);
2190                 bRet = TryWritePart2(hDrive, StartSector);
2191         }
2192
2193         if (!bRet)
2194     {
2195                 if (pPhyDrive->PartStyle == 0)
2196                 {
2197                         ForceMBR = TRUE;
2198
2199                         Log("Try write failed, now delete partition 2 for MBR...");
2200                         CHECK_CLOSE_HANDLE(hDrive);
2201
2202                         Log("Now delete partition 2...");
2203                         VDS_DeleteVtoyEFIPartition(pPhyDrive->PhyDrive);
2204
2205                         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2206                         if (hDrive == INVALID_HANDLE_VALUE)
2207                         {
2208                                 Log("Failed to GetPhysicalHandle for write.");
2209                                 rc = 1;
2210                                 goto End;
2211                         }
2212                 }
2213                 else
2214                 {
2215                         Log("TryWritePart2 failed ....");
2216                         rc = 1;
2217                         goto End;
2218                 }
2219     }
2220
2221     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
2222
2223     Log("Write Ventoy to disk ............................ ");
2224     if (0 != FormatPart2Fat(hDrive, StartSector))
2225     {
2226         rc = 1;
2227         goto End;
2228     }
2229
2230     if (hVolume != INVALID_HANDLE_VALUE)
2231     {
2232         bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2233         Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2234         CHECK_CLOSE_HANDLE(hVolume);
2235     }
2236
2237     Log("Updating Boot Image ............................. ");
2238     if (WriteGrubStage1ToPhyDrive(hDrive, pPhyDrive->PartStyle) != 0)
2239     {
2240         rc = 1;
2241         goto End;
2242     }
2243
2244     //write reserved data
2245     SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2246     bRet = WriteFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2247     Log("Write resv data ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2248
2249     // Boot Image
2250     VentoyGetLocalBootImg(&BootImg);
2251
2252     // Use Old UUID
2253     memcpy(BootImg.BootCode + 0x180, MBR.BootCode + 0x180, 16);
2254     if (pPhyDrive->PartStyle)
2255     {
2256         BootImg.BootCode[92] = 0x22;
2257     }
2258
2259     if (ForceMBR == FALSE && memcmp(BootImg.BootCode, MBR.BootCode, 440) == 0)
2260     {
2261         Log("Boot image has no difference, no need to write.");
2262     }
2263     else
2264     {
2265         Log("Boot image need to write %u.", ForceMBR);
2266
2267         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2268
2269         memcpy(MBR.BootCode, BootImg.BootCode, 440);
2270         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2271         Log("Write Boot Image ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2272     }
2273
2274     if (pPhyDrive->PartStyle == 0)
2275     {
2276         if (0x00 == MBR.PartTbl[0].Active && 0x80 == MBR.PartTbl[1].Active)
2277         {
2278             Log("Need to chage 1st partition active and 2nd partition inactive.");
2279
2280             MBR.PartTbl[0].Active = 0x80;
2281             MBR.PartTbl[1].Active = 0x00;
2282
2283             SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2284             bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2285             Log("Write NEW MBR ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2286         }
2287     }
2288
2289         if (CleanDisk)
2290         {
2291                 if (!WriteBackupDataToDisk(hDrive, 4 * 512, pBackup + 4 * 512, SIZE_2MB - 4 * 512))
2292                 {
2293                         bWriteBack = FALSE;
2294                 }
2295
2296                 if (!WriteBackupDataToDisk(hDrive, 0, pBackup, 4 * 512))
2297                 {
2298                         bWriteBack = FALSE;
2299                 }
2300
2301                 free(pBackup);
2302
2303                 if (bWriteBack)
2304                 {
2305                         Log("Write success, now delete %s", BackBinFile);
2306                         DeleteFileA(BackBinFile);
2307                 }
2308         }
2309
2310     //Refresh Drive Layout
2311     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2312
2313 End:
2314
2315     if (rc == 0)
2316     {
2317         Log("OK");
2318     }
2319     else
2320     {
2321         FindProcessOccupyDisk(hDrive, pPhyDrive);
2322     }
2323
2324     CHECK_CLOSE_HANDLE(hDrive);
2325
2326     if (Esp2Basic)
2327     {
2328                 Log("Recover GPT partition type to basic");
2329                 VDS_ChangeVtoyEFI2Basic(pPhyDrive->PhyDrive, StartSector * 512);
2330     }
2331
2332         if (pPhyDrive->PartStyle == 1)
2333         {
2334                 if (ChangeAttr || ((pPhyDrive->Part2GPTAttr >> 56) != 0xC0))
2335                 {
2336                         Log("Change EFI partition attr %u <0x%llx> to <0x%llx>", ChangeAttr, pPhyDrive->Part2GPTAttr, 0xC000000000000001ULL);
2337                         if (VDS_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, 0xC000000000000001ULL))
2338                         {
2339                                 Log("Change EFI partition attr success");
2340                                 pPhyDrive->Part2GPTAttr = 0xC000000000000001ULL;
2341                         }
2342                         else
2343                         {
2344                                 Log("Change EFI partition attr failed");
2345                         }
2346                 }
2347         }
2348
2349     if (pGptInfo)
2350     {
2351         free(pGptInfo);
2352     }
2353
2354     return rc;
2355 }
2356
2357