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     Ret = f_mkfs(TEXT("0:"), &Option, 0, 8 * 1024 * 1024);
1131     if (FR_OK == Ret)
1132     {
1133         Log("Formatting Part1 exFAT success");
1134         return 0;
1135     }
1136     else
1137     {
1138         Log("Formatting Part1 exFAT failed");
1139         return 1;
1140     }
1141 }
1142
1143
1144
1145 int ClearVentoyFromPhyDrive(HWND hWnd, PHY_DRIVE_INFO *pPhyDrive, char *pDrvLetter)
1146 {
1147     int i;
1148     int rc = 0;
1149     int state = 0;
1150     HANDLE hDrive;
1151     DWORD dwSize;
1152     BOOL bRet;
1153     CHAR MountDrive;
1154     CHAR DriveName[] = "?:\\";
1155     CHAR DriveLetters[MAX_PATH] = { 0 };
1156     LARGE_INTEGER liCurrentPosition;
1157     char *pTmpBuf = NULL;
1158     MBR_HEAD MBR;
1159
1160     *pDrvLetter = 0;
1161
1162     Log("ClearVentoyFromPhyDrive PhyDrive%d <<%s %s %dGB>>",
1163         pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1164         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1165
1166     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1167
1168     Log("Lock disk for clean ............................. ");
1169
1170     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1171     if (hDrive == INVALID_HANDLE_VALUE)
1172     {
1173         Log("Failed to open physical disk");
1174         return 1;
1175     }
1176
1177     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1178
1179     if (DriveLetters[0] == 0)
1180     {
1181         Log("No drive letter was assigned...");
1182         DriveName[0] = GetFirstUnusedDriveLetter();
1183         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1184     }
1185     else
1186     {
1187         // Unmount all mounted volumes that belong to this drive
1188         // Do it in reverse so that we always end on the first volume letter
1189         for (i = (int)strlen(DriveLetters); i > 0; i--)
1190         {
1191             DriveName[0] = DriveLetters[i - 1];
1192             bRet = DeleteVolumeMountPointA(DriveName);
1193             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1194         }
1195     }
1196
1197     MountDrive = DriveName[0];
1198     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1199
1200     // It kind of blows, but we have to relinquish access to the physical drive
1201     // for VDS to be able to delete the partitions that reside on it...
1202     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1203     CHECK_CLOSE_HANDLE(hDrive);
1204
1205     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1206
1207     if (!VDS_DeleteAllPartitions(pPhyDrive->PhyDrive))
1208     {
1209         Log("Notice: Could not delete partitions: %u", GetLastError());
1210     }
1211
1212     Log("Deleting all partitions ......................... OK");
1213
1214     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1215
1216     Log("Lock disk for write ............................. ");
1217     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1218     if (hDrive == INVALID_HANDLE_VALUE)
1219     {
1220         Log("Failed to GetPhysicalHandle for write.");
1221         rc = 1;
1222         goto End;
1223     }
1224
1225     // clear first and last 2MB space
1226     pTmpBuf = malloc(SIZE_2MB);
1227     if (!pTmpBuf)
1228     {
1229         Log("Failed to alloc memory.");
1230         rc = 1;
1231         goto End;
1232     }
1233     memset(pTmpBuf, 0, SIZE_2MB);
1234
1235     SET_FILE_POS(512);
1236     bRet = WriteFile(hDrive, pTmpBuf, SIZE_2MB - 512, &dwSize, NULL);
1237     Log("Write fisrt 1MB ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1238     if (!bRet)
1239     {
1240         rc = 1;
1241         goto End;
1242     }
1243
1244     SET_FILE_POS(pPhyDrive->SizeInBytes - SIZE_2MB);
1245     bRet = WriteFile(hDrive, pTmpBuf, SIZE_2MB, &dwSize, NULL);
1246     Log("Write 2nd 1MB ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1247     if (!bRet)
1248     {
1249         rc = 1;
1250         goto End;
1251     }
1252
1253     SET_FILE_POS(0);
1254
1255     if (pPhyDrive->SizeInBytes > 2199023255552ULL)
1256     {
1257         VTOY_GPT_INFO *pGptInfo;
1258         VTOY_GPT_HDR BackupHead;
1259         LARGE_INTEGER liCurrentPosition;
1260
1261         pGptInfo = (VTOY_GPT_INFO *)pTmpBuf;
1262
1263         VentoyFillWholeGpt(pPhyDrive->SizeInBytes, pGptInfo);
1264
1265         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1266         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1267         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1268         {
1269             rc = 1;
1270             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1271             goto End;
1272         }
1273
1274         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1275         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1276         {
1277             rc = 1;
1278             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1279             goto End;
1280         }
1281
1282         SET_FILE_POS(0);
1283         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1284         {
1285             rc = 1;
1286             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1287             goto End;
1288         }
1289
1290         Log("Write GPT Info OK ...");
1291     }
1292     else
1293     {
1294         bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
1295         Log("Read MBR ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1296         if (!bRet)
1297         {
1298             rc = 1;
1299             goto End;
1300         }
1301
1302         //clear boot code and partition table (reserved disk signature)
1303         memset(MBR.BootCode, 0, 440);
1304         memset(MBR.PartTbl, 0, sizeof(MBR.PartTbl));
1305
1306         VentoyFillMBRLocation(pPhyDrive->SizeInBytes, 2048, (UINT32)(pPhyDrive->SizeInBytes / 512 - 2048), MBR.PartTbl);
1307
1308         MBR.PartTbl[0].Active = 0x00; // bootable
1309         MBR.PartTbl[0].FsFlag = 0x07; // exFAT/NTFS/HPFS
1310
1311         SET_FILE_POS(0);
1312         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
1313         Log("Write MBR ret:%d size:%u err:%d", bRet, dwSize, LASTERR);
1314         if (!bRet)
1315         {
1316             rc = 1;
1317             goto End;
1318         }
1319     }
1320
1321     Log("Clear Ventoy successfully finished");
1322
1323         //Refresh Drive Layout
1324         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1325
1326 End:
1327
1328     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1329
1330     if (pTmpBuf)
1331     {
1332         free(pTmpBuf);
1333     }
1334
1335     if (rc == 0)
1336     {
1337         Log("Mounting Ventoy Partition ....................... ");
1338         Sleep(1000);
1339
1340         state = 0;
1341         memset(DriveLetters, 0, sizeof(DriveLetters));
1342         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1343         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1344
1345         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1346         {
1347             DriveName[0] = DriveLetters[i];
1348             Log("%s is ventoy part1, already mounted", DriveName);
1349             state = 1;
1350         }
1351
1352         if (state != 1)
1353         {
1354             Log("need to mount ventoy part1...");
1355             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, 2048, DriveLetters, sizeof(DriveLetters), FALSE))
1356             {
1357                 DriveName[0] = MountDrive;
1358                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1359                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1360
1361                 *pDrvLetter = MountDrive;
1362             }
1363             else
1364             {
1365                 Log("Failed to find ventoy volume");
1366             }
1367         }
1368
1369         Log("OK\n");
1370     }
1371     else
1372     {
1373         FindProcessOccupyDisk(hDrive, pPhyDrive);
1374     }
1375
1376     CHECK_CLOSE_HANDLE(hDrive);
1377     return rc;
1378 }
1379
1380 int InstallVentoy2FileImage(PHY_DRIVE_INFO *pPhyDrive, int PartStyle)
1381 {
1382     int i;
1383     int rc = 1;
1384     int Len = 0;
1385     int dataLen = 0;
1386     UINT size = 0;
1387     UINT segnum = 0;
1388     UINT32 chksum = 0;
1389     UINT64 data_offset = 0;
1390     UINT64 Part2StartSector = 0;
1391     UINT64 Part1StartSector = 0;
1392     UINT64 Part1SectorCount = 0;
1393     UINT8 *pData = NULL;
1394     UINT8 *pBkGptPartTbl = NULL;
1395     BYTE *ImgBuf = NULL;
1396     MBR_HEAD *pMBR = NULL;
1397     VTSI_FOOTER *pImgFooter = NULL;
1398     VTSI_SEGMENT *pSegment = NULL;
1399     VTOY_GPT_INFO *pGptInfo = NULL;
1400     VTOY_GPT_HDR *pBkGptHdr = NULL;
1401     FILE *fp = NULL;
1402
1403     Log("InstallVentoy2FileImage %s PhyDrive%d <<%s %s %dGB>>",
1404         PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1405         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1406
1407     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1408
1409     size = SIZE_1MB + VENTOY_EFI_PART_SIZE + 33 * 512 + VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT) + sizeof(VTSI_FOOTER);
1410
1411     pData = (UINT8 *)malloc(size);
1412     if (!pData)
1413     {
1414         Log("malloc image buffer failed %d.", size);
1415         goto End;
1416     }
1417
1418     pImgFooter = (VTSI_FOOTER *)(pData + size - sizeof(VTSI_FOOTER));
1419     pSegment = (VTSI_SEGMENT *)((UINT8 *)pImgFooter - VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT));
1420     memset(pImgFooter, 0, sizeof(VTSI_FOOTER));
1421     memset(pSegment, 0, VTSI_IMG_MAX_SEG * sizeof(VTSI_SEGMENT));
1422
1423     PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START);
1424
1425     Log("Writing Boot Image ............................. ");
1426     if (ReadWholeFileToBuf(VENTOY_FILE_STG1_IMG, 0, (void **)&ImgBuf, &Len))
1427     {
1428         Log("Failed to read stage1 img");
1429         goto End;
1430     }
1431
1432     unxz(ImgBuf, Len, NULL, NULL, pData, &dataLen, unxz_error);
1433     SAFE_FREE(ImgBuf);
1434
1435     Log("decompress %s len:%d", VENTOY_FILE_STG1_IMG, dataLen);
1436
1437     if (PartStyle)
1438     {
1439         pData[500] = 35;//update blocklist
1440         memmove(pData + 34 * 512, pData, SIZE_1MB - 512 * 34);
1441         memset(pData, 0, 34 * 512);
1442
1443         pGptInfo = (VTOY_GPT_INFO *)pData;
1444         memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1445         VentoyFillGpt(pPhyDrive->SizeInBytes, pGptInfo);
1446
1447         pBkGptPartTbl = pData + SIZE_1MB + VENTOY_EFI_PART_SIZE;
1448         memset(pBkGptPartTbl, 0, 33 * 512);
1449
1450         memcpy(pBkGptPartTbl, pGptInfo->PartTbl, 32 * 512);
1451         pBkGptHdr = (VTOY_GPT_HDR *)(pBkGptPartTbl + 32 * 512);
1452         VentoyFillBackupGptHead(pGptInfo, pBkGptHdr);
1453
1454         Part1StartSector = pGptInfo->PartTbl[0].StartLBA;
1455         Part1SectorCount = pGptInfo->PartTbl[0].LastLBA - Part1StartSector + 1;
1456         Part2StartSector = pGptInfo->PartTbl[1].StartLBA;
1457
1458         Log("Write GPT Info OK ...");
1459     }
1460     else
1461     {
1462         memmove(pData + 512, pData, SIZE_1MB - 512);
1463         memset(pData, 0, 512);
1464
1465         pMBR = (MBR_HEAD *)pData;
1466         VentoyFillMBR(pPhyDrive->SizeInBytes, pMBR, PartStyle);
1467         Part1StartSector = pMBR->PartTbl[0].StartSectorId;
1468         Part1SectorCount = pMBR->PartTbl[0].SectorCount;
1469         Part2StartSector = pMBR->PartTbl[1].StartSectorId;
1470
1471         Log("Write MBR OK ...");
1472     }
1473
1474     Log("Writing EFI part Image ............................. ");
1475     rc = ReadWholeFileToBuf(VENTOY_FILE_DISK_IMG, 0, (void **)&ImgBuf, &Len);
1476     if (rc)
1477     {
1478         Log("Failed to read img file %p %u", ImgBuf, Len);
1479         goto End;
1480     }
1481
1482     PROGRESS_BAR_SET_POS(PT_WRITE_VENTOY_START + 28);
1483     memset(g_part_img_buf, 0, sizeof(g_part_img_buf));
1484     unxz(ImgBuf, Len, NULL, NULL, pData + SIZE_1MB, &dataLen, unxz_error);
1485     if (dataLen == Len)
1486     {
1487         Log("decompress finished success");
1488         g_part_img_buf[0] = pData + SIZE_1MB;
1489
1490         VentoyProcSecureBoot(g_SecureBoot);
1491     }
1492     else
1493     {
1494         Log("decompress finished failed");
1495         goto End;
1496     }
1497
1498     fopen_s(&fp, "VentoySparseImg.vtsi", "wb+");
1499     if (!fp)
1500     {
1501         Log("Failed to create Ventoy img file");
1502         goto End;
1503     }
1504
1505     Log("Writing stage1 data ............................. ");
1506
1507     fwrite(pData, 1, SIZE_1MB, fp);
1508
1509     pSegment[0].disk_start_sector = 0;
1510     pSegment[0].sector_num = SIZE_1MB / 512;
1511     pSegment[0].data_offset = data_offset;
1512     data_offset += pSegment[0].sector_num * 512;
1513
1514     disk_io_set_param(INVALID_HANDLE_VALUE, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1515     disk_io_set_imghook(fp, pSegment + 1, VTSI_IMG_MAX_SEG - 1, data_offset);
1516
1517     Log("Formatting part1 exFAT ...");
1518     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1519     {
1520         Log("FormatPart1exFAT failed.");
1521         disk_io_reset_imghook(&segnum, &data_offset);
1522         goto End;
1523     }
1524
1525     disk_io_reset_imghook(&segnum, &data_offset);
1526     segnum++;
1527
1528     Log("current segment number:%d dataoff:%ld", segnum, (long)data_offset);
1529
1530     //write data
1531     Log("Writing part2 data ............................. ");
1532     fwrite(pData + SIZE_1MB, 1, VENTOY_EFI_PART_SIZE, fp);
1533     pSegment[segnum].disk_start_sector = Part2StartSector;
1534     pSegment[segnum].sector_num = VENTOY_EFI_PART_SIZE / 512;
1535     pSegment[segnum].data_offset = data_offset;
1536     data_offset += pSegment[segnum].sector_num * 512;
1537     segnum++;
1538
1539     if (PartStyle)
1540     {
1541         Log("Writing backup gpt table ............................. ");
1542         fwrite(pBkGptPartTbl, 1, 33 * 512, fp);
1543         pSegment[segnum].disk_start_sector = pPhyDrive->SizeInBytes / 512 - 33;
1544         pSegment[segnum].sector_num = 33;
1545         pSegment[segnum].data_offset = data_offset;
1546         data_offset += pSegment[segnum].sector_num * 512;
1547         segnum++;
1548     }
1549
1550     Log("Writing segment metadata ............................. ");
1551
1552     for (i = 0; i < (int)segnum; i++)
1553     {
1554         Log("SEG[%d]:  PhySector:%llu SectorNum:%llu DataOffset:%llu(sector:%llu)", i, pSegment[i].disk_start_sector, pSegment[i].sector_num,
1555             pSegment[i].data_offset, pSegment[i].data_offset / 512);
1556     }
1557
1558     dataLen = segnum * sizeof(VTSI_SEGMENT);
1559     fwrite(pSegment, 1, dataLen, fp);
1560
1561     if (dataLen % 512)
1562     {
1563         //pData + SIZE_1MB - 8192 is a temp data buffer with zero
1564         fwrite(pData + SIZE_1MB - 8192, 1, 512 - (dataLen % 512), fp);
1565     }
1566
1567     //Fill footer
1568     pImgFooter->magic = VTSI_IMG_MAGIC;
1569     pImgFooter->version = 1;
1570     pImgFooter->disk_size = pPhyDrive->SizeInBytes;
1571     memcpy(&pImgFooter->disk_signature, pPhyDrive->MBR.BootCode + 0x1b8, 4);
1572     pImgFooter->segment_num = segnum;
1573     pImgFooter->segment_offset = data_offset;
1574
1575     for (i = 0, chksum = 0; i < (int)(segnum * sizeof(VTSI_SEGMENT)); i++)
1576     {
1577         chksum += *((UINT8 *)pSegment + i);
1578     }
1579     pImgFooter->segment_chksum = ~chksum;
1580
1581     for (i = 0, chksum = 0; i < sizeof(VTSI_FOOTER); i++)
1582     {
1583         chksum += *((UINT8 *)pImgFooter + i);
1584     }
1585     pImgFooter->foot_chksum = ~chksum;
1586
1587     Log("Writing footer segnum(%u)  segoffset(%llu) ......................", segnum, data_offset);
1588     Log("disk_size=%llu disk_signature=%lx segment_offset=%llu", pImgFooter->disk_size, pImgFooter->disk_signature, pImgFooter->segment_offset);
1589
1590     fwrite(pImgFooter, 1, sizeof(VTSI_FOOTER), fp);
1591     fclose(fp);
1592
1593     Log("Writing Ventoy image file finished, the file size should be %llu .", data_offset + 512 + ((dataLen + 511) / 512 * 512));
1594
1595     rc = 0;
1596
1597 End:
1598
1599     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1600
1601     Log("retcode:%d\n", rc);
1602
1603     SAFE_FREE(pData);
1604     SAFE_FREE(ImgBuf);
1605
1606     return rc;
1607 }
1608
1609
1610 int InstallVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int PartStyle, int TryId)
1611 {
1612     int i;
1613     int rc = 0;
1614     int state = 0;
1615     HANDLE hDrive;
1616     DWORD dwSize;
1617     BOOL bRet;
1618     CHAR MountDrive;
1619     CHAR DriveName[] = "?:\\";
1620     CHAR DriveLetters[MAX_PATH] = { 0 };
1621     MBR_HEAD MBR;
1622     VTOY_GPT_INFO *pGptInfo = NULL;
1623     UINT64 Part1StartSector = 0;
1624     UINT64 Part1SectorCount = 0;
1625     UINT64 Part2StartSector = 0;
1626
1627         Log("#####################################################");
1628     Log("InstallVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
1629         PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1630         GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1631         Log("#####################################################");
1632
1633     if (PartStyle)
1634     {
1635         pGptInfo = malloc(sizeof(VTOY_GPT_INFO));
1636         memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
1637     }
1638
1639     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1640
1641     if (PartStyle)
1642     {
1643         VentoyFillGpt(pPhyDrive->SizeInBytes, pGptInfo);
1644         Part1StartSector = pGptInfo->PartTbl[0].StartLBA;
1645         Part1SectorCount = pGptInfo->PartTbl[0].LastLBA - Part1StartSector + 1;
1646         Part2StartSector = pGptInfo->PartTbl[1].StartLBA;
1647     }
1648     else
1649     {
1650         VentoyFillMBR(pPhyDrive->SizeInBytes, &MBR, PartStyle);
1651         Part1StartSector = MBR.PartTbl[0].StartSectorId;
1652         Part1SectorCount = MBR.PartTbl[0].SectorCount;
1653         Part2StartSector = MBR.PartTbl[1].StartSectorId;
1654     }
1655
1656     Log("Lock disk for clean ............................. ");
1657
1658     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
1659     if (hDrive == INVALID_HANDLE_VALUE)
1660     {
1661         Log("Failed to open physical disk");
1662         free(pGptInfo);
1663         return 1;
1664     }
1665
1666     GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1667
1668     if (DriveLetters[0] == 0)
1669     {
1670         Log("No drive letter was assigned...");
1671         DriveName[0] = GetFirstUnusedDriveLetter();
1672         Log("GetFirstUnusedDriveLetter %C: ...", DriveName[0]);
1673     }
1674     else
1675     {
1676         // Unmount all mounted volumes that belong to this drive
1677         // Do it in reverse so that we always end on the first volume letter
1678         for (i = (int)strlen(DriveLetters); i > 0; i--)
1679         {
1680             DriveName[0] = DriveLetters[i - 1];
1681             bRet = DeleteVolumeMountPointA(DriveName);
1682             Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, GetLastError());
1683         }
1684     }
1685
1686     MountDrive = DriveName[0];
1687     Log("Will use '%C:' as volume mountpoint", DriveName[0]);
1688
1689     // It kind of blows, but we have to relinquish access to the physical drive
1690     // for VDS to be able to delete the partitions that reside on it...
1691     DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
1692     CHECK_CLOSE_HANDLE(hDrive);
1693
1694     PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
1695
1696     if (!VDS_DeleteAllPartitions(pPhyDrive->PhyDrive))
1697     {
1698         Log("Notice: Could not delete partitions: 0x%x, but we continue.", GetLastError());
1699     }
1700
1701     Log("Deleting all partitions ......................... OK");
1702
1703     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
1704
1705     Log("Lock disk for write ............................. ");
1706     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1707     if (hDrive == INVALID_HANDLE_VALUE)
1708     {
1709         Log("Failed to GetPhysicalHandle for write.");
1710         rc = 1;
1711         goto End;
1712     }
1713
1714     //Refresh Drive Layout
1715     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1716
1717     disk_io_set_param(hDrive, Part1StartSector + Part1SectorCount);// include the 2048 sector gap
1718
1719     PROGRESS_BAR_SET_POS(PT_FORMAT_PART1);
1720
1721     if (PartStyle == 1 && pPhyDrive->PartStyle == 0)
1722     {
1723         Log("Wait for format part1 ...");
1724         Sleep(1000 * 5);
1725     }
1726
1727     Log("Formatting part1 exFAT ...");
1728     if (0 != FormatPart1exFAT(pPhyDrive->SizeInBytes))
1729     {
1730         Log("FormatPart1exFAT failed.");
1731         rc = 1;
1732         goto End;
1733     }
1734
1735     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1736     Log("Writing part2 FAT img ...");
1737
1738     if (0 != FormatPart2Fat(hDrive, Part2StartSector))
1739     {
1740         Log("FormatPart2Fat failed.");
1741         rc = 1;
1742         goto End;
1743     }
1744
1745     PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
1746     Log("Writing Boot Image ............................. ");
1747     if (WriteGrubStage1ToPhyDrive(hDrive, PartStyle) != 0)
1748     {
1749         Log("WriteGrubStage1ToPhyDrive failed.");
1750         rc = 1;
1751         goto End;
1752     }
1753
1754     PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
1755     Log("Writing Partition Table ........................ ");
1756     SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
1757
1758     if (PartStyle)
1759     {
1760         VTOY_GPT_HDR BackupHead;
1761         LARGE_INTEGER liCurrentPosition;
1762
1763         SET_FILE_POS(pPhyDrive->SizeInBytes - 512);
1764         VentoyFillBackupGptHead(pGptInfo, &BackupHead);
1765         if (!WriteFile(hDrive, &BackupHead, sizeof(VTOY_GPT_HDR), &dwSize, NULL))
1766         {
1767             rc = 1;
1768             Log("Write GPT Backup Head Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1769             goto End;
1770         }
1771
1772         SET_FILE_POS(pPhyDrive->SizeInBytes - 512 * 33);
1773         if (!WriteFile(hDrive, pGptInfo->PartTbl, sizeof(pGptInfo->PartTbl), &dwSize, NULL))
1774         {
1775             rc = 1;
1776             Log("Write GPT Backup Part Table Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1777             goto End;
1778         }
1779
1780         SET_FILE_POS(0);
1781         if (!WriteFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL))
1782         {
1783             rc = 1;
1784             Log("Write GPT Info Failed, dwSize:%u (%u) ErrCode:%u", dwSize, sizeof(VTOY_GPT_INFO), GetLastError());
1785             goto End;
1786         }
1787
1788         Log("Write GPT Info OK ...");
1789         memcpy(&(pPhyDrive->MBR), &(pGptInfo->MBR), 512);
1790     }
1791     else
1792     {
1793         if (!WriteFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL))
1794         {
1795             rc = 1;
1796             Log("Write MBR Failed, dwSize:%u ErrCode:%u", dwSize, GetLastError());
1797             goto End;
1798         }
1799         Log("Write MBR OK ...");
1800         memcpy(&(pPhyDrive->MBR), &MBR, 512);
1801     }
1802
1803     //Refresh Drive Layout
1804     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1805
1806 End:
1807
1808     PROGRESS_BAR_SET_POS(PT_MOUNT_VOLUME);
1809
1810     if (rc == 0)
1811     {
1812         Log("Mounting Ventoy Partition ....................... ");
1813         Sleep(1000);
1814
1815         state = 0;
1816         memset(DriveLetters, 0, sizeof(DriveLetters));
1817         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
1818         Log("Logical drive letter after write ventoy: <%s>", DriveLetters);
1819
1820         for (i = 0; i < sizeof(DriveLetters) && DriveLetters[i]; i++)
1821         {
1822             DriveName[0] = DriveLetters[i];
1823             if (IsVentoyLogicalDrive(DriveName[0]))
1824             {
1825                 Log("%s is ventoy part2, delete mountpoint", DriveName);
1826                 DeleteVolumeMountPointA(DriveName);
1827             }
1828             else
1829             {
1830                 Log("%s is ventoy part1, already mounted", DriveName);
1831                 state = 1;
1832             }
1833         }
1834
1835         if (state != 1)
1836         {
1837             Log("need to mount ventoy part1...");
1838
1839             if (0 == GetVentoyVolumeName(pPhyDrive->PhyDrive, Part1StartSector, DriveLetters, sizeof(DriveLetters), FALSE))
1840             {
1841                 DriveName[0] = MountDrive;
1842                 bRet = SetVolumeMountPointA(DriveName, DriveLetters);
1843                 Log("SetVolumeMountPoint <%s> <%s> bRet:%u code:%u", DriveName, DriveLetters, bRet, GetLastError());
1844             }
1845             else
1846             {
1847                 Log("Failed to find ventoy volume");
1848             }
1849         }
1850         Log("OK\n");
1851     }
1852     else
1853     {
1854         FindProcessOccupyDisk(hDrive, pPhyDrive);
1855     }
1856
1857     if (pGptInfo)
1858     {
1859         free(pGptInfo);
1860     }
1861
1862     CHECK_CLOSE_HANDLE(hDrive);
1863     return rc;
1864 }
1865
1866
1867 int PartitionResizeForVentoy(PHY_DRIVE_INFO *pPhyDrive)
1868 {
1869         int i, j;
1870         int rc = 1;
1871         int PhyDrive;
1872         int PartStyle;
1873         INT64 ReservedValue;
1874         UINT64 RecudeBytes;
1875         GUID Guid;
1876         MBR_HEAD MBR;
1877         VTOY_GPT_INFO *pGPT;
1878         MBR_HEAD *pMBR;
1879         DWORD dwSize = 0;
1880         VTOY_GPT_HDR BackupHead;
1881         HANDLE hDrive = INVALID_HANDLE_VALUE;
1882         GUID ZeroGuid = { 0 };
1883         static GUID WindowsDataPartType = { 0xebd0a0a2, 0xb9e5, 0x4433, { 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7 } };
1884         static GUID EspPartType = { 0xc12a7328, 0xf81f, 0x11d2, { 0xba, 0x4b, 0x00, 0xa0, 0xc9, 0x3e, 0xc9, 0x3b } };
1885         static GUID BiosGrubPartType = { 0x21686148, 0x6449, 0x6e6f, { 0x74, 0x4e, 0x65, 0x65, 0x64, 0x45, 0x46, 0x49 } };
1886
1887         Log("#####################################################");
1888         Log("PartitionResizeForVentoy PhyDrive%d <<%s %s %dGB>>",
1889                 pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
1890                 GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
1891         Log("#####################################################");
1892
1893         pGPT = &(pPhyDrive->Gpt);
1894         pMBR = &(pPhyDrive->Gpt.MBR);
1895         Log("Disksize:%llu Part2Start:%llu", pPhyDrive->SizeInBytes, pPhyDrive->ResizePart2StartSector * 512);
1896
1897         if (pMBR->PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
1898         {
1899                 PartStyle = 1;
1900         }
1901         else
1902         {
1903                 PartStyle = 0;
1904         }
1905
1906         PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
1907
1908         RecudeBytes = VENTOY_EFI_PART_SIZE;
1909         ReservedValue = GetReservedSpaceInMB();
1910         if (ReservedValue > 0)
1911         {
1912                 Log("Reduce add reserved space %lldMB", (LONGLONG)ReservedValue);
1913                 RecudeBytes += (UINT64)(ReservedValue * SIZE_1MB);
1914         }
1915
1916
1917         if (pPhyDrive->ResizeNoShrink == FALSE)
1918         {
1919                 Log("Need to shrink the volume");
1920                 if (VDS_ShrinkVolume(pPhyDrive->ResizeVolumeGuid, RecudeBytes))
1921                 {
1922                         Log("Shrink volume success, now check again");
1923
1924                         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1925                         if (hDrive == INVALID_HANDLE_VALUE)
1926                         {
1927                                 Log("Failed to GetPhysicalHandle for update.");
1928                                 goto End;
1929                         }
1930
1931                         //Refresh Drive Layout
1932                         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
1933
1934                         CHECK_CLOSE_HANDLE(hDrive);
1935
1936
1937                         if (PartResizePreCheck(NULL) && pPhyDrive->ResizeNoShrink)
1938                         {
1939                                 Log("Recheck after Shrink volume success");
1940                                 Log("After shrink Disksize:%llu Part2Start:%llu", pPhyDrive->SizeInBytes, pPhyDrive->ResizePart2StartSector * 512);
1941                         }
1942                         else
1943                         {
1944                                 Log("Recheck after Shrink volume failed %u", pPhyDrive->ResizeNoShrink);
1945                                 goto End;
1946                         }
1947                 }
1948                 else
1949                 {
1950                         Log("Shrink volume failed");
1951                         goto End;
1952                 }
1953         }
1954
1955
1956         //Now try write data
1957         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
1958         if (hDrive == INVALID_HANDLE_VALUE)
1959         {
1960                 Log("Failed to GetPhysicalHandle for update.");
1961                 goto End;
1962         }
1963
1964
1965         //Write partition 2 data
1966         PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
1967         if (0 != FormatPart2Fat(hDrive, pPhyDrive->ResizePart2StartSector))
1968         {
1969                 Log("FormatPart2Fat failed.");
1970                 goto End;
1971         }
1972
1973         //Write grub stage2 gap
1974         PROGRESS_BAR_SET_POS(PT_WRITE_STG1_IMG);
1975         Log("Writing Boot Image ............................. ");
1976         if (WriteGrubStage1ToPhyDrive(hDrive, PartStyle) != 0)
1977         {
1978                 Log("WriteGrubStage1ToPhyDrive failed.");
1979                 goto End;
1980         }
1981
1982
1983         //Write partition table
1984         PROGRESS_BAR_SET_POS(PT_WRITE_PART_TABLE);
1985         Log("Writing partition table ............................. ");
1986
1987         VentoyGetLocalBootImg(&MBR);
1988         CoCreateGuid(&Guid);
1989         memcpy(MBR.BootCode + 0x180, &Guid, 16);
1990         memcpy(pMBR->BootCode, MBR.BootCode, 440);
1991
1992         if (PartStyle == 0)
1993         {
1994                 for (i = 1; i < 4; i++)
1995                 {
1996                         if (pMBR->PartTbl[i].SectorCount == 0)
1997                         {
1998                                 break;
1999                         }
2000                 }
2001
2002                 if (i >= 4)
2003                 {
2004                         Log("Can not find MBR free partition table");
2005                         goto End;
2006                 }
2007
2008                 for (j = i - 1; j > 0; j--)
2009                 {
2010                         Log("Move MBR partition table %d --> %d", j + 1, j + 2);
2011                         memcpy(pMBR->PartTbl + (j + 1), pMBR->PartTbl + j, sizeof(PART_TABLE));
2012                 }
2013
2014                 VentoyFillMBRLocation(pPhyDrive->SizeInBytes, (UINT32)pPhyDrive->ResizePart2StartSector, VENTOY_EFI_PART_SIZE / 512, pMBR->PartTbl + 1);
2015                 pMBR->PartTbl[0].Active = 0x80; // bootable
2016                 pMBR->PartTbl[1].Active = 0x00;
2017                 pMBR->PartTbl[1].FsFlag = 0xEF; // EFI System Partition
2018
2019                 if (!WriteDataToPhyDisk(hDrive, 0, pMBR, 512))
2020                 {
2021                         Log("Legacy BIOS write MBR failed");
2022                         goto End;
2023                 }
2024         }
2025         else
2026         {
2027                 for (i = 1; i < 128; i++)
2028                 {
2029                         if (memcmp(&(pGPT->PartTbl[i].PartGuid), &ZeroGuid, sizeof(GUID)) == 0)
2030                         {
2031                                 break;
2032                         }
2033                 }
2034
2035                 if (i >= 128)
2036                 {
2037                         Log("Can not find GPT free partition table");
2038                         goto End;
2039                 }
2040
2041                 for (j = i - 1; j > 0; j--)
2042                 {
2043                         Log("Move GPT partition table %d --> %d", j + 1, j + 2);
2044                         memcpy(pGPT->PartTbl + (j + 1), pGPT->PartTbl + j, sizeof(VTOY_GPT_PART_TBL));
2045                 }
2046
2047
2048                 pMBR->BootCode[92] = 0x22;
2049
2050                 // to fix windows issue
2051                 memcpy(&(pGPT->PartTbl[1].PartType), &WindowsDataPartType, sizeof(GUID));
2052                 CoCreateGuid(&(pGPT->PartTbl[1].PartGuid));
2053
2054                 pGPT->PartTbl[1].StartLBA = pGPT->PartTbl[0].LastLBA + 1;
2055                 pGPT->PartTbl[1].LastLBA = pGPT->PartTbl[1].StartLBA + VENTOY_EFI_PART_SIZE / 512 - 1;
2056                 pGPT->PartTbl[1].Attr = 0xC000000000000001ULL;
2057                 memcpy(pGPT->PartTbl[1].Name, L"VTOYEFI", 7 * 2);
2058
2059                 //Update CRC
2060                 pGPT->Head.PartTblCrc = VentoyCrc32(pGPT->PartTbl, sizeof(pGPT->PartTbl));
2061                 pGPT->Head.Crc = VentoyCrc32(&(pGPT->Head), pGPT->Head.Length);
2062
2063                 Log("pGPT->Head.EfiStartLBA=%llu", (ULONGLONG)pGPT->Head.EfiStartLBA);
2064                 Log("pGPT->Head.EfiBackupLBA=%llu", (ULONGLONG)pGPT->Head.EfiBackupLBA);
2065
2066                 VentoyFillBackupGptHead(pGPT, &BackupHead);
2067                 if (!WriteDataToPhyDisk(hDrive, pGPT->Head.EfiBackupLBA * 512, &BackupHead, 512))
2068                 {
2069                         Log("UEFI write backup head failed");
2070                         goto End;
2071                 }
2072
2073                 if (!WriteDataToPhyDisk(hDrive, (pGPT->Head.EfiBackupLBA - 32) * 512, pGPT->PartTbl, 512 * 32))
2074                 {
2075                         Log("UEFI write backup partition table failed");
2076                         goto End;
2077                 }
2078
2079                 if (!WriteDataToPhyDisk(hDrive, 0, pGPT, 512 * 34))
2080                 {
2081                         Log("UEFI write MBR & Main partition table failed");
2082                         goto End;
2083                 }
2084         }
2085
2086
2087
2088         //Refresh Drive Layout
2089         DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2090
2091         //We must close handle here, because it will block the refresh bellow
2092         CHECK_CLOSE_HANDLE(hDrive);
2093
2094         Sleep(2000);
2095
2096         //Refresh disk list
2097         PhyDrive = pPhyDrive->PhyDrive;
2098
2099         Log("#### Now Refresh PhyDrive ####");
2100         Ventoy2DiskDestroy();
2101         Ventoy2DiskInit();
2102
2103         pPhyDrive = GetPhyDriveInfoByPhyDrive(PhyDrive);
2104         if (pPhyDrive)
2105         {
2106                 if (pPhyDrive->VentoyVersion[0] == 0)
2107                 {
2108                         Log("After process the Ventoy version is still invalid");
2109                         goto End;
2110                 }
2111
2112                 Log("### Ventoy non-destructive installation successfully finished <%s>", pPhyDrive->VentoyVersion);
2113         }
2114         else
2115         {
2116                 Log("### Ventoy non-destructive installation successfully finished <not found>");
2117         }
2118
2119         InitComboxCtrl(g_DialogHwnd, PhyDrive);
2120
2121         rc = 0;
2122
2123 End:
2124         CHECK_CLOSE_HANDLE(hDrive);
2125         return rc;
2126 }
2127
2128
2129 static BOOL DiskCheckWriteAccess(HANDLE hDrive)
2130 {
2131         DWORD dwSize;
2132         BOOL ret = FALSE;
2133         BOOL bRet = FALSE;
2134         BYTE Buffer[512];
2135         LARGE_INTEGER liCurPosition;
2136         LARGE_INTEGER liNewPosition;
2137
2138         liCurPosition.QuadPart = 2039 * 512;
2139         liNewPosition.QuadPart = 0;
2140         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2141                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2142         {
2143                 Log("SetFilePointer1 Failed %u", LASTERR);
2144                 goto out;
2145         }
2146
2147
2148         dwSize = 0;
2149         ret = ReadFile(hDrive, Buffer, 512, &dwSize, NULL);
2150         if ((!ret) || (dwSize != 512))
2151         {
2152                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2153                 goto out;
2154         }
2155
2156
2157         liCurPosition.QuadPart = 2039 * 512;
2158         liNewPosition.QuadPart = 0;
2159         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2160                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2161         {
2162                 Log("SetFilePointer2 Failed %u", LASTERR);
2163                 goto out;
2164         }
2165
2166         dwSize = 0;
2167         ret = WriteFile(hDrive, Buffer, 512, &dwSize, NULL);
2168         if ((!ret) || dwSize != 512)
2169         {
2170                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
2171                 goto out;
2172         }
2173
2174         bRet = TRUE;
2175
2176 out:
2177
2178         return bRet;
2179 }
2180
2181 static BOOL BackupDataBeforeCleanDisk(int PhyDrive, UINT64 DiskSize, BYTE **pBackup)
2182 {
2183         DWORD dwSize;
2184         DWORD dwStatus;
2185         BOOL Return = FALSE;
2186         BOOL ret = FALSE;
2187         BYTE *backup = NULL;
2188         UINT64 offset;
2189         HANDLE hDrive = INVALID_HANDLE_VALUE;
2190         LARGE_INTEGER liCurPosition;
2191         LARGE_INTEGER liNewPosition;
2192         VTOY_GPT_INFO *pGPT = NULL;
2193
2194         Log("BackupDataBeforeCleanDisk %d", PhyDrive);
2195
2196         // step1: check write access
2197         hDrive = GetPhysicalHandle(PhyDrive, TRUE, TRUE, FALSE);
2198         if (hDrive == INVALID_HANDLE_VALUE)
2199         {
2200                 Log("Failed to GetPhysicalHandle for write.");
2201                 goto out;
2202         }
2203
2204         if (DiskCheckWriteAccess(hDrive))
2205         {
2206                 Log("DiskCheckWriteAccess success");
2207                 CHECK_CLOSE_HANDLE(hDrive);
2208         }
2209         else
2210         {
2211                 Log("DiskCheckWriteAccess failed");
2212                 goto out;
2213         }
2214
2215         //step2 backup 4MB data
2216         backup = malloc(SIZE_1MB * 4);
2217         if (!backup)
2218         {
2219                 goto out;
2220         }
2221
2222         hDrive = GetPhysicalHandle(PhyDrive, FALSE, FALSE, FALSE);
2223         if (hDrive == INVALID_HANDLE_VALUE)
2224         {
2225                 goto out;
2226         }
2227
2228         //read first 2MB
2229         dwStatus = SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2230         if (dwStatus != 0)
2231         {
2232                 goto out;
2233         }
2234
2235         dwSize = 0;
2236         ret = ReadFile(hDrive, backup, SIZE_2MB, &dwSize, NULL);
2237         if ((!ret) || (dwSize != SIZE_2MB))
2238         {
2239                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2240                 goto out;
2241         }
2242
2243         pGPT = (VTOY_GPT_INFO *)backup;
2244         offset = pGPT->Head.EfiBackupLBA * 512;
2245         if (offset >= (DiskSize - SIZE_2MB) && offset < DiskSize)
2246         {
2247                 Log("EFI partition table check success");
2248         }
2249         else
2250         {
2251                 Log("Backup EFI LBA not in last 2MB range: %llu", pGPT->Head.EfiBackupLBA);
2252                 goto out;
2253         }
2254
2255         //read last 2MB
2256         liCurPosition.QuadPart = DiskSize - SIZE_2MB;
2257         liNewPosition.QuadPart = 0;
2258         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2259                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2260         {
2261                 goto out;
2262         }
2263
2264         dwSize = 0;
2265         ret = ReadFile(hDrive, backup + SIZE_2MB, SIZE_2MB, &dwSize, NULL);
2266         if ((!ret) || (dwSize != SIZE_2MB))
2267         {
2268                 Log("Failed to read %d %u 0x%x", ret, dwSize, LASTERR);
2269                 goto out;
2270         }
2271
2272         *pBackup = backup;
2273         backup = NULL; //For don't free later
2274         Return = TRUE;
2275
2276 out:
2277         CHECK_CLOSE_HANDLE(hDrive);
2278         if (backup)
2279                 free(backup);
2280
2281         return Return;
2282 }
2283
2284
2285 static BOOL WriteBackupDataToDisk(HANDLE hDrive, UINT64 Offset, BYTE *Data, DWORD Length)
2286 {
2287         DWORD dwSize = 0;
2288         BOOL ret = FALSE;
2289         LARGE_INTEGER liCurPosition;
2290         LARGE_INTEGER liNewPosition;
2291
2292         Log("WriteBackupDataToDisk %llu %p %u", Offset, Data, Length);
2293
2294         liCurPosition.QuadPart = Offset;
2295         liNewPosition.QuadPart = 0;
2296         if (0 == SetFilePointerEx(hDrive, liCurPosition, &liNewPosition, FILE_BEGIN) ||
2297                 liNewPosition.QuadPart != liCurPosition.QuadPart)
2298         {
2299                 return FALSE;
2300         }
2301
2302         ret = WriteFile(hDrive, Data, Length, &dwSize, NULL);
2303         if ((!ret) || dwSize != Length)
2304         {
2305                 Log("Failed to write %d %u %u", ret, dwSize, LASTERR);
2306                 return FALSE;
2307         }
2308
2309         Log("WriteBackupDataToDisk %llu %p %u success", Offset, Data, Length);
2310         return TRUE;
2311 }
2312
2313
2314 int UpdateVentoy2PhyDrive(PHY_DRIVE_INFO *pPhyDrive, int TryId)
2315 {
2316         int i;
2317         int rc = 0;
2318         int MaxRetry = 4;
2319         BOOL ForceMBR = FALSE;
2320         BOOL Esp2Basic = FALSE;
2321         BOOL ChangeAttr = FALSE;
2322         BOOL CleanDisk = FALSE;
2323         BOOL bWriteBack = TRUE;
2324         HANDLE hVolume;
2325         HANDLE hDrive;
2326         DWORD Status;
2327         DWORD dwSize;
2328         BOOL bRet;
2329         CHAR DriveName[] = "?:\\";
2330         CHAR DriveLetters[MAX_PATH] = { 0 };
2331         CHAR BackBinFile[MAX_PATH];
2332         UINT64 StartSector;
2333         UINT64 ReservedMB = 0;
2334         MBR_HEAD BootImg;
2335         MBR_HEAD MBR;
2336         BYTE *pBackup = NULL;
2337         VTOY_GPT_INFO *pGptInfo = NULL;
2338         UINT8 ReservedData[4096];
2339
2340         Log("#####################################################");
2341         Log("UpdateVentoy2PhyDrive try%d %s PhyDrive%d <<%s %s %dGB>>", TryId,
2342                 pPhyDrive->PartStyle ? "GPT" : "MBR", pPhyDrive->PhyDrive, pPhyDrive->VendorId, pPhyDrive->ProductId,
2343                 GetHumanReadableGBSize(pPhyDrive->SizeInBytes));
2344         Log("#####################################################");
2345
2346         PROGRESS_BAR_SET_POS(PT_LOCK_FOR_CLEAN);
2347
2348         Log("Lock disk for umount ............................ ");
2349
2350         hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, FALSE, FALSE);
2351         if (hDrive == INVALID_HANDLE_VALUE)
2352         {
2353                 Log("Failed to open physical disk");
2354                 return 1;
2355         }
2356
2357         if (pPhyDrive->PartStyle)
2358         {
2359                 pGptInfo = malloc(sizeof(VTOY_GPT_INFO));
2360                 if (!pGptInfo)
2361                 {
2362                         return 1;
2363                 }
2364
2365                 memset(pGptInfo, 0, sizeof(VTOY_GPT_INFO));
2366
2367                 // Read GPT Info
2368                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2369                 ReadFile(hDrive, pGptInfo, sizeof(VTOY_GPT_INFO), &dwSize, NULL);
2370
2371                 //MBR will be used to compare with local boot image
2372                 memcpy(&MBR, &pGptInfo->MBR, sizeof(MBR_HEAD));
2373
2374                 StartSector = pGptInfo->PartTbl[1].StartLBA;
2375                 Log("GPT StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2376
2377                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512) - 33) / 2048;
2378                 Log("GPT Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2379         }
2380         else
2381         {
2382                 // Read MBR
2383                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2384                 ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
2385
2386                 StartSector = MBR.PartTbl[1].StartSectorId;
2387                 Log("MBR StartSector in PartTbl:%llu", (ULONGLONG)StartSector);
2388
2389                 ReservedMB = (pPhyDrive->SizeInBytes / 512 - (StartSector + VENTOY_EFI_PART_SIZE / 512)) / 2048;
2390                 Log("MBR Reserved Disk Space:%llu MB", (ULONGLONG)ReservedMB);
2391         }
2392
2393         //Read Reserved Data
2394         SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2395         ReadFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2396
2397         GetLettersBelongPhyDrive(pPhyDrive->PhyDrive, DriveLetters, sizeof(DriveLetters));
2398
2399         if (DriveLetters[0] == 0)
2400         {
2401                 Log("No drive letter was assigned...");
2402         }
2403         else
2404         {
2405                 // Unmount all mounted volumes that belong to this drive
2406                 // Do it in reverse so that we always end on the first volume letter
2407                 for (i = (int)strlen(DriveLetters); i > 0; i--)
2408                 {
2409                         DriveName[0] = DriveLetters[i - 1];
2410                         if (IsVentoyLogicalDrive(DriveName[0]))
2411                         {
2412                                 Log("%s is ventoy logical drive", DriveName);
2413                                 bRet = DeleteVolumeMountPointA(DriveName);
2414                                 Log("Delete mountpoint %s ret:%u code:%u", DriveName, bRet, LASTERR);
2415                                 break;
2416                         }
2417                 }
2418         }
2419
2420         // It kind of blows, but we have to relinquish access to the physical drive
2421         // for VDS to be able to delete the partitions that reside on it...
2422         DeviceIoControl(hDrive, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2423         CHECK_CLOSE_HANDLE(hDrive);
2424
2425         if (pPhyDrive->PartStyle == 1)
2426         {
2427                 Log("TryId=%d EFI GPT partition type is 0x%llx", TryId, pPhyDrive->Part2GPTAttr);
2428                 PROGRESS_BAR_SET_POS(PT_DEL_ALL_PART);
2429
2430                 if (TryId == 1)
2431                 {
2432                         Log("Change GPT partition type to ESP");
2433                         if (VDS_ChangeVtoyEFI2ESP(pPhyDrive->PhyDrive, StartSector * 512))
2434                         {
2435                                 Esp2Basic = TRUE;
2436                                 Sleep(3000);
2437                         }
2438                 }
2439                 else if (TryId == 2)
2440                 {
2441                         Log("Change GPT partition attribute");
2442                         if (VDS_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, 0x8000000000000001))
2443                         {
2444                                 ChangeAttr = TRUE;
2445                                 Sleep(2000);
2446                         }
2447                 }
2448                 else if (TryId == 3)
2449                 {
2450                         Log("Clean disk GPT partition table");
2451                         if (BackupDataBeforeCleanDisk(pPhyDrive->PhyDrive, pPhyDrive->SizeInBytes, &pBackup))
2452                         {
2453                                 sprintf_s(BackBinFile, sizeof(BackBinFile), ".\\ventoy\\phydrive%d_%u_%d.bin",
2454                                         pPhyDrive->PhyDrive, GetCurrentProcessId(), g_backup_bin_index++);
2455                                 SaveBufToFile(BackBinFile, pBackup, 4 * SIZE_1MB);
2456                                 Log("Save backup data to %s", BackBinFile);
2457
2458                                 Log("Success to backup data before clean");
2459                                 CleanDisk = TRUE;
2460                                 if (!VDS_CleanDisk(pPhyDrive->PhyDrive))
2461                                 {
2462                                         Sleep(3000);
2463                                         DSPT_CleanDisk(pPhyDrive->PhyDrive);
2464                                 }
2465                                 Sleep(3000);
2466                         }
2467                         else
2468                         {
2469                                 Log("Failed to backup data before clean");
2470                         }
2471                 }
2472         }
2473
2474     PROGRESS_BAR_SET_POS(PT_LOCK_FOR_WRITE);
2475
2476     Log("Lock disk for update ............................ ");
2477     hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2478     if (hDrive == INVALID_HANDLE_VALUE)
2479     {
2480         Log("Failed to GetPhysicalHandle for write.");
2481         rc = 1;
2482         goto End;
2483     }
2484
2485     PROGRESS_BAR_SET_POS(PT_LOCK_VOLUME);
2486
2487     Log("Lock volume for update .......................... ");
2488     hVolume = INVALID_HANDLE_VALUE;
2489
2490         //If we change VTOYEFI to ESP, it can not have s volume name, so don't try to get it.
2491         if (CleanDisk)
2492         {
2493                 //writeback the last 2MB
2494                 if (!WriteBackupDataToDisk(hDrive, pPhyDrive->SizeInBytes - SIZE_2MB, pBackup + SIZE_2MB, SIZE_2MB))
2495                 {
2496                         bWriteBack = FALSE;
2497                 }
2498
2499                 //write the first 2MB except parttable
2500                 if (!WriteBackupDataToDisk(hDrive, 34 * 512, pBackup + 34 * 512, SIZE_2MB - 34 * 512))
2501                 {
2502                         bWriteBack = FALSE;
2503                 }
2504
2505                 Status = ERROR_NOT_FOUND;
2506         }
2507         else if (Esp2Basic)
2508         {
2509                 Status = ERROR_NOT_FOUND;
2510         }
2511         else
2512         {
2513                 for (i = 0; i < MaxRetry; i++)
2514                 {
2515                         Status = GetVentoyVolumeName(pPhyDrive->PhyDrive, StartSector, DriveLetters, sizeof(DriveLetters), TRUE);
2516                         if (ERROR_SUCCESS == Status)
2517                         {
2518                                 break;
2519                         }
2520                         else
2521                         {
2522                                 Log("==== Volume not found, wait and retry %d... ====", i);
2523                                 Sleep(2);
2524                         }
2525                 }
2526         }
2527
2528     if (ERROR_SUCCESS == Status)
2529     {
2530         Log("Now lock and dismount volume <%s>", DriveLetters);
2531
2532         for (i = 0; i < MaxRetry; i++)
2533         {
2534             hVolume = CreateFileA(DriveLetters,
2535                 GENERIC_READ | GENERIC_WRITE,
2536                 FILE_SHARE_READ,
2537                 NULL,
2538                 OPEN_EXISTING,
2539                 FILE_ATTRIBUTE_NORMAL | FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
2540                 NULL);
2541
2542             if (hVolume == INVALID_HANDLE_VALUE)
2543             {
2544                 Log("Failed to create file volume, errcode:%u, wait and retry ...", LASTERR);
2545                 Sleep(2000);
2546             }
2547             else
2548             {
2549                 break;
2550             }
2551         }
2552
2553         if (hVolume == INVALID_HANDLE_VALUE)
2554         {
2555             Log("Failed to create file volume, errcode:%u", LASTERR);
2556         }
2557         else
2558         {
2559             bRet = DeviceIoControl(hVolume, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2560             Log("FSCTL_LOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2561
2562             bRet = DeviceIoControl(hVolume, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2563             Log("FSCTL_DISMOUNT_VOLUME bRet:%u code:%u", bRet, LASTERR);
2564         }
2565     }
2566     else if (ERROR_NOT_FOUND == Status)
2567     {
2568         Log("Volume not found, maybe not supported");
2569     }
2570     else
2571     {
2572         rc = 1;
2573         goto End;
2574     }
2575
2576         bRet = TryWritePart2(hDrive, StartSector);
2577         if (FALSE == bRet && Esp2Basic)
2578         {
2579                 Log("TryWritePart2 agagin ...");
2580                 Sleep(3000);
2581                 bRet = TryWritePart2(hDrive, StartSector);
2582         }
2583
2584         if (!bRet)
2585     {
2586                 if (pPhyDrive->PartStyle == 0)
2587                 {
2588                         if (DiskCheckWriteAccess(hDrive))
2589                         {
2590                                 Log("MBR DiskCheckWriteAccess success");
2591
2592                                 ForceMBR = TRUE;
2593
2594                                 Log("Try write failed, now delete partition 2 for MBR...");
2595                                 CHECK_CLOSE_HANDLE(hDrive);
2596
2597                                 Log("Now delete partition 2...");
2598                                 VDS_DeleteVtoyEFIPartition(pPhyDrive->PhyDrive);
2599
2600                                 hDrive = GetPhysicalHandle(pPhyDrive->PhyDrive, TRUE, TRUE, FALSE);
2601                                 if (hDrive == INVALID_HANDLE_VALUE)
2602                                 {
2603                                         Log("Failed to GetPhysicalHandle for write.");
2604                                         rc = 1;
2605                                         goto End;
2606                                 }
2607                         }
2608                         else
2609                         {
2610                                 Log("MBR DiskCheckWriteAccess failed");
2611                         }
2612                 }
2613                 else
2614                 {
2615                         Log("TryWritePart2 failed ....");
2616                         rc = 1;
2617                         goto End;
2618                 }
2619     }
2620
2621     PROGRESS_BAR_SET_POS(PT_FORMAT_PART2);
2622
2623     Log("Write Ventoy to disk ............................ ");
2624     if (0 != FormatPart2Fat(hDrive, StartSector))
2625     {
2626         rc = 1;
2627         goto End;
2628     }
2629
2630     if (hVolume != INVALID_HANDLE_VALUE)
2631     {
2632         bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2633         Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2634         CHECK_CLOSE_HANDLE(hVolume);
2635     }
2636
2637     Log("Updating Boot Image ............................. ");
2638     if (WriteGrubStage1ToPhyDrive(hDrive, pPhyDrive->PartStyle) != 0)
2639     {
2640         rc = 1;
2641         goto End;
2642     }
2643
2644     //write reserved data
2645     SetFilePointer(hDrive, 512 * 2040, NULL, FILE_BEGIN);
2646     bRet = WriteFile(hDrive, ReservedData, sizeof(ReservedData), &dwSize, NULL);
2647     Log("Write resv data ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2648
2649     // Boot Image
2650     VentoyGetLocalBootImg(&BootImg);
2651
2652     // Use Old UUID
2653     memcpy(BootImg.BootCode + 0x180, MBR.BootCode + 0x180, 16);
2654     if (pPhyDrive->PartStyle)
2655     {
2656         BootImg.BootCode[92] = 0x22;
2657     }
2658
2659     if (ForceMBR == FALSE && memcmp(BootImg.BootCode, MBR.BootCode, 440) == 0)
2660     {
2661         Log("Boot image has no difference, no need to write.");
2662     }
2663     else
2664     {
2665         Log("Boot image need to write %u.", ForceMBR);
2666
2667         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2668
2669         memcpy(MBR.BootCode, BootImg.BootCode, 440);
2670         bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2671         Log("Write Boot Image ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2672     }
2673
2674     if (pPhyDrive->PartStyle == 0)
2675     {
2676         if (0x00 == MBR.PartTbl[0].Active && 0x80 == MBR.PartTbl[1].Active)
2677         {
2678             Log("Need to chage 1st partition active and 2nd partition inactive.");
2679
2680             MBR.PartTbl[0].Active = 0x80;
2681             MBR.PartTbl[1].Active = 0x00;
2682
2683             SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
2684             bRet = WriteFile(hDrive, &MBR, 512, &dwSize, NULL);
2685             Log("Write NEW MBR ret:%u dwSize:%u Error:%u", bRet, dwSize, LASTERR);
2686         }
2687     }
2688
2689         if (CleanDisk)
2690         {
2691                 if (!WriteBackupDataToDisk(hDrive, 0, pBackup, 34 * 512))
2692                 {
2693                         bWriteBack = FALSE;
2694                 }
2695
2696                 free(pBackup);
2697
2698                 if (bWriteBack)
2699                 {
2700                         Log("Write backup data success, now delete %s", BackBinFile);
2701                         DeleteFileA(BackBinFile);
2702                 }
2703                 else
2704                 {
2705                         Log("Write backup data failed");
2706                 }
2707
2708                 Sleep(1000);
2709         }
2710
2711     //Refresh Drive Layout
2712     DeviceIoControl(hDrive, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, NULL, 0, &dwSize, NULL);
2713
2714 End:
2715
2716         if (hVolume != INVALID_HANDLE_VALUE)
2717         {
2718                 bRet = DeviceIoControl(hVolume, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0, &dwSize, NULL);
2719                 Log("FSCTL_UNLOCK_VOLUME bRet:%u code:%u", bRet, LASTERR);
2720                 CHECK_CLOSE_HANDLE(hVolume);
2721         }
2722
2723     if (rc == 0)
2724     {
2725         Log("OK");
2726     }
2727     else
2728     {
2729         FindProcessOccupyDisk(hDrive, pPhyDrive);
2730     }
2731
2732     CHECK_CLOSE_HANDLE(hDrive);
2733
2734     if (Esp2Basic)
2735     {
2736                 Log("Recover GPT partition type to basic");
2737                 VDS_ChangeVtoyEFI2Basic(pPhyDrive->PhyDrive, StartSector * 512);
2738     }
2739
2740         if (pPhyDrive->PartStyle == 1)
2741         {
2742                 if (ChangeAttr || ((pPhyDrive->Part2GPTAttr >> 56) != 0xC0))
2743                 {
2744                         Log("Change EFI partition attr %u <0x%llx> to <0x%llx>", ChangeAttr, pPhyDrive->Part2GPTAttr, 0xC000000000000001ULL);
2745                         if (VDS_ChangeVtoyEFIAttr(pPhyDrive->PhyDrive, 0xC000000000000001ULL))
2746                         {
2747                                 Log("Change EFI partition attr success");
2748                                 pPhyDrive->Part2GPTAttr = 0xC000000000000001ULL;
2749                         }
2750                         else
2751                         {
2752                                 Log("Change EFI partition attr failed");
2753                         }
2754                 }
2755         }
2756
2757     if (pGptInfo)
2758     {
2759         free(pGptInfo);
2760     }
2761
2762     return rc;
2763 }
2764
2765