Ventoy2Disk/Ventoy2Disk/PhyDrive.c

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