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