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