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