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 && prop_array[i].ullOffset == 2048*512)
 234                         {
 235                             Log("Skip this partition...");
 236                             continue;
 237                         }
 238
 239
 240                         hr = IVdsAdvancedDisk_DeletePartition(pAdvancedDisk, prop_array[i].ullOffset, TRUE, TRUE);
 241                         if (hr != S_OK) {
 242                             r = FALSE;
 243                             VDS_SET_ERROR(hr);
 244                             Log("Could not delete partitions: %u", LASTERR);
 245                         }
 246                     }
 247                     r = TRUE;
 248                 }
 249                 else {
 250                     Log("No partition to delete on disk '%S'", diskprop.pwszName);
 251                     r = TRUE;
 252                 }
 253                 CoTaskMemFree(prop_array);
 254
 255 #if 0
 256                 // Issue a Clean while we're at it
 257                 HRESULT hr2 = E_FAIL;
 258                 ULONG completed;
 259                 IVdsAsync* pAsync;
 260                 hr = IVdsAdvancedDisk_Clean(pAdvancedDisk, TRUE, FALSE, FALSE, &pAsync);
 261                 while (SUCCEEDED(hr)) {
 262                     if (IS_ERROR(FormatStatus)) {
 263                         IVdsAsync_Cancel(pAsync);
 264                         break;
 265                     }
 266                     hr = IVdsAsync_QueryStatus(pAsync, &hr2, &completed);
 267                     if (SUCCEEDED(hr)) {
 268                         hr = hr2;
 269                         if (hr == S_OK)
 270                             break;
 271                         if (hr == VDS_E_OPERATION_PENDING)
 272                             hr = S_OK;
 273                     }
 274                     Sleep(500);
 275                 }
 276                 if (hr != S_OK) {
 277                     VDS_SET_ERROR(hr);
 278                     Log("Could not clean disk: %s", LASTERR);
 279                 }
 280 #endif
 281                 IVdsAdvancedDisk_Release(pAdvancedDisk);
 282                 goto out;
 283             }
 284         }
 285     }
 286
 287 out:
 288     return r;
 289 }
 290
 291
 292 static DWORD GetVentoyVolumeName(int PhyDrive, UINT32 StartSectorId, CHAR *NameBuf, UINT32 BufLen, BOOL DelSlash)
 293 {
 294     size_t len;
 295     BOOL bRet;
 296     DWORD dwSize;
 297     HANDLE hDrive;
 298     HANDLE hVolume;
 299     UINT64 PartOffset;
 300     DWORD Status = ERROR_NOT_FOUND;
 301     DISK_EXTENT *pExtents = NULL;
 302     CHAR VolumeName[MAX_PATH] = { 0 };
 303     VOLUME_DISK_EXTENTS DiskExtents;
 304
 305     PartOffset = 512ULL * StartSectorId;
 306
 307     Log("GetVentoyVolumeName PhyDrive %d PartOffset:%llu", PhyDrive, (ULONGLONG)PartOffset);
 308
 309     hVolume = FindFirstVolumeA(VolumeName, sizeof(VolumeName));
 310     if (hVolume == INVALID_HANDLE_VALUE)
 311     {
 312         return 1;
 313     }
 314
 315     do {
 316
 317         len = strlen(VolumeName);
 318         Log("Find volume:%s", VolumeName);
 319
 320         VolumeName[len - 1] = 0;
 321
 322         hDrive = CreateFileA(VolumeName, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
 323         if (hDrive == INVALID_HANDLE_VALUE)
 324         {
 325             continue;
 326         }
 327
 328         bRet = DeviceIoControl(hDrive,
 329             IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 330             NULL,
 331             0,
 332             &DiskExtents,
 333             (DWORD)(sizeof(DiskExtents)),
 334             (LPDWORD)&dwSize,
 335             NULL);
 336
 337         Log("IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS bRet:%u code:%u", bRet, LASTERR);
 338         Log("NumberOfDiskExtents:%u DiskNumber:%u", DiskExtents.NumberOfDiskExtents, DiskExtents.Extents[0].DiskNumber);
 339
 340         if (bRet && DiskExtents.NumberOfDiskExtents == 1)
 341         {
 342             pExtents = DiskExtents.Extents;
 343
 344             Log("This volume DiskNumber:%u offset:%llu", pExtents->DiskNumber, (ULONGLONG)pExtents->StartingOffset.QuadPart);
 345             if ((int)pExtents->DiskNumber == PhyDrive && pExtents->StartingOffset.QuadPart == PartOffset)
 346             {
 347                 Log("This volume match");
 348
 349                 if (!DelSlash)
 350                 {
 351                     VolumeName[len - 1] = '\\';
 352                 }
 353
 354                 sprintf_s(NameBuf, BufLen, "%s", VolumeName);
 355                 Status = ERROR_SUCCESS;
 356                 CloseHandle(hDrive);
 357                 break;
 358             }
 359         }
 360
 361         CloseHandle(hDrive);
 362     } while (FindNextVolumeA(hVolume, VolumeName, sizeof(VolumeName)));
 363
 364     FindVolumeClose(hVolume);
 365
 366     Log("GetVentoyVolumeName return %u", Status);
 367     return Status;
 368 }
 369
 370 static int GetLettersBelongPhyDrive(int PhyDrive, char *DriveLetters, size_t Length)
 371 {
 372     int n = 0;
 373     DWORD DataSize = 0;
 374     CHAR *Pos = NULL;
 375     CHAR *StringBuf = NULL;
 376
 377     DataSize = GetLogicalDriveStringsA(0, NULL);
 378     StringBuf = (CHAR *)malloc(DataSize + 1);
 379     if (StringBuf == NULL)
 380     {
 381         return 1;
 382     }
 383
 384     GetLogicalDriveStringsA(DataSize, StringBuf);
 385
 386     for (Pos = StringBuf; *Pos; Pos += strlen(Pos) + 1)
 387     {
 388         if (n < (int)Length && PhyDrive == GetPhyDriveByLogicalDrive(Pos[0]))
 389         {
 390             Log("%C: is belong to phydrive%d", Pos[0], PhyDrive);
 391             DriveLetters[n++] = Pos[0];
 392         }
 393     }
 394
 395     free(StringBuf);
 396     return 0;
 397 }
 398
 399 static HANDLE GetPhysicalHandle(int Drive, BOOLEAN bLockDrive, BOOLEAN bWriteAccess, BOOLEAN bWriteShare)
 400 {
 401     int i;
 402     DWORD dwSize;
 403     DWORD LastError;
 404     UINT64 EndTime;
 405     HANDLE hDrive = INVALID_HANDLE_VALUE;
 406     CHAR PhyDrive[128];
 407     CHAR DevPath[MAX_PATH] = { 0 };
 408
 409     safe_sprintf(PhyDrive, "\\\\.\\PhysicalDrive%d", Drive);
 410
 411     if (0 == QueryDosDeviceA(PhyDrive + 4, DevPath, sizeof(DevPath)))
 412     {
 413         Log("QueryDosDeviceA failed error:%u", GetLastError());
 414         strcpy_s(DevPath, sizeof(DevPath), "???");
 415     }
 416     else
 417     {
 418         Log("QueryDosDeviceA success %s", DevPath);
 419     }
 420
 421     for (i = 0; i < DRIVE_ACCESS_RETRIES; i++)
 422     {
 423         // Try without FILE_SHARE_WRITE (unless specifically requested) so that
 424         // we won't be bothered by the OS or other apps when we set up our data.
 425         // However this means we might have to wait for an access gap...
 426         // We keep FILE_SHARE_READ though, as this shouldn't hurt us any, and is
 427         // required for enumeration.
 428         hDrive = CreateFileA(PhyDrive,
 429             GENERIC_READ | (bWriteAccess ? GENERIC_WRITE : 0),
 430             FILE_SHARE_READ | (bWriteShare ? FILE_SHARE_WRITE : 0),
 431             NULL,
 432             OPEN_EXISTING,
 433             FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
 434             NULL);
 435
 436         LastError = GetLastError();
 437         Log("[%d] CreateFileA %s code:%u %p", i, PhyDrive, LastError, hDrive);
 438
 439         if (hDrive != INVALID_HANDLE_VALUE)
 440         {
 441             break;
 442         }
 443
 444         if ((LastError != ERROR_SHARING_VIOLATION) && (LastError != ERROR_ACCESS_DENIED))
 445         {
 446             break;
 447         }
 448
 449         if (i == 0)
 450         {
 451             Log("Waiting for access on %s [%s]...", PhyDrive, DevPath);
 452         }
 453         else if (!bWriteShare && (i > DRIVE_ACCESS_RETRIES / 3))
 454         {
 455             // If we can't seem to get a hold of the drive for some time, try to enable FILE_SHARE_WRITE...
 456             Log("Warning: Could not obtain exclusive rights. Retrying with write sharing enabled...");
 457             bWriteShare = TRUE;
 458
 459             // Try to report the process that is locking the drive
 460             // We also use bit 6 as a flag to indicate that SearchProcess was called.
 461             //access_mask = SearchProcess(DevPath, SEARCH_PROCESS_TIMEOUT, TRUE, TRUE, FALSE) | 0x40;
 462
 463         }
 464         Sleep(DRIVE_ACCESS_TIMEOUT / DRIVE_ACCESS_RETRIES);
 465     }
 466
 467     if (hDrive == INVALID_HANDLE_VALUE)
 468     {
 469         Log("Could not open %s %u", PhyDrive, LASTERR);
 470         goto End;
 471     }
 472
 473     if (bWriteAccess)
 474     {
 475         Log("Opened %s for %s write access", PhyDrive, bWriteShare ? "shared" : "exclusive");
 476     }
 477
 478     if (bLockDrive)
 479     {
 480         if (DeviceIoControl(hDrive, FSCTL_ALLOW_EXTENDED_DASD_IO, NULL, 0, NULL, 0, &dwSize, NULL))
 481         {
 482             Log("I/O boundary checks disabled");
 483         }
 484
 485         EndTime = GetTickCount64() + DRIVE_ACCESS_TIMEOUT;
 486
 487         do {
 488             if (DeviceIoControl(hDrive, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL))
 489             {
 490                 Log("FSCTL_LOCK_VOLUME success");
 491                 goto End;
 492             }
 493             Sleep(DRIVE_ACCESS_TIMEOUT / DRIVE_ACCESS_RETRIES);
 494         } while (GetTickCount64() < EndTime);
 495
 496         // If we reached this section, either we didn't manage to get a lock or the user cancelled
 497         Log("Could not lock access to %s %u", PhyDrive, LASTERR);
 498
 499         // See if we can report the processes are accessing the drive
 500         //if (!IS_ERROR(FormatStatus) && (access_mask == 0))
 501         //    access_mask = SearchProcess(DevPath, SEARCH_PROCESS_TIMEOUT, TRUE, TRUE, FALSE);
 502         // Try to continue if the only access rights we saw were for read-only
 503         //if ((access_mask & 0x07) != 0x01)
 504         //    safe_closehandle(hDrive);
 505
 506         CHECK_CLOSE_HANDLE(hDrive);
 507     }
 508
 509 End:
 510
 511     if (hDrive == INVALID_HANDLE_VALUE)
 512     {
 513         Log("Can get handle of %s, maybe some process control it.", DevPath);
 514     }
 515
 516     return hDrive;
 517 }
 518
 519 int GetPhyDriveByLogicalDrive(int DriveLetter)
 520 {
 521     BOOL Ret;
 522     DWORD dwSize;
 523     HANDLE Handle;
 524     VOLUME_DISK_EXTENTS DiskExtents;
 525     CHAR PhyPath[128];
 526
 527     safe_sprintf(PhyPath, "\\\\.\\%C:", (CHAR)DriveLetter);
 528
 529     Handle = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 530     if (Handle == INVALID_HANDLE_VALUE)
 531     {
 532         Log("Could not open the disk<%s>, error:%u", PhyPath, LASTERR);
 533         return -1;
 534     }
 535
 536     Ret = DeviceIoControl(Handle,
 537         IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 538         NULL,
 539         0,
 540         &DiskExtents,
 541         (DWORD)(sizeof(DiskExtents)),
 542         (LPDWORD)&dwSize,
 543         NULL);
 544
 545     if (!Ret || DiskExtents.NumberOfDiskExtents == 0)
 546     {
 547         Log("DeviceIoControl IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS failed %s, error:%u", PhyPath, LASTERR);
 548         CHECK_CLOSE_HANDLE(Handle);
 549         return -1;
 550     }
 551     CHECK_CLOSE_HANDLE(Handle);
 552
 553     Log("LogicalDrive:%s PhyDrive:%d Offset:%llu ExtentLength:%llu",
 554         PhyPath,
 555         DiskExtents.Extents[0].DiskNumber,
 556         DiskExtents.Extents[0].StartingOffset.QuadPart,
 557         DiskExtents.Extents[0].ExtentLength.QuadPart
 558         );
 559
 560     return (int)DiskExtents.Extents[0].DiskNumber;
 561 }
 562
 563 int GetAllPhysicalDriveInfo(PHY_DRIVE_INFO *pDriveList, DWORD *pDriveCount)
 564 {
 565     int i;
 566     int Count;
 567     int id;
 568     int Letter = 'A';
 569     BOOL  bRet;
 570     DWORD dwBytes;
 571     DWORD DriveCount = 0;
 572     HANDLE Handle = INVALID_HANDLE_VALUE;
 573     CHAR PhyDrive[128];
 574     PHY_DRIVE_INFO *CurDrive = pDriveList;
 575     GET_LENGTH_INFORMATION LengthInfo;
 576     STORAGE_PROPERTY_QUERY Query;
 577     STORAGE_DESCRIPTOR_HEADER DevDescHeader;
 578     STORAGE_DEVICE_DESCRIPTOR *pDevDesc;
 579     int PhyDriveId[VENTOY_MAX_PHY_DRIVE];
 580
 581     Count = GetPhysicalDriveCount();
 582
 583     for (i = 0; i < Count && i < VENTOY_MAX_PHY_DRIVE; i++)
 584     {
 585         PhyDriveId[i] = i;
 586     }
 587
 588     dwBytes = GetLogicalDrives();
 589     Log("Logical Drives: 0x%x", dwBytes);
 590     while (dwBytes)
 591     {
 592         if (dwBytes & 0x01)
 593         {
 594             id = GetPhyDriveByLogicalDrive(Letter);
 595             Log("%C --> %d", Letter, id);
 596             if (id >= 0)
 597             {
 598                 for (i = 0; i < Count; i++)
 599                 {
 600                     if (PhyDriveId[i] == id)
 601                     {
 602                         break;
 603                     }
 604                 }
 605
 606                 if (i >= Count)
 607                 {
 608                     Log("Add phy%d to list", i);
 609                     PhyDriveId[Count] = id;
 610                     Count++;
 611                 }
 612             }
 613         }
 614
 615         Letter++;
 616         dwBytes >>= 1;
 617     }
 618
 619     for (i = 0; i < Count && DriveCount < VENTOY_MAX_PHY_DRIVE; i++)
 620     {
 621         CHECK_CLOSE_HANDLE(Handle);
 622
 623         safe_sprintf(PhyDrive, "\\\\.\\PhysicalDrive%d", PhyDriveId[i]);
 624         Handle = CreateFileA(PhyDrive, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
 625         Log("Create file Handle:%p %s status:%u", Handle, PhyDrive, LASTERR);
 626
 627         if (Handle == INVALID_HANDLE_VALUE)
 628         {
 629             continue;
 630         }
 631
 632         bRet = DeviceIoControl(Handle,
 633                                IOCTL_DISK_GET_LENGTH_INFO, NULL,
 634                                0,
 635                                &LengthInfo,
 636                                sizeof(LengthInfo),
 637                                &dwBytes,
 638                                NULL);
 639         if (!bRet)
 640         {
 641             Log("DeviceIoControl IOCTL_DISK_GET_LENGTH_INFO failed error:%u", LASTERR);
 642             continue;
 643         }
 644
 645         Log("PHYSICALDRIVE%d size %llu bytes", i, (ULONGLONG)LengthInfo.Length.QuadPart);
 646
 647         Query.PropertyId = StorageDeviceProperty;
 648         Query.QueryType = PropertyStandardQuery;
 649
 650         bRet = DeviceIoControl(Handle,
 651                                IOCTL_STORAGE_QUERY_PROPERTY,
 652                                &Query,
 653                                sizeof(Query),
 654                                &DevDescHeader,
 655                                sizeof(STORAGE_DESCRIPTOR_HEADER),
 656                                &dwBytes,
 657                                NULL);
 658         if (!bRet)
 659         {
 660             Log("DeviceIoControl1 error:%u dwBytes:%u", LASTERR, dwBytes);
 661             continue;
 662         }
 663
 664         if (DevDescHeader.Size < sizeof(STORAGE_DEVICE_DESCRIPTOR))
 665         {
 666             Log("Invalid DevDescHeader.Size:%u", DevDescHeader.Size);
 667             continue;
 668         }
 669
 670         pDevDesc = (STORAGE_DEVICE_DESCRIPTOR *)malloc(DevDescHeader.Size);
 671         if (!pDevDesc)
 672         {
 673             Log("failed to malloc error:%u len:%u", LASTERR, DevDescHeader.Size);
 674             continue;
 675         }
 676
 677         bRet = DeviceIoControl(Handle,
 678                                IOCTL_STORAGE_QUERY_PROPERTY,
 679                                &Query,
 680                                sizeof(Query),
 681                                pDevDesc,
 682                                DevDescHeader.Size,
 683                                &dwBytes,
 684                                NULL);
 685         if (!bRet)
 686         {
 687             Log("DeviceIoControl2 error:%u dwBytes:%u", LASTERR, dwBytes);
 688             free(pDevDesc);
 689             continue;
 690         }
 691
 692         CurDrive->PhyDrive = i;
 693         CurDrive->SizeInBytes = LengthInfo.Length.QuadPart;
 694         CurDrive->DeviceType = pDevDesc->DeviceType;
 695         CurDrive->RemovableMedia = pDevDesc->RemovableMedia;
 696         CurDrive->BusType = pDevDesc->BusType;
 697
 698         if (pDevDesc->VendorIdOffset)
 699         {
 700             safe_strcpy(CurDrive->VendorId, (char *)pDevDesc + pDevDesc->VendorIdOffset);
 701             TrimString(CurDrive->VendorId);
 702         }
 703
 704         if (pDevDesc->ProductIdOffset)
 705         {
 706             safe_strcpy(CurDrive->ProductId, (char *)pDevDesc + pDevDesc->ProductIdOffset);
 707             TrimString(CurDrive->ProductId);
 708         }
 709
 710         if (pDevDesc->ProductRevisionOffset)
 711         {
 712             safe_strcpy(CurDrive->ProductRev, (char *)pDevDesc + pDevDesc->ProductRevisionOffset);
 713             TrimString(CurDrive->ProductRev);
 714         }
 715
 716         if (pDevDesc->SerialNumberOffset)
 717         {
 718             safe_strcpy(CurDrive->SerialNumber, (char *)pDevDesc + pDevDesc->SerialNumberOffset);
 719             TrimString(CurDrive->SerialNumber);
 720         }
 721
 722         CurDrive++;
 723         DriveCount++;
 724
 725         free(pDevDesc);
 726
 727         CHECK_CLOSE_HANDLE(Handle);
 728     }
 729
 730     for (i = 0, CurDrive = pDriveList; i < (int)DriveCount; i++, CurDrive++)
 731     {
 732         Log("PhyDrv:%d BusType:%-4s Removable:%u Size:%dGB(%llu) Name:%s %s",
 733             CurDrive->PhyDrive, GetBusTypeString(CurDrive->BusType), CurDrive->RemovableMedia,
 734             GetHumanReadableGBSize(CurDrive->SizeInBytes), CurDrive->SizeInBytes,
 735             CurDrive->VendorId, CurDrive->ProductId);
 736     }
 737
 738     *pDriveCount = DriveCount;
 739
 740     return 0;
 741 }
 742
 743
 744 static HANDLE g_FatPhyDrive;
 745 static UINT64 g_Part2StartSec;
 746 static int GetVentoyVersionFromFatFile(CHAR *VerBuf, size_t BufLen)
 747 {
 748     int rc = 1;
 749     int size = 0;
 750     char *buf = NULL;
 751     void *flfile = NULL;
 752
 753     flfile = fl_fopen("/grub/grub.cfg", "rb");
 754     if (flfile)
 755     {
 756         fl_fseek(flfile, 0, SEEK_END);
 757         size = (int)fl_ftell(flfile);
 758
 759         fl_fseek(flfile, 0, SEEK_SET);
 760
 761         buf = (char *)malloc(size + 1);
 762         if (buf)
 763         {
 764             fl_fread(buf, 1, size, flfile);
 765             buf[size] = 0;
 766
 767             rc = 0;
 768             sprintf_s(VerBuf, BufLen, "%s", ParseVentoyVersionFromString(buf));
 769             free(buf);
 770         }
 771
 772         fl_fclose(flfile);
 773     }
 774
 775     return rc;
 776 }
 777
 778 static int VentoyFatDiskRead(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 779 {
 780     DWORD dwSize;
 781     BOOL bRet;
 782     DWORD ReadSize;
 783     LARGE_INTEGER liCurrentPosition;
 784
 785     liCurrentPosition.QuadPart = Sector + g_Part2StartSec;
 786     liCurrentPosition.QuadPart *= 512;
 787     SetFilePointerEx(g_FatPhyDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 788
 789     ReadSize = (DWORD)(SectorCount * 512);
 790
 791     bRet = ReadFile(g_FatPhyDrive, Buffer, ReadSize, &dwSize, NULL);
 792     if (bRet == FALSE || dwSize != ReadSize)
 793     {
 794         Log("ReadFile error bRet:%u WriteSize:%u dwSize:%u ErrCode:%u\n", bRet, ReadSize, dwSize, LASTERR);
 795     }
 796
 797     return 1;
 798 }
 799
 800
 801 int GetVentoyVerInPhyDrive(const PHY_DRIVE_INFO *pDriveInfo, CHAR *VerBuf, size_t BufLen)
 802 {
 803     int rc = 0;
 804     HANDLE hDrive;
 805
 806     hDrive = GetPhysicalHandle(pDriveInfo->PhyDrive, FALSE, FALSE, FALSE);
 807     if (hDrive == INVALID_HANDLE_VALUE)
 808     {
 809         return 1;
 810     }
 811
 812     g_FatPhyDrive = hDrive;
 813     g_Part2StartSec = (pDriveInfo->SizeInBytes - VENTOY_EFI_PART_SIZE) / 512;
 814
 815     Log("Parse FAT fs...");
 816
 817     fl_init();
 818
 819     if (0 == fl_attach_media(VentoyFatDiskRead, NULL))
 820     {
 821         rc = GetVentoyVersionFromFatFile(VerBuf, BufLen);
 822     }
 823     else
 824     {
 825         rc = 1;
 826     }
 827
 828     fl_shutdown();
 829
 830     CHECK_CLOSE_HANDLE(hDrive);
 831
 832     return rc;
 833 }
 834
 835
 836
 837
 838
 839 static unsigned int g_disk_unxz_len = 0;
 840 static BYTE *g_part_img_pos = NULL;
 841 static BYTE *g_part_img_buf[VENTOY_EFI_PART_SIZE / SIZE_1MB];
 842
 843
 844 static int VentoyFatMemRead(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 845 {
 846         uint32 i;
 847         uint32 offset;
 848         BYTE *MbBuf = NULL;
 849
 850         for (i = 0; i < SectorCount; i++)
 851         {
 852                 offset = (Sector + i) * 512;
 853
 854                 if (g_part_img_buf[1] == NULL)
 855                 {
 856                         MbBuf = g_part_img_buf[0] + offset;
 857                         memcpy(Buffer + i * 512, MbBuf, 512);
 858                 }
 859                 else
 860                 {
 861                         MbBuf = g_part_img_buf[offset / SIZE_1MB];
 862                         memcpy(Buffer + i * 512, MbBuf + (offset % SIZE_1MB), 512);
 863                 }
 864         }
 865
 866         return 1;
 867 }
 868
 869
 870 static int VentoyFatMemWrite(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 871 {
 872         uint32 i;
 873         uint32 offset;
 874         BYTE *MbBuf = NULL;
 875
 876         for (i = 0; i < SectorCount; i++)
 877         {
 878                 offset = (Sector + i) * 512;
 879
 880                 if (g_part_img_buf[1] == NULL)
 881                 {
 882                         MbBuf = g_part_img_buf[0] + offset;
 883                         memcpy(MbBuf, Buffer + i * 512, 512);
 884                 }
 885                 else
 886                 {
 887                         MbBuf = g_part_img_buf[offset / SIZE_1MB];
 888                         memcpy(MbBuf + (offset % SIZE_1MB), Buffer + i * 512, 512);
 889                 }
 890         }
 891
 892         return 1;
 893 }
 894
 895 int VentoyProcSecureBoot(BOOL SecureBoot)
 896 {
 897         int rc = 0;
 898         int size;
 899         char *filebuf = NULL;
 900         void *file = NULL;
 901
 902         Log("VentoyProcSecureBoot %d ...", SecureBoot);
 903
 904         if (SecureBoot)
 905         {
 906                 Log("Secure boot is enabled ...");
 907                 return 0;
 908         }
 909
 910         fl_init();
 911
 912         if (0 == fl_attach_media(VentoyFatMemRead, VentoyFatMemWrite))
 913         {
 914                 file = fl_fopen("/EFI/BOOT/grubx64_real.efi", "rb");
 915                 Log("Open ventoy efi file %p ", file);
 916                 if (file)
 917                 {
 918                         fl_fseek(file, 0, SEEK_END);
 919                         size = (int)fl_ftell(file);
 920                         fl_fseek(file, 0, SEEK_SET);
 921
 922                         Log("ventoy efi file size %d ...", size);
 923
 924                         filebuf = (char *)malloc(size);
 925                         if (filebuf)
 926                         {
 927                                 fl_fread(filebuf, 1, size, file);
 928                         }
 929
 930                         fl_fclose(file);
 931
 932                         Log("Now delete all efi files ...");
 933                         fl_remove("/EFI/BOOT/BOOTX64.EFI");
 934                         fl_remove("/EFI/BOOT/grubx64.efi");
 935                         fl_remove("/EFI/BOOT/grubx64_real.efi");
 936                         fl_remove("/EFI/BOOT/MokManager.efi");
 937
 938                         file = fl_fopen("/EFI/BOOT/BOOTX64.EFI", "wb");
 939                         Log("Open bootx64 efi file %p ", file);
 940                         if (file)
 941                         {
 942                                 if (filebuf)
 943                                 {
 944                                         fl_fwrite(filebuf, 1, size, file);
 945                                 }
 946
 947                                 fl_fflush(file);
 948                                 fl_fclose(file);
 949                         }
 950
 951                         if (filebuf)
 952                         {
 953                                 free(filebuf);
 954                         }
 955                 }
 956         }
 957         else
 958         {
 959                 rc = 1;
 960         }
 961
 962         fl_shutdown();
 963
 964         return rc;
 965 }
 966
 967
 968
 969 static int disk_xz_flush(void *src, unsigned int size)
 970 {
 971     unsigned int i;
 972     BYTE *buf = (BYTE *)src;
 973
 974     for (i = 0; i < size; i++)
 975     {
 976         *g_part_img_pos = *buf++;
 977
 978         g_disk_unxz_len++;
 979         if ((g_disk_unxz_len % SIZE_1MB) == 0)
 980         {
 981             g_part_img_pos = g_part_img_buf[g_disk_unxz_len / SIZE_1MB];
 982         }
 983         else
 984         {
 985             g_part_img_pos++;
 986         }
 987     }
 988
 989     return (int)size;
 990 }
 991
 992 static void unxz_error(char *x)
 993 {
 994     Log("%s", x);
 995 }
 996
 997 static BOOL TryWritePart2(HANDLE hDrive, UINT64 StartSectorId)
 998 {
 999     BOOL bRet;
1000     DWORD TrySize = 16 * 1024;
1001     DWORD dwSize;
1002     BYTE *Buffer = NULL;
1003     unsigned char *data = NULL;
1004     LARGE_INTEGER liCurrentPosition;
1005
1006     liCurrentPosition.QuadPart = StartSectorId * 512;
1007     SetFilePointerEx(hDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
1008
1009     Buffer = malloc(TrySize);
1010
1011     bRet = WriteFile(hDrive, Buffer, TrySize, &dwSize, NULL);
1012
1013     free(Buffer);
1014
1015     Log("Try write part2 bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
1016
1017     if (bRet && dwSize == TrySize)
1018     {
1019         return TRUE;
1020     }
1021
1022     return FALSE;
1023 }
1024
1025 static int FormatPart2Fat(HANDLE hDrive, UINT64 StartSectorId)
1026 {
1027     int i;
1028     int rc = 0;
1029     int len = 0;
1030     int writelen = 0;
1031     int partwrite = 0;
1032     DWORD dwSize = 0;
1033     BOOL bRet;
1034     unsigned char *data = NULL;
1035     LARGE_INTEGER liCurrentPosition;
1036
1037     Log("FormatPart2Fat ...");
1038
1039     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&data, &len);
1040     if (rc)
1041     {
1042         Log("Failed to read img file %p %u", data, len);
1043         return 1;
1044     }
1045
1046     liCurrentPosition.QuadPart = StartSectorId * 512;
1047     SetFilePointerEx(hDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
1048
1049     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
1050
1051     g_part_img_buf[0] = (BYTE *)malloc(VENTOY_EFI_PART_SIZE);
1052     if (g_part_img_buf[0])
1053     {
1054         Log("Malloc whole img buffer success, now decompress ...");
1055         unxz(data, len, NULL, NULL, g_part_img_buf[0], &writelen, unxz_error);
1056
1057         if (len == writelen)
1058         {
1059             Log("decompress finished success");
1060
1061                         VentoyProcSecureBoot(g_SecureBoot);
1062
1063             for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1064             {
1065                 dwSize = 0;
1066                 bRet = WriteFile(hDrive, g_part_img_buf[0] + i * SIZE_1MB, SIZE_1MB, &dwSize, NULL);
1067                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
1068
1069                 if (!bRet)
1070                 {
1071                     rc = 1;
1072                     goto End;
1073                 }
1074
1075                 PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START + i);
1076             }
1077         }
1078         else
1079         {
1080             rc = 1;
1081             Log("decompress finished failed");
1082             goto End;
1083         }
1084     }
1085     else
1086     {
1087         Log("Failed to malloc whole img size %u, now split it", VENTOY_EFI_PART_SIZE);
1088
1089         partwrite = 1;
1090         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1091         {
1092             g_part_img_buf[i] = (BYTE *)malloc(SIZE_1MB);
1093             if (g_part_img_buf[i] == NULL)
1094             {
1095                 rc = 1;
1096                 goto End;
1097             }
1098         }
1099
1100         Log("Malloc part img buffer success, now decompress ...");
1101
1102         g_part_img_pos = g_part_img_buf[0];
1103
1104         unxz(data, len, NULL, disk_xz_flush, NULL, NULL, unxz_error);
1105
1106         if (g_disk_unxz_len == VENTOY_EFI_PART_SIZE)
1107         {
1108             Log("decompress finished success");
1109
1110                         VentoyProcSecureBoot(g_SecureBoot);
1111
1112             for (int i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1113             {
1114                 dwSize = 0;
1115                 bRet = WriteFile(hDrive, g_part_img_buf[i], SIZE_1MB, &dwSize, NULL);
1116                 Log("Write part data bRet:%u dwSize:%u code:%u", bRet, dwSize, LASTERR);
1117
1118                 if (!bRet)
1119                 {
1120                     rc = 1;
1121                     goto End;
1122                 }
1123
1124                 PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START + i);
1125             }
1126         }
1127         else
1128         {
1129             rc = 1;
1130             Log("decompress finished failed");
1131             goto End;
1132         }
1133     }
1134
1135 End:
1136
1137     if (data) free(data);
1138
1139     if (partwrite)
1140     {
1141         for (i = 0; i < VENTOY_EFI_PART_SIZE / SIZE_1MB; i++)
1142         {
1143             if (g_part_img_buf[i]) free(g_part_img_buf[i]);
1144         }
1145     }
1146     else
1147     {
1148         if (g_part_img_buf[0]) free(g_part_img_buf[0]);
1149     }
1150
1151     return rc;
1152 }
1153
1154 static int WriteGrubStage1ToPhyDrive(HANDLE hDrive)
1155 {
1156     int Len = 0;
1157     int readLen = 0;
1158     BOOL bRet;
1159     DWORD dwSize;
1160     BYTE *ImgBuf = NULL;
1161     BYTE *RawBuf = NULL;
1162
1163     Log("WriteGrubStage1ToPhyDrive ...");
1164
1165     RawBuf = (BYTE *)malloc(SIZE_1MB);
1166     if (!RawBuf)
1167     {
1168         return 1;
1169     }
1170
1171     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1172     {
1173         Log("Failed to read stage1 img");
1174         free(RawBuf);
1175         return 1;
1176     }
1177
1178     unxz(ImgBuf, Len, NULL, NULL, RawBuf, &readLen, unxz_error);
1179
1180     SetFilePointer(hDrive, 512, NULL, FILE_BEGIN);
1181
1182     bRet = WriteFile(hDrive, RawBuf, SIZE_1MB - 512, &dwSize, NULL);
1183     Log("WriteFile Ret:%u dwSize:%u ErrCode:%u", bRet, dwSize, GetLastError());
1184
1185     free(RawBuf);
1186     free(ImgBuf);
1187     return 0;
1188 }
1189
1190
1191
1192 static int FormatPart1exFAT(UINT64 DiskSizeBytes)
1193 {
1194     MKFS_PARM Option;
1195     FRESULT Ret;
1196     FATFS fs;
1197
1198     Option.fmt = FM_EXFAT;
1199     Option.n_fat = 1;
1200     Option.align = 8;
1201     Option.n_root = 1;
1202
1203     // < 32GB select 32KB as cluster size
1204     // > 32GB select 128KB as cluster size
1205     if (DiskSizeBytes / 1024 / 1024 / 1024 <= 32)
1206     {
1207         Option.au_size = 32768;
1208     }
1209     else
1210     {
1211         Option.au_size = 131072;
1212     }
1213
1214     Log("Formatting Part1 exFAT ...");
1215
1216     Ret = f_mkfs(TEXT("0:"), &Option, 0, 8 * 1024 * 1024);
1217
1218     if (FR_OK == Ret)
1219     {
1220         Log("Formatting Part1 exFAT success");
1221
1222         Ret = f_mount(&fs, TEXT("0:"), 1);
1223         Log("mount part %d", Ret);
1224
1225         if (FR_OK == Ret)
1226         {
1227             Ret = f_setlabel(TEXT("Ventoy"));
1228             Log("f_setlabel %d", Ret);
1229
1230             Ret = f_mount(0, TEXT("0:"), 1);
1231             Log("umount part %d", Ret);
1232         }
1233
1234         return 0;
1235     }
1236     else
1237     {
1238         Log("Formatting Part1 exFAT failed");
1239         return 1;
1240     }
1241 }
1242
1243
1244 int InstallVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive)
1245 {
1246     int i;
1247     int rc = 0;
1248     int state = 0;
1249     HANDLE hDrive;
1250     DWORD dwSize;
1251     BOOL bRet;
1252     CHAR MountDrive;
1253     CHAR DriveName[] = "?:\\";
1254     CHAR DriveLetters[MAX_PATH] = { 0 };
1255     MBR_HEAD MBR;
1256
1257     Log("InstallVentoy2PhyDrive PhyDrive%d <<%s %s %dGB>>",
1258         pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1259         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1260
1261     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1262
1263     VentoyFillMBR(pPhyDrive->SizeInBytes, &MBR);
1264
1265     Log("Lock disk for clean ............................. ");
1266
1267     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1268     if (hDrive == INVALID_HANDLE_VALUE)
1269     {
1270         Log("Failed to open physical disk");
1271         return 1;
1272     }
1273
1274     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1275
1276     if (DriveLetters[0] == 0)
1277     {
1278         Log("No drive letter was assigned...");
1279         DriveName[0] = GetFirstUnusedDriveLetter();
1280         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1281     }
1282     else
1283     {
1284         // Unmount all mounted volumes that belong to this drive
1285         // Do it in reverse so that we always end on the first volume letter
1286         for (i = (int)strlen(DriveLetters); i > 0; i--)
1287         {
1288             DriveName[0] = DriveLetters[i - 1];
1289             bRet = DeleteVolumeMountPointA(DriveName);
1290             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1291         }
1292     }
1293
1294     MountDrive = DriveName[0];
1295     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1296
1297     // It kind of blows, but we have to relinquish access to the physical drive
1298     // for VDS to be able to delete the partitions that reside on it...
1299     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1300     CHECK_CLOSE_HANDLE(hDrive);
1301
1302     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1303
1304     if (!DeletePartitions(pPhyDrive->PhyDrive, FALSE))
1305     {
1306         Log("Notice: Could not delete partitions: %u", GetLastError());
1307     }
1308
1309     Log("Deleting all partitions ......................... OK");
1310
1311     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1312
1313     Log("Lock disk for write ............................. ");
1314     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1315     if (hDrive == INVALID_HANDLE_VALUE)
1316     {
1317         Log("Failed to GetPhysicalHandle for write.");
1318         rc = 1;
1319         goto End;
1320     }
1321
1322     //Refresh Drive Layout
1323     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1324
1325     disk_io_set_param(hDrive, MBR.PartTbl[0].StartSectorId + MBR.PartTbl[0].SectorCount);
1326
1327     PROGRESS_BAR_SET_POS(PT_FORMAT_PART1);
1328
1329     Log("Formatting part1 exFAT ...");
1330     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1331     {
1332         rc = 1;
1333         goto End;
1334     }
1335
1336     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1337     Log("Writing part2 FAT img ...");
1338     if (0 != FormatPart2Fat(hDrive, MBR.PartTbl[1].StartSectorId))
1339     {
1340         rc = 1;
1341         goto End;
1342     }
1343
1344     PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
1345     Log("Writting Boot Image ............................. ");
1346     if (WriteGrubStage1ToPhyDrive(hDrive) != 0)
1347     {
1348         rc = 1;
1349         goto End;
1350     }
1351
1352
1353     PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
1354     Log("Writting Partition Table ........................ ");
1355     SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1356     if (!WriteFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL))
1357     {
1358         rc = 1;
1359         Log("Write MBR Failed, dwSize:%u ErrCode:%u", dwSize, GetLastError());
1360         goto End;
1361     }
1362
1363     Log("Write MBR OK ...");
1364
1365     //Refresh Drive Layout
1366     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1367
1368 End:
1369     CHECK_CLOSE_HANDLE(hDrive);
1370
1371     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1372     Log("Mounting Ventoy Partition ....................... ");
1373     Sleep(1000);
1374
1375     state = 0;
1376     memset(DriveLetters, 0, sizeof(DriveLetters));
1377     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1378     Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1379
1380     for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1381     {
1382         DriveName[0] = DriveLetters[i];
1383         if (IsVentoyLogicalDrive(DriveName[0]))
1384         {
1385             Log("%s is ventoy part2, delete mountpoint", DriveName);
1386             DeleteVolumeMountPointA(DriveName);
1387         }
1388         else
1389         {
1390             Log("%s is ventoy part1, already mounted", DriveName);
1391             state = 1;
1392         }
1393     }
1394
1395     if (state != 1)
1396     {
1397         Log("need to mount ventoy part1...");
1398         if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, MBR.PartTbl[0].StartSectorId, DriveLetters, sizeof(DriveLetters), FALSE))
1399         {
1400             DriveName[0] = MountDrive;
1401             bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1402             Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1403         }
1404         else
1405         {
1406             Log("Failed to find ventoy volume");
1407         }
1408     }
1409
1410     Log("OK\n");
1411
1412     return rc;
1413 }
1414
1415 int UpdateVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive)
1416 {
1417     int i;
1418     int rc = 0;
1419     BOOL ForceMBR = FALSE;
1420     HANDLE hVolume;
1421     HANDLE hDrive;
1422     DWORD Status;
1423     DWORD dwSize;
1424     BOOL bRet;
1425     CHAR DriveName[] = "?:\\";
1426     CHAR DriveLetters[MAX_PATH] = { 0 };
1427     UINT32 StartSector;
1428     MBR_HEAD BootImg;
1429     MBR_HEAD MBR;
1430
1431     StartSector = (UINT32)(pPhyDrive->SizeInBytes / 512 - VENTOY_EFI_PART_SIZE / 512);
1432
1433     Log("UpdateVentoy2PhyDrive PhyDrive%d <<%s %s %dGB>>",
1434         pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1435         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1436
1437     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1438
1439     Log("Lock disk for umount ............................ ");
1440
1441     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1442     if (hDrive == INVALID_HANDLE_VALUE)
1443     {
1444         Log("Failed to open physical disk");
1445         return 1;
1446     }
1447
1448     // Read MBR
1449     ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
1450
1451     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1452
1453     if (DriveLetters[0] == 0)
1454     {
1455         Log("No drive letter was assigned...");
1456     }
1457     else
1458     {
1459         // Unmount all mounted volumes that belong to this drive
1460         // Do it in reverse so that we always end on the first volume letter
1461         for (i = (int)strlen(DriveLetters); i > 0; i--)
1462         {
1463             DriveName[0] = DriveLetters[i - 1];
1464             if (IsVentoyLogicalDrive(DriveName[0]))
1465             {
1466                 Log("%s is ventoy logical drive", DriveName);
1467                 bRet = DeleteVolumeMountPointA(DriveName);
1468                 Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, LASTERR);
1469                 break;
1470             }
1471         }
1472     }
1473
1474     // It kind of blows, but we have to relinquish access to the physical drive
1475     // for VDS to be able to delete the partitions that reside on it...
1476     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1477     CHECK_CLOSE_HANDLE(hDrive);
1478
1479     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1480
1481     Log("Lock disk for update ............................ ");
1482     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1483     if (hDrive == INVALID_HANDLE_VALUE)
1484     {
1485         Log("Failed to GetPhysicalHandle for write.");
1486         rc = 1;
1487         goto End;
1488     }
1489
1490     PROGRESS_BAR_SET_POS(PT_LOCK_VOLUME);
1491
1492     Log("Lock volume for update .......................... ");
1493     hVolume = INVALID_HANDLE_VALUE;
1494     Status = GetVentoyVolumeName(pPhyDrive->PhyDrive, MBR.PartTbl[1].StartSectorId, DriveLetters, sizeof(DriveLetters), TRUE);
1495     if (ERROR_SUCCESS == Status)
1496     {
1497         Log("Now lock and dismount volume <%s>", DriveLetters);
1498         hVolume = CreateFileA(DriveLetters,
1499             GENERIC_READ | GENERIC_WRITE,
1500             FILE_SHARE_READ,
1501             NULL,
1502             OPEN_EXISTING,
1503             FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
1504             NULL);
1505
1506         if (hVolume == INVALID_HANDLE_VALUE)
1507         {
1508             Log("Failed to create file volume, errcode:%u", LASTERR);
1509             rc = 1;
1510             goto End;
1511         }
1512
1513         bRet = DeviceIoControl(hVolume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1514         Log("FSCTL_LOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
1515
1516         bRet = DeviceIoControl(hVolume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1517         Log("FSCTL_DISMOUNT_VOLUME bRet:%u code:%u", bRet, LASTERR);
1518     }
1519     else if (ERROR_NOT_FOUND == Status)
1520     {
1521         Log("Volume not found, maybe not supported");
1522     }
1523     else
1524     {
1525         rc = 1;
1526         goto End;
1527     }
1528
1529
1530     if (!TryWritePart2(hDrive, StartSector))
1531     {
1532         ForceMBR = TRUE;
1533         Log("Try write failed, now delete partition 2...");
1534
1535         CHECK_CLOSE_HANDLE(hDrive);
1536
1537         Log("Now delete partition 2...");
1538         DeletePartitions(pPhyDrive->PhyDrive, TRUE);
1539
1540         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1541         if (hDrive == INVALID_HANDLE_VALUE)
1542         {
1543             Log("Failed to GetPhysicalHandle for write.");
1544             rc = 1;
1545             goto End;
1546         }
1547     }
1548
1549
1550     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1551
1552     Log("Write Ventoy to disk ............................ ");
1553     if (0 != FormatPart2Fat(hDrive, StartSector))
1554     {
1555         rc = 1;
1556         goto End;
1557     }
1558
1559     if (hVolume != INVALID_HANDLE_VALUE)
1560     {
1561         bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1562         Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
1563         CHECK_CLOSE_HANDLE(hVolume);
1564     }
1565
1566     Log("Updating Boot Image ............................. ");
1567     if (WriteGrubStage1ToPhyDrive(hDrive) != 0)
1568     {
1569         rc = 1;
1570         goto End;
1571     }
1572
1573     // Boot Image
1574     VentoyGetLocalBootImg(&BootImg);
1575
1576     // Use Old UUID
1577     memcpy(BootImg.BootCode + 0x180, MBR.BootCode + 0x180, 16);
1578
1579     if (ForceMBR == FALSE && memcmp(BootImg.BootCode, MBR.BootCode, 440) == 0)
1580     {
1581         Log("Boot image has no difference, no need to write.");
1582     }
1583     else
1584     {
1585         Log("Boot image need to write %u.", ForceMBR);
1586
1587         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1588
1589         memcpy(MBR.BootCode, BootImg.BootCode, 440);
1590         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
1591         Log("Write Boot Image ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
1592     }
1593
1594     if (0x00 == MBR.PartTbl[0].Active && 0x80 == MBR.PartTbl[1].Active)
1595     {
1596         Log("Need to chage 1st partition active and 2nd partition inactive.");
1597
1598         MBR.PartTbl[0].Active = 0x80;
1599         MBR.PartTbl[1].Active = 0x00;
1600
1601         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1602         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
1603         Log("Write NEW MBR ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
1604     }
1605
1606     //Refresh Drive Layout
1607     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1608
1609 End:
1610     CHECK_CLOSE_HANDLE(hDrive);
1611
1612     return rc;
1613 }
1614