Ventoy2Disk/Ventoy2Disk/PhyDrive.c

   1 /******************************************************************************
   2  * PhyDrive.c
   3  *
   4  * Copyright (c) 2020, longpanda <admin@ventoy.net>
   5  *
   6  * This program is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU General Public License as
   8  * published by the Free Software Foundation; either version 3 of the
   9  * License, or (at your option) any later version.
  10  *
  11  * This program is distributed in the hope that it will be useful, but
  12  * WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, see <http://www.gnu.org/licenses/>.
  18  *
  19  */
  20
  21 #include <Windows.h>
  22 #include <winternl.h>
  23 #include <commctrl.h>
  24 #include <initguid.h>
  25 #include <vds.h>
  26 #include "resource.h"
  27 #include "Language.h"
  28 #include "Ventoy2Disk.h"
  29 #include "fat_filelib.h"
  30 #include "ff.h"
  31
  32 #define VDS_SET_ERROR SetLastError
  33 #define IVdsServiceLoader_LoadService(This, pwszMachineName, ppService) (This)->lpVtbl->LoadService(This, pwszMachineName, ppService)
  34 #define IVdsServiceLoader_Release(This) (This)->lpVtbl->Release(This)
  35 #define IVdsService_QueryProviders(This, masks, ppEnum) (This)->lpVtbl->QueryProviders(This, masks, ppEnum)
  36 #define IVdsService_WaitForServiceReady(This) ((This)->lpVtbl->WaitForServiceReady(This))
  37 #define IVdsService_CleanupObsoleteMountPoints(This) ((This)->lpVtbl->CleanupObsoleteMountPoints(This))
  38 #define IVdsService_Refresh(This) ((This)->lpVtbl->Refresh(This))
  39 #define IVdsService_Reenumerate(This) ((This)->lpVtbl->Reenumerate(This))
  40 #define IVdsSwProvider_QueryInterface(This, riid, ppvObject) (This)->lpVtbl->QueryInterface(This, riid, ppvObject)
  41 #define IVdsProvider_Release(This) (This)->lpVtbl->Release(This)
  42 #define IVdsSwProvider_QueryPacks(This, ppEnum) (This)->lpVtbl->QueryPacks(This, ppEnum)
  43 #define IVdsSwProvider_Release(This) (This)->lpVtbl->Release(This)
  44 #define IVdsPack_QueryDisks(This, ppEnum) (This)->lpVtbl->QueryDisks(This, ppEnum)
  45 #define IVdsDisk_GetProperties(This, pDiskProperties) (This)->lpVtbl->GetProperties(This, pDiskProperties)
  46 #define IVdsDisk_Release(This) (This)->lpVtbl->Release(This)
  47 #define IVdsDisk_QueryInterface(This, riid, ppvObject) (This)->lpVtbl->QueryInterface(This, riid, ppvObject)
  48 #define IVdsAdvancedDisk_QueryPartitions(This, ppPartitionPropArray, plNumberOfPartitions) (This)->lpVtbl->QueryPartitions(This, ppPartitionPropArray, plNumberOfPartitions)
  49 #define IVdsAdvancedDisk_DeletePartition(This, ullOffset, bForce, bForceProtected) (This)->lpVtbl->DeletePartition(This, ullOffset, bForce, bForceProtected)
  50 #define IVdsAdvancedDisk_Clean(This, bForce, bForceOEM, bFullClean, ppAsync) (This)->lpVtbl->Clean(This, bForce, bForceOEM, bFullClean, ppAsync)
  51 #define IVdsAdvancedDisk_Release(This) (This)->lpVtbl->Release(This)
  52 #define IEnumVdsObject_Next(This, celt, ppObjectArray, pcFetched) (This)->lpVtbl->Next(This, celt, ppObjectArray, pcFetched)
  53 #define IVdsPack_QueryVolumes(This, ppEnum) (This)->lpVtbl->QueryVolumes(This, ppEnum)
  54 #define IVdsVolume_QueryInterface(This, riid, ppvObject) (This)->lpVtbl->QueryInterface(This, riid, ppvObject)
  55 #define IVdsVolume_Release(This) (This)->lpVtbl->Release(This)
  56 #define IVdsVolumeMF3_QueryVolumeGuidPathnames(This, pwszPathArray, pulNumberOfPaths) (This)->lpVtbl->QueryVolumeGuidPathnames(This,pwszPathArray,pulNumberOfPaths)
  57 #define IVdsVolumeMF3_FormatEx2(This, pwszFileSystemTypeName, usFileSystemRevision, ulDesiredUnitAllocationSize, pwszLabel, Options, ppAsync) (This)->lpVtbl->FormatEx2(This, pwszFileSystemTypeName, usFileSystemRevision, ulDesiredUnitAllocationSize, pwszLabel, Options, ppAsync)
  58 #define IVdsVolumeMF3_Release(This) (This)->lpVtbl->Release(This)
  59 #define IVdsVolume_GetProperties(This, pVolumeProperties) (This)->lpVtbl->GetProperties(This,pVolumeProperties)
  60 #define IVdsAsync_Cancel(This) (This)->lpVtbl->Cancel(This)
  61 #define IVdsAsync_QueryStatus(This,pHrResult,pulPercentCompleted) (This)->lpVtbl->QueryStatus(This,pHrResult,pulPercentCompleted)
  62 #define IVdsAsync_Wait(This,pHrResult,pAsyncOut) (This)->lpVtbl->Wait(This,pHrResult,pAsyncOut)
  63 #define IVdsAsync_Release(This) (This)->lpVtbl->Release(This)
  64
  65 #define IUnknown_QueryInterface(This, a, b) (This)->lpVtbl->QueryInterface(This,a,b)
  66 #define IUnknown_Release(This) (This)->lpVtbl->Release(This)
  67
  68 /*
  69 * Delete all the partitions from a disk, using VDS
  70 * 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
  71 */
  72 BOOL DeletePartitions(DWORD DriveIndex, BOOL OnlyPart2)
  73 {
  74     BOOL r = FALSE;
  75     HRESULT hr;
  76     ULONG ulFetched;
  77     wchar_t wPhysicalName[48];
  78     IVdsServiceLoader *pLoader;
  79     IVdsService *pService;
  80     IEnumVdsObject *pEnum;
  81     IUnknown *pUnk;
  82
  83     swprintf_s(wPhysicalName, ARRAYSIZE(wPhysicalName), L"\\\\?\\PhysicalDrive%lu", DriveIndex);
  84
  85     // Initialize COM
  86     CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);
  87     CoInitializeSecurity(NULL, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_CONNECT,
  88         RPC_C_IMP_LEVEL_IMPERSONATE, NULL, 0, NULL);
  89
  90     // Create a VDS Loader Instance
  91     hr = CoCreateInstance(&CLSID_VdsLoader, NULL, CLSCTX_LOCAL_SERVER | CLSCTX_REMOTE_SERVER,
  92         &IID_IVdsServiceLoader, (void **)&pLoader);
  93     if (hr != S_OK) {
  94         VDS_SET_ERROR(hr);
  95         Log("Could not create VDS Loader Instance: %u", LASTERR);
  96         goto out;
  97     }
  98
  99     // Load the VDS Service
 100     hr = IVdsServiceLoader_LoadService(pLoader, L"", &pService);
 101     IVdsServiceLoader_Release(pLoader);
 102     if (hr != S_OK) {
 103         VDS_SET_ERROR(hr);
 104         Log("Could not load VDS Service: %u", LASTERR);
 105         goto out;
 106     }
 107
 108     // Wait for the Service to become ready if needed
 109     hr = IVdsService_WaitForServiceReady(pService);
 110     if (hr != S_OK) {
 111         VDS_SET_ERROR(hr);
 112         Log("VDS Service is not ready: %u", LASTERR);
 113         goto out;
 114     }
 115
 116     // Query the VDS Service Providers
 117     hr = IVdsService_QueryProviders(pService, VDS_QUERY_SOFTWARE_PROVIDERS, &pEnum);
 118     if (hr != S_OK) {
 119         VDS_SET_ERROR(hr);
 120         Log("Could not query VDS Service Providers: %u", LASTERR);
 121         goto out;
 122     }
 123
 124     while (IEnumVdsObject_Next(pEnum, 1, &pUnk, &ulFetched) == S_OK) {
 125         IVdsProvider *pProvider;
 126         IVdsSwProvider *pSwProvider;
 127         IEnumVdsObject *pEnumPack;
 128         IUnknown *pPackUnk;
 129
 130         // Get VDS Provider
 131         hr = IUnknown_QueryInterface(pUnk, &IID_IVdsProvider, (void **)&pProvider);
 132         IUnknown_Release(pUnk);
 133         if (hr != S_OK) {
 134             VDS_SET_ERROR(hr);
 135             Log("Could not get VDS Provider: %u", LASTERR);
 136             goto out;
 137         }
 138
 139         // Get VDS Software Provider
 140         hr = IVdsSwProvider_QueryInterface(pProvider, &IID_IVdsSwProvider, (void **)&pSwProvider);
 141         IVdsProvider_Release(pProvider);
 142         if (hr != S_OK) {
 143             VDS_SET_ERROR(hr);
 144             Log("Could not get VDS Software Provider: %u", LASTERR);
 145             goto out;
 146         }
 147
 148         // Get VDS Software Provider Packs
 149         hr = IVdsSwProvider_QueryPacks(pSwProvider, &pEnumPack);
 150         IVdsSwProvider_Release(pSwProvider);
 151         if (hr != S_OK) {
 152             VDS_SET_ERROR(hr);
 153             Log("Could not get VDS Software Provider Packs: %u", LASTERR);
 154             goto out;
 155         }
 156
 157         // Enumerate Provider Packs
 158         while (IEnumVdsObject_Next(pEnumPack, 1, &pPackUnk, &ulFetched) == S_OK) {
 159             IVdsPack *pPack;
 160             IEnumVdsObject *pEnumDisk;
 161             IUnknown *pDiskUnk;
 162
 163             hr = IUnknown_QueryInterface(pPackUnk, &IID_IVdsPack, (void **)&pPack);
 164             IUnknown_Release(pPackUnk);
 165             if (hr != S_OK) {
 166                 VDS_SET_ERROR(hr);
 167                 Log("Could not query VDS Software Provider Pack: %u", LASTERR);
 168                 goto out;
 169             }
 170
 171             // Use the pack interface to access the disks
 172             hr = IVdsPack_QueryDisks(pPack, &pEnumDisk);
 173             if (hr != S_OK) {
 174                 VDS_SET_ERROR(hr);
 175                 Log("Could not query VDS disks: %u", LASTERR);
 176                 goto out;
 177             }
 178
 179             // List disks
 180             while (IEnumVdsObject_Next(pEnumDisk, 1, &pDiskUnk, &ulFetched) == S_OK) {
 181                 VDS_DISK_PROP diskprop;
 182                 VDS_PARTITION_PROP* prop_array;
 183                 LONG i, prop_array_size;
 184                 IVdsDisk *pDisk;
 185                 IVdsAdvancedDisk *pAdvancedDisk;
 186
 187                 // Get the disk interface.
 188                 hr = IUnknown_QueryInterface(pDiskUnk, &IID_IVdsDisk, (void **)&pDisk);
 189                 if (hr != S_OK) {
 190                     VDS_SET_ERROR(hr);
 191                     Log("Could not query VDS Disk Interface: %u", LASTERR);
 192                     goto out;
 193                 }
 194
 195                 // Get the disk properties
 196                 hr = IVdsDisk_GetProperties(pDisk, &diskprop);
 197                 if (hr != S_OK) {
 198                     VDS_SET_ERROR(hr);
 199                     Log("Could not query VDS Disk Properties: %u", LASTERR);
 200                     goto out;
 201                 }
 202
 203                 // Isolate the disk we want
 204                 if (_wcsicmp(wPhysicalName, diskprop.pwszName) != 0) {
 205                     IVdsDisk_Release(pDisk);
 206                     continue;
 207                 }
 208
 209                 // Instantiate the AdvanceDisk interface for our disk.
 210                 hr = IVdsDisk_QueryInterface(pDisk, &IID_IVdsAdvancedDisk, (void **)&pAdvancedDisk);
 211                 IVdsDisk_Release(pDisk);
 212                 if (hr != S_OK) {
 213                     VDS_SET_ERROR(hr);
 214                     Log("Could not access VDS Advanced Disk interface: %u", LASTERR);
 215                     goto out;
 216                 }
 217
 218                 // Query the partition data, so we can get the start offset, which we need for deletion
 219                 hr = IVdsAdvancedDisk_QueryPartitions(pAdvancedDisk, &prop_array, &prop_array_size);
 220                 if (hr == S_OK) {
 221                     Log("Deleting ALL partition(s) from disk '%S':", diskprop.pwszName);
 222                     // Now go through each partition
 223                     for (i = 0; i < prop_array_size; i++) {
 224
 225                         Log("* Partition %d (offset: %lld, size: %llu)", prop_array[i].ulPartitionNumber,
 226                             prop_array[i].ullOffset, (ULONGLONG)prop_array[i].ullSize);
 227
 228                         if (OnlyPart2 && prop_array[i].ullOffset == 2048*512)
 229                         {
 230                             Log("Skip this partition...");
 231                             continue;
 232                         }
 233
 234
 235                         hr = IVdsAdvancedDisk_DeletePartition(pAdvancedDisk, prop_array[i].ullOffset, TRUE, TRUE);
 236                         if (hr != S_OK) {
 237                             r = FALSE;
 238                             VDS_SET_ERROR(hr);
 239                             Log("Could not delete partitions: %u", LASTERR);
 240                         }
 241                     }
 242                     r = TRUE;
 243                 }
 244                 else {
 245                     Log("No partition to delete on disk '%S'", diskprop.pwszName);
 246                     r = TRUE;
 247                 }
 248                 CoTaskMemFree(prop_array);
 249
 250 #if 0
 251                 // Issue a Clean while we're at it
 252                 HRESULT hr2 = E_FAIL;
 253                 ULONG completed;
 254                 IVdsAsync* pAsync;
 255                 hr = IVdsAdvancedDisk_Clean(pAdvancedDisk, TRUE, FALSE, FALSE, &pAsync);
 256                 while (SUCCEEDED(hr)) {
 257                     if (IS_ERROR(FormatStatus)) {
 258                         IVdsAsync_Cancel(pAsync);
 259                         break;
 260                     }
 261                     hr = IVdsAsync_QueryStatus(pAsync, &hr2, &completed);
 262                     if (SUCCEEDED(hr)) {
 263                         hr = hr2;
 264                         if (hr == S_OK)
 265                             break;
 266                         if (hr == VDS_E_OPERATION_PENDING)
 267                             hr = S_OK;
 268                     }
 269                     Sleep(500);
 270                 }
 271                 if (hr != S_OK) {
 272                     VDS_SET_ERROR(hr);
 273                     Log("Could not clean disk: %s", LASTERR);
 274                 }
 275 #endif
 276                 IVdsAdvancedDisk_Release(pAdvancedDisk);
 277                 goto out;
 278             }
 279         }
 280     }
 281
 282 out:
 283     return r;
 284 }
 285
 286
 287 static DWORD GetVentoyVolumeName(int PhyDrive, UINT32 StartSectorId, CHAR *NameBuf, UINT32 BufLen, BOOL DelSlash)
 288 {
 289     size_t len;
 290     BOOL bRet;
 291     DWORD dwSize;
 292     HANDLE hDrive;
 293     HANDLE hVolume;
 294     UINT64 PartOffset;
 295     DWORD Status = ERROR_NOT_FOUND;
 296     DISK_EXTENT *pExtents = NULL;
 297     CHAR VolumeName[MAX_PATH] = { 0 };
 298     VOLUME_DISK_EXTENTS DiskExtents;
 299
 300     PartOffset = 512ULL * StartSectorId;
 301
 302     Log("GetVentoyVolumeName PhyDrive %d PartOffset:%llu", PhyDrive, (ULONGLONG)PartOffset);
 303
 304     hVolume = FindFirstVolumeA(VolumeName, sizeof(VolumeName));
 305     if (hVolume == INVALID_HANDLE_VALUE)
 306     {
 307         return 1;
 308     }
 309
 310     do {
 311
 312         len = strlen(VolumeName);
 313         Log("Find volume:%s", VolumeName);
 314
 315         VolumeName[len - 1] = 0;
 316
 317         hDrive = CreateFileA(VolumeName, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
 318         if (hDrive == INVALID_HANDLE_VALUE)
 319         {
 320             continue;
 321         }
 322
 323         bRet = DeviceIoControl(hDrive,
 324             IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 325             NULL,
 326             0,
 327             &DiskExtents,
 328             (DWORD)(sizeof(DiskExtents)),
 329             (LPDWORD)&dwSize,
 330             NULL);
 331
 332         Log("IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS bRet:%u code:%u", bRet, LASTERR);
 333         Log("NumberOfDiskExtents:%u DiskNumber:%u", DiskExtents.NumberOfDiskExtents, DiskExtents.Extents[0].DiskNumber);
 334
 335         if (bRet && DiskExtents.NumberOfDiskExtents == 1)
 336         {
 337             pExtents = DiskExtents.Extents;
 338
 339             Log("This volume DiskNumber:%u offset:%llu", pExtents->DiskNumber, (ULONGLONG)pExtents->StartingOffset.QuadPart);
 340             if ((int)pExtents->DiskNumber == PhyDrive && pExtents->StartingOffset.QuadPart == PartOffset)
 341             {
 342                 Log("This volume match");
 343
 344                 if (!DelSlash)
 345                 {
 346                     VolumeName[len - 1] = '\\';
 347                 }
 348
 349                 sprintf_s(NameBuf, BufLen, "%s", VolumeName);
 350                 Status = ERROR_SUCCESS;
 351                 CloseHandle(hDrive);
 352                 break;
 353             }
 354         }
 355
 356         CloseHandle(hDrive);
 357     } while (FindNextVolumeA(hVolume, VolumeName, sizeof(VolumeName)));
 358
 359     FindVolumeClose(hVolume);
 360
 361     Log("GetVentoyVolumeName return %u", Status);
 362     return Status;
 363 }
 364
 365 static int GetLettersBelongPhyDrive(int PhyDrive, char *DriveLetters, size_t Length)
 366 {
 367     int n = 0;
 368     DWORD DataSize = 0;
 369     CHAR *Pos = NULL;
 370     CHAR *StringBuf = NULL;
 371
 372     DataSize = GetLogicalDriveStringsA(0, NULL);
 373     StringBuf = (CHAR *)malloc(DataSize + 1);
 374     if (StringBuf == NULL)
 375     {
 376         return 1;
 377     }
 378
 379     GetLogicalDriveStringsA(DataSize, StringBuf);
 380
 381     for (Pos = StringBuf; *Pos; Pos += strlen(Pos) + 1)
 382     {
 383         if (n < (int)Length && PhyDrive == GetPhyDriveByLogicalDrive(Pos[0]))
 384         {
 385             Log("%C: is belong to phydrive%d", Pos[0], PhyDrive);
 386             DriveLetters[n++] = Pos[0];
 387         }
 388     }
 389
 390     free(StringBuf);
 391     return 0;
 392 }
 393
 394 static HANDLE GetPhysicalHandle(int Drive, BOOLEAN bLockDrive, BOOLEAN bWriteAccess, BOOLEAN bWriteShare)
 395 {
 396     int i;
 397     DWORD dwSize;
 398     DWORD LastError;
 399     UINT64 EndTime;
 400     HANDLE hDrive = INVALID_HANDLE_VALUE;
 401     CHAR PhyDrive[128];
 402     CHAR DevPath[MAX_PATH] = { 0 };
 403
 404     safe_sprintf(PhyDrive, "\\\\.\\PhysicalDrive%d", Drive);
 405
 406     if (0 == QueryDosDeviceA(PhyDrive + 4, DevPath, sizeof(DevPath)))
 407     {
 408         Log("QueryDosDeviceA failed error:%u", GetLastError());
 409         strcpy_s(DevPath, sizeof(DevPath), "???");
 410     }
 411     else
 412     {
 413         Log("QueryDosDeviceA success %s", DevPath);
 414     }
 415
 416     for (i = 0; i < DRIVE_ACCESS_RETRIES; i++)
 417     {
 418         // Try without FILE_SHARE_WRITE (unless specifically requested) so that
 419         // we won't be bothered by the OS or other apps when we set up our data.
 420         // However this means we might have to wait for an access gap...
 421         // We keep FILE_SHARE_READ though, as this shouldn't hurt us any, and is
 422         // required for enumeration.
 423         hDrive = CreateFileA(PhyDrive,
 424             GENERIC_READ | (bWriteAccess ? GENERIC_WRITE : 0),
 425             FILE_SHARE_READ | (bWriteShare ? FILE_SHARE_WRITE : 0),
 426             NULL,
 427             OPEN_EXISTING,
 428             FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
 429             NULL);
 430
 431         LastError = GetLastError();
 432         Log("[%d] CreateFileA %s code:%u %p", i, PhyDrive, LastError, hDrive);
 433
 434         if (hDrive != INVALID_HANDLE_VALUE)
 435         {
 436             break;
 437         }
 438
 439         if ((LastError != ERROR_SHARING_VIOLATION) && (LastError != ERROR_ACCESS_DENIED))
 440         {
 441             break;
 442         }
 443
 444         if (i == 0)
 445         {
 446             Log("Waiting for access on %s [%s]...", PhyDrive, DevPath);
 447         }
 448         else if (!bWriteShare && (i > DRIVE_ACCESS_RETRIES / 3))
 449         {
 450             // If we can't seem to get a hold of the drive for some time, try to enable FILE_SHARE_WRITE...
 451             Log("Warning: Could not obtain exclusive rights. Retrying with write sharing enabled...");
 452             bWriteShare = TRUE;
 453
 454             // Try to report the process that is locking the drive
 455             // We also use bit 6 as a flag to indicate that SearchProcess was called.
 456             //access_mask = SearchProcess(DevPath, SEARCH_PROCESS_TIMEOUT, TRUE, TRUE, FALSE) | 0x40;
 457
 458         }
 459         Sleep(DRIVE_ACCESS_TIMEOUT / DRIVE_ACCESS_RETRIES);
 460     }
 461
 462     if (hDrive == INVALID_HANDLE_VALUE)
 463     {
 464         Log("Could not open %s %u", PhyDrive, LASTERR);
 465         goto End;
 466     }
 467
 468     if (bWriteAccess)
 469     {
 470         Log("Opened %s for %s write access", PhyDrive, bWriteShare ? "shared" : "exclusive");
 471     }
 472
 473     if (bLockDrive)
 474     {
 475         if (DeviceIoControl(hDrive, FSCTL_ALLOW_EXTENDED_DASD_IO, NULL, 0, NULL, 0, &dwSize, NULL))
 476         {
 477             Log("I/O boundary checks disabled");
 478         }
 479
 480         EndTime = GetTickCount64() + DRIVE_ACCESS_TIMEOUT;
 481
 482         do {
 483             if (DeviceIoControl(hDrive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL))
 484             {
 485                 Log("FSCTL_LOCK_VOLUME success");
 486                 goto End;
 487             }
 488             Sleep(DRIVE_ACCESS_TIMEOUT / DRIVE_ACCESS_RETRIES);
 489         } while (GetTickCount64() < EndTime);
 490
 491         // If we reached this section, either we didn't manage to get a lock or the user cancelled
 492         Log("Could not lock access to %s %u", PhyDrive, LASTERR);
 493
 494         // See if we can report the processes are accessing the drive
 495         //if (!IS_ERROR(FormatStatus) && (access_mask == 0))
 496         //    access_mask = SearchProcess(DevPath, SEARCH_PROCESS_TIMEOUT, TRUE, TRUE, FALSE);
 497         // Try to continue if the only access rights we saw were for read-only
 498         //if ((access_mask & 0x07) != 0x01)
 499         //    safe_closehandle(hDrive);
 500
 501         CHECK_CLOSE_HANDLE(hDrive);
 502     }
 503
 504 End:
 505
 506     if (hDrive == INVALID_HANDLE_VALUE)
 507     {
 508         Log("Can get handle of %s, maybe some process control it.", DevPath);
 509     }
 510
 511     return hDrive;
 512 }
 513
 514 int GetPhyDriveByLogicalDrive(int DriveLetter)
 515 {
 516     BOOL Ret;
 517     DWORD dwSize;
 518     HANDLE Handle;
 519     VOLUME_DISK_EXTENTS DiskExtents;
 520     CHAR PhyPath[128];
 521
 522     safe_sprintf(PhyPath, "\\\\.\\%C:", (CHAR)DriveLetter);
 523
 524     Handle = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 525     if (Handle == INVALID_HANDLE_VALUE)
 526     {
 527         Log("Could not open the disk<%s>, error:%u", PhyPath, LASTERR);
 528         return -1;
 529     }
 530
 531     Ret = DeviceIoControl(Handle,
 532         IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 533         NULL,
 534         0,
 535         &DiskExtents,
 536         (DWORD)(sizeof(DiskExtents)),
 537         (LPDWORD)&dwSize,
 538         NULL);
 539
 540     if (!Ret || DiskExtents.NumberOfDiskExtents == 0)
 541     {
 542         Log("DeviceIoControl IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS failed %s, error:%u", PhyPath, LASTERR);
 543         CHECK_CLOSE_HANDLE(Handle);
 544         return -1;
 545     }
 546     CHECK_CLOSE_HANDLE(Handle);
 547
 548     Log("LogicalDrive:%s PhyDrive:%d Offset:%llu ExtentLength:%llu",
 549         PhyPath,
 550         DiskExtents.Extents[0].DiskNumber,
 551         DiskExtents.Extents[0].StartingOffset.QuadPart,
 552         DiskExtents.Extents[0].ExtentLength.QuadPart
 553         );
 554
 555     return (int)DiskExtents.Extents[0].DiskNumber;
 556 }
 557
 558 int GetAllPhysicalDriveInfo(PHY_DRIVE_INFO *pDriveList, DWORD *pDriveCount)
 559 {
 560     int i;
 561     int Count;
 562     int id;
 563     int Letter = 'A';
 564     BOOL  bRet;
 565     DWORD dwBytes;
 566     DWORD DriveCount = 0;
 567     HANDLE Handle = INVALID_HANDLE_VALUE;
 568     CHAR PhyDrive[128];
 569     PHY_DRIVE_INFO *CurDrive = pDriveList;
 570     GET_LENGTH_INFORMATION LengthInfo;
 571     STORAGE_PROPERTY_QUERY Query;
 572     STORAGE_DESCRIPTOR_HEADER DevDescHeader;
 573     STORAGE_DEVICE_DESCRIPTOR *pDevDesc;
 574     int PhyDriveId[VENTOY_MAX_PHY_DRIVE];
 575
 576     Count = GetPhysicalDriveCount();
 577
 578     for (i = 0; i < Count && i < VENTOY_MAX_PHY_DRIVE; i++)
 579     {
 580         PhyDriveId[i] = i;
 581     }
 582
 583     dwBytes = GetLogicalDrives();
 584     Log("Logical Drives: 0x%x", dwBytes);
 585     while (dwBytes)
 586     {
 587         if (dwBytes & 0x01)
 588         {
 589             id = GetPhyDriveByLogicalDrive(Letter);
 590             Log("%C --> %d", Letter, id);
 591             if (id >= 0)
 592             {
 593                 for (i = 0; i < Count; i++)
 594                 {
 595                     if (PhyDriveId[i] == id)
 596                     {
 597                         break;
 598                     }
 599                 }
 600
 601                 if (i >= Count)
 602                 {
 603                     Log("Add phy%d to list", i);
 604                     PhyDriveId[Count] = id;
 605                     Count++;
 606                 }
 607             }
 608         }
 609
 610         Letter++;
 611         dwBytes >>= 1;
 612     }
 613
 614     for (i = 0; i < Count && DriveCount < VENTOY_MAX_PHY_DRIVE; i++)
 615     {
 616         CHECK_CLOSE_HANDLE(Handle);
 617
 618         safe_sprintf(PhyDrive, "\\\\.\\PhysicalDrive%d", PhyDriveId[i]);
 619         Handle = CreateFileA(PhyDrive, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
 620         Log("Create file Handle:%p %s status:%u", Handle, PhyDrive, LASTERR);
 621
 622         if (Handle == INVALID_HANDLE_VALUE)
 623         {
 624             continue;
 625         }
 626
 627         bRet = DeviceIoControl(Handle,
 628                                IOCTL_DISK_GET_LENGTH_INFO, NULL,
 629                                0,
 630                                &LengthInfo,
 631                                sizeof(LengthInfo),
 632                                &dwBytes,
 633                                NULL);
 634         if (!bRet)
 635         {
 636             Log("DeviceIoControl IOCTL_DISK_GET_LENGTH_INFO failed error:%u", LASTERR);
 637             continue;
 638         }
 639
 640         Log("PHYSICALDRIVE%d size %llu bytes", i, (ULONGLONG)LengthInfo.Length.QuadPart);
 641
 642         Query.PropertyId = StorageDeviceProperty;
 643         Query.QueryType = PropertyStandardQuery;
 644
 645         bRet = DeviceIoControl(Handle,
 646                                IOCTL_STORAGE_QUERY_PROPERTY,
 647                                &Query,
 648                                sizeof(Query),
 649                                &DevDescHeader,
 650                                sizeof(STORAGE_DESCRIPTOR_HEADER),
 651                                &dwBytes,
 652                                NULL);
 653         if (!bRet)
 654         {
 655             Log("DeviceIoControl1 error:%u dwBytes:%u", LASTERR, dwBytes);
 656             continue;
 657         }
 658
 659         pDevDesc = (STORAGE_DEVICE_DESCRIPTOR *)malloc(DevDescHeader.Size);
 660         if (!pDevDesc)
 661         {
 662             Log("failed to malloc error:%u len:%u", LASTERR, DevDescHeader.Size);
 663             continue;
 664         }
 665
 666         bRet = DeviceIoControl(Handle,
 667                                IOCTL_STORAGE_QUERY_PROPERTY,
 668                                &Query,
 669                                sizeof(Query),
 670                                pDevDesc,
 671                                DevDescHeader.Size,
 672                                &dwBytes,
 673                                NULL);
 674         if (!bRet)
 675         {
 676             Log("DeviceIoControl2 error:%u dwBytes:%u", LASTERR, dwBytes);
 677             free(pDevDesc);
 678             continue;
 679         }
 680
 681         CurDrive->PhyDrive = i;
 682         CurDrive->SizeInBytes = LengthInfo.Length.QuadPart;
 683         CurDrive->DeviceType = pDevDesc->DeviceType;
 684         CurDrive->RemovableMedia = pDevDesc->RemovableMedia;
 685         CurDrive->BusType = pDevDesc->BusType;
 686
 687         if (pDevDesc->VendorIdOffset)
 688         {
 689             safe_strcpy(CurDrive->VendorId, (char *)pDevDesc + pDevDesc->VendorIdOffset);
 690             TrimString(CurDrive->VendorId);
 691         }
 692
 693         if (pDevDesc->ProductIdOffset)
 694         {
 695             safe_strcpy(CurDrive->ProductId, (char *)pDevDesc + pDevDesc->ProductIdOffset);
 696             TrimString(CurDrive->ProductId);
 697         }
 698
 699         if (pDevDesc->ProductRevisionOffset)
 700         {
 701             safe_strcpy(CurDrive->ProductRev, (char *)pDevDesc + pDevDesc->ProductRevisionOffset);
 702             TrimString(CurDrive->ProductRev);
 703         }
 704
 705         if (pDevDesc->SerialNumberOffset)
 706         {
 707             safe_strcpy(CurDrive->SerialNumber, (char *)pDevDesc + pDevDesc->SerialNumberOffset);
 708             TrimString(CurDrive->SerialNumber);
 709         }
 710
 711         CurDrive++;
 712         DriveCount++;
 713
 714         free(pDevDesc);
 715
 716         CHECK_CLOSE_HANDLE(Handle);
 717     }
 718
 719     for (i = 0, CurDrive = pDriveList; i < (int)DriveCount; i++, CurDrive++)
 720     {
 721         Log("PhyDrv:%d BusType:%-4s Removable:%u Size:%dGB(%llu) Name:%s %s",
 722             CurDrive->PhyDrive, GetBusTypeString(CurDrive->BusType), CurDrive->RemovableMedia,
 723             GetHumanReadableGBSize(CurDrive->SizeInBytes), CurDrive->SizeInBytes,
 724             CurDrive->VendorId, CurDrive->ProductId);
 725     }
 726
 727     *pDriveCount = DriveCount;
 728
 729     return 0;
 730 }
 731
 732
 733 static HANDLE g_FatPhyDrive;
 734 static UINT64 g_Part2StartSec;
 735 static int GetVentoyVersionFromFatFile(CHAR *VerBuf, size_t BufLen)
 736 {
 737     int rc = 1;
 738     int size = 0;
 739     char *buf = NULL;
 740     void *flfile = NULL;
 741
 742     flfile = fl_fopen("/grub/grub.cfg", "rb");
 743     if (flfile)
 744     {
 745         fl_fseek(flfile, 0, SEEK_END);
 746         size = (int)fl_ftell(flfile);
 747
 748         fl_fseek(flfile, 0, SEEK_SET);
 749
 750         buf = (char *)malloc(size + 1);
 751         if (buf)
 752         {
 753             fl_fread(buf, 1, size, flfile);
 754             buf[size] = 0;
 755
 756             rc = 0;
 757             sprintf_s(VerBuf, BufLen, "%s", ParseVentoyVersionFromString(buf));
 758             free(buf);
 759         }
 760
 761         fl_fclose(flfile);
 762     }
 763
 764     return rc;
 765 }
 766
 767 static int VentoyFatDiskRead(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 768 {
 769     DWORD dwSize;
 770     BOOL bRet;
 771     DWORD ReadSize;
 772     LARGE_INTEGER liCurrentPosition;
 773
 774     liCurrentPosition.QuadPart = Sector + g_Part2StartSec;
 775     liCurrentPosition.QuadPart *= 512;
 776     SetFilePointerEx(g_FatPhyDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 777
 778     ReadSize = (DWORD)(SectorCount * 512);
 779
 780     bRet = ReadFile(g_FatPhyDrive, Buffer, ReadSize, &dwSize, NULL);
 781     if (bRet == FALSE || dwSize != ReadSize)
 782     {
 783         Log("ReadFile error bRet:%u WriteSize:%u dwSize:%u ErrCode:%u\n", bRet, ReadSize, dwSize, LASTERR);
 784     }
 785
 786     return 1;
 787 }
 788
 789
 790 int GetVentoyVerInPhyDrive(const PHY_DRIVE_INFO *pDriveInfo, CHAR *VerBuf, size_t BufLen)
 791 {
 792     int rc = 0;
 793     HANDLE hDrive;
 794
 795     hDrive = GetPhysicalHandle(pDriveInfo->PhyDrive, FALSE, FALSE, FALSE);
 796     if (hDrive == INVALID_HANDLE_VALUE)
 797     {
 798         return 1;
 799     }
 800
 801     g_FatPhyDrive = hDrive;
 802     g_Part2StartSec = (pDriveInfo->SizeInBytes - VENTOY_EFI_PART_SIZE) / 512;
 803
 804     Log("Parse FAT fs...");
 805
 806     fl_init();
 807
 808     if (0 == fl_attach_media(VentoyFatDiskRead, NULL))
 809     {
 810         rc = GetVentoyVersionFromFatFile(VerBuf, BufLen);
 811     }
 812     else
 813     {
 814         rc = 1;
 815     }
 816
 817     fl_shutdown();
 818
 819     CHECK_CLOSE_HANDLE(hDrive);
 820
 821     return rc;
 822 }
 823
 824
 825
 826
 827 static unsigned int g_disk_unxz_len = 0;
 828 static BYTE *g_part_img_pos = NULL;
 829 static BYTE *g_part_img_buf[VENTOY_EFI_PART_SIZE / SIZE_1MB];
 830
 831 static int disk_xz_flush(void *src, unsigned int size)
 832 {
 833     unsigned int i;
 834     BYTE *buf = (BYTE *)src;
 835
 836     for (i = 0; i < size; i++)
 837     {
 838         *g_part_img_pos = *buf++;
 839
 840         g_disk_unxz_len++;
 841         if ((g_disk_unxz_len % SIZE_1MB) == 0)
 842         {
 843             g_part_img_pos = g_part_img_buf[g_disk_unxz_len / SIZE_1MB];
 844         }
 845         else
 846         {
 847             g_part_img_pos++;
 848         }
 849     }
 850
 851     return (int)size;
 852 }
 853
 854 static void unxz_error(char *x)
 855 {
 856     Log("%s", x);
 857 }
 858
 859 static BOOL TryWritePart2(HANDLE hDrive, UINT64 StartSectorId)
 860 {
 861     BOOL bRet;
 862     DWORD TrySize = 16 * 1024;
 863     DWORD dwSize;
 864     BYTE *Buffer = NULL;
 865     unsigned char *data = NULL;
 866     LARGE_INTEGER liCurrentPosition;
 867
 868     liCurrentPosition.QuadPart = StartSectorId * 512;
 869     SetFilePointerEx(hDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 870
 871     Buffer = malloc(TrySize);
 872
 873     bRet = WriteFile(hDrive, Buffer, TrySize, &dwSize, NULL);
 874
 875     free(Buffer);
 876
 877     Log("Try write part2 bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 878
 879     if (bRet && dwSize == TrySize)
 880     {
 881         return TRUE;
 882     }
 883
 884     return FALSE;
 885 }
 886
 887 static int FormatPart2Fat(HANDLE hDrive, UINT64 StartSectorId)
 888 {
 889     int i;
 890     int rc = 0;
 891     int len = 0;
 892     int writelen = 0;
 893     int partwrite = 0;
 894     DWORD dwSize = 0;
 895     BOOL bRet;
 896     unsigned char *data = NULL;
 897     LARGE_INTEGER liCurrentPosition;
 898
 899     Log("FormatPart2Fat ...");
 900
 901     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&data, &len);
 902     if (rc)
 903     {
 904         Log("Failed to read img file %p %u", data, len);
 905         return 1;
 906     }
 907
 908     liCurrentPosition.QuadPart = StartSectorId * 512;
 909     SetFilePointerEx(hDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 910
 911     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
 912
 913     g_part_img_buf[0] = (BYTE *)malloc(VENTOY_EFI_PART_SIZE);
 914     if (g_part_img_buf[0])
 915     {
 916         Log("Malloc whole img buffer success, now decompress ...");
 917         unxz(data, len, NULL, NULL, g_part_img_buf[0], &writelen, unxz_error);
 918
 919         if (len == writelen)
 920         {
 921             Log("decompress finished success");
 922             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 923             {
 924                 dwSize = 0;
 925                 bRet = WriteFile(hDrive, g_part_img_buf[0] + i * SIZE_1MB, SIZE_1MB, &dwSize, NULL);
 926                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 927
 928                 if (!bRet)
 929                 {
 930                     rc = 1;
 931                     goto End;
 932                 }
 933
 934                 PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START + i);
 935             }
 936         }
 937         else
 938         {
 939             rc = 1;
 940             Log("decompress finished failed");
 941             goto End;
 942         }
 943     }
 944     else
 945     {
 946         Log("Failed to malloc whole img size %u, now split it", VENTOY_EFI_PART_SIZE);
 947
 948         partwrite = 1;
 949         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 950         {
 951             g_part_img_buf[i] = (BYTE *)malloc(SIZE_1MB);
 952             if (g_part_img_buf[i] == NULL)
 953             {
 954                 rc = 1;
 955                 goto End;
 956             }
 957         }
 958
 959         Log("Malloc part img buffer success, now decompress ...");
 960
 961         g_part_img_pos = g_part_img_buf[0];
 962
 963         unxz(data, len, NULL, disk_xz_flush, NULL, NULL, unxz_error);
 964
 965         if (g_disk_unxz_len == VENTOY_EFI_PART_SIZE)
 966         {
 967             Log("decompress finished success");
 968             for (int i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 969             {
 970                 dwSize = 0;
 971                 bRet = WriteFile(hDrive, g_part_img_buf[i], SIZE_1MB, &dwSize, NULL);
 972                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
 973
 974                 if (!bRet)
 975                 {
 976                     rc = 1;
 977                     goto End;
 978                 }
 979
 980                 PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START + i);
 981             }
 982         }
 983         else
 984         {
 985             rc = 1;
 986             Log("decompress finished failed");
 987             goto End;
 988         }
 989     }
 990
 991 End:
 992
 993     if (data) free(data);
 994
 995     if (partwrite)
 996     {
 997         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
 998         {
 999             if (g_part_img_buf[i]) free(g_part_img_buf[i]);
1000         }
1001     }
1002     else
1003     {
1004         if (g_part_img_buf[0]) free(g_part_img_buf[0]);
1005     }
1006
1007     return rc;
1008 }
1009
1010 static int WriteGrubStage1ToPhyDrive(HANDLE hDrive)
1011 {
1012     int Len = 0;
1013     int readLen = 0;
1014     BOOL bRet;
1015     DWORD dwSize;
1016     BYTE *ImgBuf = NULL;
1017     BYTE *RawBuf = NULL;
1018
1019     Log("WriteGrubStage1ToPhyDrive ...");
1020
1021     RawBuf = (BYTE *)malloc(SIZE_1MB);
1022     if (!RawBuf)
1023     {
1024         return 1;
1025     }
1026
1027     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1028     {
1029         Log("Failed to read stage1 img");
1030         free(RawBuf);
1031         return 1;
1032     }
1033
1034     unxz(ImgBuf, Len, NULL, NULL, RawBuf, &readLen, unxz_error);
1035
1036     SetFilePointer(hDrive, 512, NULL, FILE_BEGIN);
1037
1038     bRet = WriteFile(hDrive, RawBuf, SIZE_1MB - 512, &dwSize, NULL);
1039     Log("WriteFile Ret:%u dwSize:%u ErrCode:%u", bRet, dwSize, GetLastError());
1040
1041     free(RawBuf);
1042     free(ImgBuf);
1043     return 0;
1044 }
1045
1046
1047
1048 static int FormatPart1exFAT(UINT64 DiskSizeBytes)
1049 {
1050     MKFS_PARM Option;
1051     FRESULT Ret;
1052     FATFS fs;
1053
1054     Option.fmt = FM_EXFAT;
1055     Option.n_fat = 1;
1056     Option.align = 8;
1057     Option.n_root = 1;
1058
1059     // < 32GB select 32KB as cluster size
1060     // > 32GB select 128KB as cluster size
1061     if (DiskSizeBytes / 1024 / 1024 / 1024 <= 32)
1062     {
1063         Option.au_size = 32768;
1064     }
1065     else
1066     {
1067         Option.au_size = 131072;
1068     }
1069
1070     Log("Formatting Part1 exFAT ...");
1071
1072     Ret = f_mkfs(TEXT("0:"), &Option, 0, 8 * 1024 * 1024);
1073
1074     if (FR_OK == Ret)
1075     {
1076         Log("Formatting Part1 exFAT success");
1077
1078         Ret = f_mount(&fs, TEXT("0:"), 1);
1079         Log("mount part %d", Ret);
1080
1081         if (FR_OK == Ret)
1082         {
1083             Ret = f_setlabel(TEXT("Ventoy"));
1084             Log("f_setlabel %d", Ret);
1085
1086             Ret = f_mount(0, TEXT("0:"), 1);
1087             Log("umount part %d", Ret);
1088         }
1089
1090         return 0;
1091     }
1092     else
1093     {
1094         Log("Formatting Part1 exFAT failed");
1095         return 1;
1096     }
1097 }
1098
1099
1100 int InstallVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive)
1101 {
1102     int i;
1103     int rc = 0;
1104     int state = 0;
1105     HANDLE hDrive;
1106     DWORD dwSize;
1107     BOOL bRet;
1108     CHAR MountDrive;
1109     CHAR DriveName[] = "?:\\";
1110     CHAR DriveLetters[MAX_PATH] = { 0 };
1111     MBR_HEAD MBR;
1112
1113     Log("InstallVentoy2PhyDrive PhyDrive%d <<%s %s %dGB>>",
1114         pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1115         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1116
1117     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1118
1119     VentoyFillMBR(pPhyDrive->SizeInBytes, &MBR);
1120
1121     Log("Lock disk for clean ............................. ");
1122
1123     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1124     if (hDrive == INVALID_HANDLE_VALUE)
1125     {
1126         Log("Failed to open physical disk");
1127         return 1;
1128     }
1129
1130     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1131
1132     if (DriveLetters[0] == 0)
1133     {
1134         Log("No drive letter was assigned...");
1135         DriveName[0] = GetFirstUnusedDriveLetter();
1136         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1137     }
1138     else
1139     {
1140         // Unmount all mounted volumes that belong to this drive
1141         // Do it in reverse so that we always end on the first volume letter
1142         for (i = (int)strlen(DriveLetters); i > 0; i--)
1143         {
1144             DriveName[0] = DriveLetters[i - 1];
1145             bRet = DeleteVolumeMountPointA(DriveName);
1146             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1147         }
1148     }
1149
1150     MountDrive = DriveName[0];
1151     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1152
1153     // It kind of blows, but we have to relinquish access to the physical drive
1154     // for VDS to be able to delete the partitions that reside on it...
1155     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1156     CHECK_CLOSE_HANDLE(hDrive);
1157
1158     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1159
1160     if (!DeletePartitions(pPhyDrive->PhyDrive, FALSE))
1161     {
1162         Log("Notice: Could not delete partitions: %u", GetLastError());
1163     }
1164
1165     Log("Deleting all partitions ......................... OK");
1166
1167     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1168
1169     Log("Lock disk for write ............................. ");
1170     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1171     if (hDrive == INVALID_HANDLE_VALUE)
1172     {
1173         Log("Failed to GetPhysicalHandle for write.");
1174         rc = 1;
1175         goto End;
1176     }
1177
1178     //Refresh Drive Layout
1179     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1180
1181     disk_io_set_param(hDrive, MBR.PartTbl[0].StartSectorId + MBR.PartTbl[0].SectorCount);
1182
1183     PROGRESS_BAR_SET_POS(PT_FORMAT_PART1);
1184
1185     Log("Formatting part1 exFAT ...");
1186     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1187     {
1188         rc = 1;
1189         goto End;
1190     }
1191
1192     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1193     Log("Writing part2 FAT img ...");
1194     if (0 != FormatPart2Fat(hDrive, MBR.PartTbl[1].StartSectorId))
1195     {
1196         rc = 1;
1197         goto End;
1198     }
1199
1200     PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
1201     Log("Writting Boot Image ............................. ");
1202     if (WriteGrubStage1ToPhyDrive(hDrive) != 0)
1203     {
1204         rc = 1;
1205         goto End;
1206     }
1207
1208
1209     PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
1210     Log("Writting Partition Table ........................ ");
1211     SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1212     if (!WriteFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL))
1213     {
1214         rc = 1;
1215         Log("Write MBR Failed, dwSize:%u ErrCode:%u", dwSize, GetLastError());
1216         goto End;
1217     }
1218
1219     Log("Write MBR OK ...");
1220
1221     //Refresh Drive Layout
1222     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1223
1224 End:
1225     CHECK_CLOSE_HANDLE(hDrive);
1226
1227     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1228     Log("Mounting Ventoy Partition ....................... ");
1229     Sleep(1000);
1230
1231     state = 0;
1232     memset(DriveLetters, 0, sizeof(DriveLetters));
1233     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1234     Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1235
1236     for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1237     {
1238         DriveName[0] = DriveLetters[i];
1239         if (IsVentoyLogicalDrive(DriveName[0]))
1240         {
1241             Log("%s is ventoy part2, delete mountpoint", DriveName);
1242             DeleteVolumeMountPointA(DriveName);
1243         }
1244         else
1245         {
1246             Log("%s is ventoy part1, already mounted", DriveName);
1247             state = 1;
1248         }
1249     }
1250
1251     if (state != 1)
1252     {
1253         Log("need to mount ventoy part1...");
1254         if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, MBR.PartTbl[0].StartSectorId, DriveLetters, sizeof(DriveLetters), FALSE))
1255         {
1256             DriveName[0] = MountDrive;
1257             bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1258             Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1259         }
1260         else
1261         {
1262             Log("Failed to find ventoy volume");
1263         }
1264     }
1265
1266     Log("OK\n");
1267
1268     return rc;
1269 }
1270
1271 int UpdateVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive)
1272 {
1273     int i;
1274     int rc = 0;
1275     BOOL ForceMBR = FALSE;
1276     HANDLE hVolume;
1277     HANDLE hDrive;
1278     DWORD Status;
1279     DWORD dwSize;
1280     BOOL bRet;
1281     CHAR DriveName[] = "?:\\";
1282     CHAR DriveLetters[MAX_PATH] = { 0 };
1283     UINT32 StartSector;
1284     MBR_HEAD BootImg;
1285     MBR_HEAD MBR;
1286
1287     StartSector = (UINT32)(pPhyDrive->SizeInBytes / 512 - VENTOY_EFI_PART_SIZE / 512);
1288
1289     Log("UpdateVentoy2PhyDrive PhyDrive%d <<%s %s %dGB>>",
1290         pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1291         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1292
1293     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1294
1295     Log("Lock disk for umount ............................ ");
1296
1297     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1298     if (hDrive == INVALID_HANDLE_VALUE)
1299     {
1300         Log("Failed to open physical disk");
1301         return 1;
1302     }
1303
1304     // Read MBR
1305     ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
1306
1307     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1308
1309     if (DriveLetters[0] == 0)
1310     {
1311         Log("No drive letter was assigned...");
1312     }
1313     else
1314     {
1315         // Unmount all mounted volumes that belong to this drive
1316         // Do it in reverse so that we always end on the first volume letter
1317         for (i = (int)strlen(DriveLetters); i > 0; i--)
1318         {
1319             DriveName[0] = DriveLetters[i - 1];
1320             if (IsVentoyLogicalDrive(DriveName[0]))
1321             {
1322                 Log("%s is ventoy logical drive", DriveName);
1323                 bRet = DeleteVolumeMountPointA(DriveName);
1324                 Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, LASTERR);
1325                 break;
1326             }
1327         }
1328     }
1329
1330     // It kind of blows, but we have to relinquish access to the physical drive
1331     // for VDS to be able to delete the partitions that reside on it...
1332     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1333     CHECK_CLOSE_HANDLE(hDrive);
1334
1335     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1336
1337     Log("Lock disk for update ............................ ");
1338     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1339     if (hDrive == INVALID_HANDLE_VALUE)
1340     {
1341         Log("Failed to GetPhysicalHandle for write.");
1342         rc = 1;
1343         goto End;
1344     }
1345
1346     PROGRESS_BAR_SET_POS(PT_LOCK_VOLUME);
1347
1348     Log("Lock volume for update .......................... ");
1349     hVolume = INVALID_HANDLE_VALUE;
1350     Status = GetVentoyVolumeName(pPhyDrive->PhyDrive, MBR.PartTbl[1].StartSectorId, DriveLetters, sizeof(DriveLetters), TRUE);
1351     if (ERROR_SUCCESS == Status)
1352     {
1353         Log("Now lock and dismount volume <%s>", DriveLetters);
1354         hVolume = CreateFileA(DriveLetters,
1355             GENERIC_READ | GENERIC_WRITE,
1356             FILE_SHARE_READ,
1357             NULL,
1358             OPEN_EXISTING,
1359             FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
1360             NULL);
1361
1362         if (hVolume == INVALID_HANDLE_VALUE)
1363         {
1364             Log("Failed to create file volume, errcode:%u", LASTERR);
1365             rc = 1;
1366             goto End;
1367         }
1368
1369         bRet = DeviceIoControl(hVolume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1370         Log("FSCTL_LOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
1371
1372         bRet = DeviceIoControl(hVolume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1373         Log("FSCTL_DISMOUNT_VOLUME bRet:%u code:%u", bRet, LASTERR);
1374     }
1375     else if (ERROR_NOT_FOUND == Status)
1376     {
1377         Log("Volume not found, maybe not supported");
1378     }
1379     else
1380     {
1381         rc = 1;
1382         goto End;
1383     }
1384
1385
1386     if (!TryWritePart2(hDrive, StartSector))
1387     {
1388         ForceMBR = TRUE;
1389         Log("Try write failed, now delete partition 2...");
1390
1391         CHECK_CLOSE_HANDLE(hDrive);
1392
1393         Log("Now delete partition 2...");
1394         DeletePartitions(pPhyDrive->PhyDrive, TRUE);
1395
1396         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1397         if (hDrive == INVALID_HANDLE_VALUE)
1398         {
1399             Log("Failed to GetPhysicalHandle for write.");
1400             rc = 1;
1401             goto End;
1402         }
1403     }
1404
1405
1406     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1407
1408     Log("Write Ventoy to disk ............................ ");
1409     if (0 != FormatPart2Fat(hDrive, StartSector))
1410     {
1411         rc = 1;
1412         goto End;
1413     }
1414
1415     if (hVolume != INVALID_HANDLE_VALUE)
1416     {
1417         bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1418         Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
1419         CHECK_CLOSE_HANDLE(hVolume);
1420     }
1421
1422     Log("Updating Boot Image ............................. ");
1423     if (WriteGrubStage1ToPhyDrive(hDrive) != 0)
1424     {
1425         rc = 1;
1426         goto End;
1427     }
1428
1429     // Boot Image
1430     VentoyGetLocalBootImg(&BootImg);
1431
1432     // Use Old UUID
1433     memcpy(BootImg.BootCode + 0x180, MBR.BootCode + 0x180, 16);
1434
1435     if (ForceMBR == FALSE && memcmp(BootImg.BootCode, MBR.BootCode, 440) == 0)
1436     {
1437         Log("Boot image has no difference, no need to write.");
1438     }
1439     else
1440     {
1441         Log("Boot image need to write %u.", ForceMBR);
1442
1443         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1444
1445         memcpy(MBR.BootCode, BootImg.BootCode, 440);
1446         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
1447         Log("Write Boot Image ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
1448     }
1449
1450     //Refresh Drive Layout
1451     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1452
1453 End:
1454     CHECK_CLOSE_HANDLE(hDrive);
1455
1456     return rc;
1457 }
1458