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             fl_remove("/EFI/BOOT/grub.efi");
 735
 736                         file = fl_fopen("/EFI/BOOT/BOOTX64.EFI", "wb");
 737                         Log("Open bootx64 efi file %p ", file);
 738                         if (file)
 739                         {
 740                                 if (filebuf)
 741                                 {
 742                                         fl_fwrite(filebuf, 1, size, file);
 743                                 }
 744
 745                                 fl_fflush(file);
 746                                 fl_fclose(file);
 747                         }
 748
 749                         if (filebuf)
 750                         {
 751                                 free(filebuf);
 752                         }
 753                 }
 754
 755         file = fl_fopen("/EFI/BOOT/grubia32_real.efi", "rb");
 756         Log("Open ventoy efi file %p ", file);
 757         if (file)
 758         {
 759             fl_fseek(file, 0, SEEK_END);
 760             size = (int)fl_ftell(file);
 761             fl_fseek(file, 0, SEEK_SET);
 762
 763             Log("ventoy efi file size %d ...", size);
 764
 765             filebuf = (char *)malloc(size);
 766             if (filebuf)
 767             {
 768                 fl_fread(filebuf, 1, size, file);
 769             }
 770
 771             fl_fclose(file);
 772
 773             Log("Now delete all efi files ...");
 774             fl_remove("/EFI/BOOT/BOOTIA32.EFI");
 775             fl_remove("/EFI/BOOT/grubia32.efi");
 776             fl_remove("/EFI/BOOT/grubia32_real.efi");
 777             fl_remove("/EFI/BOOT/mmia32.efi");
 778
 779             file = fl_fopen("/EFI/BOOT/BOOTIA32.EFI", "wb");
 780             Log("Open bootia32 efi file %p ", file);
 781             if (file)
 782             {
 783                 if (filebuf)
 784                 {
 785                     fl_fwrite(filebuf, 1, size, file);
 786                 }
 787
 788                 fl_fflush(file);
 789                 fl_fclose(file);
 790             }
 791
 792             if (filebuf)
 793             {
 794                 free(filebuf);
 795             }
 796         }
 797
 798         }
 799         else
 800         {
 801                 rc = 1;
 802         }
 803
 804         fl_shutdown();
 805
 806         return rc;
 807 }
 808
 809
 810
 811 static int disk_xz_flush(void *src, unsigned int size)
 812 {
 813     unsigned int i;
 814     BYTE *buf = (BYTE *)src;
 815
 816     for (i = 0; i < size; i++)
 817     {
 818         *g_part_img_pos = *buf++;
 819
 820         g_disk_unxz_len++;
 821         if ((g_disk_unxz_len % SIZE_1MB) == 0)
 822         {
 823             g_part_img_pos = g_part_img_buf[g_disk_unxz_len / SIZE_1MB];
 824         }
 825         else
 826         {
 827             g_part_img_pos++;
 828         }
 829     }
 830
 831     return (int)size;
 832 }
 833
 834 static void unxz_error(char *x)
 835 {
 836     Log("%s", x);
 837 }
 838
 839 static BOOL TryWritePart2(HANDLE hDrive, UINT64 StartSectorId)
 840 {
 841     BOOL bRet;
 842     DWORD TrySize = 16 * 1024;
 843     DWORD dwSize;
 844     BYTE *Buffer = NULL;
 845     unsigned char *data = NULL;
 846     LARGE_INTEGER liCurrentPosition;
 847
 848     liCurrentPosition.QuadPart = StartSectorId * 512;
 849     SetFilePointerEx(hDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 850
 851     Buffer = malloc(TrySize);
 852
 853     bRet = WriteFile(hDrive, Buffer, TrySize, &dwSize, NULL);
 854
 855     free(Buffer);
 856
 857     Log("Try write part2 bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 858
 859     if (bRet && dwSize == TrySize)
 860     {
 861         return TRUE;
 862     }
 863
 864     return FALSE;
 865 }
 866
 867 static int FormatPart2Fat(HANDLE hDrive, UINT64 StartSectorId)
 868 {
 869     int i;
 870     int rc = 0;
 871     int len = 0;
 872     int writelen = 0;
 873     int partwrite = 0;
 874     int Pos = PT_WRITE_VENTOY_START;
 875     DWORD dwSize = 0;
 876     BOOL bRet;
 877     unsigned char *data = NULL;
 878     LARGE_INTEGER liCurrentPosition;
 879         LARGE_INTEGER liNewPosition;
 880     BYTE *CheckBuf = NULL;
 881
 882         Log("FormatPart2Fat %llu...", (ULONGLONG)StartSectorId);
 883
 884     CheckBuf = malloc(SIZE_1MB);
 885     if (!CheckBuf)
 886     {
 887         Log("Failed to malloc check buf");
 888         return 1;
 889     }
 890
 891     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&data, &len);
 892     if (rc)
 893     {
 894         Log("Failed to read img file %p %u", data, len);
 895         free(CheckBuf);
 896         return 1;
 897     }
 898
 899     liCurrentPosition.QuadPart = StartSectorId * 512;
 900     SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
 901
 902     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
 903
 904     g_part_img_buf[0] = (BYTE *)malloc(VENTOY_EFI_PART_SIZE);
 905     if (g_part_img_buf[0])
 906     {
 907         Log("Malloc whole img buffer success, now decompress ...");
 908         unxz(data, len, NULL, NULL, g_part_img_buf[0], &writelen, unxz_error);
 909
 910         if (len == writelen)
 911         {
 912             Log("decompress finished success");
 913
 914                         VentoyProcSecureBoot(g_SecureBoot);
 915
 916             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 917             {
 918                 dwSize = 0;
 919                                 bRet = WriteFile(hDrive, g_part_img_buf[0] + i * SIZE_1MB, SIZE_1MB, &dwSize, NULL);
 920                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 921
 922                 if (!bRet)
 923                 {
 924                     rc = 1;
 925                     goto End;
 926                 }
 927
 928                 PROGRESS_BAR_SET_POS(Pos);
 929                 if (i % 2 == 0)
 930                 {
 931                     Pos++;
 932                 }
 933             }
 934
 935             //Read and check the data
 936             liCurrentPosition.QuadPart = StartSectorId * 512;
 937             SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
 938
 939             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 940             {
 941                 bRet = ReadFile(hDrive, CheckBuf, SIZE_1MB, &dwSize, NULL);
 942                 Log("Read part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 943
 944                 if (!bRet || memcmp(CheckBuf, g_part_img_buf[0] + i * SIZE_1MB, SIZE_1MB))
 945                 {
 946                     Log("### [Check Fail] The data write and read does not match");
 947                     rc = 1;
 948                     goto End;
 949                 }
 950
 951                 PROGRESS_BAR_SET_POS(Pos);
 952                 if (i % 2 == 0)
 953                 {
 954                     Pos++;
 955                 }
 956             }
 957         }
 958         else
 959         {
 960             rc = 1;
 961             Log("decompress finished failed");
 962             goto End;
 963         }
 964     }
 965     else
 966     {
 967         Log("Failed to malloc whole img size %u, now split it", VENTOY_EFI_PART_SIZE);
 968
 969         partwrite = 1;
 970         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 971         {
 972             g_part_img_buf[i] = (BYTE *)malloc(SIZE_1MB);
 973             if (g_part_img_buf[i] == NULL)
 974             {
 975                 rc = 1;
 976                 goto End;
 977             }
 978         }
 979
 980         Log("Malloc part img buffer success, now decompress ...");
 981
 982         g_part_img_pos = g_part_img_buf[0];
 983
 984         unxz(data, len, NULL, disk_xz_flush, NULL, NULL, unxz_error);
 985
 986         if (g_disk_unxz_len == VENTOY_EFI_PART_SIZE)
 987         {
 988             Log("decompress finished success");
 989
 990                         VentoyProcSecureBoot(g_SecureBoot);
 991
 992             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 993             {
 994                 dwSize = 0;
 995                 bRet = WriteFile(hDrive, g_part_img_buf[i], SIZE_1MB, &dwSize, NULL);
 996                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 997
 998                 if (!bRet)
 999                 {
1000                     rc = 1;
1001                     goto End;
1002                 }
1003
1004                 PROGRESS_BAR_SET_POS(Pos);
1005                 if (i % 2 == 0)
1006                 {
1007                     Pos++;
1008                 }
1009             }
1010
1011             //Read and check the data
1012             liCurrentPosition.QuadPart = StartSectorId * 512;
1013             SetFilePointerEx(hDrive, liCurrentPosition, &liNewPosition, FILE_BEGIN);
1014
1015             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1016             {
1017                 bRet = ReadFile(hDrive, CheckBuf, SIZE_1MB, &dwSize, NULL);
1018                 Log("Read part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
1019
1020                 if (!bRet || memcmp(CheckBuf, g_part_img_buf[i], SIZE_1MB))
1021                 {
1022                     Log("### [Check Fail] The data write and read does not match");
1023                     rc = 1;
1024                     goto End;
1025                 }
1026
1027                 PROGRESS_BAR_SET_POS(Pos);
1028                 if (i % 2 == 0)
1029                 {
1030                     Pos++;
1031                 }
1032             }
1033         }
1034         else
1035         {
1036             rc = 1;
1037             Log("decompress finished failed");
1038             goto End;
1039         }
1040     }
1041
1042 End:
1043
1044     if (data) free(data);
1045     if (CheckBuf)free(CheckBuf);
1046
1047     if (partwrite)
1048     {
1049         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1050         {
1051             if (g_part_img_buf[i]) free(g_part_img_buf[i]);
1052         }
1053     }
1054     else
1055     {
1056         if (g_part_img_buf[0]) free(g_part_img_buf[0]);
1057     }
1058
1059     return rc;
1060 }
1061
1062 static int WriteGrubStage1ToPhyDrive(HANDLE hDrive, int PartStyle)
1063 {
1064     int Len = 0;
1065     int readLen = 0;
1066     BOOL bRet;
1067     DWORD dwSize;
1068     BYTE *ImgBuf = NULL;
1069     BYTE *RawBuf = NULL;
1070
1071     Log("WriteGrubStage1ToPhyDrive ...");
1072
1073     RawBuf = (BYTE *)malloc(SIZE_1MB);
1074     if (!RawBuf)
1075     {
1076         return 1;
1077     }
1078
1079     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1080     {
1081         Log("Failed to read stage1 img");
1082         free(RawBuf);
1083         return 1;
1084     }
1085
1086     unxz(ImgBuf, Len, NULL, NULL, RawBuf, &readLen, unxz_error);
1087
1088     if (PartStyle)
1089     {
1090         Log("Write GPT stage1 ...");
1091         RawBuf[500] = 35;//update blocklist
1092         SetFilePointer(hDrive, 512 * 34, NULL, FILE_BEGIN);
1093         bRet = WriteFile(hDrive, RawBuf, SIZE_1MB - 512 * 34, &dwSize, NULL);
1094     }
1095     else
1096     {
1097         Log("Write MBR stage1 ...");
1098         SetFilePointer(hDrive, 512, NULL, FILE_BEGIN);
1099         bRet = WriteFile(hDrive, RawBuf, SIZE_1MB - 512, &dwSize, NULL);
1100     }
1101
1102     Log("WriteFile Ret:%u dwSize:%u ErrCode:%u", bRet, dwSize, GetLastError());
1103
1104     free(RawBuf);
1105     free(ImgBuf);
1106     return 0;
1107 }
1108
1109
1110 static int FormatPart1LargeFAT32(UINT64 DiskSizeBytes, int CluserSize)
1111 {
1112     MKFS_PARM Option;
1113     FRESULT Ret;
1114     FATFS FS;
1115
1116     Option.fmt = FM_FAT32;
1117     Option.n_fat = 1;
1118     Option.align = 8;
1119     Option.n_root = 1;
1120
1121     if (CluserSize == 0)
1122     {
1123         // < 32GB select 32KB as cluster size
1124         // > 32GB select 128KB as cluster size
1125         if (DiskSizeBytes / 1024 / 1024 / 1024 <= 32)
1126         {
1127             Option.au_size = 32768;
1128         }
1129         else
1130         {
1131             Option.au_size = 131072;
1132         }
1133     }
1134     else
1135     {
1136         Option.au_size = CluserSize;
1137     }
1138
1139     Log("Formatting Part1 large FAT32 ClusterSize:%u(%uKB) ...", CluserSize, CluserSize / 1024);
1140
1141     disk_io_reset_write_error();
1142
1143     Ret = f_mkfs(TEXT("0:"), &Option, 0, 8 * 1024 * 1024);
1144     if (FR_OK == Ret)
1145     {
1146         if (disk_io_is_write_error())
1147         {
1148             Log("Formatting Part1 large FAT32 failed, write error.");
1149             return 1;
1150         }
1151
1152         Log("Formatting Part1 large FAT32 success, now set label");
1153
1154         Ret = f_mount(&FS, TEXT("0:"), 1);
1155         if (FR_OK == Ret)
1156         {
1157             Log("f_mount SUCCESS");
1158             Ret = f_setlabel(TEXT("0:Ventoy"));
1159             if (FR_OK == Ret)
1160             {
1161                 Log("f_setlabel SUCCESS");
1162                 Ret = f_unmount(TEXT("0:"));
1163                 Log("f_unmount %d %s", Ret, (FR_OK == Ret) ? "SUCCESS" : "FAILED");
1164             }
1165             else
1166             {
1167                 Log("f_setlabel failed %d", Ret);
1168             }
1169         }
1170         else
1171         {
1172             Log("f_mount failed %d", Ret);
1173         }
1174
1175         return 0;
1176     }
1177     else
1178     {
1179         Log("Formatting Part1 large FAT32 failed");
1180         return 1;
1181     }
1182 }
1183
1184 static int FormatPart1exFAT(UINT64 DiskSizeBytes)
1185 {
1186     MKFS_PARM Option;
1187     FRESULT Ret;
1188
1189     Option.fmt = FM_EXFAT;
1190     Option.n_fat = 1;
1191     Option.align = 8;
1192     Option.n_root = 1;
1193
1194     // < 32GB select 32KB as cluster size
1195     // > 32GB select 128KB as cluster size
1196     if (DiskSizeBytes / 1024 / 1024 / 1024 <= 32)
1197     {
1198         Option.au_size = 32768;
1199     }
1200     else
1201     {
1202         Option.au_size = 131072;
1203     }
1204
1205     Log("Formatting Part1 exFAT ...");
1206
1207         disk_io_reset_write_error();
1208
1209     Ret = f_mkfs(TEXT("0:"), &Option, 0, 8 * 1024 * 1024);
1210     if (FR_OK == Ret)
1211     {
1212                 if (disk_io_is_write_error())
1213                 {
1214                         Log("Formatting Part1 exFAT failed, write error.");
1215                         return 1;
1216                 }
1217
1218         Log("Formatting Part1 exFAT success");
1219         return 0;
1220     }
1221     else
1222     {
1223         Log("Formatting Part1 exFAT failed");
1224         return 1;
1225     }
1226 }
1227
1228 static int ZeroPart1FileSystem(HANDLE hDrive, UINT64 Part2StartSector)
1229 {
1230     int i;
1231     DWORD dwSize = 0;
1232     LARGE_INTEGER liCurPos;
1233     LARGE_INTEGER liNewPos;
1234     CHAR TmpBuffer[1024] = { 0 };
1235
1236     liCurPos.QuadPart = VENTOY_PART1_START_SECTOR * 512;
1237     liNewPos.QuadPart = 0;
1238     if (0 == SetFilePointerEx(hDrive, liCurPos, &liNewPos, FILE_BEGIN) ||
1239         liNewPos.QuadPart != liCurPos.QuadPart)
1240     {
1241         Log("SetFilePointerEx Failed %u %llu %llu", LASTERR, (ULONGLONG)liCurPos.QuadPart, (ULONGLONG)liNewPos.QuadPart);
1242         return 1;
1243     }
1244
1245     for (i = 0; i < 1024; i++)
1246     {
1247         WriteFile(hDrive, TmpBuffer, 1024, &dwSize, NULL);
1248     }
1249
1250     liCurPos.QuadPart = (Part2StartSector * 512) - (1024 * 1024);
1251     liNewPos.QuadPart = 0;
1252     if (0 == SetFilePointerEx(hDrive, liCurPos, &liNewPos, FILE_BEGIN) ||
1253         liNewPos.QuadPart != liCurPos.QuadPart)
1254     {
1255         Log("SetFilePointerEx Failed %u %llu %llu", LASTERR, (ULONGLONG)liCurPos.QuadPart, (ULONGLONG)liNewPos.QuadPart);
1256         return 1;
1257     }
1258
1259     for (i = 0; i < 1024; i++)
1260     {
1261         WriteFile(hDrive, TmpBuffer, 1024, &dwSize, NULL);
1262     }
1263
1264     Log("Zero Part1 SUCCESS");
1265     return 0;
1266 }
1267
1268 int ClearVentoyFromPhyDrive(HWND hWnd, PHY_DRIVE_INFO *pPhyDrive, char *pDrvLetter)
1269 {
1270     int i;
1271     int rc = 0;
1272     int state = 0;
1273     HANDLE hDrive;
1274     DWORD dwSize;
1275     BOOL bRet;
1276     CHAR MountDrive;
1277     CHAR DriveName[] = "?:\\";
1278     CHAR DriveLetters[MAX_PATH] = { 0 };
1279     LARGE_INTEGER liCurrentPosition;
1280     char *pTmpBuf = NULL;
1281     MBR_HEAD MBR;
1282
1283     *pDrvLetter = 0;
1284
1285     Log("ClearVentoyFromPhyDrive PhyDrive%d <<%s %s %dGB>>",
1286         pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1287         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1288
1289     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1290
1291     Log("Lock disk for clean ............................. ");
1292
1293     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1294     if (hDrive == INVALID_HANDLE_VALUE)
1295     {
1296         Log("Failed to open physical disk");
1297         return 1;
1298     }
1299
1300     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1301
1302     if (DriveLetters[0] == 0)
1303     {
1304         Log("No drive letter was assigned...");
1305         DriveName[0] = GetFirstUnusedDriveLetter();
1306         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1307     }
1308     else
1309     {
1310         // Unmount all mounted volumes that belong to this drive
1311         // Do it in reverse so that we always end on the first volume letter
1312         for (i = (int)strlen(DriveLetters); i > 0; i--)
1313         {
1314             DriveName[0] = DriveLetters[i - 1];
1315             bRet = DeleteVolumeMountPointA(DriveName);
1316             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1317         }
1318     }
1319
1320     MountDrive = DriveName[0];
1321     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1322
1323     // It kind of blows, but we have to relinquish access to the physical drive
1324     // for VDS to be able to delete the partitions that reside on it...
1325     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1326     CHECK_CLOSE_HANDLE(hDrive);
1327
1328     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1329
1330     if (!VDS_DeleteAllPartitions(pPhyDrive->PhyDrive))
1331     {
1332         Log("Notice: Could not delete partitions: %u", GetLastError());
1333     }
1334
1335     Log("Deleting all partitions ......................... OK");
1336
1337     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1338
1339     Log("Lock disk for write ............................. ");
1340     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1341     if (hDrive == INVALID_HANDLE_VALUE)
1342     {
1343         Log("Failed to GetPhysicalHandle for write.");
1344         rc = 1;
1345         goto End;
1346     }
1347
1348     // clear first and last 2MB space
1349     pTmpBuf = malloc(SIZE_2MB);
1350     if (!pTmpBuf)
1351     {
1352         Log("Failed to alloc memory.");
1353         rc = 1;
1354         goto End;
1355     }
1356     memset(pTmpBuf, 0, SIZE_2MB);
1357
1358     SET_FILE_POS(512);
1359     bRet = WriteFile(hDrive, pTmpBuf, SIZE_2MB - 512, &dwSize, NULL);
1360     Log("Write fisrt 1MB ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1361     if (!bRet)
1362     {
1363         rc = 1;
1364         goto End;
1365     }
1366
1367     SET_FILE_POS(pPhyDrive->SizeInBytes - SIZE_2MB);
1368     bRet = WriteFile(hDrive, pTmpBuf, SIZE_2MB, &dwSize, NULL);
1369     Log("Write 2nd 1MB ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1370     if (!bRet)
1371     {
1372         rc = 1;
1373         goto End;
1374     }
1375
1376     SET_FILE_POS(0);
1377
1378     if (pPhyDrive->SizeInBytes > 2199023255552ULL)
1379     {
1380         VTOY_GPT_INFO *pGptInfo;
1381         VTOY_GPT_HDR BackupHead;
1382         LARGE_INTEGER liCurrentPosition;
1383
1384         pGptInfo = (VTOY_GPT_INFO *)pTmpBuf;
1385
1386         VentoyFillWholeGpt(pPhyDrive->SizeInBytes, pGptInfo);
1387
1388         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1389         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1390         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1391         {
1392             rc = 1;
1393             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1394             goto End;
1395         }
1396
1397         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1398         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1399         {
1400             rc = 1;
1401             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1402             goto End;
1403         }
1404
1405         SET_FILE_POS(0);
1406         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1407         {
1408             rc = 1;
1409             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1410             goto End;
1411         }
1412
1413         Log("Write GPT Info OK ...");
1414     }
1415     else
1416     {
1417         bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
1418         Log("Read MBR ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1419         if (!bRet)
1420         {
1421             rc = 1;
1422             goto End;
1423         }
1424
1425         //clear boot code and partition table (reserved disk signature)
1426         memset(MBR.BootCode, 0, 440);
1427         memset(MBR.PartTbl, 0, sizeof(MBR.PartTbl));
1428
1429         VentoyFillMBRLocation(pPhyDrive->SizeInBytes, 2048, (UINT32)(pPhyDrive->SizeInBytes / 512 - 2048), MBR.PartTbl);
1430
1431         MBR.PartTbl[0].Active = 0x00; // bootable
1432         MBR.PartTbl[0].FsFlag = 0x07; // exFAT/NTFS/HPFS
1433
1434         SET_FILE_POS(0);
1435         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
1436         Log("Write MBR ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1437         if (!bRet)
1438         {
1439             rc = 1;
1440             goto End;
1441         }
1442     }
1443
1444     Log("Clear Ventoy successfully finished");
1445
1446         //Refresh Drive Layout
1447         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1448
1449 End:
1450
1451     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1452     PROGRESS_BAR_SET_POS(PT_REFORMAT_FINISH);
1453
1454     if (pTmpBuf)
1455     {
1456         free(pTmpBuf);
1457     }
1458
1459     if (rc == 0)
1460     {
1461         Log("Mounting Ventoy Partition ....................... ");
1462         Sleep(1000);
1463
1464         state = 0;
1465         memset(DriveLetters, 0, sizeof(DriveLetters));
1466         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1467         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1468
1469         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1470         {
1471             DriveName[0] = DriveLetters[i];
1472             Log("%s is ventoy part1, already mounted", DriveName);
1473             state = 1;
1474         }
1475
1476         if (state != 1)
1477         {
1478             Log("need to mount ventoy part1...");
1479             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, 2048, DriveLetters, sizeof(DriveLetters), FALSE))
1480             {
1481                 DriveName[0] = MountDrive;
1482                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1483                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1484
1485                 *pDrvLetter = MountDrive;
1486             }
1487             else
1488             {
1489                 Log("Failed to find ventoy volume");
1490             }
1491         }
1492
1493         Log("OK\n");
1494     }
1495     else
1496     {
1497         FindProcessOccupyDisk(hDrive, pPhyDrive);
1498     }
1499
1500     CHECK_CLOSE_HANDLE(hDrive);
1501     return rc;
1502 }
1503
1504 int InstallVentoy2FileImage(PHY_DRIVE_INFO *pPhyDrive, int PartStyle)
1505 {
1506     int i;
1507     int rc = 1;
1508     int Len = 0;
1509     int dataLen = 0;
1510     UINT size = 0;
1511     UINT segnum = 0;
1512     UINT32 chksum = 0;
1513     UINT64 data_offset = 0;
1514     UINT64 Part2StartSector = 0;
1515     UINT64 Part1StartSector = 0;
1516     UINT64 Part1SectorCount = 0;
1517     UINT8 *pData = NULL;
1518     UINT8 *pBkGptPartTbl = NULL;
1519     BYTE *ImgBuf = NULL;
1520     MBR_HEAD *pMBR = NULL;
1521     VTSI_FOOTER *pImgFooter = NULL;
1522     VTSI_SEGMENT *pSegment = NULL;
1523     VTOY_GPT_INFO *pGptInfo = NULL;
1524     VTOY_GPT_HDR *pBkGptHdr = NULL;
1525     FILE *fp = NULL;
1526
1527     Log("InstallVentoy2FileImage %s PhyDrive%d <<%s %s %dGB>>",
1528         PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1529         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1530
1531     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1532
1533     size = SIZE_1MB + VENTOY_EFI_PART_SIZE + 33 * 512 + VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT) + sizeof(VTSI_FOOTER);
1534
1535     pData = (UINT8 *)malloc(size);
1536     if (!pData)
1537     {
1538         Log("malloc image buffer failed %d.", size);
1539         goto End;
1540     }
1541
1542     pImgFooter = (VTSI_FOOTER *)(pData + size - sizeof(VTSI_FOOTER));
1543     pSegment = (VTSI_SEGMENT *)((UINT8 *)pImgFooter - VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT));
1544     memset(pImgFooter, 0, sizeof(VTSI_FOOTER));
1545     memset(pSegment, 0, VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT));
1546
1547     PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START);
1548
1549     Log("Writing Boot Image ............................. ");
1550     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1551     {
1552         Log("Failed to read stage1 img");
1553         goto End;
1554     }
1555
1556     unxz(ImgBuf, Len, NULL, NULL, pData, &dataLen, unxz_error);
1557     SAFE_FREE(ImgBuf);
1558
1559     Log("decompress %s len:%d", VENTOY_FILE_STG1_IMG, dataLen);
1560
1561     if (PartStyle)
1562     {
1563         pData[500] = 35;//update blocklist
1564         memmove(pData + 34 * 512, pData, SIZE_1MB - 512 * 34);
1565         memset(pData, 0, 34 * 512);
1566
1567         pGptInfo = (VTOY_GPT_INFO *)pData;
1568         memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1569         VentoyFillGpt(pPhyDrive->SizeInBytes, pGptInfo);
1570
1571         pBkGptPartTbl = pData + SIZE_1MB + VENTOY_EFI_PART_SIZE;
1572         memset(pBkGptPartTbl, 0, 33 * 512);
1573
1574         memcpy(pBkGptPartTbl, pGptInfo->PartTbl, 32 * 512);
1575         pBkGptHdr = (VTOY_GPT_HDR *)(pBkGptPartTbl + 32 * 512);
1576         VentoyFillBackupGptHead(pGptInfo, pBkGptHdr);
1577
1578         Part1StartSector = pGptInfo->PartTbl[0].StartLBA;
1579         Part1SectorCount = pGptInfo->PartTbl[0].LastLBA - Part1StartSector + 1;
1580         Part2StartSector = pGptInfo->PartTbl[1].StartLBA;
1581
1582         Log("Write GPT Info OK ...");
1583     }
1584     else
1585     {
1586         memmove(pData + 512, pData, SIZE_1MB - 512);
1587         memset(pData, 0, 512);
1588
1589         pMBR = (MBR_HEAD *)pData;
1590         VentoyFillMBR(pPhyDrive->SizeInBytes, pMBR, PartStyle);
1591         Part1StartSector = pMBR->PartTbl[0].StartSectorId;
1592         Part1SectorCount = pMBR->PartTbl[0].SectorCount;
1593         Part2StartSector = pMBR->PartTbl[1].StartSectorId;
1594
1595         Log("Write MBR OK ...");
1596     }
1597
1598     Log("Writing EFI part Image ............................. ");
1599     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&ImgBuf, &Len);
1600     if (rc)
1601     {
1602         Log("Failed to read img file %p %u", ImgBuf, Len);
1603         goto End;
1604     }
1605
1606     PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START + 28);
1607     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
1608     unxz(ImgBuf, Len, NULL, NULL, pData + SIZE_1MB, &dataLen, unxz_error);
1609     if (dataLen == Len)
1610     {
1611         Log("decompress finished success");
1612         g_part_img_buf[0] = pData + SIZE_1MB;
1613
1614         VentoyProcSecureBoot(g_SecureBoot);
1615     }
1616     else
1617     {
1618         Log("decompress finished failed");
1619         goto End;
1620     }
1621
1622     fopen_s(&fp, "VentoySparseImg.vtsi", "wb+");
1623     if (!fp)
1624     {
1625         Log("Failed to create Ventoy img file");
1626         goto End;
1627     }
1628
1629     Log("Writing stage1 data ............................. ");
1630
1631     fwrite(pData, 1, SIZE_1MB, fp);
1632
1633     pSegment[0].disk_start_sector = 0;
1634     pSegment[0].sector_num = SIZE_1MB / 512;
1635     pSegment[0].data_offset = data_offset;
1636     data_offset += pSegment[0].sector_num * 512;
1637
1638     disk_io_set_param(INVALID_HANDLE_VALUE, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1639     disk_io_set_imghook(fp, pSegment + 1, VTSI_IMG_MAX_SEG - 1, data_offset);
1640
1641     Log("Formatting part1 exFAT ...");
1642     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1643     {
1644         Log("FormatPart1exFAT failed.");
1645         disk_io_reset_imghook(&segnum, &data_offset);
1646         goto End;
1647     }
1648
1649     disk_io_reset_imghook(&segnum, &data_offset);
1650     segnum++;
1651
1652     Log("current segment number:%d dataoff:%ld", segnum, (long)data_offset);
1653
1654     //write data
1655     Log("Writing part2 data ............................. ");
1656     fwrite(pData + SIZE_1MB, 1, VENTOY_EFI_PART_SIZE, fp);
1657     pSegment[segnum].disk_start_sector = Part2StartSector;
1658     pSegment[segnum].sector_num = VENTOY_EFI_PART_SIZE / 512;
1659     pSegment[segnum].data_offset = data_offset;
1660     data_offset += pSegment[segnum].sector_num * 512;
1661     segnum++;
1662
1663     if (PartStyle)
1664     {
1665         Log("Writing backup gpt table ............................. ");
1666         fwrite(pBkGptPartTbl, 1, 33 * 512, fp);
1667         pSegment[segnum].disk_start_sector = pPhyDrive->SizeInBytes / 512 - 33;
1668         pSegment[segnum].sector_num = 33;
1669         pSegment[segnum].data_offset = data_offset;
1670         data_offset += pSegment[segnum].sector_num * 512;
1671         segnum++;
1672     }
1673
1674     Log("Writing segment metadata ............................. ");
1675
1676     for (i = 0; i < (int)segnum; i++)
1677     {
1678         Log("SEG[%d]:  PhySector:%llu SectorNum:%llu DataOffset:%llu(sector:%llu)", i, pSegment[i].disk_start_sector, pSegment[i].sector_num,
1679             pSegment[i].data_offset, pSegment[i].data_offset / 512);
1680     }
1681
1682     dataLen = segnum * sizeof(VTSI_SEGMENT);
1683     fwrite(pSegment, 1, dataLen, fp);
1684
1685     if (dataLen % 512)
1686     {
1687         //pData + SIZE_1MB - 8192 is a temp data buffer with zero
1688         fwrite(pData + SIZE_1MB - 8192, 1, 512 - (dataLen % 512), fp);
1689     }
1690
1691     //Fill footer
1692     pImgFooter->magic = VTSI_IMG_MAGIC;
1693     pImgFooter->version = 1;
1694     pImgFooter->disk_size = pPhyDrive->SizeInBytes;
1695     memcpy(&pImgFooter->disk_signature, pPhyDrive->MBR.BootCode + 0x1b8, 4);
1696     pImgFooter->segment_num = segnum;
1697     pImgFooter->segment_offset = data_offset;
1698
1699     for (i = 0, chksum = 0; i < (int)(segnum * sizeof(VTSI_SEGMENT)); i++)
1700     {
1701         chksum += *((UINT8 *)pSegment + i);
1702     }
1703     pImgFooter->segment_chksum = ~chksum;
1704
1705     for (i = 0, chksum = 0; i < sizeof(VTSI_FOOTER); i++)
1706     {
1707         chksum += *((UINT8 *)pImgFooter + i);
1708     }
1709     pImgFooter->foot_chksum = ~chksum;
1710
1711     Log("Writing footer segnum(%u)  segoffset(%llu) ......................", segnum, data_offset);
1712     Log("disk_size=%llu disk_signature=%lx segment_offset=%llu", pImgFooter->disk_size, pImgFooter->disk_signature, pImgFooter->segment_offset);
1713
1714     fwrite(pImgFooter, 1, sizeof(VTSI_FOOTER), fp);
1715     fclose(fp);
1716
1717     Log("Writing Ventoy image file finished, the file size should be %llu .", data_offset + 512 + ((dataLen + 511) / 512 * 512));
1718
1719     rc = 0;
1720
1721 End:
1722
1723     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1724     PROGRESS_BAR_SET_POS(PT_REFORMAT_FINISH);
1725
1726     Log("retcode:%d\n", rc);
1727
1728     SAFE_FREE(pData);
1729     SAFE_FREE(ImgBuf);
1730
1731     return rc;
1732 }
1733
1734
1735 int InstallVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int PartStyle, int TryId)
1736 {
1737     int i;
1738     int rc = 0;
1739     int state = 0;
1740     BOOL ReformatOK;
1741     HANDLE hDrive;
1742     DWORD dwSize;
1743     BOOL bRet;
1744     CHAR MountDrive;
1745     CHAR DriveName[] = "?:\\";
1746     CHAR DriveLetters[MAX_PATH] = { 0 };
1747     MBR_HEAD MBR;
1748     VTOY_GPT_INFO *pGptInfo = NULL;
1749     UINT64 Part1StartSector = 0;
1750     UINT64 Part1SectorCount = 0;
1751     UINT64 Part2StartSector = 0;
1752     BOOL LargeFAT32 = FALSE;
1753     BOOL DefaultExFAT = FALSE;
1754
1755         Log("#####################################################");
1756     Log("InstallVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
1757         PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1758         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1759         Log("#####################################################");
1760
1761     if (PartStyle)
1762     {
1763         pGptInfo = malloc(sizeof(VTOY_GPT_INFO));
1764         memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1765     }
1766
1767     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1768
1769     if (PartStyle)
1770     {
1771         VentoyFillGpt(pPhyDrive->SizeInBytes, pGptInfo);
1772         Part1StartSector = pGptInfo->PartTbl[0].StartLBA;
1773         Part1SectorCount = pGptInfo->PartTbl[0].LastLBA - Part1StartSector + 1;
1774         Part2StartSector = pGptInfo->PartTbl[1].StartLBA;
1775     }
1776     else
1777     {
1778         VentoyFillMBR(pPhyDrive->SizeInBytes, &MBR, PartStyle);
1779         Part1StartSector = MBR.PartTbl[0].StartSectorId;
1780         Part1SectorCount = MBR.PartTbl[0].SectorCount;
1781         Part2StartSector = MBR.PartTbl[1].StartSectorId;
1782     }
1783
1784     Log("Lock disk for clean ............................. ");
1785
1786     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1787     if (hDrive == INVALID_HANDLE_VALUE)
1788     {
1789         Log("Failed to open physical disk");
1790         free(pGptInfo);
1791         return 1;
1792     }
1793
1794     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1795
1796     if (DriveLetters[0] == 0)
1797     {
1798         Log("No drive letter was assigned...");
1799         DriveName[0] = GetFirstUnusedDriveLetter();
1800         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1801     }
1802     else
1803     {
1804         // Unmount all mounted volumes that belong to this drive
1805         // Do it in reverse so that we always end on the first volume letter
1806         for (i = (int)strlen(DriveLetters); i > 0; i--)
1807         {
1808             DriveName[0] = DriveLetters[i - 1];
1809             bRet = DeleteVolumeMountPointA(DriveName);
1810             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1811         }
1812     }
1813
1814     MountDrive = DriveName[0];
1815     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1816
1817     // It kind of blows, but we have to relinquish access to the physical drive
1818     // for VDS to be able to delete the partitions that reside on it...
1819     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1820     CHECK_CLOSE_HANDLE(hDrive);
1821
1822     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1823
1824     if (!VDS_DeleteAllPartitions(pPhyDrive->PhyDrive))
1825     {
1826         Log("Notice: Could not delete partitions: 0x%x, but we continue.", GetLastError());
1827     }
1828
1829     Log("Deleting all partitions ......................... OK");
1830
1831     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1832
1833     Log("Lock disk for write ............................. ");
1834     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1835     if (hDrive == INVALID_HANDLE_VALUE)
1836     {
1837         Log("Failed to GetPhysicalHandle for write.");
1838         rc = 1;
1839         goto End;
1840     }
1841
1842     //Refresh Drive Layout
1843     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1844
1845     disk_io_set_param(hDrive, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1846
1847     PROGRESS_BAR_SET_POS(PT_FORMAT_PART1);
1848
1849     if (PartStyle == 1 && pPhyDrive->PartStyle == 0)
1850     {
1851         Log("Wait for format part1 ...");
1852         Sleep(1000 * 5);
1853     }
1854
1855     if (GetVentoyFsType() == VTOY_FS_FAT32 && (Part1SectorCount * 512 >= FAT32_MAX_LIMIT))
1856     {
1857         Log("Formatting part1 large FAT32 ...");
1858         LargeFAT32 = TRUE;
1859         if (0 != FormatPart1LargeFAT32(pPhyDrive->SizeInBytes, GetClusterSize()))
1860         {
1861             Log("FormatPart1LargeFAT32 failed.");
1862             rc = 1;
1863             goto End;
1864         }
1865     }
1866     else if (GetVentoyFsType() == VTOY_FS_EXFAT && GetClusterSize() == 0)
1867     {
1868         Log("Formatting part1 exFAT ...");
1869         DefaultExFAT = TRUE;
1870         if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1871         {
1872             Log("FormatPart1exFAT failed.");
1873             rc = 1;
1874             goto End;
1875         }
1876     }
1877     else
1878     {
1879         Log("Zero part1 file system ...");
1880         if (0 != ZeroPart1FileSystem(hDrive, Part2StartSector))
1881         {
1882             Log("ZeroPart1FileSystem failed.");
1883             rc = 1;
1884             goto End;
1885         }
1886     }
1887
1888     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1889     Log("Writing part2 FAT img ...");
1890
1891     if (0 != FormatPart2Fat(hDrive, Part2StartSector))
1892     {
1893         Log("FormatPart2Fat failed.");
1894         rc = 1;
1895         goto End;
1896     }
1897
1898     PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
1899     Log("Writing Boot Image ............................. ");
1900     if (WriteGrubStage1ToPhyDrive(hDrive, PartStyle) != 0)
1901     {
1902         Log("WriteGrubStage1ToPhyDrive failed.");
1903         rc = 1;
1904         goto End;
1905     }
1906
1907     PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
1908     Log("Writing Partition Table ........................ ");
1909     SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1910
1911     if (PartStyle)
1912     {
1913         VTOY_GPT_HDR BackupHead;
1914         LARGE_INTEGER liCurrentPosition;
1915
1916         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1917         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1918         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1919         {
1920             rc = 1;
1921             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1922             goto End;
1923         }
1924
1925         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1926         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1927         {
1928             rc = 1;
1929             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1930             goto End;
1931         }
1932
1933         SET_FILE_POS(0);
1934         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1935         {
1936             rc = 1;
1937             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1938             goto End;
1939         }
1940
1941         Log("Write GPT Info OK ...");
1942         memcpy(&(pPhyDrive->MBR), &(pGptInfo->MBR), 512);
1943     }
1944     else
1945     {
1946         if (!WriteFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL))
1947         {
1948             rc = 1;
1949             Log("Write MBR Failed, dwSize:%u ErrCode:%u", dwSize, GetLastError());
1950             goto End;
1951         }
1952         Log("Write MBR OK ...");
1953         memcpy(&(pPhyDrive->MBR), &MBR, 512);
1954     }
1955
1956     //Refresh Drive Layout
1957     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1958
1959 End:
1960
1961     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1962
1963     if (rc == 0)
1964     {
1965         Log("Mounting Ventoy Partition ....................... ");
1966         Sleep(1000);
1967
1968         state = 0;
1969         memset(DriveLetters, 0, sizeof(DriveLetters));
1970         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1971         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1972
1973         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1974         {
1975             DriveName[0] = DriveLetters[i];
1976             if (IsVentoyLogicalDrive(DriveName[0]))
1977             {
1978                 Log("%s is ventoy part2, delete mountpoint", DriveName);
1979                 DeleteVolumeMountPointA(DriveName);
1980             }
1981             else
1982             {
1983                 Log("%s is ventoy part1, already mounted", DriveName);
1984                 MountDrive = DriveName[0];
1985                 state = 1;
1986             }
1987         }
1988
1989         if (state != 1)
1990         {
1991             Log("need to mount ventoy part1...");
1992
1993             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, Part1StartSector, DriveLetters, sizeof(DriveLetters), FALSE))
1994             {
1995                 DriveName[0] = MountDrive;
1996                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1997                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1998
1999                 if (bRet)
2000                 {
2001                     state = 1;
2002                 }
2003             }
2004             else
2005             {
2006                 Log("Failed to find ventoy volume");
2007             }
2008         }
2009
2010         // close handle, or it will deny reformat
2011         Log("Close handle ...");
2012         CHECK_CLOSE_HANDLE(hDrive);
2013
2014         ReformatOK = TRUE;
2015
2016         if (state)
2017         {
2018             if (LargeFAT32)
2019             {
2020                 Log("No need to reformat for large FAT32");
2021                 pPhyDrive->VentoyFsClusterSize = GetVolumeClusterSize(MountDrive);
2022             }
2023             else if (DefaultExFAT)
2024             {
2025                 Log("No need to reformat for default exfat");
2026                 pPhyDrive->VentoyFsClusterSize = GetVolumeClusterSize(MountDrive);
2027             }
2028             else
2029             {
2030                 bRet = DISK_FormatVolume(MountDrive, GetVentoyFsType(), Part1SectorCount * 512);
2031                 for (i = 0; bRet == FALSE && i < 2; i++)
2032                 {
2033                     Log("Wait and retry reformat ...");
2034                     Sleep(1000);
2035                     bRet = DISK_FormatVolume(MountDrive, GetVentoyFsType(), Part1SectorCount * 512);
2036                 }
2037
2038                 if (bRet)
2039                 {
2040                     Log("Reformat %C:\\ to %s SUCCESS", MountDrive, GetVentoyFsName());
2041                     pPhyDrive->VentoyFsClusterSize = GetVolumeClusterSize(MountDrive);
2042
2043                     if ((GetVentoyFsType() != VTOY_FS_UDF) && (pPhyDrive->VentoyFsClusterSize < 2048))
2044                     {
2045                         for (i = 0; i < 10; i++)
2046                         {
2047                             Log("### Invalid cluster size %d ###", pPhyDrive->VentoyFsClusterSize);
2048                         }
2049                     }
2050                 }
2051                 else
2052                 {
2053                     ReformatOK = FALSE;
2054                     Log("Reformat %C:\\ to %s FAILED", MountDrive, GetVentoyFsName());
2055                 }
2056             }
2057         }
2058         else
2059         {
2060             Log("Can not reformat %s to %s", DriveName, GetVentoyFsName());
2061         }
2062
2063         if (!ReformatOK)
2064         {
2065             Log("Format to exfat with built-in algorithm");
2066
2067             hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2068             if (hDrive == INVALID_HANDLE_VALUE)
2069             {
2070                 Log("Failed to GetPhysicalHandle for write.");
2071             }
2072             else
2073             {
2074                 if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
2075                 {
2076                     Log("FormatPart1exFAT SUCCESS.");
2077                 }
2078                 else
2079                 {
2080                     Log("FormatPart1exFAT FAILED.");
2081                 }
2082
2083                 CHECK_CLOSE_HANDLE(hDrive);
2084             }
2085         }
2086
2087         Log("OK\n");
2088     }
2089     else
2090     {
2091                 PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
2092
2093         FindProcessOccupyDisk(hDrive, pPhyDrive);
2094
2095                 if (!VDS_IsLastAvaliable())
2096                 {
2097                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2098                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2099                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2100                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2101                         Log("###### [Error:] Virtual Disk Service (VDS) Unavailable ######");
2102                 }
2103
2104         CHECK_CLOSE_HANDLE(hDrive);
2105     }
2106
2107     if (pGptInfo)
2108     {
2109         free(pGptInfo);
2110     }
2111
2112     return rc;
2113 }
2114
2115
2116 int PartitionResizeForVentoy(PHY_DRIVE_INFO *pPhyDrive)
2117 {
2118         int i, j;
2119         int rc = 1;
2120         int PhyDrive;
2121         int PartStyle;
2122         INT64 ReservedValue;
2123         UINT64 RecudeBytes;
2124         GUID Guid;
2125         MBR_HEAD MBR;
2126         VTOY_GPT_INFO *pGPT;
2127         MBR_HEAD *pMBR;
2128         DWORD dwSize = 0;
2129         VTOY_GPT_HDR BackupHead;
2130         HANDLE hDrive = INVALID_HANDLE_VALUE;
2131         GUID ZeroGuid = { 0 };
2132         static GUID WindowsDataPartType = { 0xebd0a0a2, 0xb9e5, 0x4433, { 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7 } };
2133         static GUID EspPartType = { 0xc12a7328, 0xf81f, 0x11d2, { 0xba, 0x4b, 0x00, 0xa0, 0xc9, 0x3e, 0xc9, 0x3b } };
2134         static GUID BiosGrubPartType = { 0x21686148, 0x6449, 0x6e6f, { 0x74, 0x4e, 0x65, 0x65, 0x64, 0x45, 0x46, 0x49 } };
2135
2136         Log("#####################################################");
2137         Log("PartitionResizeForVentoy PhyDrive%d <<%s %s %dGB>>",
2138                 pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
2139                 GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
2140         Log("#####################################################");
2141
2142         pGPT = &(pPhyDrive->Gpt);
2143         pMBR = &(pPhyDrive->Gpt.MBR);
2144         Log("Disksize:%llu Part2Start:%llu", pPhyDrive->SizeInBytes, pPhyDrive->ResizePart2StartSector * 512);
2145
2146         if (pMBR->PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
2147         {
2148                 PartStyle = 1;
2149         }
2150         else
2151         {
2152                 PartStyle = 0;
2153         }
2154
2155         PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
2156
2157         RecudeBytes = VENTOY_EFI_PART_SIZE;
2158         ReservedValue = GetReservedSpaceInMB();
2159         if (ReservedValue > 0)
2160         {
2161                 Log("Reduce add reserved space %lldMB", (LONGLONG)ReservedValue);
2162                 RecudeBytes += (UINT64)(ReservedValue * SIZE_1MB);
2163         }
2164
2165
2166         if (pPhyDrive->ResizeNoShrink == FALSE)
2167         {
2168                 Log("Need to shrink the volume");
2169                 if (DISK_ShrinkVolume(pPhyDrive->PhyDrive, pPhyDrive->ResizeVolumeGuid, pPhyDrive->Part1DriveLetter, pPhyDrive->ResizeOldPart1Size, RecudeBytes))
2170                 {
2171                         Log("Shrink volume success, now check again");
2172
2173                         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2174                         if (hDrive == INVALID_HANDLE_VALUE)
2175                         {
2176                                 Log("Failed to GetPhysicalHandle for update.");
2177                                 goto End;
2178                         }
2179
2180                         //Refresh Drive Layout
2181                         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2182
2183                         CHECK_CLOSE_HANDLE(hDrive);
2184
2185
2186                         if (PartResizePreCheck(NULL) && pPhyDrive->ResizeNoShrink)
2187                         {
2188                                 Log("Recheck after Shrink volume success");
2189                                 Log("After shrink Disksize:%llu Part2Start:%llu", pPhyDrive->SizeInBytes, pPhyDrive->ResizePart2StartSector * 512);
2190                         }
2191                         else
2192                         {
2193                                 Log("Recheck after Shrink volume failed %u", pPhyDrive->ResizeNoShrink);
2194                                 goto End;
2195                         }
2196                 }
2197                 else
2198                 {
2199                         Log("Shrink volume failed");
2200                         goto End;
2201                 }
2202         }
2203
2204
2205         //Now try write data
2206         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2207         if (hDrive == INVALID_HANDLE_VALUE)
2208         {
2209                 Log("Failed to GetPhysicalHandle for update.");
2210                 goto End;
2211         }
2212
2213
2214         //Write partition 2 data
2215         PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
2216         if (0 != FormatPart2Fat(hDrive, pPhyDrive->ResizePart2StartSector))
2217         {
2218                 Log("FormatPart2Fat failed.");
2219                 goto End;
2220         }
2221
2222         //Write grub stage2 gap
2223         PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
2224         Log("Writing Boot Image ............................. ");
2225         if (WriteGrubStage1ToPhyDrive(hDrive, PartStyle) != 0)
2226         {
2227                 Log("WriteGrubStage1ToPhyDrive failed.");
2228                 goto End;
2229         }
2230
2231
2232         //Write partition table
2233         PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
2234         Log("Writing partition table ............................. ");
2235
2236         VentoyGetLocalBootImg(&MBR);
2237         CoCreateGuid(&Guid);
2238         memcpy(MBR.BootCode + 0x180, &Guid, 16);
2239         memcpy(pMBR->BootCode, MBR.BootCode, 440);
2240
2241         if (PartStyle == 0)
2242         {
2243                 for (i = 1; i < 4; i++)
2244                 {
2245                         if (pMBR->PartTbl[i].SectorCount == 0)
2246                         {
2247                                 break;
2248                         }
2249                 }
2250
2251                 if (i >= 4)
2252                 {
2253                         Log("Can not find MBR free partition table");
2254                         goto End;
2255                 }
2256
2257                 for (j = i - 1; j > 0; j--)
2258                 {
2259                         Log("Move MBR partition table %d --> %d", j + 1, j + 2);
2260                         memcpy(pMBR->PartTbl + (j + 1), pMBR->PartTbl + j, sizeof(PART_TABLE));
2261                 }
2262
2263         memset(pMBR->PartTbl + 1, 0, sizeof(PART_TABLE));
2264                 VentoyFillMBRLocation(pPhyDrive->SizeInBytes, (UINT32)pPhyDrive->ResizePart2StartSector, VENTOY_EFI_PART_SIZE / 512, pMBR->PartTbl + 1);
2265                 pMBR->PartTbl[0].Active = 0x80; // bootable
2266                 pMBR->PartTbl[1].Active = 0x00;
2267                 pMBR->PartTbl[1].FsFlag = 0xEF; // EFI System Partition
2268
2269                 if (!WriteDataToPhyDisk(hDrive, 0, pMBR, 512))
2270                 {
2271                         Log("Legacy BIOS write MBR failed");
2272                         goto End;
2273                 }
2274         }
2275         else
2276         {
2277                 for (i = 1; i < 128; i++)
2278                 {
2279                         if (memcmp(&(pGPT->PartTbl[i].PartGuid), &ZeroGuid, sizeof(GUID)) == 0)
2280                         {
2281                                 break;
2282                         }
2283                 }
2284
2285                 if (i >= 128)
2286                 {
2287                         Log("Can not find GPT free partition table");
2288                         goto End;
2289                 }
2290
2291                 for (j = i - 1; j > 0; j--)
2292                 {
2293                         Log("Move GPT partition table %d --> %d", j + 1, j + 2);
2294                         memcpy(pGPT->PartTbl + (j + 1), pGPT->PartTbl + j, sizeof(VTOY_GPT_PART_TBL));
2295                 }
2296
2297
2298                 pMBR->BootCode[92] = 0x22;
2299
2300                 // to fix windows issue
2301         memset(pGPT->PartTbl + 1, 0, sizeof(VTOY_GPT_PART_TBL));
2302                 memcpy(&(pGPT->PartTbl[1].PartType), &WindowsDataPartType, sizeof(GUID));
2303                 CoCreateGuid(&(pGPT->PartTbl[1].PartGuid));
2304
2305                 pGPT->PartTbl[1].StartLBA = pGPT->PartTbl[0].LastLBA + 1;
2306                 pGPT->PartTbl[1].LastLBA = pGPT->PartTbl[1].StartLBA + VENTOY_EFI_PART_SIZE / 512 - 1;
2307                 pGPT->PartTbl[1].Attr = 0xC000000000000001ULL;
2308                 memcpy(pGPT->PartTbl[1].Name, L"VTOYEFI", 7 * 2);
2309
2310                 //Update CRC
2311                 pGPT->Head.PartTblCrc = VentoyCrc32(pGPT->PartTbl, sizeof(pGPT->PartTbl));
2312         pGPT->Head.Crc = 0;
2313                 pGPT->Head.Crc = VentoyCrc32(&(pGPT->Head), pGPT->Head.Length);
2314
2315                 Log("pGPT->Head.EfiStartLBA=%llu", (ULONGLONG)pGPT->Head.EfiStartLBA);
2316                 Log("pGPT->Head.EfiBackupLBA=%llu", (ULONGLONG)pGPT->Head.EfiBackupLBA);
2317
2318                 VentoyFillBackupGptHead(pGPT, &BackupHead);
2319                 if (!WriteDataToPhyDisk(hDrive, pGPT->Head.EfiBackupLBA * 512, &BackupHead, 512))
2320                 {
2321                         Log("UEFI write backup head failed");
2322                         goto End;
2323                 }
2324
2325                 if (!WriteDataToPhyDisk(hDrive, (pGPT->Head.EfiBackupLBA - 32) * 512, pGPT->PartTbl, 512 * 32))
2326                 {
2327                         Log("UEFI write backup partition table failed");
2328                         goto End;
2329                 }
2330
2331                 if (!WriteDataToPhyDisk(hDrive, 0, pGPT, 512 * 34))
2332                 {
2333                         Log("UEFI write MBR & Main partition table failed");
2334                         goto End;
2335                 }
2336         }
2337
2338
2339
2340         //Refresh Drive Layout
2341         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2342
2343         //We must close handle here, because it will block the refresh bellow
2344         CHECK_CLOSE_HANDLE(hDrive);
2345
2346         Sleep(2000);
2347
2348         //Refresh disk list
2349         PhyDrive = pPhyDrive->PhyDrive;
2350
2351         Log("#### Now Refresh PhyDrive ####");
2352         Ventoy2DiskDestroy();
2353         Ventoy2DiskInit();
2354
2355         pPhyDrive = GetPhyDriveInfoByPhyDrive(PhyDrive);
2356         if (pPhyDrive)
2357         {
2358                 if (pPhyDrive->VentoyVersion[0] == 0)
2359                 {
2360                         Log("After process the Ventoy version is still invalid");
2361                         goto End;
2362                 }
2363
2364                 Log("### Ventoy non-destructive installation successfully finished <%s>", pPhyDrive->VentoyVersion);
2365         }
2366         else
2367         {
2368                 Log("### Ventoy non-destructive installation successfully finished <not found>");
2369         }
2370
2371         InitComboxCtrl(g_DialogHwnd, PhyDrive);
2372
2373         rc = 0;
2374
2375 End:
2376         CHECK_CLOSE_HANDLE(hDrive);
2377         return rc;
2378 }
2379
2380
2381 static BOOL DiskCheckWriteAccess(HANDLE hDrive)
2382 {
2383         DWORD dwSize;
2384         BOOL ret = FALSE;
2385         BOOL bRet = FALSE;
2386         BYTE Buffer[512];
2387         LARGE_INTEGER liCurPosition;
2388         LARGE_INTEGER liNewPosition;
2389
2390         liCurPosition.QuadPart = 2039 * 512;
2391         liNewPosition.QuadPart = 0;
2392         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2393                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2394         {
2395                 Log("SetFilePointer1 Failed %u", LASTERR);
2396                 goto out;
2397         }
2398
2399
2400         dwSize = 0;
2401         ret = ReadFile(hDrive, Buffer, 512, &dwSize, NULL);
2402         if ((!ret) || (dwSize != 512))
2403         {
2404                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2405                 goto out;
2406         }
2407
2408
2409         liCurPosition.QuadPart = 2039 * 512;
2410         liNewPosition.QuadPart = 0;
2411         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2412                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2413         {
2414                 Log("SetFilePointer2 Failed %u", LASTERR);
2415                 goto out;
2416         }
2417
2418         dwSize = 0;
2419         ret = WriteFile(hDrive, Buffer, 512, &dwSize, NULL);
2420         if ((!ret) || dwSize != 512)
2421         {
2422                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
2423                 goto out;
2424         }
2425
2426         bRet = TRUE;
2427
2428 out:
2429
2430         return bRet;
2431 }
2432
2433 static BOOL BackupDataBeforeCleanDisk(int PhyDrive, UINT64 DiskSize, BYTE **pBackup)
2434 {
2435         DWORD dwSize;
2436         DWORD dwStatus;
2437         BOOL Return = FALSE;
2438         BOOL ret = FALSE;
2439         BYTE *backup = NULL;
2440         UINT64 offset;
2441         HANDLE hDrive = INVALID_HANDLE_VALUE;
2442         LARGE_INTEGER liCurPosition;
2443         LARGE_INTEGER liNewPosition;
2444         VTOY_GPT_INFO *pGPT = NULL;
2445
2446         Log("BackupDataBeforeCleanDisk %d", PhyDrive);
2447
2448         // step1: check write access
2449         hDrive = GetPhysicalHandle(PhyDrive, TRUE, TRUE, FALSE);
2450         if (hDrive == INVALID_HANDLE_VALUE)
2451         {
2452                 Log("Failed to GetPhysicalHandle for write.");
2453                 goto out;
2454         }
2455
2456         if (DiskCheckWriteAccess(hDrive))
2457         {
2458                 Log("DiskCheckWriteAccess success");
2459                 CHECK_CLOSE_HANDLE(hDrive);
2460         }
2461         else
2462         {
2463                 Log("DiskCheckWriteAccess failed");
2464                 goto out;
2465         }
2466
2467         //step2 backup 4MB data
2468         backup = malloc(SIZE_1MB * 4);
2469         if (!backup)
2470         {
2471                 goto out;
2472         }
2473
2474         hDrive = GetPhysicalHandle(PhyDrive, FALSE, FALSE, FALSE);
2475         if (hDrive == INVALID_HANDLE_VALUE)
2476         {
2477                 goto out;
2478         }
2479
2480         //read first 2MB
2481         dwStatus = SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2482         if (dwStatus != 0)
2483         {
2484                 goto out;
2485         }
2486
2487         dwSize = 0;
2488         ret = ReadFile(hDrive, backup, SIZE_2MB, &dwSize, NULL);
2489         if ((!ret) || (dwSize != SIZE_2MB))
2490         {
2491                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2492                 goto out;
2493         }
2494
2495         pGPT = (VTOY_GPT_INFO *)backup;
2496         offset = pGPT->Head.EfiBackupLBA * 512;
2497         if (offset >= (DiskSize - SIZE_2MB) && offset < DiskSize)
2498         {
2499                 Log("EFI partition table check success");
2500         }
2501         else
2502         {
2503                 Log("Backup EFI LBA not in last 2MB range: %llu", pGPT->Head.EfiBackupLBA);
2504                 goto out;
2505         }
2506
2507         //read last 2MB
2508         liCurPosition.QuadPart = DiskSize - SIZE_2MB;
2509         liNewPosition.QuadPart = 0;
2510         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2511                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2512         {
2513                 goto out;
2514         }
2515
2516         dwSize = 0;
2517         ret = ReadFile(hDrive, backup + SIZE_2MB, SIZE_2MB, &dwSize, NULL);
2518         if ((!ret) || (dwSize != SIZE_2MB))
2519         {
2520                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2521                 goto out;
2522         }
2523
2524         *pBackup = backup;
2525         backup = NULL; //For don't free later
2526         Return = TRUE;
2527
2528 out:
2529         CHECK_CLOSE_HANDLE(hDrive);
2530         if (backup)
2531                 free(backup);
2532
2533         return Return;
2534 }
2535
2536
2537 static BOOL WriteBackupDataToDisk(HANDLE hDrive, UINT64 Offset, BYTE *Data, DWORD Length)
2538 {
2539         DWORD dwSize = 0;
2540         BOOL ret = FALSE;
2541         LARGE_INTEGER liCurPosition;
2542         LARGE_INTEGER liNewPosition;
2543
2544         Log("WriteBackupDataToDisk %llu %p %u", Offset, Data, Length);
2545
2546         liCurPosition.QuadPart = Offset;
2547         liNewPosition.QuadPart = 0;
2548         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2549                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2550         {
2551                 return FALSE;
2552         }
2553
2554         ret = WriteFile(hDrive, Data, Length, &dwSize, NULL);
2555         if ((!ret) || dwSize != Length)
2556         {
2557                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
2558                 return FALSE;
2559         }
2560
2561         Log("WriteBackupDataToDisk %llu %p %u success", Offset, Data, Length);
2562         return TRUE;
2563 }
2564
2565
2566 int UpdateVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int TryId)
2567 {
2568         int i;
2569         int rc = 0;
2570         int MaxRetry = 4;
2571         BOOL ForceMBR = FALSE;
2572         BOOL Esp2Basic = FALSE;
2573         BOOL ChangeAttr = FALSE;
2574         BOOL CleanDisk = FALSE;
2575         BOOL DelEFI = FALSE;
2576         BOOL bWriteBack = TRUE;
2577         HANDLE hVolume;
2578         HANDLE hDrive;
2579         DWORD Status;
2580         DWORD dwSize;
2581         BOOL bRet;
2582         CHAR DriveName[] = "?:\\";
2583         CHAR DriveLetters[MAX_PATH] = { 0 };
2584         CHAR BackBinFile[MAX_PATH];
2585         UINT64 StartSector;
2586         UINT64 ReservedMB = 0;
2587         MBR_HEAD BootImg;
2588         MBR_HEAD MBR;
2589         BYTE *pBackup = NULL;
2590         VTOY_GPT_INFO *pGptInfo = NULL;
2591         VTOY_GPT_INFO *pGptBkup = NULL;
2592         UINT8 ReservedData[4096];
2593
2594         Log("#####################################################");
2595         Log("UpdateVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
2596                 pPhyDrive->PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
2597                 GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
2598         Log("#####################################################");
2599
2600         PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
2601
2602         Log("Lock disk for umount ............................ ");
2603
2604         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
2605         if (hDrive == INVALID_HANDLE_VALUE)
2606         {
2607                 Log("Failed to open physical disk");
2608                 return 1;
2609         }
2610
2611         if (pPhyDrive->PartStyle)
2612         {
2613                 pGptInfo = malloc(2 * sizeof(VTOY_GPT_INFO));
2614                 if (!pGptInfo)
2615                 {
2616                         return 1;
2617                 }
2618
2619                 memset(pGptInfo, 0, 2 * sizeof(VTOY_GPT_INFO));
2620                 pGptBkup = pGptInfo + 1;
2621
2622                 // Read GPT Info
2623                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2624                 ReadFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL);
2625                 memcpy(pGptBkup, pGptInfo, sizeof(VTOY_GPT_INFO));
2626
2627                 //MBR will be used to compare with local boot image
2628                 memcpy(&MBR, &pGptInfo->MBR, sizeof(MBR_HEAD));
2629
2630                 StartSector = pGptInfo->PartTbl[1].StartLBA;
2631                 Log("GPT StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2632
2633                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512) - 33) / 2048;
2634                 Log("GPT Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2635         }
2636         else
2637         {
2638                 // Read MBR
2639                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2640                 ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
2641
2642                 StartSector = MBR.PartTbl[1].StartSectorId;
2643                 Log("MBR StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2644
2645                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512)) / 2048;
2646                 Log("MBR Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2647         }
2648
2649         //Read Reserved Data
2650         SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2651         ReadFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2652
2653         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
2654
2655         if (DriveLetters[0] == 0)
2656         {
2657                 Log("No drive letter was assigned...");
2658         }
2659         else
2660         {
2661                 // Unmount all mounted volumes that belong to this drive
2662                 // Do it in reverse so that we always end on the first volume letter
2663                 for (i = (int)strlen(DriveLetters); i > 0; i--)
2664                 {
2665                         DriveName[0] = DriveLetters[i - 1];
2666                         if (IsVentoyLogicalDrive(DriveName[0]))
2667                         {
2668                                 Log("%s is ventoy logical drive", DriveName);
2669                                 bRet = DeleteVolumeMountPointA(DriveName);
2670                                 Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, LASTERR);
2671                                 break;
2672                         }
2673                 }
2674         }
2675
2676         // It kind of blows, but we have to relinquish access to the physical drive
2677         // for VDS to be able to delete the partitions that reside on it...
2678         DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2679         CHECK_CLOSE_HANDLE(hDrive);
2680
2681         if (pPhyDrive->PartStyle == 1)
2682         {
2683                 Log("TryId=%d EFI GPT partition type is 0x%llx", TryId, pPhyDrive->Part2GPTAttr);
2684                 PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
2685
2686                 if (TryId == 1)
2687                 {
2688                         Log("Change GPT partition type to ESP");
2689                         if (DISK_ChangeVtoyEFI2ESP(pPhyDrive->PhyDrive, StartSector * 512ULL))
2690                         {
2691                                 Esp2Basic = TRUE;
2692                                 Sleep(3000);
2693                         }
2694                 }
2695                 else if (TryId == 2)
2696                 {
2697                         Log("Change GPT partition attribute");
2698                         if (DISK_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, StartSector * 512ULL, 0x8000000000000001))
2699                         {
2700                                 ChangeAttr = TRUE;
2701                                 Sleep(2000);
2702                         }
2703                 }
2704                 else if (TryId == 3)
2705                 {
2706                         DISK_DeleteVtoyEFIPartition(pPhyDrive->PhyDrive, StartSector * 512ULL);
2707                         DelEFI = TRUE;
2708                 }
2709                 else if (TryId == 4)
2710                 {
2711                         Log("Clean disk GPT partition table");
2712                         if (BackupDataBeforeCleanDisk(pPhyDrive->PhyDrive, pPhyDrive->SizeInBytes, &pBackup))
2713                         {
2714                                 sprintf_s(BackBinFile, sizeof(BackBinFile), ".\\ventoy\\phydrive%d_%u_%d.bin",
2715                                         pPhyDrive->PhyDrive, GetCurrentProcessId(), g_backup_bin_index++);
2716                                 SaveBufToFile(BackBinFile, pBackup, 4 * SIZE_1MB);
2717                                 Log("Save backup data to %s", BackBinFile);
2718
2719                                 Log("Success to backup data before clean");
2720                                 CleanDisk = TRUE;
2721                                 DISK_CleanDisk(pPhyDrive->PhyDrive);
2722                                 Sleep(3000);
2723                         }
2724                         else
2725                         {
2726                                 Log("Failed to backup data before clean");
2727                         }
2728                 }
2729         }
2730
2731     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
2732
2733     Log("Lock disk for update ............................ ");
2734     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2735     if (hDrive == INVALID_HANDLE_VALUE)
2736     {
2737         Log("Failed to GetPhysicalHandle for write.");
2738         rc = 1;
2739         goto End;
2740     }
2741
2742     PROGRESS_BAR_SET_POS(PT_LOCK_VOLUME);
2743
2744     Log("Lock volume for update .......................... ");
2745     hVolume = INVALID_HANDLE_VALUE;
2746
2747         //If we change VTOYEFI to ESP, it can not have s volume name, so don't try to get it.
2748         if (CleanDisk)
2749         {
2750                 //writeback the last 2MB
2751                 if (!WriteBackupDataToDisk(hDrive, pPhyDrive->SizeInBytes - SIZE_2MB, pBackup + SIZE_2MB, SIZE_2MB))
2752                 {
2753                         bWriteBack = FALSE;
2754                 }
2755
2756                 //write the first 2MB except parttable
2757                 if (!WriteBackupDataToDisk(hDrive, 34 * 512, pBackup + 34 * 512, SIZE_2MB - 34 * 512))
2758                 {
2759                         bWriteBack = FALSE;
2760                 }
2761
2762                 Status = ERROR_NOT_FOUND;
2763         }
2764         else if (DelEFI)
2765         {
2766                 Status = ERROR_NOT_FOUND;
2767         }
2768         else if (Esp2Basic)
2769         {
2770                 Status = ERROR_NOT_FOUND;
2771         }
2772         else
2773         {
2774                 for (i = 0; i < MaxRetry; i++)
2775                 {
2776                         Status = GetVentoyVolumeName(pPhyDrive->PhyDrive, StartSector, DriveLetters, sizeof(DriveLetters), TRUE);
2777                         if (ERROR_SUCCESS == Status)
2778                         {
2779                                 break;
2780                         }
2781                         else
2782                         {
2783                                 Log("==== Volume not found, wait and retry %d... ====", i);
2784                                 Sleep(2);
2785                         }
2786                 }
2787         }
2788
2789     if (ERROR_SUCCESS == Status)
2790     {
2791         Log("Now lock and dismount volume <%s>", DriveLetters);
2792
2793         for (i = 0; i < MaxRetry; i++)
2794         {
2795             hVolume = CreateFileA(DriveLetters,
2796                 GENERIC_READ | GENERIC_WRITE,
2797                 FILE_SHARE_READ,
2798                 NULL,
2799                 OPEN_EXISTING,
2800                 FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
2801                 NULL);
2802
2803             if (hVolume == INVALID_HANDLE_VALUE)
2804             {
2805                 Log("Failed to create file volume, errcode:%u, wait and retry ...", LASTERR);
2806                 Sleep(2000);
2807             }
2808             else
2809             {
2810                 break;
2811             }
2812         }
2813
2814         if (hVolume == INVALID_HANDLE_VALUE)
2815         {
2816             Log("Failed to create file volume, errcode:%u", LASTERR);
2817         }
2818         else
2819         {
2820             bRet = DeviceIoControl(hVolume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2821             Log("FSCTL_LOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2822
2823             bRet = DeviceIoControl(hVolume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2824             Log("FSCTL_DISMOUNT_VOLUME bRet:%u code:%u", bRet, LASTERR);
2825         }
2826     }
2827     else if (ERROR_NOT_FOUND == Status)
2828     {
2829         Log("Volume not found, maybe not supported");
2830     }
2831     else
2832     {
2833         rc = 1;
2834         goto End;
2835     }
2836
2837         bRet = TryWritePart2(hDrive, StartSector);
2838         if (FALSE == bRet && Esp2Basic)
2839         {
2840                 Log("TryWritePart2 agagin ...");
2841                 Sleep(3000);
2842                 bRet = TryWritePart2(hDrive, StartSector);
2843         }
2844
2845         if (!bRet)
2846     {
2847                 if (pPhyDrive->PartStyle == 0)
2848                 {
2849                         if (DiskCheckWriteAccess(hDrive))
2850                         {
2851                                 Log("MBR DiskCheckWriteAccess success");
2852
2853                                 ForceMBR = TRUE;
2854
2855                                 Log("Try write failed, now delete partition 2 for MBR...");
2856                                 CHECK_CLOSE_HANDLE(hDrive);
2857
2858                                 Log("Now delete partition 2...");
2859                                 DISK_DeleteVtoyEFIPartition(pPhyDrive->PhyDrive, StartSector * 512ULL);
2860
2861                                 hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2862                                 if (hDrive == INVALID_HANDLE_VALUE)
2863                                 {
2864                                         Log("Failed to GetPhysicalHandle for write.");
2865                                         rc = 1;
2866                                         goto End;
2867                                 }
2868                         }
2869                         else
2870                         {
2871                                 Log("MBR DiskCheckWriteAccess failed");
2872                         }
2873                 }
2874                 else
2875                 {
2876                         Log("TryWritePart2 failed ....");
2877                         rc = 1;
2878                         goto End;
2879                 }
2880     }
2881
2882     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
2883
2884     Log("Write Ventoy to disk ............................ ");
2885     if (0 != FormatPart2Fat(hDrive, StartSector))
2886     {
2887         rc = 1;
2888         goto End;
2889     }
2890
2891     if (hVolume != INVALID_HANDLE_VALUE)
2892     {
2893         bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2894         Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2895         CHECK_CLOSE_HANDLE(hVolume);
2896     }
2897
2898     Log("Updating Boot Image ............................. ");
2899     if (WriteGrubStage1ToPhyDrive(hDrive, pPhyDrive->PartStyle) != 0)
2900     {
2901         rc = 1;
2902         goto End;
2903     }
2904
2905     //write reserved data
2906     SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2907     bRet = WriteFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2908     Log("Write resv data ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2909
2910     // Boot Image
2911     VentoyGetLocalBootImg(&BootImg);
2912
2913     // Use Old UUID
2914     memcpy(BootImg.BootCode + 0x180, MBR.BootCode + 0x180, 16);
2915     if (pPhyDrive->PartStyle)
2916     {
2917         BootImg.BootCode[92] = 0x22;
2918     }
2919
2920     if (ForceMBR == FALSE && memcmp(BootImg.BootCode, MBR.BootCode, 440) == 0)
2921     {
2922         Log("Boot image has no difference, no need to write.");
2923     }
2924     else
2925     {
2926         Log("Boot image need to write %u.", ForceMBR);
2927
2928         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2929
2930         memcpy(MBR.BootCode, BootImg.BootCode, 440);
2931         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2932         Log("Write Boot Image ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2933     }
2934
2935     if (pPhyDrive->PartStyle == 0)
2936     {
2937         if (0x00 == MBR.PartTbl[0].Active && 0x80 == MBR.PartTbl[1].Active)
2938         {
2939             Log("Need to chage 1st partition active and 2nd partition inactive.");
2940
2941             MBR.PartTbl[0].Active = 0x80;
2942             MBR.PartTbl[1].Active = 0x00;
2943
2944             SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2945             bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2946             Log("Write NEW MBR ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2947         }
2948     }
2949
2950         if (CleanDisk)
2951         {
2952                 if (!WriteBackupDataToDisk(hDrive, 0, pBackup, 34 * 512))
2953                 {
2954                         bWriteBack = FALSE;
2955                 }
2956
2957                 free(pBackup);
2958
2959                 if (bWriteBack)
2960                 {
2961                         Log("Write backup data success, now delete %s", BackBinFile);
2962                         DeleteFileA(BackBinFile);
2963                 }
2964                 else
2965                 {
2966                         Log("Write backup data failed");
2967                 }
2968
2969                 Sleep(1000);
2970         }
2971         else if (DelEFI)
2972         {
2973                 VTOY_GPT_HDR BackupHdr;
2974
2975                 VentoyFillBackupGptHead(pGptBkup, &BackupHdr);
2976                 if (!WriteBackupDataToDisk(hDrive, 512 * pGptBkup->Head.EfiBackupLBA, (BYTE*)(&BackupHdr), 512))
2977                 {
2978                         bWriteBack = FALSE;
2979                 }
2980
2981                 if (!WriteBackupDataToDisk(hDrive, 512 * (pGptBkup->Head.EfiBackupLBA - 32), (BYTE*)(pGptBkup->PartTbl), 32 * 512))
2982                 {
2983                         bWriteBack = FALSE;
2984                 }
2985
2986                 if (!WriteBackupDataToDisk(hDrive, 512, (BYTE*)pGptBkup + 512, 33 * 512))
2987                 {
2988                         bWriteBack = FALSE;
2989                 }
2990
2991                 if (bWriteBack)
2992                 {
2993                         Log("Write backup partition table success");
2994                 }
2995                 else
2996                 {
2997                         Log("Write backup partition table failed");
2998                 }
2999
3000                 Sleep(1000);
3001         }
3002
3003     //Refresh Drive Layout
3004     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
3005
3006 End:
3007
3008         if (hVolume != INVALID_HANDLE_VALUE)
3009         {
3010                 bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
3011                 Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
3012                 CHECK_CLOSE_HANDLE(hVolume);
3013         }
3014
3015     if (rc == 0)
3016     {
3017         Log("OK");
3018     }
3019     else
3020     {
3021                 PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
3022         FindProcessOccupyDisk(hDrive, pPhyDrive);
3023     }
3024
3025     CHECK_CLOSE_HANDLE(hDrive);
3026
3027     if (Esp2Basic)
3028     {
3029                 Log("Recover GPT partition type to basic");
3030                 DISK_ChangeVtoyEFI2Basic(pPhyDrive->PhyDrive, StartSector * 512);
3031     }
3032
3033         if (pPhyDrive->PartStyle == 1)
3034         {
3035                 if (ChangeAttr || ((pPhyDrive->Part2GPTAttr >> 56) != 0xC0))
3036                 {
3037                         Log("Change EFI partition attr %u <0x%llx> to <0x%llx>", ChangeAttr, pPhyDrive->Part2GPTAttr, 0xC000000000000001ULL);
3038                         if (DISK_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, StartSector * 512ULL, 0xC000000000000001ULL))
3039                         {
3040                                 Log("Change EFI partition attr success");
3041                                 pPhyDrive->Part2GPTAttr = 0xC000000000000001ULL;
3042                         }
3043                         else
3044                         {
3045                                 Log("Change EFI partition attr failed");
3046                         }
3047                 }
3048         }
3049
3050     if (pGptInfo)
3051     {
3052         free(pGptInfo);
3053     }
3054
3055     return rc;
3056 }
3057
3058