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
  37 static BOOL WriteDataToPhyDisk(HANDLE hDrive, UINT64 Offset, VOID *buffer, DWORD len)
  38 {
  39         BOOL bRet;
  40         DWORD dwSize = 0;
  41         LARGE_INTEGER liCurPosition;
  42         LARGE_INTEGER liNewPosition;
  43
  44         liCurPosition.QuadPart = (LONGLONG)Offset;
  45         liNewPosition.QuadPart = 0;
  46         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
  47                 liNewPosition.QuadPart != liCurPosition.QuadPart)
  48         {
  49                 Log("SetFilePointerEx Failed %u", LASTERR);
  50                 return FALSE;
  51         }
  52
  53         bRet = WriteFile(hDrive, buffer, len, &dwSize, NULL);
  54         if (bRet == FALSE || dwSize != len)
  55         {
  56                 Log("Write file error %u %u", dwSize, LASTERR);
  57                 return FALSE;
  58         }
  59
  60         return TRUE;
  61 }
  62
  63
  64 static DWORD GetVentoyVolumeName(int PhyDrive, UINT64 StartSectorId, CHAR *NameBuf, UINT32 BufLen, BOOL DelSlash)
  65 {
  66     size_t len;
  67     BOOL bRet;
  68     DWORD dwSize;
  69     HANDLE hDrive;
  70     HANDLE hVolume;
  71     UINT64 PartOffset;
  72     DWORD Status = ERROR_NOT_FOUND;
  73     DISK_EXTENT *pExtents = NULL;
  74     CHAR VolumeName[MAX_PATH] = { 0 };
  75     VOLUME_DISK_EXTENTS DiskExtents;
  76
  77     PartOffset = 512ULL * StartSectorId;
  78
  79         Log("GetVentoyVolumeName PhyDrive %d SectorStart:%llu PartOffset:%llu", PhyDrive, (ULONGLONG)StartSectorId, (ULONGLONG)PartOffset);
  80
  81     hVolume = FindFirstVolumeA(VolumeName, sizeof(VolumeName));
  82     if (hVolume == INVALID_HANDLE_VALUE)
  83     {
  84         return 1;
  85     }
  86
  87     do {
  88
  89         len = strlen(VolumeName);
  90         Log("Find volume:%s", VolumeName);
  91
  92         VolumeName[len - 1] = 0;
  93
  94         hDrive = CreateFileA(VolumeName, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
  95         if (hDrive == INVALID_HANDLE_VALUE)
  96         {
  97             continue;
  98         }
  99
 100         bRet = DeviceIoControl(hDrive,
 101             IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 102             NULL,
 103             0,
 104             &DiskExtents,
 105             (DWORD)(sizeof(DiskExtents)),
 106             (LPDWORD)&dwSize,
 107             NULL);
 108
 109         Log("IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS bRet:%u code:%u", bRet, LASTERR);
 110         Log("NumberOfDiskExtents:%u DiskNumber:%u", DiskExtents.NumberOfDiskExtents, DiskExtents.Extents[0].DiskNumber);
 111
 112         if (bRet && DiskExtents.NumberOfDiskExtents == 1)
 113         {
 114             pExtents = DiskExtents.Extents;
 115
 116             Log("This volume DiskNumber:%u offset:%llu", pExtents->DiskNumber, (ULONGLONG)pExtents->StartingOffset.QuadPart);
 117             if ((int)pExtents->DiskNumber == PhyDrive && pExtents->StartingOffset.QuadPart == PartOffset)
 118             {
 119                 Log("This volume match");
 120
 121                 if (!DelSlash)
 122                 {
 123                     VolumeName[len - 1] = '\\';
 124                 }
 125
 126                 sprintf_s(NameBuf, BufLen, "%s", VolumeName);
 127                 Status = ERROR_SUCCESS;
 128                 CloseHandle(hDrive);
 129                 break;
 130             }
 131         }
 132
 133         CloseHandle(hDrive);
 134     } while (FindNextVolumeA(hVolume, VolumeName, sizeof(VolumeName)));
 135
 136     FindVolumeClose(hVolume);
 137
 138     Log("GetVentoyVolumeName return %u", Status);
 139     return Status;
 140 }
 141
 142 static int GetLettersBelongPhyDrive(int PhyDrive, char *DriveLetters, size_t Length)
 143 {
 144     int n = 0;
 145     DWORD DataSize = 0;
 146     CHAR *Pos = NULL;
 147     CHAR *StringBuf = NULL;
 148
 149     DataSize = GetLogicalDriveStringsA(0, NULL);
 150     StringBuf = (CHAR *)malloc(DataSize + 1);
 151     if (StringBuf == NULL)
 152     {
 153         return 1;
 154     }
 155
 156     GetLogicalDriveStringsA(DataSize, StringBuf);
 157
 158     for (Pos = StringBuf; *Pos; Pos += strlen(Pos) + 1)
 159     {
 160         if (n < (int)Length && PhyDrive == GetPhyDriveByLogicalDrive(Pos[0], NULL))
 161         {
 162             Log("%C: is belong to phydrive%d", Pos[0], PhyDrive);
 163             DriveLetters[n++] = Pos[0];
 164         }
 165     }
 166
 167     free(StringBuf);
 168     return 0;
 169 }
 170
 171 HANDLE GetPhysicalHandle(int Drive, BOOLEAN bLockDrive, BOOLEAN bWriteAccess, BOOLEAN bWriteShare)
 172 {
 173     int i;
 174     DWORD dwSize;
 175     DWORD LastError;
 176     UINT64 EndTime;
 177     HANDLE hDrive = INVALID_HANDLE_VALUE;
 178     CHAR PhyDrive[128];
 179     CHAR DevPath[MAX_PATH] = { 0 };
 180
 181     safe_sprintf(PhyDrive, "\\\\.\\PhysicalDrive%d", Drive);
 182
 183     if (0 == QueryDosDeviceA(PhyDrive + 4, DevPath, sizeof(DevPath)))
 184     {
 185         Log("QueryDosDeviceA failed error:%u", GetLastError());
 186         strcpy_s(DevPath, sizeof(DevPath), "???");
 187     }
 188     else
 189     {
 190         Log("QueryDosDeviceA success %s", DevPath);
 191     }
 192
 193     for (i = 0; i < DRIVE_ACCESS_RETRIES; i++)
 194     {
 195         // Try without FILE_SHARE_WRITE (unless specifically requested) so that
 196         // we won't be bothered by the OS or other apps when we set up our data.
 197         // However this means we might have to wait for an access gap...
 198         // We keep FILE_SHARE_READ though, as this shouldn't hurt us any, and is
 199         // required for enumeration.
 200         hDrive = CreateFileA(PhyDrive,
 201             GENERIC_READ | (bWriteAccess ? GENERIC_WRITE : 0),
 202             FILE_SHARE_READ | (bWriteShare ? FILE_SHARE_WRITE : 0),
 203             NULL,
 204             OPEN_EXISTING,
 205             FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
 206             NULL);
 207
 208         LastError = GetLastError();
 209         Log("[%d] CreateFileA %s code:%u %p", i, PhyDrive, LastError, hDrive);
 210
 211         if (hDrive != INVALID_HANDLE_VALUE)
 212         {
 213             break;
 214         }
 215
 216         if ((LastError != ERROR_SHARING_VIOLATION) && (LastError != ERROR_ACCESS_DENIED))
 217         {
 218             break;
 219         }
 220
 221         if (i == 0)
 222         {
 223             Log("Waiting for access on %s [%s]...", PhyDrive, DevPath);
 224         }
 225         else if (!bWriteShare && (i > DRIVE_ACCESS_RETRIES / 3))
 226         {
 227             // If we can't seem to get a hold of the drive for some time, try to enable FILE_SHARE_WRITE...
 228             Log("Warning: Could not obtain exclusive rights. Retrying with write sharing enabled...");
 229             bWriteShare = TRUE;
 230
 231             // Try to report the process that is locking the drive
 232             // We also use bit 6 as a flag to indicate that SearchProcess was called.
 233             //access_mask = SearchProcess(DevPath, SEARCH_PROCESS_TIMEOUT, TRUE, TRUE, FALSE) | 0x40;
 234
 235         }
 236         Sleep(DRIVE_ACCESS_TIMEOUT / DRIVE_ACCESS_RETRIES);
 237     }
 238
 239     if (hDrive == INVALID_HANDLE_VALUE)
 240     {
 241         Log("Could not open %s %u", PhyDrive, LASTERR);
 242         goto End;
 243     }
 244
 245     if (bWriteAccess)
 246     {
 247         Log("Opened %s for %s write access", PhyDrive, bWriteShare ? "shared" : "exclusive");
 248     }
 249
 250     if (bLockDrive)
 251     {
 252         if (DeviceIoControl(hDrive, FSCTL_ALLOW_EXTENDED_DASD_IO, NULL, 0, NULL, 0, &dwSize, NULL))
 253         {
 254             Log("I/O boundary checks disabled");
 255         }
 256
 257         EndTime = GetTickCount64() + DRIVE_ACCESS_TIMEOUT;
 258
 259         do {
 260             if (DeviceIoControl(hDrive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL))
 261             {
 262                 Log("FSCTL_LOCK_VOLUME success");
 263                 goto End;
 264             }
 265             Sleep(DRIVE_ACCESS_TIMEOUT / DRIVE_ACCESS_RETRIES);
 266         } while (GetTickCount64() < EndTime);
 267
 268         // If we reached this section, either we didn't manage to get a lock or the user cancelled
 269         Log("Could not lock access to %s %u", PhyDrive, LASTERR);
 270
 271         // See if we can report the processes are accessing the drive
 272         //if (!IS_ERROR(FormatStatus) && (access_mask == 0))
 273         //    access_mask = SearchProcess(DevPath, SEARCH_PROCESS_TIMEOUT, TRUE, TRUE, FALSE);
 274         // Try to continue if the only access rights we saw were for read-only
 275         //if ((access_mask & 0x07) != 0x01)
 276         //    safe_closehandle(hDrive);
 277
 278         CHECK_CLOSE_HANDLE(hDrive);
 279     }
 280
 281 End:
 282
 283     if (hDrive == INVALID_HANDLE_VALUE)
 284     {
 285         Log("Can get handle of %s, maybe some process control it.", DevPath);
 286     }
 287
 288     return hDrive;
 289 }
 290
 291 int GetPhyDriveByLogicalDrive(int DriveLetter, UINT64 *Offset)
 292 {
 293     BOOL Ret = FALSE;
 294     DWORD dwSize = 0;
 295     HANDLE Handle = INVALID_HANDLE_VALUE;
 296     VOLUME_DISK_EXTENTS DiskExtents;
 297     CHAR PhyPath[128];
 298
 299     safe_sprintf(PhyPath, "\\\\.\\%C:", (CHAR)DriveLetter);
 300
 301     Handle = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 302     if (Handle == INVALID_HANDLE_VALUE)
 303     {
 304         Log("Could not open the disk<%s>, error:%u", PhyPath, LASTERR);
 305         return -1;
 306     }
 307
 308     memset(&DiskExtents, 0, sizeof(DiskExtents));
 309     Ret = DeviceIoControl(Handle,
 310         IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 311         NULL,
 312         0,
 313         &DiskExtents,
 314         (DWORD)(sizeof(DiskExtents)),
 315         (LPDWORD)&dwSize,
 316         NULL);
 317
 318     if (!Ret || DiskExtents.NumberOfDiskExtents == 0)
 319     {
 320         Log("DeviceIoControl IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS failed %s, error:%u", PhyPath, LASTERR);
 321         CHECK_CLOSE_HANDLE(Handle);
 322         return -1;
 323     }
 324     CHECK_CLOSE_HANDLE(Handle);
 325
 326     Log("LogicalDrive:%s PhyDrive:%d Num:%d Offset:%llu ExtentLength:%llu",
 327         PhyPath,
 328         DiskExtents.Extents[0].DiskNumber,
 329         DiskExtents.NumberOfDiskExtents,
 330         DiskExtents.Extents[0].StartingOffset.QuadPart,
 331         DiskExtents.Extents[0].ExtentLength.QuadPart
 332         );
 333
 334         if (Offset)
 335         {
 336                 *Offset = (UINT64)(DiskExtents.Extents[0].StartingOffset.QuadPart);
 337         }
 338
 339     return (int)DiskExtents.Extents[0].DiskNumber;
 340 }
 341
 342 int GetAllPhysicalDriveInfo(PHY_DRIVE_INFO *pDriveList, DWORD *pDriveCount)
 343 {
 344     int i;
 345     int Count;
 346     int id;
 347     int Letter = 'A';
 348     BOOL  bRet;
 349     DWORD dwBytes;
 350     DWORD DriveCount = 0;
 351     HANDLE Handle = INVALID_HANDLE_VALUE;
 352     CHAR PhyDrive[128];
 353     PHY_DRIVE_INFO *CurDrive = pDriveList;
 354     GET_LENGTH_INFORMATION LengthInfo;
 355     STORAGE_PROPERTY_QUERY Query;
 356     STORAGE_DESCRIPTOR_HEADER DevDescHeader;
 357     STORAGE_DEVICE_DESCRIPTOR *pDevDesc;
 358     int PhyDriveId[VENTOY_MAX_PHY_DRIVE];
 359
 360     Count = GetPhysicalDriveCount();
 361
 362     for (i = 0; i < Count && i < VENTOY_MAX_PHY_DRIVE; i++)
 363     {
 364         PhyDriveId[i] = i;
 365     }
 366
 367     dwBytes = GetLogicalDrives();
 368     Log("Logical Drives: 0x%x", dwBytes);
 369     while (dwBytes)
 370     {
 371         if (dwBytes & 0x01)
 372         {
 373                         id = GetPhyDriveByLogicalDrive(Letter, NULL);
 374             Log("%C --> %d", Letter, id);
 375             if (id >= 0)
 376             {
 377                 for (i = 0; i < Count; i++)
 378                 {
 379                     if (PhyDriveId[i] == id)
 380                     {
 381                         break;
 382                     }
 383                 }
 384
 385                 if (i >= Count)
 386                 {
 387                     Log("Add phy%d to list", i);
 388                     PhyDriveId[Count] = id;
 389                     Count++;
 390                 }
 391             }
 392         }
 393
 394         Letter++;
 395         dwBytes >>= 1;
 396     }
 397
 398     for (i = 0; i < Count && DriveCount < VENTOY_MAX_PHY_DRIVE; i++)
 399     {
 400         CHECK_CLOSE_HANDLE(Handle);
 401
 402         safe_sprintf(PhyDrive, "\\\\.\\PhysicalDrive%d", PhyDriveId[i]);
 403         Handle = CreateFileA(PhyDrive, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
 404         Log("Create file Handle:%p %s status:%u", Handle, PhyDrive, LASTERR);
 405
 406         if (Handle == INVALID_HANDLE_VALUE)
 407         {
 408             continue;
 409         }
 410
 411         bRet = DeviceIoControl(Handle,
 412                                IOCTL_DISK_GET_LENGTH_INFO, NULL,
 413                                0,
 414                                &LengthInfo,
 415                                sizeof(LengthInfo),
 416                                &dwBytes,
 417                                NULL);
 418         if (!bRet)
 419         {
 420             Log("DeviceIoControl IOCTL_DISK_GET_LENGTH_INFO failed error:%u", LASTERR);
 421             continue;
 422         }
 423
 424         Log("PHYSICALDRIVE%d size %llu bytes", i, (ULONGLONG)LengthInfo.Length.QuadPart);
 425
 426         Query.PropertyId = StorageDeviceProperty;
 427         Query.QueryType = PropertyStandardQuery;
 428
 429         bRet = DeviceIoControl(Handle,
 430                                IOCTL_STORAGE_QUERY_PROPERTY,
 431                                &Query,
 432                                sizeof(Query),
 433                                &DevDescHeader,
 434                                sizeof(STORAGE_DESCRIPTOR_HEADER),
 435                                &dwBytes,
 436                                NULL);
 437         if (!bRet)
 438         {
 439             Log("DeviceIoControl1 error:%u dwBytes:%u", LASTERR, dwBytes);
 440             continue;
 441         }
 442
 443         if (DevDescHeader.Size < sizeof(STORAGE_DEVICE_DESCRIPTOR))
 444         {
 445             Log("Invalid DevDescHeader.Size:%u", DevDescHeader.Size);
 446             continue;
 447         }
 448
 449         pDevDesc = (STORAGE_DEVICE_DESCRIPTOR *)malloc(DevDescHeader.Size);
 450         if (!pDevDesc)
 451         {
 452             Log("failed to malloc error:%u len:%u", LASTERR, DevDescHeader.Size);
 453             continue;
 454         }
 455
 456         bRet = DeviceIoControl(Handle,
 457                                IOCTL_STORAGE_QUERY_PROPERTY,
 458                                &Query,
 459                                sizeof(Query),
 460                                pDevDesc,
 461                                DevDescHeader.Size,
 462                                &dwBytes,
 463                                NULL);
 464         if (!bRet)
 465         {
 466             Log("DeviceIoControl2 error:%u dwBytes:%u", LASTERR, dwBytes);
 467             free(pDevDesc);
 468             continue;
 469         }
 470
 471         CurDrive->PhyDrive = i;
 472         CurDrive->SizeInBytes = LengthInfo.Length.QuadPart;
 473         CurDrive->DeviceType = pDevDesc->DeviceType;
 474         CurDrive->RemovableMedia = pDevDesc->RemovableMedia;
 475         CurDrive->BusType = pDevDesc->BusType;
 476
 477         if (pDevDesc->VendorIdOffset)
 478         {
 479             safe_strcpy(CurDrive->VendorId, (char *)pDevDesc + pDevDesc->VendorIdOffset);
 480             TrimString(CurDrive->VendorId);
 481         }
 482
 483         if (pDevDesc->ProductIdOffset)
 484         {
 485             safe_strcpy(CurDrive->ProductId, (char *)pDevDesc + pDevDesc->ProductIdOffset);
 486             TrimString(CurDrive->ProductId);
 487         }
 488
 489         if (pDevDesc->ProductRevisionOffset)
 490         {
 491             safe_strcpy(CurDrive->ProductRev, (char *)pDevDesc + pDevDesc->ProductRevisionOffset);
 492             TrimString(CurDrive->ProductRev);
 493         }
 494
 495         if (pDevDesc->SerialNumberOffset)
 496         {
 497             safe_strcpy(CurDrive->SerialNumber, (char *)pDevDesc + pDevDesc->SerialNumberOffset);
 498             TrimString(CurDrive->SerialNumber);
 499         }
 500
 501         CurDrive++;
 502         DriveCount++;
 503
 504         free(pDevDesc);
 505
 506         CHECK_CLOSE_HANDLE(Handle);
 507     }
 508
 509     for (i = 0, CurDrive = pDriveList; i < (int)DriveCount; i++, CurDrive++)
 510     {
 511         Log("PhyDrv:%d BusType:%-4s Removable:%u Size:%dGB(%llu) Name:%s %s",
 512             CurDrive->PhyDrive, GetBusTypeString(CurDrive->BusType), CurDrive->RemovableMedia,
 513             GetHumanReadableGBSize(CurDrive->SizeInBytes), CurDrive->SizeInBytes,
 514             CurDrive->VendorId, CurDrive->ProductId);
 515     }
 516
 517     *pDriveCount = DriveCount;
 518
 519     return 0;
 520 }
 521
 522
 523 static HANDLE g_FatPhyDrive;
 524 static UINT64 g_Part2StartSec;
 525 static int GetVentoyVersionFromFatFile(CHAR *VerBuf, size_t BufLen)
 526 {
 527     int rc = 1;
 528     int size = 0;
 529     char *buf = NULL;
 530     void *flfile = NULL;
 531
 532     flfile = fl_fopen("/grub/grub.cfg", "rb");
 533     if (flfile)
 534     {
 535         fl_fseek(flfile, 0, SEEK_END);
 536         size = (int)fl_ftell(flfile);
 537
 538         fl_fseek(flfile, 0, SEEK_SET);
 539
 540         buf = (char *)malloc(size + 1);
 541         if (buf)
 542         {
 543             fl_fread(buf, 1, size, flfile);
 544             buf[size] = 0;
 545
 546             rc = 0;
 547             sprintf_s(VerBuf, BufLen, "%s", ParseVentoyVersionFromString(buf));
 548             free(buf);
 549         }
 550
 551         fl_fclose(flfile);
 552     }
 553
 554     return rc;
 555 }
 556
 557 static int VentoyFatDiskRead(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 558 {
 559     DWORD dwSize;
 560     BOOL bRet;
 561     DWORD ReadSize;
 562     LARGE_INTEGER liCurrentPosition;
 563
 564     liCurrentPosition.QuadPart = Sector + g_Part2StartSec;
 565     liCurrentPosition.QuadPart *= 512;
 566     SetFilePointerEx(g_FatPhyDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 567
 568     ReadSize = (DWORD)(SectorCount * 512);
 569
 570     bRet = ReadFile(g_FatPhyDrive, Buffer, ReadSize, &dwSize, NULL);
 571     if (bRet == FALSE || dwSize != ReadSize)
 572     {
 573         Log("ReadFile error bRet:%u WriteSize:%u dwSize:%u ErrCode:%u\n", bRet, ReadSize, dwSize, LASTERR);
 574     }
 575
 576     return 1;
 577 }
 578
 579
 580 int GetVentoyVerInPhyDrive(const PHY_DRIVE_INFO *pDriveInfo, UINT64 Part2StartSector, CHAR *VerBuf, size_t BufLen, BOOL *pSecureBoot)
 581 {
 582     int rc = 0;
 583     HANDLE hDrive;
 584     void *flfile;
 585
 586     hDrive = GetPhysicalHandle(pDriveInfo->PhyDrive, FALSE, FALSE, FALSE);
 587     if (hDrive == INVALID_HANDLE_VALUE)
 588     {
 589         return 1;
 590     }
 591
 592     g_FatPhyDrive = hDrive;
 593         g_Part2StartSec = Part2StartSector;
 594
 595     Log("Parse FAT fs...");
 596
 597     fl_init();
 598
 599     if (0 == fl_attach_media(VentoyFatDiskRead, NULL))
 600     {
 601         Log("attach media success...");
 602         rc = GetVentoyVersionFromFatFile(VerBuf, BufLen);
 603     }
 604     else
 605     {
 606         Log("attach media failed...");
 607         rc = 1;
 608     }
 609
 610     Log("GetVentoyVerInPhyDrive rc=%d...", rc);
 611     if (rc == 0)
 612     {
 613         Log("VentoyVerInPhyDrive %d is <%s>...", pDriveInfo->PhyDrive, VerBuf);
 614
 615         flfile = fl_fopen("/EFI/BOOT/grubx64_real.efi", "rb");
 616         if (flfile)
 617         {
 618             *pSecureBoot = TRUE;
 619             fl_fclose(flfile);
 620         }
 621     }
 622
 623     fl_shutdown();
 624
 625     CHECK_CLOSE_HANDLE(hDrive);
 626
 627     return rc;
 628 }
 629
 630
 631
 632
 633
 634 static unsigned int g_disk_unxz_len = 0;
 635 static BYTE *g_part_img_pos = NULL;
 636 static BYTE *g_part_img_buf[VENTOY_EFI_PART_SIZE / SIZE_1MB];
 637
 638
 639 static int VentoyFatMemRead(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 640 {
 641         uint32 i;
 642         uint32 offset;
 643         BYTE *MbBuf = NULL;
 644
 645         for (i = 0; i < SectorCount; i++)
 646         {
 647                 offset = (Sector + i) * 512;
 648
 649                 if (g_part_img_buf[1] == NULL)
 650                 {
 651                         MbBuf = g_part_img_buf[0] + offset;
 652                         memcpy(Buffer + i * 512, MbBuf, 512);
 653                 }
 654                 else
 655                 {
 656                         MbBuf = g_part_img_buf[offset / SIZE_1MB];
 657                         memcpy(Buffer + i * 512, MbBuf + (offset % SIZE_1MB), 512);
 658                 }
 659         }
 660
 661         return 1;
 662 }
 663
 664
 665 static int VentoyFatMemWrite(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 666 {
 667         uint32 i;
 668         uint32 offset;
 669         BYTE *MbBuf = NULL;
 670
 671         for (i = 0; i < SectorCount; i++)
 672         {
 673                 offset = (Sector + i) * 512;
 674
 675                 if (g_part_img_buf[1] == NULL)
 676                 {
 677                         MbBuf = g_part_img_buf[0] + offset;
 678                         memcpy(MbBuf, Buffer + i * 512, 512);
 679                 }
 680                 else
 681                 {
 682                         MbBuf = g_part_img_buf[offset / SIZE_1MB];
 683                         memcpy(MbBuf + (offset % SIZE_1MB), Buffer + i * 512, 512);
 684                 }
 685         }
 686
 687         return 1;
 688 }
 689
 690 int VentoyProcSecureBoot(BOOL SecureBoot)
 691 {
 692         int rc = 0;
 693         int size;
 694         char *filebuf = NULL;
 695         void *file = NULL;
 696
 697         Log("VentoyProcSecureBoot %d ...", SecureBoot);
 698
 699         if (SecureBoot)
 700         {
 701                 Log("Secure boot is enabled ...");
 702                 return 0;
 703         }
 704
 705         fl_init();
 706
 707         if (0 == fl_attach_media(VentoyFatMemRead, VentoyFatMemWrite))
 708         {
 709                 file = fl_fopen("/EFI/BOOT/grubx64_real.efi", "rb");
 710                 Log("Open ventoy efi file %p ", file);
 711                 if (file)
 712                 {
 713                         fl_fseek(file, 0, SEEK_END);
 714                         size = (int)fl_ftell(file);
 715                         fl_fseek(file, 0, SEEK_SET);
 716
 717                         Log("ventoy efi file size %d ...", size);
 718
 719                         filebuf = (char *)malloc(size);
 720                         if (filebuf)
 721                         {
 722                                 fl_fread(filebuf, 1, size, file);
 723                         }
 724
 725                         fl_fclose(file);
 726
 727                         Log("Now delete all efi files ...");
 728                         fl_remove("/EFI/BOOT/BOOTX64.EFI");
 729                         fl_remove("/EFI/BOOT/grubx64.efi");
 730                         fl_remove("/EFI/BOOT/grubx64_real.efi");
 731                         fl_remove("/EFI/BOOT/MokManager.efi");
 732                         fl_remove("/EFI/BOOT/mmx64.efi");
 733             fl_remove("/ENROLL_THIS_KEY_IN_MOKMANAGER.cer");
 734
 735                         file = fl_fopen("/EFI/BOOT/BOOTX64.EFI", "wb");
 736                         Log("Open bootx64 efi file %p ", file);
 737                         if (file)
 738                         {
 739                                 if (filebuf)
 740                                 {
 741                                         fl_fwrite(filebuf, 1, size, file);
 742                                 }
 743
 744                                 fl_fflush(file);
 745                                 fl_fclose(file);
 746                         }
 747
 748                         if (filebuf)
 749                         {
 750                                 free(filebuf);
 751                         }
 752                 }
 753
 754         file = fl_fopen("/EFI/BOOT/grubia32_real.efi", "rb");
 755         Log("Open ventoy efi file %p ", file);
 756         if (file)
 757         {
 758             fl_fseek(file, 0, SEEK_END);
 759             size = (int)fl_ftell(file);
 760             fl_fseek(file, 0, SEEK_SET);
 761
 762             Log("ventoy efi file size %d ...", size);
 763
 764             filebuf = (char *)malloc(size);
 765             if (filebuf)
 766             {
 767                 fl_fread(filebuf, 1, size, file);
 768             }
 769
 770             fl_fclose(file);
 771
 772             Log("Now delete all efi files ...");
 773             fl_remove("/EFI/BOOT/BOOTIA32.EFI");
 774             fl_remove("/EFI/BOOT/grubia32.efi");
 775             fl_remove("/EFI/BOOT/grubia32_real.efi");
 776             fl_remove("/EFI/BOOT/mmia32.efi");
 777
 778             file = fl_fopen("/EFI/BOOT/BOOTIA32.EFI", "wb");
 779             Log("Open bootia32 efi file %p ", file);
 780             if (file)
 781             {
 782                 if (filebuf)
 783                 {
 784                     fl_fwrite(filebuf, 1, size, file);
 785                 }
 786
 787                 fl_fflush(file);
 788                 fl_fclose(file);
 789             }
 790
 791             if (filebuf)
 792             {
 793                 free(filebuf);
 794             }
 795         }
 796
 797         }
 798         else
 799         {
 800                 rc = 1;
 801         }
 802
 803         fl_shutdown();
 804
 805         return rc;
 806 }
 807
 808
 809
 810 static int disk_xz_flush(void *src, unsigned int size)
 811 {
 812     unsigned int i;
 813     BYTE *buf = (BYTE *)src;
 814
 815     for (i = 0; i < size; i++)
 816     {
 817         *g_part_img_pos = *buf++;
 818
 819         g_disk_unxz_len++;
 820         if ((g_disk_unxz_len % SIZE_1MB) == 0)
 821         {
 822             g_part_img_pos = g_part_img_buf[g_disk_unxz_len / SIZE_1MB];
 823         }
 824         else
 825         {
 826             g_part_img_pos++;
 827         }
 828     }
 829
 830     return (int)size;
 831 }
 832
 833 static void unxz_error(char *x)
 834 {
 835     Log("%s", x);
 836 }
 837
 838 static BOOL TryWritePart2(HANDLE hDrive, UINT64 StartSectorId)
 839 {
 840     BOOL bRet;
 841     DWORD TrySize = 16 * 1024;
 842     DWORD dwSize;
 843     BYTE *Buffer = NULL;
 844     unsigned char *data = NULL;
 845     LARGE_INTEGER liCurrentPosition;
 846
 847     liCurrentPosition.QuadPart = StartSectorId * 512;
 848     SetFilePointerEx(hDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 849
 850     Buffer = malloc(TrySize);
 851
 852     bRet = WriteFile(hDrive, Buffer, TrySize, &dwSize, NULL);
 853
 854     free(Buffer);
 855
 856     Log("Try write part2 bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 857
 858     if (bRet && dwSize == TrySize)
 859     {
 860         return TRUE;
 861     }
 862
 863     return FALSE;
 864 }
 865
 866 static int FormatPart2Fat(HANDLE hDrive, UINT64 StartSectorId)
 867 {
 868     int i;
 869     int rc = 0;
 870     int len = 0;
 871     int writelen = 0;
 872     int partwrite = 0;
 873     int Pos = PT_WRITE_VENTOY_START;
 874     DWORD dwSize = 0;
 875     BOOL bRet;
 876     unsigned char *data = NULL;
 877     LARGE_INTEGER liCurrentPosition;
 878         LARGE_INTEGER liNewPosition;
 879     BYTE *CheckBuf = NULL;
 880
 881         Log("FormatPart2Fat %llu...", (ULONGLONG)StartSectorId);
 882
 883     CheckBuf = malloc(SIZE_1MB);
 884     if (!CheckBuf)
 885     {
 886         Log("Failed to malloc check buf");
 887         return 1;
 888     }
 889
 890     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&data, &len);
 891     if (rc)
 892     {
 893         Log("Failed to read img file %p %u", data, len);
 894         free(CheckBuf);
 895         return 1;
 896     }
 897
 898     liCurrentPosition.QuadPart = StartSectorId * 512;
 899     SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
 900
 901     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
 902
 903     g_part_img_buf[0] = (BYTE *)malloc(VENTOY_EFI_PART_SIZE);
 904     if (g_part_img_buf[0])
 905     {
 906         Log("Malloc whole img buffer success, now decompress ...");
 907         unxz(data, len, NULL, NULL, g_part_img_buf[0], &writelen, unxz_error);
 908
 909         if (len == writelen)
 910         {
 911             Log("decompress finished success");
 912
 913                         VentoyProcSecureBoot(g_SecureBoot);
 914
 915             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 916             {
 917                 dwSize = 0;
 918                                 bRet = WriteFile(hDrive, g_part_img_buf[0] + i * SIZE_1MB, SIZE_1MB, &dwSize, NULL);
 919                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 920
 921                 if (!bRet)
 922                 {
 923                     rc = 1;
 924                     goto End;
 925                 }
 926
 927                 PROGRESS_BAR_SET_POS(Pos);
 928                 if (i % 2 == 0)
 929                 {
 930                     Pos++;
 931                 }
 932             }
 933
 934             //Read and check the data
 935             liCurrentPosition.QuadPart = StartSectorId * 512;
 936             SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
 937
 938             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 939             {
 940                 bRet = ReadFile(hDrive, CheckBuf, SIZE_1MB, &dwSize, NULL);
 941                 Log("Read part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 942
 943                 if (!bRet || memcmp(CheckBuf, g_part_img_buf[0] + i * SIZE_1MB, SIZE_1MB))
 944                 {
 945                     Log("### [Check Fail] The data write and read does not match");
 946                     rc = 1;
 947                     goto End;
 948                 }
 949
 950                 PROGRESS_BAR_SET_POS(Pos);
 951                 if (i % 2 == 0)
 952                 {
 953                     Pos++;
 954                 }
 955             }
 956         }
 957         else
 958         {
 959             rc = 1;
 960             Log("decompress finished failed");
 961             goto End;
 962         }
 963     }
 964     else
 965     {
 966         Log("Failed to malloc whole img size %u, now split it", VENTOY_EFI_PART_SIZE);
 967
 968         partwrite = 1;
 969         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 970         {
 971             g_part_img_buf[i] = (BYTE *)malloc(SIZE_1MB);
 972             if (g_part_img_buf[i] == NULL)
 973             {
 974                 rc = 1;
 975                 goto End;
 976             }
 977         }
 978
 979         Log("Malloc part img buffer success, now decompress ...");
 980
 981         g_part_img_pos = g_part_img_buf[0];
 982
 983         unxz(data, len, NULL, disk_xz_flush, NULL, NULL, unxz_error);
 984
 985         if (g_disk_unxz_len == VENTOY_EFI_PART_SIZE)
 986         {
 987             Log("decompress finished success");
 988
 989                         VentoyProcSecureBoot(g_SecureBoot);
 990
 991             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 992             {
 993                 dwSize = 0;
 994                 bRet = WriteFile(hDrive, g_part_img_buf[i], SIZE_1MB, &dwSize, NULL);
 995                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 996
 997                 if (!bRet)
 998                 {
 999                     rc = 1;
1000                     goto End;
1001                 }
1002
1003                 PROGRESS_BAR_SET_POS(Pos);
1004                 if (i % 2 == 0)
1005                 {
1006                     Pos++;
1007                 }
1008             }
1009
1010             //Read and check the data
1011             liCurrentPosition.QuadPart = StartSectorId * 512;
1012             SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
1013
1014             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1015             {
1016                 bRet = ReadFile(hDrive, CheckBuf, SIZE_1MB, &dwSize, NULL);
1017                 Log("Read part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
1018
1019                 if (!bRet || memcmp(CheckBuf, g_part_img_buf[i], SIZE_1MB))
1020                 {
1021                     Log("### [Check Fail] The data write and read does not match");
1022                     rc = 1;
1023                     goto End;
1024                 }
1025
1026                 PROGRESS_BAR_SET_POS(Pos);
1027                 if (i % 2 == 0)
1028                 {
1029                     Pos++;
1030                 }
1031             }
1032         }
1033         else
1034         {
1035             rc = 1;
1036             Log("decompress finished failed");
1037             goto End;
1038         }
1039     }
1040
1041 End:
1042
1043     if (data) free(data);
1044     if (CheckBuf)free(CheckBuf);
1045
1046     if (partwrite)
1047     {
1048         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1049         {
1050             if (g_part_img_buf[i]) free(g_part_img_buf[i]);
1051         }
1052     }
1053     else
1054     {
1055         if (g_part_img_buf[0]) free(g_part_img_buf[0]);
1056     }
1057
1058     return rc;
1059 }
1060
1061 static int WriteGrubStage1ToPhyDrive(HANDLE hDrive, int PartStyle)
1062 {
1063     int Len = 0;
1064     int readLen = 0;
1065     BOOL bRet;
1066     DWORD dwSize;
1067     BYTE *ImgBuf = NULL;
1068     BYTE *RawBuf = NULL;
1069
1070     Log("WriteGrubStage1ToPhyDrive ...");
1071
1072     RawBuf = (BYTE *)malloc(SIZE_1MB);
1073     if (!RawBuf)
1074     {
1075         return 1;
1076     }
1077
1078     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1079     {
1080         Log("Failed to read stage1 img");
1081         free(RawBuf);
1082         return 1;
1083     }
1084
1085     unxz(ImgBuf, Len, NULL, NULL, RawBuf, &readLen, unxz_error);
1086
1087     if (PartStyle)
1088     {
1089         Log("Write GPT stage1 ...");
1090         RawBuf[500] = 35;//update blocklist
1091         SetFilePointer(hDrive, 512 * 34, NULL, FILE_BEGIN);
1092         bRet = WriteFile(hDrive, RawBuf, SIZE_1MB - 512 * 34, &dwSize, NULL);
1093     }
1094     else
1095     {
1096         Log("Write MBR stage1 ...");
1097         SetFilePointer(hDrive, 512, NULL, FILE_BEGIN);
1098         bRet = WriteFile(hDrive, RawBuf, SIZE_1MB - 512, &dwSize, NULL);
1099     }
1100
1101     Log("WriteFile Ret:%u dwSize:%u ErrCode:%u", bRet, dwSize, GetLastError());
1102
1103     free(RawBuf);
1104     free(ImgBuf);
1105     return 0;
1106 }
1107
1108
1109 static int FormatPart1LargeFAT32(UINT64 DiskSizeBytes, int CluserSize)
1110 {
1111     MKFS_PARM Option;
1112     FRESULT Ret;
1113     FATFS FS;
1114
1115     Option.fmt = FM_FAT32;
1116     Option.n_fat = 1;
1117     Option.align = 8;
1118     Option.n_root = 1;
1119
1120     if (CluserSize == 0)
1121     {
1122         // < 32GB select 32KB as cluster size
1123         // > 32GB select 128KB as cluster size
1124         if (DiskSizeBytes / 1024 / 1024 / 1024 <= 32)
1125         {
1126             Option.au_size = 32768;
1127         }
1128         else
1129         {
1130             Option.au_size = 131072;
1131         }
1132     }
1133     else
1134     {
1135         Option.au_size = CluserSize;
1136     }
1137
1138     Log("Formatting Part1 large FAT32 ClusterSize:%u(%uKB) ...", CluserSize, CluserSize / 1024);
1139
1140     disk_io_reset_write_error();
1141
1142     Ret = f_mkfs(TEXT("0:"), &Option, 0, 8 * 1024 * 1024);
1143     if (FR_OK == Ret)
1144     {
1145         if (disk_io_is_write_error())
1146         {
1147             Log("Formatting Part1 large FAT32 failed, write error.");
1148             return 1;
1149         }
1150
1151         Log("Formatting Part1 large FAT32 success, now set label");
1152
1153         Ret = f_mount(&FS, TEXT("0:"), 1);
1154         if (FR_OK == Ret)
1155         {
1156             Log("f_mount SUCCESS");
1157             Ret = f_setlabel(TEXT("0:Ventoy"));
1158             if (FR_OK == Ret)
1159             {
1160                 Log("f_setlabel SUCCESS");
1161                 Ret = f_unmount(TEXT("0:"));
1162                 Log("f_unmount %d %s", Ret, (FR_OK == Ret) ? "SUCCESS" : "FAILED");
1163             }
1164             else
1165             {
1166                 Log("f_setlabel failed %d", Ret);
1167             }
1168         }
1169         else
1170         {
1171             Log("f_mount failed %d", Ret);
1172         }
1173
1174         return 0;
1175     }
1176     else
1177     {
1178         Log("Formatting Part1 large FAT32 failed");
1179         return 1;
1180     }
1181 }
1182
1183 static int FormatPart1exFAT(UINT64 DiskSizeBytes)
1184 {
1185     MKFS_PARM Option;
1186     FRESULT Ret;
1187
1188     Option.fmt = FM_EXFAT;
1189     Option.n_fat = 1;
1190     Option.align = 8;
1191     Option.n_root = 1;
1192
1193     // < 32GB select 32KB as cluster size
1194     // > 32GB select 128KB as cluster size
1195     if (DiskSizeBytes / 1024 / 1024 / 1024 <= 32)
1196     {
1197         Option.au_size = 32768;
1198     }
1199     else
1200     {
1201         Option.au_size = 131072;
1202     }
1203
1204     Log("Formatting Part1 exFAT ...");
1205
1206         disk_io_reset_write_error();
1207
1208     Ret = f_mkfs(TEXT("0:"), &Option, 0, 8 * 1024 * 1024);
1209     if (FR_OK == Ret)
1210     {
1211                 if (disk_io_is_write_error())
1212                 {
1213                         Log("Formatting Part1 exFAT failed, write error.");
1214                         return 1;
1215                 }
1216
1217         Log("Formatting Part1 exFAT success");
1218         return 0;
1219     }
1220     else
1221     {
1222         Log("Formatting Part1 exFAT failed");
1223         return 1;
1224     }
1225 }
1226
1227
1228
1229 int ClearVentoyFromPhyDrive(HWND hWnd, PHY_DRIVE_INFO *pPhyDrive, char *pDrvLetter)
1230 {
1231     int i;
1232     int rc = 0;
1233     int state = 0;
1234     HANDLE hDrive;
1235     DWORD dwSize;
1236     BOOL bRet;
1237     CHAR MountDrive;
1238     CHAR DriveName[] = "?:\\";
1239     CHAR DriveLetters[MAX_PATH] = { 0 };
1240     LARGE_INTEGER liCurrentPosition;
1241     char *pTmpBuf = NULL;
1242     MBR_HEAD MBR;
1243
1244     *pDrvLetter = 0;
1245
1246     Log("ClearVentoyFromPhyDrive PhyDrive%d <<%s %s %dGB>>",
1247         pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1248         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1249
1250     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1251
1252     Log("Lock disk for clean ............................. ");
1253
1254     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1255     if (hDrive == INVALID_HANDLE_VALUE)
1256     {
1257         Log("Failed to open physical disk");
1258         return 1;
1259     }
1260
1261     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1262
1263     if (DriveLetters[0] == 0)
1264     {
1265         Log("No drive letter was assigned...");
1266         DriveName[0] = GetFirstUnusedDriveLetter();
1267         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1268     }
1269     else
1270     {
1271         // Unmount all mounted volumes that belong to this drive
1272         // Do it in reverse so that we always end on the first volume letter
1273         for (i = (int)strlen(DriveLetters); i > 0; i--)
1274         {
1275             DriveName[0] = DriveLetters[i - 1];
1276             bRet = DeleteVolumeMountPointA(DriveName);
1277             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1278         }
1279     }
1280
1281     MountDrive = DriveName[0];
1282     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1283
1284     // It kind of blows, but we have to relinquish access to the physical drive
1285     // for VDS to be able to delete the partitions that reside on it...
1286     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1287     CHECK_CLOSE_HANDLE(hDrive);
1288
1289     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1290
1291     if (!VDS_DeleteAllPartitions(pPhyDrive->PhyDrive))
1292     {
1293         Log("Notice: Could not delete partitions: %u", GetLastError());
1294     }
1295
1296     Log("Deleting all partitions ......................... OK");
1297
1298     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1299
1300     Log("Lock disk for write ............................. ");
1301     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1302     if (hDrive == INVALID_HANDLE_VALUE)
1303     {
1304         Log("Failed to GetPhysicalHandle for write.");
1305         rc = 1;
1306         goto End;
1307     }
1308
1309     // clear first and last 2MB space
1310     pTmpBuf = malloc(SIZE_2MB);
1311     if (!pTmpBuf)
1312     {
1313         Log("Failed to alloc memory.");
1314         rc = 1;
1315         goto End;
1316     }
1317     memset(pTmpBuf, 0, SIZE_2MB);
1318
1319     SET_FILE_POS(512);
1320     bRet = WriteFile(hDrive, pTmpBuf, SIZE_2MB - 512, &dwSize, NULL);
1321     Log("Write fisrt 1MB ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1322     if (!bRet)
1323     {
1324         rc = 1;
1325         goto End;
1326     }
1327
1328     SET_FILE_POS(pPhyDrive->SizeInBytes - SIZE_2MB);
1329     bRet = WriteFile(hDrive, pTmpBuf, SIZE_2MB, &dwSize, NULL);
1330     Log("Write 2nd 1MB ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1331     if (!bRet)
1332     {
1333         rc = 1;
1334         goto End;
1335     }
1336
1337     SET_FILE_POS(0);
1338
1339     if (pPhyDrive->SizeInBytes > 2199023255552ULL)
1340     {
1341         VTOY_GPT_INFO *pGptInfo;
1342         VTOY_GPT_HDR BackupHead;
1343         LARGE_INTEGER liCurrentPosition;
1344
1345         pGptInfo = (VTOY_GPT_INFO *)pTmpBuf;
1346
1347         VentoyFillWholeGpt(pPhyDrive->SizeInBytes, pGptInfo);
1348
1349         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1350         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1351         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1352         {
1353             rc = 1;
1354             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1355             goto End;
1356         }
1357
1358         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1359         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1360         {
1361             rc = 1;
1362             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1363             goto End;
1364         }
1365
1366         SET_FILE_POS(0);
1367         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1368         {
1369             rc = 1;
1370             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1371             goto End;
1372         }
1373
1374         Log("Write GPT Info OK ...");
1375     }
1376     else
1377     {
1378         bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
1379         Log("Read MBR ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1380         if (!bRet)
1381         {
1382             rc = 1;
1383             goto End;
1384         }
1385
1386         //clear boot code and partition table (reserved disk signature)
1387         memset(MBR.BootCode, 0, 440);
1388         memset(MBR.PartTbl, 0, sizeof(MBR.PartTbl));
1389
1390         VentoyFillMBRLocation(pPhyDrive->SizeInBytes, 2048, (UINT32)(pPhyDrive->SizeInBytes / 512 - 2048), MBR.PartTbl);
1391
1392         MBR.PartTbl[0].Active = 0x00; // bootable
1393         MBR.PartTbl[0].FsFlag = 0x07; // exFAT/NTFS/HPFS
1394
1395         SET_FILE_POS(0);
1396         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
1397         Log("Write MBR ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1398         if (!bRet)
1399         {
1400             rc = 1;
1401             goto End;
1402         }
1403     }
1404
1405     Log("Clear Ventoy successfully finished");
1406
1407         //Refresh Drive Layout
1408         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1409
1410 End:
1411
1412     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1413     PROGRESS_BAR_SET_POS(PT_REFORMAT_FINISH);
1414
1415     if (pTmpBuf)
1416     {
1417         free(pTmpBuf);
1418     }
1419
1420     if (rc == 0)
1421     {
1422         Log("Mounting Ventoy Partition ....................... ");
1423         Sleep(1000);
1424
1425         state = 0;
1426         memset(DriveLetters, 0, sizeof(DriveLetters));
1427         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1428         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1429
1430         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1431         {
1432             DriveName[0] = DriveLetters[i];
1433             Log("%s is ventoy part1, already mounted", DriveName);
1434             state = 1;
1435         }
1436
1437         if (state != 1)
1438         {
1439             Log("need to mount ventoy part1...");
1440             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, 2048, DriveLetters, sizeof(DriveLetters), FALSE))
1441             {
1442                 DriveName[0] = MountDrive;
1443                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1444                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1445
1446                 *pDrvLetter = MountDrive;
1447             }
1448             else
1449             {
1450                 Log("Failed to find ventoy volume");
1451             }
1452         }
1453
1454         Log("OK\n");
1455     }
1456     else
1457     {
1458         FindProcessOccupyDisk(hDrive, pPhyDrive);
1459     }
1460
1461     CHECK_CLOSE_HANDLE(hDrive);
1462     return rc;
1463 }
1464
1465 int InstallVentoy2FileImage(PHY_DRIVE_INFO *pPhyDrive, int PartStyle)
1466 {
1467     int i;
1468     int rc = 1;
1469     int Len = 0;
1470     int dataLen = 0;
1471     UINT size = 0;
1472     UINT segnum = 0;
1473     UINT32 chksum = 0;
1474     UINT64 data_offset = 0;
1475     UINT64 Part2StartSector = 0;
1476     UINT64 Part1StartSector = 0;
1477     UINT64 Part1SectorCount = 0;
1478     UINT8 *pData = NULL;
1479     UINT8 *pBkGptPartTbl = NULL;
1480     BYTE *ImgBuf = NULL;
1481     MBR_HEAD *pMBR = NULL;
1482     VTSI_FOOTER *pImgFooter = NULL;
1483     VTSI_SEGMENT *pSegment = NULL;
1484     VTOY_GPT_INFO *pGptInfo = NULL;
1485     VTOY_GPT_HDR *pBkGptHdr = NULL;
1486     FILE *fp = NULL;
1487
1488     Log("InstallVentoy2FileImage %s PhyDrive%d <<%s %s %dGB>>",
1489         PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1490         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1491
1492     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1493
1494     size = SIZE_1MB + VENTOY_EFI_PART_SIZE + 33 * 512 + VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT) + sizeof(VTSI_FOOTER);
1495
1496     pData = (UINT8 *)malloc(size);
1497     if (!pData)
1498     {
1499         Log("malloc image buffer failed %d.", size);
1500         goto End;
1501     }
1502
1503     pImgFooter = (VTSI_FOOTER *)(pData + size - sizeof(VTSI_FOOTER));
1504     pSegment = (VTSI_SEGMENT *)((UINT8 *)pImgFooter - VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT));
1505     memset(pImgFooter, 0, sizeof(VTSI_FOOTER));
1506     memset(pSegment, 0, VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT));
1507
1508     PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START);
1509
1510     Log("Writing Boot Image ............................. ");
1511     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1512     {
1513         Log("Failed to read stage1 img");
1514         goto End;
1515     }
1516
1517     unxz(ImgBuf, Len, NULL, NULL, pData, &dataLen, unxz_error);
1518     SAFE_FREE(ImgBuf);
1519
1520     Log("decompress %s len:%d", VENTOY_FILE_STG1_IMG, dataLen);
1521
1522     if (PartStyle)
1523     {
1524         pData[500] = 35;//update blocklist
1525         memmove(pData + 34 * 512, pData, SIZE_1MB - 512 * 34);
1526         memset(pData, 0, 34 * 512);
1527
1528         pGptInfo = (VTOY_GPT_INFO *)pData;
1529         memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1530         VentoyFillGpt(pPhyDrive->SizeInBytes, pGptInfo);
1531
1532         pBkGptPartTbl = pData + SIZE_1MB + VENTOY_EFI_PART_SIZE;
1533         memset(pBkGptPartTbl, 0, 33 * 512);
1534
1535         memcpy(pBkGptPartTbl, pGptInfo->PartTbl, 32 * 512);
1536         pBkGptHdr = (VTOY_GPT_HDR *)(pBkGptPartTbl + 32 * 512);
1537         VentoyFillBackupGptHead(pGptInfo, pBkGptHdr);
1538
1539         Part1StartSector = pGptInfo->PartTbl[0].StartLBA;
1540         Part1SectorCount = pGptInfo->PartTbl[0].LastLBA - Part1StartSector + 1;
1541         Part2StartSector = pGptInfo->PartTbl[1].StartLBA;
1542
1543         Log("Write GPT Info OK ...");
1544     }
1545     else
1546     {
1547         memmove(pData + 512, pData, SIZE_1MB - 512);
1548         memset(pData, 0, 512);
1549
1550         pMBR = (MBR_HEAD *)pData;
1551         VentoyFillMBR(pPhyDrive->SizeInBytes, pMBR, PartStyle);
1552         Part1StartSector = pMBR->PartTbl[0].StartSectorId;
1553         Part1SectorCount = pMBR->PartTbl[0].SectorCount;
1554         Part2StartSector = pMBR->PartTbl[1].StartSectorId;
1555
1556         Log("Write MBR OK ...");
1557     }
1558
1559     Log("Writing EFI part Image ............................. ");
1560     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&ImgBuf, &Len);
1561     if (rc)
1562     {
1563         Log("Failed to read img file %p %u", ImgBuf, Len);
1564         goto End;
1565     }
1566
1567     PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START + 28);
1568     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
1569     unxz(ImgBuf, Len, NULL, NULL, pData + SIZE_1MB, &dataLen, unxz_error);
1570     if (dataLen == Len)
1571     {
1572         Log("decompress finished success");
1573         g_part_img_buf[0] = pData + SIZE_1MB;
1574
1575         VentoyProcSecureBoot(g_SecureBoot);
1576     }
1577     else
1578     {
1579         Log("decompress finished failed");
1580         goto End;
1581     }
1582
1583     fopen_s(&fp, "VentoySparseImg.vtsi", "wb+");
1584     if (!fp)
1585     {
1586         Log("Failed to create Ventoy img file");
1587         goto End;
1588     }
1589
1590     Log("Writing stage1 data ............................. ");
1591
1592     fwrite(pData, 1, SIZE_1MB, fp);
1593
1594     pSegment[0].disk_start_sector = 0;
1595     pSegment[0].sector_num = SIZE_1MB / 512;
1596     pSegment[0].data_offset = data_offset;
1597     data_offset += pSegment[0].sector_num * 512;
1598
1599     disk_io_set_param(INVALID_HANDLE_VALUE, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1600     disk_io_set_imghook(fp, pSegment + 1, VTSI_IMG_MAX_SEG - 1, data_offset);
1601
1602     Log("Formatting part1 exFAT ...");
1603     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1604     {
1605         Log("FormatPart1exFAT failed.");
1606         disk_io_reset_imghook(&segnum, &data_offset);
1607         goto End;
1608     }
1609
1610     disk_io_reset_imghook(&segnum, &data_offset);
1611     segnum++;
1612
1613     Log("current segment number:%d dataoff:%ld", segnum, (long)data_offset);
1614
1615     //write data
1616     Log("Writing part2 data ............................. ");
1617     fwrite(pData + SIZE_1MB, 1, VENTOY_EFI_PART_SIZE, fp);
1618     pSegment[segnum].disk_start_sector = Part2StartSector;
1619     pSegment[segnum].sector_num = VENTOY_EFI_PART_SIZE / 512;
1620     pSegment[segnum].data_offset = data_offset;
1621     data_offset += pSegment[segnum].sector_num * 512;
1622     segnum++;
1623
1624     if (PartStyle)
1625     {
1626         Log("Writing backup gpt table ............................. ");
1627         fwrite(pBkGptPartTbl, 1, 33 * 512, fp);
1628         pSegment[segnum].disk_start_sector = pPhyDrive->SizeInBytes / 512 - 33;
1629         pSegment[segnum].sector_num = 33;
1630         pSegment[segnum].data_offset = data_offset;
1631         data_offset += pSegment[segnum].sector_num * 512;
1632         segnum++;
1633     }
1634
1635     Log("Writing segment metadata ............................. ");
1636
1637     for (i = 0; i < (int)segnum; i++)
1638     {
1639         Log("SEG[%d]:  PhySector:%llu SectorNum:%llu DataOffset:%llu(sector:%llu)", i, pSegment[i].disk_start_sector, pSegment[i].sector_num,
1640             pSegment[i].data_offset, pSegment[i].data_offset / 512);
1641     }
1642
1643     dataLen = segnum * sizeof(VTSI_SEGMENT);
1644     fwrite(pSegment, 1, dataLen, fp);
1645
1646     if (dataLen % 512)
1647     {
1648         //pData + SIZE_1MB - 8192 is a temp data buffer with zero
1649         fwrite(pData + SIZE_1MB - 8192, 1, 512 - (dataLen % 512), fp);
1650     }
1651
1652     //Fill footer
1653     pImgFooter->magic = VTSI_IMG_MAGIC;
1654     pImgFooter->version = 1;
1655     pImgFooter->disk_size = pPhyDrive->SizeInBytes;
1656     memcpy(&pImgFooter->disk_signature, pPhyDrive->MBR.BootCode + 0x1b8, 4);
1657     pImgFooter->segment_num = segnum;
1658     pImgFooter->segment_offset = data_offset;
1659
1660     for (i = 0, chksum = 0; i < (int)(segnum * sizeof(VTSI_SEGMENT)); i++)
1661     {
1662         chksum += *((UINT8 *)pSegment + i);
1663     }
1664     pImgFooter->segment_chksum = ~chksum;
1665
1666     for (i = 0, chksum = 0; i < sizeof(VTSI_FOOTER); i++)
1667     {
1668         chksum += *((UINT8 *)pImgFooter + i);
1669     }
1670     pImgFooter->foot_chksum = ~chksum;
1671
1672     Log("Writing footer segnum(%u)  segoffset(%llu) ......................", segnum, data_offset);
1673     Log("disk_size=%llu disk_signature=%lx segment_offset=%llu", pImgFooter->disk_size, pImgFooter->disk_signature, pImgFooter->segment_offset);
1674
1675     fwrite(pImgFooter, 1, sizeof(VTSI_FOOTER), fp);
1676     fclose(fp);
1677
1678     Log("Writing Ventoy image file finished, the file size should be %llu .", data_offset + 512 + ((dataLen + 511) / 512 * 512));
1679
1680     rc = 0;
1681
1682 End:
1683
1684     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1685     PROGRESS_BAR_SET_POS(PT_REFORMAT_FINISH);
1686
1687     Log("retcode:%d\n", rc);
1688
1689     SAFE_FREE(pData);
1690     SAFE_FREE(ImgBuf);
1691
1692     return rc;
1693 }
1694
1695
1696 int InstallVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int PartStyle, int TryId)
1697 {
1698     int i;
1699     int rc = 0;
1700     int state = 0;
1701     HANDLE hDrive;
1702     DWORD dwSize;
1703     BOOL bRet;
1704     CHAR MountDrive;
1705     CHAR DriveName[] = "?:\\";
1706     CHAR DriveLetters[MAX_PATH] = { 0 };
1707     MBR_HEAD MBR;
1708     VTOY_GPT_INFO *pGptInfo = NULL;
1709     UINT64 Part1StartSector = 0;
1710     UINT64 Part1SectorCount = 0;
1711     UINT64 Part2StartSector = 0;
1712     BOOL LargeFAT32 = FALSE;
1713
1714         Log("#####################################################");
1715     Log("InstallVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
1716         PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1717         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1718         Log("#####################################################");
1719
1720     if (PartStyle)
1721     {
1722         pGptInfo = malloc(sizeof(VTOY_GPT_INFO));
1723         memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1724     }
1725
1726     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1727
1728     if (PartStyle)
1729     {
1730         VentoyFillGpt(pPhyDrive->SizeInBytes, pGptInfo);
1731         Part1StartSector = pGptInfo->PartTbl[0].StartLBA;
1732         Part1SectorCount = pGptInfo->PartTbl[0].LastLBA - Part1StartSector + 1;
1733         Part2StartSector = pGptInfo->PartTbl[1].StartLBA;
1734     }
1735     else
1736     {
1737         VentoyFillMBR(pPhyDrive->SizeInBytes, &MBR, PartStyle);
1738         Part1StartSector = MBR.PartTbl[0].StartSectorId;
1739         Part1SectorCount = MBR.PartTbl[0].SectorCount;
1740         Part2StartSector = MBR.PartTbl[1].StartSectorId;
1741     }
1742
1743     Log("Lock disk for clean ............................. ");
1744
1745     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1746     if (hDrive == INVALID_HANDLE_VALUE)
1747     {
1748         Log("Failed to open physical disk");
1749         free(pGptInfo);
1750         return 1;
1751     }
1752
1753     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1754
1755     if (DriveLetters[0] == 0)
1756     {
1757         Log("No drive letter was assigned...");
1758         DriveName[0] = GetFirstUnusedDriveLetter();
1759         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1760     }
1761     else
1762     {
1763         // Unmount all mounted volumes that belong to this drive
1764         // Do it in reverse so that we always end on the first volume letter
1765         for (i = (int)strlen(DriveLetters); i > 0; i--)
1766         {
1767             DriveName[0] = DriveLetters[i - 1];
1768             bRet = DeleteVolumeMountPointA(DriveName);
1769             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1770         }
1771     }
1772
1773     MountDrive = DriveName[0];
1774     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1775
1776     // It kind of blows, but we have to relinquish access to the physical drive
1777     // for VDS to be able to delete the partitions that reside on it...
1778     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1779     CHECK_CLOSE_HANDLE(hDrive);
1780
1781     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1782
1783     if (!VDS_DeleteAllPartitions(pPhyDrive->PhyDrive))
1784     {
1785         Log("Notice: Could not delete partitions: 0x%x, but we continue.", GetLastError());
1786     }
1787
1788     Log("Deleting all partitions ......................... OK");
1789
1790     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1791
1792     Log("Lock disk for write ............................. ");
1793     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1794     if (hDrive == INVALID_HANDLE_VALUE)
1795     {
1796         Log("Failed to GetPhysicalHandle for write.");
1797         rc = 1;
1798         goto End;
1799     }
1800
1801     //Refresh Drive Layout
1802     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1803
1804     disk_io_set_param(hDrive, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1805
1806     PROGRESS_BAR_SET_POS(PT_FORMAT_PART1);
1807
1808     if (PartStyle == 1 && pPhyDrive->PartStyle == 0)
1809     {
1810         Log("Wait for format part1 ...");
1811         Sleep(1000 * 5);
1812     }
1813
1814     if (GetVentoyFsType() == VTOY_FS_FAT32 && (Part1SectorCount * 512 >= FAT32_MAX_LIMIT))
1815     {
1816         Log("Formatting part1 large FAT32 ...");
1817         LargeFAT32 = TRUE;
1818         if (0 != FormatPart1LargeFAT32(pPhyDrive->SizeInBytes, GetClusterSize()))
1819         {
1820             Log("FormatPart1LargeFAT32 failed.");
1821             rc = 1;
1822             goto End;
1823         }
1824     }
1825     else
1826     {
1827         Log("Formatting part1 exFAT ...");
1828         if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1829         {
1830             Log("FormatPart1exFAT failed.");
1831             rc = 1;
1832             goto End;
1833         }
1834     }
1835
1836
1837     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1838     Log("Writing part2 FAT img ...");
1839
1840     if (0 != FormatPart2Fat(hDrive, Part2StartSector))
1841     {
1842         Log("FormatPart2Fat failed.");
1843         rc = 1;
1844         goto End;
1845     }
1846
1847     PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
1848     Log("Writing Boot Image ............................. ");
1849     if (WriteGrubStage1ToPhyDrive(hDrive, PartStyle) != 0)
1850     {
1851         Log("WriteGrubStage1ToPhyDrive failed.");
1852         rc = 1;
1853         goto End;
1854     }
1855
1856     PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
1857     Log("Writing Partition Table ........................ ");
1858     SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1859
1860     if (PartStyle)
1861     {
1862         VTOY_GPT_HDR BackupHead;
1863         LARGE_INTEGER liCurrentPosition;
1864
1865         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1866         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1867         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1868         {
1869             rc = 1;
1870             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1871             goto End;
1872         }
1873
1874         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1875         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1876         {
1877             rc = 1;
1878             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1879             goto End;
1880         }
1881
1882         SET_FILE_POS(0);
1883         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1884         {
1885             rc = 1;
1886             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1887             goto End;
1888         }
1889
1890         Log("Write GPT Info OK ...");
1891         memcpy(&(pPhyDrive->MBR), &(pGptInfo->MBR), 512);
1892     }
1893     else
1894     {
1895         if (!WriteFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL))
1896         {
1897             rc = 1;
1898             Log("Write MBR Failed, dwSize:%u ErrCode:%u", dwSize, GetLastError());
1899             goto End;
1900         }
1901         Log("Write MBR OK ...");
1902         memcpy(&(pPhyDrive->MBR), &MBR, 512);
1903     }
1904
1905     //Refresh Drive Layout
1906     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1907
1908 End:
1909
1910     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1911
1912     if (rc == 0)
1913     {
1914         Log("Mounting Ventoy Partition ....................... ");
1915         Sleep(1000);
1916
1917         state = 0;
1918         memset(DriveLetters, 0, sizeof(DriveLetters));
1919         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1920         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1921
1922         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1923         {
1924             DriveName[0] = DriveLetters[i];
1925             if (IsVentoyLogicalDrive(DriveName[0]))
1926             {
1927                 Log("%s is ventoy part2, delete mountpoint", DriveName);
1928                 DeleteVolumeMountPointA(DriveName);
1929             }
1930             else
1931             {
1932                 Log("%s is ventoy part1, already mounted", DriveName);
1933                 MountDrive = DriveName[0];
1934                 state = 1;
1935             }
1936         }
1937
1938         if (state != 1)
1939         {
1940             Log("need to mount ventoy part1...");
1941
1942             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, Part1StartSector, DriveLetters, sizeof(DriveLetters), FALSE))
1943             {
1944                 DriveName[0] = MountDrive;
1945                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1946                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1947
1948                 if (bRet)
1949                 {
1950                     state = 1;
1951                 }
1952             }
1953             else
1954             {
1955                 Log("Failed to find ventoy volume");
1956             }
1957         }
1958
1959         if (state)
1960         {
1961             if (LargeFAT32)
1962             {
1963                 Log("No need to reformat for large FAT32");
1964                 pPhyDrive->VentoyFsClusterSize = GetVolumeClusterSize(MountDrive);
1965             }
1966             else if (DISK_FormatVolume(MountDrive, GetVentoyFsType(), Part1SectorCount * 512))
1967             {
1968                 Log("Reformat %C:\\ to %s SUCCESS", MountDrive, GetVentoyFsName());
1969                 pPhyDrive->VentoyFsClusterSize = GetVolumeClusterSize(MountDrive);
1970
1971                 if ((GetVentoyFsType() != VTOY_FS_UDF) && (pPhyDrive->VentoyFsClusterSize < 2048))
1972                 {
1973                     for (i = 0; i < 10; i++)
1974                     {
1975                         Log("### Invalid cluster size %d ###", pPhyDrive->VentoyFsClusterSize);
1976                     }
1977                 }
1978             }
1979             else
1980             {
1981                 Log("Reformat %C:\\ to %s FAILED", MountDrive, GetVentoyFsName());
1982             }
1983         }
1984         else
1985         {
1986             Log("Can not reformat %s to %s", DriveName, GetVentoyFsName());
1987         }
1988
1989         Log("OK\n");
1990     }
1991     else
1992     {
1993                 PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1994
1995         FindProcessOccupyDisk(hDrive, pPhyDrive);
1996
1997                 if (!VDS_IsLastAvaliable())
1998                 {
1999                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2000                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2001                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2002                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2003                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2004                 }
2005     }
2006
2007     if (pGptInfo)
2008     {
2009         free(pGptInfo);
2010     }
2011
2012     CHECK_CLOSE_HANDLE(hDrive);
2013     return rc;
2014 }
2015
2016
2017 int PartitionResizeForVentoy(PHY_DRIVE_INFO *pPhyDrive)
2018 {
2019         int i, j;
2020         int rc = 1;
2021         int PhyDrive;
2022         int PartStyle;
2023         INT64 ReservedValue;
2024         UINT64 RecudeBytes;
2025         GUID Guid;
2026         MBR_HEAD MBR;
2027         VTOY_GPT_INFO *pGPT;
2028         MBR_HEAD *pMBR;
2029         DWORD dwSize = 0;
2030         VTOY_GPT_HDR BackupHead;
2031         HANDLE hDrive = INVALID_HANDLE_VALUE;
2032         GUID ZeroGuid = { 0 };
2033         static GUID WindowsDataPartType = { 0xebd0a0a2, 0xb9e5, 0x4433, { 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7 } };
2034         static GUID EspPartType = { 0xc12a7328, 0xf81f, 0x11d2, { 0xba, 0x4b, 0x00, 0xa0, 0xc9, 0x3e, 0xc9, 0x3b } };
2035         static GUID BiosGrubPartType = { 0x21686148, 0x6449, 0x6e6f, { 0x74, 0x4e, 0x65, 0x65, 0x64, 0x45, 0x46, 0x49 } };
2036
2037         Log("#####################################################");
2038         Log("PartitionResizeForVentoy PhyDrive%d <<%s %s %dGB>>",
2039                 pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
2040                 GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
2041         Log("#####################################################");
2042
2043         pGPT = &(pPhyDrive->Gpt);
2044         pMBR = &(pPhyDrive->Gpt.MBR);
2045         Log("Disksize:%llu Part2Start:%llu", pPhyDrive->SizeInBytes, pPhyDrive->ResizePart2StartSector * 512);
2046
2047         if (pMBR->PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
2048         {
2049                 PartStyle = 1;
2050         }
2051         else
2052         {
2053                 PartStyle = 0;
2054         }
2055
2056         PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
2057
2058         RecudeBytes = VENTOY_EFI_PART_SIZE;
2059         ReservedValue = GetReservedSpaceInMB();
2060         if (ReservedValue > 0)
2061         {
2062                 Log("Reduce add reserved space %lldMB", (LONGLONG)ReservedValue);
2063                 RecudeBytes += (UINT64)(ReservedValue * SIZE_1MB);
2064         }
2065
2066
2067         if (pPhyDrive->ResizeNoShrink == FALSE)
2068         {
2069                 Log("Need to shrink the volume");
2070                 if (DISK_ShrinkVolume(pPhyDrive->PhyDrive, pPhyDrive->ResizeVolumeGuid, pPhyDrive->Part1DriveLetter, pPhyDrive->ResizeOldPart1Size, RecudeBytes))
2071                 {
2072                         Log("Shrink volume success, now check again");
2073
2074                         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2075                         if (hDrive == INVALID_HANDLE_VALUE)
2076                         {
2077                                 Log("Failed to GetPhysicalHandle for update.");
2078                                 goto End;
2079                         }
2080
2081                         //Refresh Drive Layout
2082                         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2083
2084                         CHECK_CLOSE_HANDLE(hDrive);
2085
2086
2087                         if (PartResizePreCheck(NULL) && pPhyDrive->ResizeNoShrink)
2088                         {
2089                                 Log("Recheck after Shrink volume success");
2090                                 Log("After shrink Disksize:%llu Part2Start:%llu", pPhyDrive->SizeInBytes, pPhyDrive->ResizePart2StartSector * 512);
2091                         }
2092                         else
2093                         {
2094                                 Log("Recheck after Shrink volume failed %u", pPhyDrive->ResizeNoShrink);
2095                                 goto End;
2096                         }
2097                 }
2098                 else
2099                 {
2100                         Log("Shrink volume failed");
2101                         goto End;
2102                 }
2103         }
2104
2105
2106         //Now try write data
2107         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2108         if (hDrive == INVALID_HANDLE_VALUE)
2109         {
2110                 Log("Failed to GetPhysicalHandle for update.");
2111                 goto End;
2112         }
2113
2114
2115         //Write partition 2 data
2116         PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
2117         if (0 != FormatPart2Fat(hDrive, pPhyDrive->ResizePart2StartSector))
2118         {
2119                 Log("FormatPart2Fat failed.");
2120                 goto End;
2121         }
2122
2123         //Write grub stage2 gap
2124         PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
2125         Log("Writing Boot Image ............................. ");
2126         if (WriteGrubStage1ToPhyDrive(hDrive, PartStyle) != 0)
2127         {
2128                 Log("WriteGrubStage1ToPhyDrive failed.");
2129                 goto End;
2130         }
2131
2132
2133         //Write partition table
2134         PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
2135         Log("Writing partition table ............................. ");
2136
2137         VentoyGetLocalBootImg(&MBR);
2138         CoCreateGuid(&Guid);
2139         memcpy(MBR.BootCode + 0x180, &Guid, 16);
2140         memcpy(pMBR->BootCode, MBR.BootCode, 440);
2141
2142         if (PartStyle == 0)
2143         {
2144                 for (i = 1; i < 4; i++)
2145                 {
2146                         if (pMBR->PartTbl[i].SectorCount == 0)
2147                         {
2148                                 break;
2149                         }
2150                 }
2151
2152                 if (i >= 4)
2153                 {
2154                         Log("Can not find MBR free partition table");
2155                         goto End;
2156                 }
2157
2158                 for (j = i - 1; j > 0; j--)
2159                 {
2160                         Log("Move MBR partition table %d --> %d", j + 1, j + 2);
2161                         memcpy(pMBR->PartTbl + (j + 1), pMBR->PartTbl + j, sizeof(PART_TABLE));
2162                 }
2163
2164         memset(pMBR->PartTbl + 1, 0, sizeof(PART_TABLE));
2165                 VentoyFillMBRLocation(pPhyDrive->SizeInBytes, (UINT32)pPhyDrive->ResizePart2StartSector, VENTOY_EFI_PART_SIZE / 512, pMBR->PartTbl + 1);
2166                 pMBR->PartTbl[0].Active = 0x80; // bootable
2167                 pMBR->PartTbl[1].Active = 0x00;
2168                 pMBR->PartTbl[1].FsFlag = 0xEF; // EFI System Partition
2169
2170                 if (!WriteDataToPhyDisk(hDrive, 0, pMBR, 512))
2171                 {
2172                         Log("Legacy BIOS write MBR failed");
2173                         goto End;
2174                 }
2175         }
2176         else
2177         {
2178                 for (i = 1; i < 128; i++)
2179                 {
2180                         if (memcmp(&(pGPT->PartTbl[i].PartGuid), &ZeroGuid, sizeof(GUID)) == 0)
2181                         {
2182                                 break;
2183                         }
2184                 }
2185
2186                 if (i >= 128)
2187                 {
2188                         Log("Can not find GPT free partition table");
2189                         goto End;
2190                 }
2191
2192                 for (j = i - 1; j > 0; j--)
2193                 {
2194                         Log("Move GPT partition table %d --> %d", j + 1, j + 2);
2195                         memcpy(pGPT->PartTbl + (j + 1), pGPT->PartTbl + j, sizeof(VTOY_GPT_PART_TBL));
2196                 }
2197
2198
2199                 pMBR->BootCode[92] = 0x22;
2200
2201                 // to fix windows issue
2202         memset(pGPT->PartTbl + 1, 0, sizeof(VTOY_GPT_PART_TBL));
2203                 memcpy(&(pGPT->PartTbl[1].PartType), &WindowsDataPartType, sizeof(GUID));
2204                 CoCreateGuid(&(pGPT->PartTbl[1].PartGuid));
2205
2206                 pGPT->PartTbl[1].StartLBA = pGPT->PartTbl[0].LastLBA + 1;
2207                 pGPT->PartTbl[1].LastLBA = pGPT->PartTbl[1].StartLBA + VENTOY_EFI_PART_SIZE / 512 - 1;
2208                 pGPT->PartTbl[1].Attr = 0xC000000000000001ULL;
2209                 memcpy(pGPT->PartTbl[1].Name, L"VTOYEFI", 7 * 2);
2210
2211                 //Update CRC
2212                 pGPT->Head.PartTblCrc = VentoyCrc32(pGPT->PartTbl, sizeof(pGPT->PartTbl));
2213         pGPT->Head.Crc = 0;
2214                 pGPT->Head.Crc = VentoyCrc32(&(pGPT->Head), pGPT->Head.Length);
2215
2216                 Log("pGPT->Head.EfiStartLBA=%llu", (ULONGLONG)pGPT->Head.EfiStartLBA);
2217                 Log("pGPT->Head.EfiBackupLBA=%llu", (ULONGLONG)pGPT->Head.EfiBackupLBA);
2218
2219                 VentoyFillBackupGptHead(pGPT, &BackupHead);
2220                 if (!WriteDataToPhyDisk(hDrive, pGPT->Head.EfiBackupLBA * 512, &BackupHead, 512))
2221                 {
2222                         Log("UEFI write backup head failed");
2223                         goto End;
2224                 }
2225
2226                 if (!WriteDataToPhyDisk(hDrive, (pGPT->Head.EfiBackupLBA - 32) * 512, pGPT->PartTbl, 512 * 32))
2227                 {
2228                         Log("UEFI write backup partition table failed");
2229                         goto End;
2230                 }
2231
2232                 if (!WriteDataToPhyDisk(hDrive, 0, pGPT, 512 * 34))
2233                 {
2234                         Log("UEFI write MBR & Main partition table failed");
2235                         goto End;
2236                 }
2237         }
2238
2239
2240
2241         //Refresh Drive Layout
2242         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2243
2244         //We must close handle here, because it will block the refresh bellow
2245         CHECK_CLOSE_HANDLE(hDrive);
2246
2247         Sleep(2000);
2248
2249         //Refresh disk list
2250         PhyDrive = pPhyDrive->PhyDrive;
2251
2252         Log("#### Now Refresh PhyDrive ####");
2253         Ventoy2DiskDestroy();
2254         Ventoy2DiskInit();
2255
2256         pPhyDrive = GetPhyDriveInfoByPhyDrive(PhyDrive);
2257         if (pPhyDrive)
2258         {
2259                 if (pPhyDrive->VentoyVersion[0] == 0)
2260                 {
2261                         Log("After process the Ventoy version is still invalid");
2262                         goto End;
2263                 }
2264
2265                 Log("### Ventoy non-destructive installation successfully finished <%s>", pPhyDrive->VentoyVersion);
2266         }
2267         else
2268         {
2269                 Log("### Ventoy non-destructive installation successfully finished <not found>");
2270         }
2271
2272         InitComboxCtrl(g_DialogHwnd, PhyDrive);
2273
2274         rc = 0;
2275
2276 End:
2277         CHECK_CLOSE_HANDLE(hDrive);
2278         return rc;
2279 }
2280
2281
2282 static BOOL DiskCheckWriteAccess(HANDLE hDrive)
2283 {
2284         DWORD dwSize;
2285         BOOL ret = FALSE;
2286         BOOL bRet = FALSE;
2287         BYTE Buffer[512];
2288         LARGE_INTEGER liCurPosition;
2289         LARGE_INTEGER liNewPosition;
2290
2291         liCurPosition.QuadPart = 2039 * 512;
2292         liNewPosition.QuadPart = 0;
2293         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2294                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2295         {
2296                 Log("SetFilePointer1 Failed %u", LASTERR);
2297                 goto out;
2298         }
2299
2300
2301         dwSize = 0;
2302         ret = ReadFile(hDrive, Buffer, 512, &dwSize, NULL);
2303         if ((!ret) || (dwSize != 512))
2304         {
2305                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2306                 goto out;
2307         }
2308
2309
2310         liCurPosition.QuadPart = 2039 * 512;
2311         liNewPosition.QuadPart = 0;
2312         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2313                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2314         {
2315                 Log("SetFilePointer2 Failed %u", LASTERR);
2316                 goto out;
2317         }
2318
2319         dwSize = 0;
2320         ret = WriteFile(hDrive, Buffer, 512, &dwSize, NULL);
2321         if ((!ret) || dwSize != 512)
2322         {
2323                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
2324                 goto out;
2325         }
2326
2327         bRet = TRUE;
2328
2329 out:
2330
2331         return bRet;
2332 }
2333
2334 static BOOL BackupDataBeforeCleanDisk(int PhyDrive, UINT64 DiskSize, BYTE **pBackup)
2335 {
2336         DWORD dwSize;
2337         DWORD dwStatus;
2338         BOOL Return = FALSE;
2339         BOOL ret = FALSE;
2340         BYTE *backup = NULL;
2341         UINT64 offset;
2342         HANDLE hDrive = INVALID_HANDLE_VALUE;
2343         LARGE_INTEGER liCurPosition;
2344         LARGE_INTEGER liNewPosition;
2345         VTOY_GPT_INFO *pGPT = NULL;
2346
2347         Log("BackupDataBeforeCleanDisk %d", PhyDrive);
2348
2349         // step1: check write access
2350         hDrive = GetPhysicalHandle(PhyDrive, TRUE, TRUE, FALSE);
2351         if (hDrive == INVALID_HANDLE_VALUE)
2352         {
2353                 Log("Failed to GetPhysicalHandle for write.");
2354                 goto out;
2355         }
2356
2357         if (DiskCheckWriteAccess(hDrive))
2358         {
2359                 Log("DiskCheckWriteAccess success");
2360                 CHECK_CLOSE_HANDLE(hDrive);
2361         }
2362         else
2363         {
2364                 Log("DiskCheckWriteAccess failed");
2365                 goto out;
2366         }
2367
2368         //step2 backup 4MB data
2369         backup = malloc(SIZE_1MB * 4);
2370         if (!backup)
2371         {
2372                 goto out;
2373         }
2374
2375         hDrive = GetPhysicalHandle(PhyDrive, FALSE, FALSE, FALSE);
2376         if (hDrive == INVALID_HANDLE_VALUE)
2377         {
2378                 goto out;
2379         }
2380
2381         //read first 2MB
2382         dwStatus = SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2383         if (dwStatus != 0)
2384         {
2385                 goto out;
2386         }
2387
2388         dwSize = 0;
2389         ret = ReadFile(hDrive, backup, SIZE_2MB, &dwSize, NULL);
2390         if ((!ret) || (dwSize != SIZE_2MB))
2391         {
2392                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2393                 goto out;
2394         }
2395
2396         pGPT = (VTOY_GPT_INFO *)backup;
2397         offset = pGPT->Head.EfiBackupLBA * 512;
2398         if (offset >= (DiskSize - SIZE_2MB) && offset < DiskSize)
2399         {
2400                 Log("EFI partition table check success");
2401         }
2402         else
2403         {
2404                 Log("Backup EFI LBA not in last 2MB range: %llu", pGPT->Head.EfiBackupLBA);
2405                 goto out;
2406         }
2407
2408         //read last 2MB
2409         liCurPosition.QuadPart = DiskSize - SIZE_2MB;
2410         liNewPosition.QuadPart = 0;
2411         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2412                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2413         {
2414                 goto out;
2415         }
2416
2417         dwSize = 0;
2418         ret = ReadFile(hDrive, backup + SIZE_2MB, SIZE_2MB, &dwSize, NULL);
2419         if ((!ret) || (dwSize != SIZE_2MB))
2420         {
2421                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2422                 goto out;
2423         }
2424
2425         *pBackup = backup;
2426         backup = NULL; //For don't free later
2427         Return = TRUE;
2428
2429 out:
2430         CHECK_CLOSE_HANDLE(hDrive);
2431         if (backup)
2432                 free(backup);
2433
2434         return Return;
2435 }
2436
2437
2438 static BOOL WriteBackupDataToDisk(HANDLE hDrive, UINT64 Offset, BYTE *Data, DWORD Length)
2439 {
2440         DWORD dwSize = 0;
2441         BOOL ret = FALSE;
2442         LARGE_INTEGER liCurPosition;
2443         LARGE_INTEGER liNewPosition;
2444
2445         Log("WriteBackupDataToDisk %llu %p %u", Offset, Data, Length);
2446
2447         liCurPosition.QuadPart = Offset;
2448         liNewPosition.QuadPart = 0;
2449         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2450                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2451         {
2452                 return FALSE;
2453         }
2454
2455         ret = WriteFile(hDrive, Data, Length, &dwSize, NULL);
2456         if ((!ret) || dwSize != Length)
2457         {
2458                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
2459                 return FALSE;
2460         }
2461
2462         Log("WriteBackupDataToDisk %llu %p %u success", Offset, Data, Length);
2463         return TRUE;
2464 }
2465
2466
2467 int UpdateVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int TryId)
2468 {
2469         int i;
2470         int rc = 0;
2471         int MaxRetry = 4;
2472         BOOL ForceMBR = FALSE;
2473         BOOL Esp2Basic = FALSE;
2474         BOOL ChangeAttr = FALSE;
2475         BOOL CleanDisk = FALSE;
2476         BOOL DelEFI = FALSE;
2477         BOOL bWriteBack = TRUE;
2478         HANDLE hVolume;
2479         HANDLE hDrive;
2480         DWORD Status;
2481         DWORD dwSize;
2482         BOOL bRet;
2483         CHAR DriveName[] = "?:\\";
2484         CHAR DriveLetters[MAX_PATH] = { 0 };
2485         CHAR BackBinFile[MAX_PATH];
2486         UINT64 StartSector;
2487         UINT64 ReservedMB = 0;
2488         MBR_HEAD BootImg;
2489         MBR_HEAD MBR;
2490         BYTE *pBackup = NULL;
2491         VTOY_GPT_INFO *pGptInfo = NULL;
2492         VTOY_GPT_INFO *pGptBkup = NULL;
2493         UINT8 ReservedData[4096];
2494
2495         Log("#####################################################");
2496         Log("UpdateVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
2497                 pPhyDrive->PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
2498                 GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
2499         Log("#####################################################");
2500
2501         PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
2502
2503         Log("Lock disk for umount ............................ ");
2504
2505         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
2506         if (hDrive == INVALID_HANDLE_VALUE)
2507         {
2508                 Log("Failed to open physical disk");
2509                 return 1;
2510         }
2511
2512         if (pPhyDrive->PartStyle)
2513         {
2514                 pGptInfo = malloc(2 * sizeof(VTOY_GPT_INFO));
2515                 if (!pGptInfo)
2516                 {
2517                         return 1;
2518                 }
2519
2520                 memset(pGptInfo, 0, 2 * sizeof(VTOY_GPT_INFO));
2521                 pGptBkup = pGptInfo + 1;
2522
2523                 // Read GPT Info
2524                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2525                 ReadFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL);
2526                 memcpy(pGptBkup, pGptInfo, sizeof(VTOY_GPT_INFO));
2527
2528                 //MBR will be used to compare with local boot image
2529                 memcpy(&MBR, &pGptInfo->MBR, sizeof(MBR_HEAD));
2530
2531                 StartSector = pGptInfo->PartTbl[1].StartLBA;
2532                 Log("GPT StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2533
2534                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512) - 33) / 2048;
2535                 Log("GPT Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2536         }
2537         else
2538         {
2539                 // Read MBR
2540                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2541                 ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
2542
2543                 StartSector = MBR.PartTbl[1].StartSectorId;
2544                 Log("MBR StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2545
2546                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512)) / 2048;
2547                 Log("MBR Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2548         }
2549
2550         //Read Reserved Data
2551         SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2552         ReadFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2553
2554         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
2555
2556         if (DriveLetters[0] == 0)
2557         {
2558                 Log("No drive letter was assigned...");
2559         }
2560         else
2561         {
2562                 // Unmount all mounted volumes that belong to this drive
2563                 // Do it in reverse so that we always end on the first volume letter
2564                 for (i = (int)strlen(DriveLetters); i > 0; i--)
2565                 {
2566                         DriveName[0] = DriveLetters[i - 1];
2567                         if (IsVentoyLogicalDrive(DriveName[0]))
2568                         {
2569                                 Log("%s is ventoy logical drive", DriveName);
2570                                 bRet = DeleteVolumeMountPointA(DriveName);
2571                                 Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, LASTERR);
2572                                 break;
2573                         }
2574                 }
2575         }
2576
2577         // It kind of blows, but we have to relinquish access to the physical drive
2578         // for VDS to be able to delete the partitions that reside on it...
2579         DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2580         CHECK_CLOSE_HANDLE(hDrive);
2581
2582         if (pPhyDrive->PartStyle == 1)
2583         {
2584                 Log("TryId=%d EFI GPT partition type is 0x%llx", TryId, pPhyDrive->Part2GPTAttr);
2585                 PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
2586
2587                 if (TryId == 1)
2588                 {
2589                         Log("Change GPT partition type to ESP");
2590                         if (DISK_ChangeVtoyEFI2ESP(pPhyDrive->PhyDrive, StartSector * 512ULL))
2591                         {
2592                                 Esp2Basic = TRUE;
2593                                 Sleep(3000);
2594                         }
2595                 }
2596                 else if (TryId == 2)
2597                 {
2598                         Log("Change GPT partition attribute");
2599                         if (DISK_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, StartSector * 512ULL, 0x8000000000000001))
2600                         {
2601                                 ChangeAttr = TRUE;
2602                                 Sleep(2000);
2603                         }
2604                 }
2605                 else if (TryId == 3)
2606                 {
2607                         DISK_DeleteVtoyEFIPartition(pPhyDrive->PhyDrive, StartSector * 512ULL);
2608                         DelEFI = TRUE;
2609                 }
2610                 else if (TryId == 4)
2611                 {
2612                         Log("Clean disk GPT partition table");
2613                         if (BackupDataBeforeCleanDisk(pPhyDrive->PhyDrive, pPhyDrive->SizeInBytes, &pBackup))
2614                         {
2615                                 sprintf_s(BackBinFile, sizeof(BackBinFile), ".\\ventoy\\phydrive%d_%u_%d.bin",
2616                                         pPhyDrive->PhyDrive, GetCurrentProcessId(), g_backup_bin_index++);
2617                                 SaveBufToFile(BackBinFile, pBackup, 4 * SIZE_1MB);
2618                                 Log("Save backup data to %s", BackBinFile);
2619
2620                                 Log("Success to backup data before clean");
2621                                 CleanDisk = TRUE;
2622                                 DISK_CleanDisk(pPhyDrive->PhyDrive);
2623                                 Sleep(3000);
2624                         }
2625                         else
2626                         {
2627                                 Log("Failed to backup data before clean");
2628                         }
2629                 }
2630         }
2631
2632     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
2633
2634     Log("Lock disk for update ............................ ");
2635     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2636     if (hDrive == INVALID_HANDLE_VALUE)
2637     {
2638         Log("Failed to GetPhysicalHandle for write.");
2639         rc = 1;
2640         goto End;
2641     }
2642
2643     PROGRESS_BAR_SET_POS(PT_LOCK_VOLUME);
2644
2645     Log("Lock volume for update .......................... ");
2646     hVolume = INVALID_HANDLE_VALUE;
2647
2648         //If we change VTOYEFI to ESP, it can not have s volume name, so don't try to get it.
2649         if (CleanDisk)
2650         {
2651                 //writeback the last 2MB
2652                 if (!WriteBackupDataToDisk(hDrive, pPhyDrive->SizeInBytes - SIZE_2MB, pBackup + SIZE_2MB, SIZE_2MB))
2653                 {
2654                         bWriteBack = FALSE;
2655                 }
2656
2657                 //write the first 2MB except parttable
2658                 if (!WriteBackupDataToDisk(hDrive, 34 * 512, pBackup + 34 * 512, SIZE_2MB - 34 * 512))
2659                 {
2660                         bWriteBack = FALSE;
2661                 }
2662
2663                 Status = ERROR_NOT_FOUND;
2664         }
2665         else if (DelEFI)
2666         {
2667                 Status = ERROR_NOT_FOUND;
2668         }
2669         else if (Esp2Basic)
2670         {
2671                 Status = ERROR_NOT_FOUND;
2672         }
2673         else
2674         {
2675                 for (i = 0; i < MaxRetry; i++)
2676                 {
2677                         Status = GetVentoyVolumeName(pPhyDrive->PhyDrive, StartSector, DriveLetters, sizeof(DriveLetters), TRUE);
2678                         if (ERROR_SUCCESS == Status)
2679                         {
2680                                 break;
2681                         }
2682                         else
2683                         {
2684                                 Log("==== Volume not found, wait and retry %d... ====", i);
2685                                 Sleep(2);
2686                         }
2687                 }
2688         }
2689
2690     if (ERROR_SUCCESS == Status)
2691     {
2692         Log("Now lock and dismount volume <%s>", DriveLetters);
2693
2694         for (i = 0; i < MaxRetry; i++)
2695         {
2696             hVolume = CreateFileA(DriveLetters,
2697                 GENERIC_READ | GENERIC_WRITE,
2698                 FILE_SHARE_READ,
2699                 NULL,
2700                 OPEN_EXISTING,
2701                 FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
2702                 NULL);
2703
2704             if (hVolume == INVALID_HANDLE_VALUE)
2705             {
2706                 Log("Failed to create file volume, errcode:%u, wait and retry ...", LASTERR);
2707                 Sleep(2000);
2708             }
2709             else
2710             {
2711                 break;
2712             }
2713         }
2714
2715         if (hVolume == INVALID_HANDLE_VALUE)
2716         {
2717             Log("Failed to create file volume, errcode:%u", LASTERR);
2718         }
2719         else
2720         {
2721             bRet = DeviceIoControl(hVolume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2722             Log("FSCTL_LOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2723
2724             bRet = DeviceIoControl(hVolume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2725             Log("FSCTL_DISMOUNT_VOLUME bRet:%u code:%u", bRet, LASTERR);
2726         }
2727     }
2728     else if (ERROR_NOT_FOUND == Status)
2729     {
2730         Log("Volume not found, maybe not supported");
2731     }
2732     else
2733     {
2734         rc = 1;
2735         goto End;
2736     }
2737
2738         bRet = TryWritePart2(hDrive, StartSector);
2739         if (FALSE == bRet && Esp2Basic)
2740         {
2741                 Log("TryWritePart2 agagin ...");
2742                 Sleep(3000);
2743                 bRet = TryWritePart2(hDrive, StartSector);
2744         }
2745
2746         if (!bRet)
2747     {
2748                 if (pPhyDrive->PartStyle == 0)
2749                 {
2750                         if (DiskCheckWriteAccess(hDrive))
2751                         {
2752                                 Log("MBR DiskCheckWriteAccess success");
2753
2754                                 ForceMBR = TRUE;
2755
2756                                 Log("Try write failed, now delete partition 2 for MBR...");
2757                                 CHECK_CLOSE_HANDLE(hDrive);
2758
2759                                 Log("Now delete partition 2...");
2760                                 DISK_DeleteVtoyEFIPartition(pPhyDrive->PhyDrive, StartSector * 512ULL);
2761
2762                                 hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2763                                 if (hDrive == INVALID_HANDLE_VALUE)
2764                                 {
2765                                         Log("Failed to GetPhysicalHandle for write.");
2766                                         rc = 1;
2767                                         goto End;
2768                                 }
2769                         }
2770                         else
2771                         {
2772                                 Log("MBR DiskCheckWriteAccess failed");
2773                         }
2774                 }
2775                 else
2776                 {
2777                         Log("TryWritePart2 failed ....");
2778                         rc = 1;
2779                         goto End;
2780                 }
2781     }
2782
2783     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
2784
2785     Log("Write Ventoy to disk ............................ ");
2786     if (0 != FormatPart2Fat(hDrive, StartSector))
2787     {
2788         rc = 1;
2789         goto End;
2790     }
2791
2792     if (hVolume != INVALID_HANDLE_VALUE)
2793     {
2794         bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2795         Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2796         CHECK_CLOSE_HANDLE(hVolume);
2797     }
2798
2799     Log("Updating Boot Image ............................. ");
2800     if (WriteGrubStage1ToPhyDrive(hDrive, pPhyDrive->PartStyle) != 0)
2801     {
2802         rc = 1;
2803         goto End;
2804     }
2805
2806     //write reserved data
2807     SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2808     bRet = WriteFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2809     Log("Write resv data ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2810
2811     // Boot Image
2812     VentoyGetLocalBootImg(&BootImg);
2813
2814     // Use Old UUID
2815     memcpy(BootImg.BootCode + 0x180, MBR.BootCode + 0x180, 16);
2816     if (pPhyDrive->PartStyle)
2817     {
2818         BootImg.BootCode[92] = 0x22;
2819     }
2820
2821     if (ForceMBR == FALSE && memcmp(BootImg.BootCode, MBR.BootCode, 440) == 0)
2822     {
2823         Log("Boot image has no difference, no need to write.");
2824     }
2825     else
2826     {
2827         Log("Boot image need to write %u.", ForceMBR);
2828
2829         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2830
2831         memcpy(MBR.BootCode, BootImg.BootCode, 440);
2832         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2833         Log("Write Boot Image ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2834     }
2835
2836     if (pPhyDrive->PartStyle == 0)
2837     {
2838         if (0x00 == MBR.PartTbl[0].Active && 0x80 == MBR.PartTbl[1].Active)
2839         {
2840             Log("Need to chage 1st partition active and 2nd partition inactive.");
2841
2842             MBR.PartTbl[0].Active = 0x80;
2843             MBR.PartTbl[1].Active = 0x00;
2844
2845             SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2846             bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2847             Log("Write NEW MBR ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2848         }
2849     }
2850
2851         if (CleanDisk)
2852         {
2853                 if (!WriteBackupDataToDisk(hDrive, 0, pBackup, 34 * 512))
2854                 {
2855                         bWriteBack = FALSE;
2856                 }
2857
2858                 free(pBackup);
2859
2860                 if (bWriteBack)
2861                 {
2862                         Log("Write backup data success, now delete %s", BackBinFile);
2863                         DeleteFileA(BackBinFile);
2864                 }
2865                 else
2866                 {
2867                         Log("Write backup data failed");
2868                 }
2869
2870                 Sleep(1000);
2871         }
2872         else if (DelEFI)
2873         {
2874                 VTOY_GPT_HDR BackupHdr;
2875
2876                 VentoyFillBackupGptHead(pGptBkup, &BackupHdr);
2877                 if (!WriteBackupDataToDisk(hDrive, 512 * pGptBkup->Head.EfiBackupLBA, (BYTE*)(&BackupHdr), 512))
2878                 {
2879                         bWriteBack = FALSE;
2880                 }
2881
2882                 if (!WriteBackupDataToDisk(hDrive, 512 * (pGptBkup->Head.EfiBackupLBA - 32), (BYTE*)(pGptBkup->PartTbl), 32 * 512))
2883                 {
2884                         bWriteBack = FALSE;
2885                 }
2886
2887                 if (!WriteBackupDataToDisk(hDrive, 512, (BYTE*)pGptBkup + 512, 33 * 512))
2888                 {
2889                         bWriteBack = FALSE;
2890                 }
2891
2892                 if (bWriteBack)
2893                 {
2894                         Log("Write backup partition table success");
2895                 }
2896                 else
2897                 {
2898                         Log("Write backup partition table failed");
2899                 }
2900
2901                 Sleep(1000);
2902         }
2903
2904     //Refresh Drive Layout
2905     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2906
2907 End:
2908
2909         if (hVolume != INVALID_HANDLE_VALUE)
2910         {
2911                 bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2912                 Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2913                 CHECK_CLOSE_HANDLE(hVolume);
2914         }
2915
2916     if (rc == 0)
2917     {
2918         Log("OK");
2919     }
2920     else
2921     {
2922                 PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
2923         FindProcessOccupyDisk(hDrive, pPhyDrive);
2924     }
2925
2926     CHECK_CLOSE_HANDLE(hDrive);
2927
2928     if (Esp2Basic)
2929     {
2930                 Log("Recover GPT partition type to basic");
2931                 DISK_ChangeVtoyEFI2Basic(pPhyDrive->PhyDrive, StartSector * 512);
2932     }
2933
2934         if (pPhyDrive->PartStyle == 1)
2935         {
2936                 if (ChangeAttr || ((pPhyDrive->Part2GPTAttr >> 56) != 0xC0))
2937                 {
2938                         Log("Change EFI partition attr %u <0x%llx> to <0x%llx>", ChangeAttr, pPhyDrive->Part2GPTAttr, 0xC000000000000001ULL);
2939                         if (DISK_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, StartSector * 512ULL, 0xC000000000000001ULL))
2940                         {
2941                                 Log("Change EFI partition attr success");
2942                                 pPhyDrive->Part2GPTAttr = 0xC000000000000001ULL;
2943                         }
2944                         else
2945                         {
2946                                 Log("Change EFI partition attr failed");
2947                         }
2948                 }
2949         }
2950
2951     if (pGptInfo)
2952     {
2953         free(pGptInfo);
2954     }
2955
2956     return rc;
2957 }
2958
2959