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, but we continue.", GetLastError());
1664     }
1665
1666     Log("Deleting all partitions ......................... OK");
1667
1668     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1669
1670     Log("Lock disk for write ............................. ");
1671     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1672     if (hDrive == INVALID_HANDLE_VALUE)
1673     {
1674         Log("Failed to GetPhysicalHandle for write.");
1675         rc = 1;
1676         goto End;
1677     }
1678
1679     //Refresh Drive Layout
1680     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1681
1682     disk_io_set_param(hDrive, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1683
1684     PROGRESS_BAR_SET_POS(PT_FORMAT_PART1);
1685
1686     if (PartStyle == 1 && pPhyDrive->PartStyle == 0)
1687     {
1688         Log("Wait for format part1 ...");
1689         Sleep(1000 * 5);
1690     }
1691
1692     Log("Formatting part1 exFAT ...");
1693     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1694     {
1695         Log("FormatPart1exFAT failed.");
1696         rc = 1;
1697         goto End;
1698     }
1699
1700     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1701     Log("Writing part2 FAT img ...");
1702
1703     if (0 != FormatPart2Fat(hDrive, Part2StartSector))
1704     {
1705         Log("FormatPart2Fat failed.");
1706         rc = 1;
1707         goto End;
1708     }
1709
1710     PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
1711     Log("Writing Boot Image ............................. ");
1712     if (WriteGrubStage1ToPhyDrive(hDrive, PartStyle) != 0)
1713     {
1714         Log("WriteGrubStage1ToPhyDrive failed.");
1715         rc = 1;
1716         goto End;
1717     }
1718
1719     PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
1720     Log("Writing Partition Table ........................ ");
1721     SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1722
1723     if (PartStyle)
1724     {
1725         VTOY_GPT_HDR BackupHead;
1726         LARGE_INTEGER liCurrentPosition;
1727
1728         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1729         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1730         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1731         {
1732             rc = 1;
1733             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1734             goto End;
1735         }
1736
1737         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1738         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1739         {
1740             rc = 1;
1741             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1742             goto End;
1743         }
1744
1745         SET_FILE_POS(0);
1746         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1747         {
1748             rc = 1;
1749             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1750             goto End;
1751         }
1752
1753         Log("Write GPT Info OK ...");
1754         memcpy(&(pPhyDrive->MBR), &(pGptInfo->MBR), 512);
1755     }
1756     else
1757     {
1758         if (!WriteFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL))
1759         {
1760             rc = 1;
1761             Log("Write MBR Failed, dwSize:%u ErrCode:%u", dwSize, GetLastError());
1762             goto End;
1763         }
1764         Log("Write MBR OK ...");
1765         memcpy(&(pPhyDrive->MBR), &MBR, 512);
1766     }
1767
1768     //Refresh Drive Layout
1769     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1770
1771 End:
1772
1773     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1774
1775     if (rc == 0)
1776     {
1777         Log("Mounting Ventoy Partition ....................... ");
1778         Sleep(1000);
1779
1780         state = 0;
1781         memset(DriveLetters, 0, sizeof(DriveLetters));
1782         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1783         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1784
1785         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1786         {
1787             DriveName[0] = DriveLetters[i];
1788             if (IsVentoyLogicalDrive(DriveName[0]))
1789             {
1790                 Log("%s is ventoy part2, delete mountpoint", DriveName);
1791                 DeleteVolumeMountPointA(DriveName);
1792             }
1793             else
1794             {
1795                 Log("%s is ventoy part1, already mounted", DriveName);
1796                 state = 1;
1797             }
1798         }
1799
1800         if (state != 1)
1801         {
1802             Log("need to mount ventoy part1...");
1803
1804             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, Part1StartSector, DriveLetters, sizeof(DriveLetters), FALSE))
1805             {
1806                 DriveName[0] = MountDrive;
1807                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1808                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1809             }
1810             else
1811             {
1812                 Log("Failed to find ventoy volume");
1813             }
1814         }
1815         Log("OK\n");
1816     }
1817     else
1818     {
1819         FindProcessOccupyDisk(hDrive, pPhyDrive);
1820     }
1821
1822     if (pGptInfo)
1823     {
1824         free(pGptInfo);
1825     }
1826
1827     CHECK_CLOSE_HANDLE(hDrive);
1828     return rc;
1829 }
1830
1831 static BOOL BackupDataBeforeCleanDisk(int PhyDrive, UINT64 DiskSize, BYTE **pBackup)
1832 {
1833         DWORD dwSize;
1834         DWORD dwStatus;
1835         BOOL Return = FALSE;
1836         BOOL ret = FALSE;
1837         BYTE *backup = NULL;
1838         UINT64 offset;
1839         BYTE Buffer[512];
1840         HANDLE hDrive = INVALID_HANDLE_VALUE;
1841         LARGE_INTEGER liCurPosition;
1842         LARGE_INTEGER liNewPosition;
1843         VTOY_GPT_INFO *pGPT = NULL;
1844
1845         Log("BackupDataBeforeCleanDisk %d", PhyDrive);
1846
1847         // step1: check write access
1848         hDrive = GetPhysicalHandle(PhyDrive, TRUE, TRUE, FALSE);
1849         if (hDrive == INVALID_HANDLE_VALUE)
1850         {
1851                 Log("Failed to GetPhysicalHandle for write.");
1852                 goto out;
1853         }
1854
1855         liCurPosition.QuadPart = 2039 * 512;
1856         liNewPosition.QuadPart = 0;
1857         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
1858                 liNewPosition.QuadPart != liCurPosition.QuadPart)
1859         {
1860                 Log("SetFilePointer1 Failed %u", LASTERR);
1861                 goto out;
1862         }
1863
1864
1865         dwSize = 0;
1866         ret = ReadFile(hDrive, Buffer, 512, &dwSize, NULL);
1867         if ((!ret) || (dwSize != 512))
1868         {
1869                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
1870                 goto out;
1871         }
1872
1873
1874         liCurPosition.QuadPart = 2039 * 512;
1875         liNewPosition.QuadPart = 0;
1876         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
1877                 liNewPosition.QuadPart != liCurPosition.QuadPart)
1878         {
1879                 Log("SetFilePointer2 Failed %u", LASTERR);
1880                 goto out;
1881         }
1882
1883         ret = WriteFile(hDrive, Buffer, 512, &dwSize, NULL);
1884         if ((!ret) || dwSize != 512)
1885         {
1886                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
1887                 goto out;
1888         }
1889
1890         CHECK_CLOSE_HANDLE(hDrive);
1891         Log("Write access check success");
1892
1893
1894         //step2 backup 4MB data
1895         backup = malloc(SIZE_1MB * 4);
1896         if (!backup)
1897         {
1898                 goto out;
1899         }
1900
1901         hDrive = GetPhysicalHandle(PhyDrive, FALSE, FALSE, FALSE);
1902         if (hDrive == INVALID_HANDLE_VALUE)
1903         {
1904                 goto out;
1905         }
1906
1907         //read first 2MB
1908         dwStatus = SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1909         if (dwStatus != 0)
1910         {
1911                 goto out;
1912         }
1913
1914         dwSize = 0;
1915         ret = ReadFile(hDrive, backup, SIZE_2MB, &dwSize, NULL);
1916         if ((!ret) || (dwSize != SIZE_2MB))
1917         {
1918                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
1919                 goto out;
1920         }
1921
1922         pGPT = (VTOY_GPT_INFO *)backup;
1923         offset = pGPT->Head.EfiBackupLBA * 512;
1924         if (offset >= (DiskSize - SIZE_2MB) && offset < DiskSize)
1925         {
1926                 Log("EFI partition table check success");
1927         }
1928         else
1929         {
1930                 Log("Backup EFI LBA not in last 2MB range: %llu", pGPT->Head.EfiBackupLBA);
1931                 goto out;
1932         }
1933
1934         //read last 2MB
1935         liCurPosition.QuadPart = DiskSize - SIZE_2MB;
1936         liNewPosition.QuadPart = 0;
1937         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
1938                 liNewPosition.QuadPart != liCurPosition.QuadPart)
1939         {
1940                 goto out;
1941         }
1942
1943         dwSize = 0;
1944         ret = ReadFile(hDrive, backup + SIZE_2MB, SIZE_2MB, &dwSize, NULL);
1945         if ((!ret) || (dwSize != SIZE_2MB))
1946         {
1947                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
1948                 goto out;
1949         }
1950
1951         *pBackup = backup;
1952         backup = NULL; //For don't free later
1953         Return = TRUE;
1954
1955 out:
1956         CHECK_CLOSE_HANDLE(hDrive);
1957         if (backup)
1958                 free(backup);
1959
1960         return Return;
1961 }
1962
1963
1964 static BOOL WriteBackupDataToDisk(HANDLE hDrive, UINT64 Offset, BYTE *Data, DWORD Length)
1965 {
1966         DWORD dwSize = 0;
1967         BOOL ret = FALSE;
1968         LARGE_INTEGER liCurPosition;
1969         LARGE_INTEGER liNewPosition;
1970
1971         Log("WriteBackupDataToDisk %llu %p %u", Offset, Data, Length);
1972
1973         liCurPosition.QuadPart = Offset;
1974         liNewPosition.QuadPart = 0;
1975         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
1976                 liNewPosition.QuadPart != liCurPosition.QuadPart)
1977         {
1978                 return FALSE;
1979         }
1980
1981         ret = WriteFile(hDrive, Data, Length, &dwSize, NULL);
1982         if ((!ret) || dwSize != Length)
1983         {
1984                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
1985                 return FALSE;
1986         }
1987
1988         Log("WriteBackupDataToDisk %llu %p %u success", Offset, Data, Length);
1989         return TRUE;
1990 }
1991
1992
1993 int UpdateVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int TryId)
1994 {
1995         int i;
1996         int rc = 0;
1997         int MaxRetry = 4;
1998         BOOL ForceMBR = FALSE;
1999         BOOL Esp2Basic = FALSE;
2000         BOOL ChangeAttr = FALSE;
2001         BOOL CleanDisk = FALSE;
2002         BOOL bWriteBack = TRUE;
2003         HANDLE hVolume;
2004         HANDLE hDrive;
2005         DWORD Status;
2006         DWORD dwSize;
2007         BOOL bRet;
2008         CHAR DriveName[] = "?:\\";
2009         CHAR DriveLetters[MAX_PATH] = { 0 };
2010         CHAR BackBinFile[MAX_PATH];
2011         UINT64 StartSector;
2012         UINT64 ReservedMB = 0;
2013         MBR_HEAD BootImg;
2014         MBR_HEAD MBR;
2015         BYTE *pBackup = NULL;
2016         VTOY_GPT_INFO *pGptInfo = NULL;
2017         UINT8 ReservedData[4096];
2018
2019         Log("UpdateVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
2020                 pPhyDrive->PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
2021                 GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
2022
2023         PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
2024
2025         Log("Lock disk for umount ............................ ");
2026
2027         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
2028         if (hDrive == INVALID_HANDLE_VALUE)
2029         {
2030                 Log("Failed to open physical disk");
2031                 return 1;
2032         }
2033
2034         if (pPhyDrive->PartStyle)
2035         {
2036                 pGptInfo = malloc(sizeof(VTOY_GPT_INFO));
2037                 if (!pGptInfo)
2038                 {
2039                         return 1;
2040                 }
2041
2042                 memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
2043
2044                 // Read GPT Info
2045                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2046                 ReadFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL);
2047
2048                 //MBR will be used to compare with local boot image
2049                 memcpy(&MBR, &pGptInfo->MBR, sizeof(MBR_HEAD));
2050
2051                 StartSector = pGptInfo->PartTbl[1].StartLBA;
2052                 Log("GPT StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2053
2054                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512) - 33) / 2048;
2055                 Log("GPT Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2056         }
2057         else
2058         {
2059                 // Read MBR
2060                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2061                 ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
2062
2063                 StartSector = MBR.PartTbl[1].StartSectorId;
2064                 Log("MBR StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2065
2066                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512)) / 2048;
2067                 Log("MBR Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2068         }
2069
2070         //Read Reserved Data
2071         SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2072         ReadFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2073
2074         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
2075
2076         if (DriveLetters[0] == 0)
2077         {
2078                 Log("No drive letter was assigned...");
2079         }
2080         else
2081         {
2082                 // Unmount all mounted volumes that belong to this drive
2083                 // Do it in reverse so that we always end on the first volume letter
2084                 for (i = (int)strlen(DriveLetters); i > 0; i--)
2085                 {
2086                         DriveName[0] = DriveLetters[i - 1];
2087                         if (IsVentoyLogicalDrive(DriveName[0]))
2088                         {
2089                                 Log("%s is ventoy logical drive", DriveName);
2090                                 bRet = DeleteVolumeMountPointA(DriveName);
2091                                 Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, LASTERR);
2092                                 break;
2093                         }
2094                 }
2095         }
2096
2097         // It kind of blows, but we have to relinquish access to the physical drive
2098         // for VDS to be able to delete the partitions that reside on it...
2099         DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2100         CHECK_CLOSE_HANDLE(hDrive);
2101
2102         if (pPhyDrive->PartStyle == 1)
2103         {
2104                 Log("TryId=%d EFI GPT partition type is 0x%llx", TryId, pPhyDrive->Part2GPTAttr);
2105                 PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
2106
2107                 if (TryId == 1)
2108                 {
2109                         Log("Change GPT partition type to ESP");
2110                         if (VDS_ChangeVtoyEFI2ESP(pPhyDrive->PhyDrive, StartSector * 512))
2111                         {
2112                                 Esp2Basic = TRUE;
2113                                 Sleep(3000);
2114                         }
2115                 }
2116                 else if (TryId == 2)
2117                 {
2118                         Log("Change GPT partition attribute");
2119                         if (VDS_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, 0x8000000000000001))
2120                         {
2121                                 ChangeAttr = TRUE;
2122                                 Sleep(2000);
2123                         }
2124                 }
2125                 else if (TryId == 3)
2126                 {
2127                         Log("Clean disk GPT partition table");
2128                         if (BackupDataBeforeCleanDisk(pPhyDrive->PhyDrive, pPhyDrive->SizeInBytes, &pBackup))
2129                         {
2130                                 sprintf_s(BackBinFile, sizeof(BackBinFile), ".\\ventoy\\phydrive%d_%u_%d.bin",
2131                                         pPhyDrive->PhyDrive, GetCurrentProcessId(), g_backup_bin_index++);
2132                                 SaveBufToFile(BackBinFile, pBackup, 4 * SIZE_1MB);
2133                                 Log("Save backup data to %s", BackBinFile);
2134
2135                                 Log("Success to backup data before clean");
2136                                 CleanDisk = TRUE;
2137                                 if (!VDS_CleanDisk(pPhyDrive->PhyDrive))
2138                                 {
2139                                         Sleep(3000);
2140                                         DSPT_CleanDisk(pPhyDrive->PhyDrive);
2141                                 }
2142                                 Sleep(3000);
2143                         }
2144                         else
2145                         {
2146                                 Log("Failed to backup data before clean");
2147                         }
2148                 }
2149         }
2150
2151     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
2152
2153     Log("Lock disk for update ............................ ");
2154     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2155     if (hDrive == INVALID_HANDLE_VALUE)
2156     {
2157         Log("Failed to GetPhysicalHandle for write.");
2158         rc = 1;
2159         goto End;
2160     }
2161
2162     PROGRESS_BAR_SET_POS(PT_LOCK_VOLUME);
2163
2164     Log("Lock volume for update .......................... ");
2165     hVolume = INVALID_HANDLE_VALUE;
2166
2167         //If we change VTOYEFI to ESP, it can not have s volume name, so don't try to get it.
2168         if (CleanDisk)
2169         {
2170                 //writeback the last 2MB
2171                 if (!WriteBackupDataToDisk(hDrive, pPhyDrive->SizeInBytes - SIZE_2MB, pBackup + SIZE_2MB, SIZE_2MB))
2172                 {
2173                         bWriteBack = FALSE;
2174                 }
2175
2176                 //write the first 2MB except parttable
2177                 if (!WriteBackupDataToDisk(hDrive, 34 * 512, pBackup + 34 * 512, SIZE_2MB - 34 * 512))
2178                 {
2179                         bWriteBack = FALSE;
2180                 }
2181
2182                 Status = ERROR_NOT_FOUND;
2183         }
2184         else if (Esp2Basic)
2185         {
2186                 Status = ERROR_NOT_FOUND;
2187         }
2188         else
2189         {
2190                 for (i = 0; i < MaxRetry; i++)
2191                 {
2192                         Status = GetVentoyVolumeName(pPhyDrive->PhyDrive, StartSector, DriveLetters, sizeof(DriveLetters), TRUE);
2193                         if (ERROR_SUCCESS == Status)
2194                         {
2195                                 break;
2196                         }
2197                         else
2198                         {
2199                                 Log("==== Volume not found, wait and retry %d... ====", i);
2200                                 Sleep(2);
2201                         }
2202                 }
2203         }
2204
2205     if (ERROR_SUCCESS == Status)
2206     {
2207         Log("Now lock and dismount volume <%s>", DriveLetters);
2208
2209         for (i = 0; i < MaxRetry; i++)
2210         {
2211             hVolume = CreateFileA(DriveLetters,
2212                 GENERIC_READ | GENERIC_WRITE,
2213                 FILE_SHARE_READ,
2214                 NULL,
2215                 OPEN_EXISTING,
2216                 FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
2217                 NULL);
2218
2219             if (hVolume == INVALID_HANDLE_VALUE)
2220             {
2221                 Log("Failed to create file volume, errcode:%u, wait and retry ...", LASTERR);
2222                 Sleep(2000);
2223             }
2224             else
2225             {
2226                 break;
2227             }
2228         }
2229
2230         if (hVolume == INVALID_HANDLE_VALUE)
2231         {
2232             Log("Failed to create file volume, errcode:%u", LASTERR);
2233         }
2234         else
2235         {
2236             bRet = DeviceIoControl(hVolume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2237             Log("FSCTL_LOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2238
2239             bRet = DeviceIoControl(hVolume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2240             Log("FSCTL_DISMOUNT_VOLUME bRet:%u code:%u", bRet, LASTERR);
2241         }
2242     }
2243     else if (ERROR_NOT_FOUND == Status)
2244     {
2245         Log("Volume not found, maybe not supported");
2246     }
2247     else
2248     {
2249         rc = 1;
2250         goto End;
2251     }
2252
2253         bRet = TryWritePart2(hDrive, StartSector);
2254         if (FALSE == bRet && Esp2Basic)
2255         {
2256                 Log("TryWritePart2 agagin ...");
2257                 Sleep(3000);
2258                 bRet = TryWritePart2(hDrive, StartSector);
2259         }
2260
2261         if (!bRet)
2262     {
2263                 if (pPhyDrive->PartStyle == 0)
2264                 {
2265                         ForceMBR = TRUE;
2266
2267                         Log("Try write failed, now delete partition 2 for MBR...");
2268                         CHECK_CLOSE_HANDLE(hDrive);
2269
2270                         Log("Now delete partition 2...");
2271                         VDS_DeleteVtoyEFIPartition(pPhyDrive->PhyDrive);
2272
2273                         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2274                         if (hDrive == INVALID_HANDLE_VALUE)
2275                         {
2276                                 Log("Failed to GetPhysicalHandle for write.");
2277                                 rc = 1;
2278                                 goto End;
2279                         }
2280                 }
2281                 else
2282                 {
2283                         Log("TryWritePart2 failed ....");
2284                         rc = 1;
2285                         goto End;
2286                 }
2287     }
2288
2289     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
2290
2291     Log("Write Ventoy to disk ............................ ");
2292     if (0 != FormatPart2Fat(hDrive, StartSector))
2293     {
2294         rc = 1;
2295         goto End;
2296     }
2297
2298     if (hVolume != INVALID_HANDLE_VALUE)
2299     {
2300         bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2301         Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2302         CHECK_CLOSE_HANDLE(hVolume);
2303     }
2304
2305     Log("Updating Boot Image ............................. ");
2306     if (WriteGrubStage1ToPhyDrive(hDrive, pPhyDrive->PartStyle) != 0)
2307     {
2308         rc = 1;
2309         goto End;
2310     }
2311
2312     //write reserved data
2313     SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2314     bRet = WriteFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2315     Log("Write resv data ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2316
2317     // Boot Image
2318     VentoyGetLocalBootImg(&BootImg);
2319
2320     // Use Old UUID
2321     memcpy(BootImg.BootCode + 0x180, MBR.BootCode + 0x180, 16);
2322     if (pPhyDrive->PartStyle)
2323     {
2324         BootImg.BootCode[92] = 0x22;
2325     }
2326
2327     if (ForceMBR == FALSE && memcmp(BootImg.BootCode, MBR.BootCode, 440) == 0)
2328     {
2329         Log("Boot image has no difference, no need to write.");
2330     }
2331     else
2332     {
2333         Log("Boot image need to write %u.", ForceMBR);
2334
2335         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2336
2337         memcpy(MBR.BootCode, BootImg.BootCode, 440);
2338         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2339         Log("Write Boot Image ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2340     }
2341
2342     if (pPhyDrive->PartStyle == 0)
2343     {
2344         if (0x00 == MBR.PartTbl[0].Active && 0x80 == MBR.PartTbl[1].Active)
2345         {
2346             Log("Need to chage 1st partition active and 2nd partition inactive.");
2347
2348             MBR.PartTbl[0].Active = 0x80;
2349             MBR.PartTbl[1].Active = 0x00;
2350
2351             SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2352             bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2353             Log("Write NEW MBR ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2354         }
2355     }
2356
2357         if (CleanDisk)
2358         {
2359                 if (!WriteBackupDataToDisk(hDrive, 0, pBackup, 34 * 512))
2360                 {
2361                         bWriteBack = FALSE;
2362                 }
2363
2364                 free(pBackup);
2365
2366                 if (bWriteBack)
2367                 {
2368                         Log("Write backup data success, now delete %s", BackBinFile);
2369                         DeleteFileA(BackBinFile);
2370                 }
2371                 else
2372                 {
2373                         Log("Write backup data failed");
2374                 }
2375
2376                 Sleep(1000);
2377         }
2378
2379     //Refresh Drive Layout
2380     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2381
2382 End:
2383
2384         if (hVolume != INVALID_HANDLE_VALUE)
2385         {
2386                 bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2387                 Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2388                 CHECK_CLOSE_HANDLE(hVolume);
2389         }
2390
2391     if (rc == 0)
2392     {
2393         Log("OK");
2394     }
2395     else
2396     {
2397         FindProcessOccupyDisk(hDrive, pPhyDrive);
2398     }
2399
2400     CHECK_CLOSE_HANDLE(hDrive);
2401
2402     if (Esp2Basic)
2403     {
2404                 Log("Recover GPT partition type to basic");
2405                 VDS_ChangeVtoyEFI2Basic(pPhyDrive->PhyDrive, StartSector * 512);
2406     }
2407
2408         if (pPhyDrive->PartStyle == 1)
2409         {
2410                 if (ChangeAttr || ((pPhyDrive->Part2GPTAttr >> 56) != 0xC0))
2411                 {
2412                         Log("Change EFI partition attr %u <0x%llx> to <0x%llx>", ChangeAttr, pPhyDrive->Part2GPTAttr, 0xC000000000000001ULL);
2413                         if (VDS_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, 0xC000000000000001ULL))
2414                         {
2415                                 Log("Change EFI partition attr success");
2416                                 pPhyDrive->Part2GPTAttr = 0xC000000000000001ULL;
2417                         }
2418                         else
2419                         {
2420                                 Log("Change EFI partition attr failed");
2421                         }
2422                 }
2423         }
2424
2425     if (pGptInfo)
2426     {
2427         free(pGptInfo);
2428     }
2429
2430     return rc;
2431 }
2432
2433