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