Ventoy2Disk/Ventoy2Disk/Utility.c

   1 /******************************************************************************
   2  * Utility.c
   3  *
   4  * Copyright (c) 2021, 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 #include <Windows.h>
  22 #include "Ventoy2Disk.h"
  23
  24 void TraceOut(const char *Fmt, ...)
  25 {
  26     va_list Arg;
  27     int Len = 0;
  28     FILE *File = NULL;
  29     char szBuf[1024];
  30
  31     va_start(Arg, Fmt);
  32     Len += vsnprintf_s(szBuf + Len, sizeof(szBuf)-Len, sizeof(szBuf)-Len, Fmt, Arg);
  33     va_end(Arg);
  34
  35     fopen_s(&File, VENTOY_FILE_LOG, "a+");
  36     if (File)
  37     {
  38         fwrite(szBuf, 1, Len, File);
  39         fclose(File);
  40     }
  41 }
  42
  43 void Log(const char *Fmt, ...)
  44 {
  45     va_list Arg;
  46     int Len = 0;
  47     FILE *File = NULL;
  48     SYSTEMTIME Sys;
  49     char szBuf[1024];
  50
  51     GetLocalTime(&Sys);
  52     Len += safe_sprintf(szBuf,
  53         "[%4d/%02d/%02d %02d:%02d:%02d.%03d] ",
  54         Sys.wYear, Sys.wMonth, Sys.wDay,
  55         Sys.wHour, Sys.wMinute, Sys.wSecond,
  56         Sys.wMilliseconds);
  57
  58     va_start(Arg, Fmt);
  59     Len += vsnprintf_s(szBuf + Len, sizeof(szBuf)-Len, sizeof(szBuf)-Len, Fmt, Arg);
  60     va_end(Arg);
  61
  62     //printf("%s\n", szBuf);
  63
  64 #if 1
  65     fopen_s(&File, VENTOY_FILE_LOG, "a+");
  66     if (File)
  67     {
  68         fwrite(szBuf, 1, Len, File);
  69         fwrite("\n", 1, 1, File);
  70         fclose(File);
  71     }
  72 #endif
  73
  74 }
  75
  76 BOOL IsPathExist(BOOL Dir, const char *Fmt, ...)
  77 {
  78     va_list Arg;
  79     HANDLE hFile;
  80     DWORD Attr;
  81     CHAR FilePath[MAX_PATH];
  82
  83     va_start(Arg, Fmt);
  84     vsnprintf_s(FilePath, sizeof(FilePath), sizeof(FilePath), Fmt, Arg);
  85     va_end(Arg);
  86
  87     hFile = CreateFileA(FilePath, FILE_READ_EA, FILE_SHARE_READ, 0, OPEN_EXISTING, 0, 0);
  88     if (INVALID_HANDLE_VALUE == hFile)
  89     {
  90         return FALSE;
  91     }
  92
  93     CloseHandle(hFile);
  94
  95     Attr = GetFileAttributesA(FilePath);
  96
  97     if (Dir)
  98     {
  99         if ((Attr & FILE_ATTRIBUTE_DIRECTORY) == 0)
 100         {
 101             return FALSE;
 102         }
 103     }
 104     else
 105     {
 106         if (Attr & FILE_ATTRIBUTE_DIRECTORY)
 107         {
 108             return FALSE;
 109         }
 110     }
 111
 112     return TRUE;
 113 }
 114
 115 int SaveBufToFile(const CHAR *FileName, const void *Buffer, int BufLen)
 116 {
 117     FILE *File = NULL;
 118     void *Data = NULL;
 119
 120     fopen_s(&File, FileName, "wb");
 121     if (File == NULL)
 122     {
 123         Log("Failed to open file %s", FileName);
 124         return 1;
 125     }
 126
 127     fwrite(Buffer, 1, BufLen, File);
 128     fclose(File);
 129     return 0;
 130 }
 131
 132 int ReadWholeFileToBuf(const CHAR *FileName, int ExtLen, void **Bufer, int *BufLen)
 133 {
 134     int FileSize;
 135     FILE *File = NULL;
 136     void *Data = NULL;
 137
 138     fopen_s(&File, FileName, "rb");
 139     if (File == NULL)
 140     {
 141         Log("Failed to open file %s", FileName);
 142         return 1;
 143     }
 144
 145     fseek(File, 0, SEEK_END);
 146     FileSize = (int)ftell(File);
 147
 148     Data = malloc(FileSize + ExtLen);
 149     if (!Data)
 150     {
 151         fclose(File);
 152         return 1;
 153     }
 154
 155     fseek(File, 0, SEEK_SET);
 156     fread(Data, 1, FileSize, File);
 157
 158     fclose(File);
 159
 160     *Bufer = Data;
 161     *BufLen = FileSize;
 162
 163     return 0;
 164 }
 165
 166 const CHAR* GetLocalVentoyVersion(void)
 167 {
 168     int rc;
 169     int FileSize;
 170     CHAR *Pos = NULL;
 171     CHAR *Buf = NULL;
 172     static CHAR LocalVersion[64] = { 0 };
 173
 174     if (LocalVersion[0] == 0)
 175     {
 176         rc = ReadWholeFileToBuf(VENTOY_FILE_VERSION, 1, (void **)&Buf, &FileSize);
 177         if (rc)
 178         {
 179             return "";
 180         }
 181         Buf[FileSize] = 0;
 182
 183         for (Pos = Buf; *Pos; Pos++)
 184         {
 185             if (*Pos == '\r' || *Pos == '\n')
 186             {
 187                 *Pos = 0;
 188                 break;
 189             }
 190         }
 191
 192         safe_sprintf(LocalVersion, "%s", Buf);
 193         free(Buf);
 194     }
 195
 196     return LocalVersion;
 197 }
 198
 199 const CHAR* ParseVentoyVersionFromString(CHAR *Buf)
 200 {
 201     CHAR *Pos = NULL;
 202     CHAR *End = NULL;
 203     static CHAR LocalVersion[64] = { 0 };
 204
 205     Pos = strstr(Buf, "VENTOY_VERSION=");
 206     if (Pos)
 207     {
 208         Pos += strlen("VENTOY_VERSION=");
 209         if (*Pos == '"')
 210         {
 211             Pos++;
 212         }
 213
 214         End = Pos;
 215         while (*End != 0 && *End != '"' && *End != '\r' && *End != '\n')
 216         {
 217             End++;
 218         }
 219
 220         *End = 0;
 221
 222         safe_sprintf(LocalVersion, "%s", Pos);
 223         return LocalVersion;
 224     }
 225
 226     return "";
 227 }
 228
 229 BOOL IsWow64(void)
 230 {
 231     typedef BOOL(WINAPI *LPFN_ISWOW64PROCESS)(HANDLE, PBOOL);
 232     LPFN_ISWOW64PROCESS fnIsWow64Process;
 233     BOOL bIsWow64 = FALSE;
 234
 235     fnIsWow64Process = (LPFN_ISWOW64PROCESS)GetProcAddress(GetModuleHandleA("kernel32"), "IsWow64Process");
 236     if (NULL != fnIsWow64Process)
 237     {
 238         fnIsWow64Process(GetCurrentProcess(), &bIsWow64);
 239     }
 240
 241     return bIsWow64;
 242 }
 243
 244 /*
 245 * Some code and functions in the file are copied from rufus.
 246 * https://github.com/pbatard/rufus
 247 */
 248
 249 /* Windows versions */
 250 enum WindowsVersion {
 251     WINDOWS_UNDEFINED = -1,
 252     WINDOWS_UNSUPPORTED = 0,
 253     WINDOWS_XP = 0x51,
 254     WINDOWS_2003 = 0x52,        // Also XP_64
 255     WINDOWS_VISTA = 0x60,       // Also Server 2008
 256     WINDOWS_7 = 0x61,           // Also Server 2008_R2
 257     WINDOWS_8 = 0x62,           // Also Server 2012
 258     WINDOWS_8_1 = 0x63,         // Also Server 2012_R2
 259     WINDOWS_10_PREVIEW1 = 0x64,
 260     WINDOWS_10 = 0xA0,          // Also Server 2016, also Server 2019
 261     WINDOWS_11 = 0xB0,          // Also Server 2022
 262     WINDOWS_MAX
 263 };
 264
 265 static const char* GetEdition(DWORD ProductType)
 266 {
 267     // From: https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-getproductinfo
 268     // These values can be found in the winnt.h header.
 269     switch (ProductType) {
 270     case 0x00000000: return ""; //  Undefined
 271     case 0x00000001: return "Ultimate";
 272     case 0x00000002: return "Home Basic";
 273     case 0x00000003: return "Home Premium";
 274     case 0x00000004: return "Enterprise";
 275     case 0x00000005: return "Home Basic N";
 276     case 0x00000006: return "Business";
 277     case 0x00000007: return "Standard Server";
 278     case 0x00000008: return "Datacenter Server";
 279     case 0x00000009: return "Smallbusiness Server";
 280     case 0x0000000A: return "Enterprise Server";
 281     case 0x0000000B: return "Starter";
 282     case 0x00000010: return "Business N";
 283     case 0x00000011: return "Web Server";
 284     case 0x00000012: return "Cluster Server";
 285     case 0x00000013: return "Home Server";
 286     case 0x0000001A: return "Home Premium N";
 287     case 0x0000001B: return "Enterprise N";
 288     case 0x0000001C: return "Ultimate N";
 289     case 0x00000022: return "Home Premium Server";
 290     case 0x0000002F: return "Starter N";
 291     case 0x00000030: return "Pro";
 292     case 0x00000031: return "Pro N";
 293     case 0x00000042: return "Starter E";
 294     case 0x00000043: return "Home Basic E";
 295     case 0x00000044: return "Premium E";
 296     case 0x00000045: return "Pro E";
 297     case 0x00000046: return "Enterprise E";
 298     case 0x00000047: return "Ultimate E";
 299     case 0x00000048: return "Enterprise Eval";
 300     case 0x00000054: return "Enterprise N Eval";
 301     case 0x00000057: return "Thin PC";
 302     case 0x0000006F: return "Core Connected";
 303     case 0x00000070: return "Pro Student";
 304     case 0x00000071: return "Core Connected N";
 305     case 0x00000072: return "Pro Student N";
 306     case 0x00000073: return "Core Connected Single Language";
 307     case 0x00000074: return "Core Connected China";
 308     case 0x00000079: return "Edu";
 309     case 0x0000007A: return "Edu N";
 310     case 0x0000007D: return "Enterprise S";
 311     case 0x0000007E: return "Enterprise S N";
 312     case 0x0000007F: return "Pro S";
 313     case 0x00000080: return "Pro S N";
 314     case 0x00000081: return "Enterprise S Eval";
 315     case 0x00000082: return "Enterprise S N Eval";
 316     case 0x0000008A: return "Pro Single Language";
 317     case 0x0000008B: return "Pro China";
 318     case 0x0000008C: return "Enterprise Subscription";
 319     case 0x0000008D: return "Enterprise Subscription N";
 320     case 0x00000095: return "Utility VM";
 321     case 0x000000A1: return "Pro Workstation";
 322     case 0x000000A2: return "Pro Workstation N";
 323     case 0x000000A4: return "Pro for Education";
 324     case 0x000000A5: return "Pro for Education N";
 325     case 0x000000AB: return "Enterprise G";     // I swear Microsoft are just making up editions...
 326     case 0x000000AC: return "Enterprise G N";
 327     case 0x000000B6: return "Core OS";
 328     case 0x000000B7: return "Cloud E";
 329     case 0x000000B8: return "Cloud E N";
 330     case 0x000000BD: return "Lite";
 331     case 0xABCDABCD: return "(Unlicensed)";
 332     default: return "(Unknown Edition)";
 333     }
 334 }
 335
 336 #define is_x64 IsWow64
 337 #define static_strcpy safe_strcpy
 338 #define REGKEY_HKCU HKEY_CURRENT_USER
 339 #define REGKEY_HKLM HKEY_LOCAL_MACHINE
 340 static int  nWindowsVersion = WINDOWS_UNDEFINED;
 341 static int  nWindowsBuildNumber = -1;
 342 static char WindowsVersionStr[128] = "";
 343
 344 /* Helpers for 32 bit registry operations */
 345
 346 /*
 347 * Read a generic registry key value. If a short key_name is used, assume that
 348 * it belongs to the application and create the app subkey if required
 349 */
 350 static __inline BOOL _GetRegistryKey(HKEY key_root, const char* key_name, DWORD reg_type,
 351     LPBYTE dest, DWORD dest_size)
 352 {
 353     const char software_prefix[] = "SOFTWARE\\";
 354     char long_key_name[MAX_PATH] = { 0 };
 355     BOOL r = FALSE;
 356     size_t i;
 357     LONG s;
 358     HKEY hSoftware = NULL, hApp = NULL;
 359     DWORD dwType = -1, dwSize = dest_size;
 360
 361     memset(dest, 0, dest_size);
 362
 363     if (key_name == NULL)
 364         return FALSE;
 365
 366     for (i = strlen(key_name); i>0; i--) {
 367         if (key_name[i] == '\\')
 368             break;
 369     }
 370
 371     if (i > 0) {
 372         // Prefix with "SOFTWARE" if needed
 373         if (_strnicmp(key_name, software_prefix, sizeof(software_prefix)-1) != 0) {
 374             if (i + sizeof(software_prefix) >= sizeof(long_key_name))
 375                 return FALSE;
 376             strcpy_s(long_key_name, sizeof(long_key_name), software_prefix);
 377             strcat_s(long_key_name, sizeof(long_key_name), key_name);
 378             long_key_name[sizeof(software_prefix)+i - 1] = 0;
 379         }
 380         else {
 381             if (i >= sizeof(long_key_name))
 382                 return FALSE;
 383             static_strcpy(long_key_name, key_name);
 384             long_key_name[i] = 0;
 385         }
 386         i++;
 387         if (RegOpenKeyExA(key_root, long_key_name, 0, KEY_READ, &hApp) != ERROR_SUCCESS) {
 388             hApp = NULL;
 389             goto out;
 390         }
 391     }
 392     else {
 393         if (RegOpenKeyExA(key_root, "SOFTWARE", 0, KEY_READ | KEY_CREATE_SUB_KEY, &hSoftware) != ERROR_SUCCESS) {
 394             hSoftware = NULL;
 395             goto out;
 396         }
 397     }
 398
 399     s = RegQueryValueExA(hApp, &key_name[i], NULL, &dwType, (LPBYTE)dest, &dwSize);
 400     // No key means default value of 0 or empty string
 401     if ((s == ERROR_FILE_NOT_FOUND) || ((s == ERROR_SUCCESS) && (dwType == reg_type) && (dwSize > 0))) {
 402         r = TRUE;
 403     }
 404 out:
 405     if (hSoftware != NULL)
 406         RegCloseKey(hSoftware);
 407     if (hApp != NULL)
 408         RegCloseKey(hApp);
 409     return r;
 410 }
 411
 412 #define GetRegistryKey32(root, key, pval) _GetRegistryKey(root, key, REG_DWORD, (LPBYTE)pval, sizeof(DWORD))
 413 static __inline INT32 ReadRegistryKey32(HKEY root, const char* key) {
 414     DWORD val;
 415     GetRegistryKey32(root, key, &val);
 416     return (INT32)val;
 417 }
 418
 419 /*
 420 * Modified from smartmontools' os_win32.cpp
 421 */
 422 void GetWindowsVersion(void)
 423 {
 424     OSVERSIONINFOEXA vi, vi2;
 425     DWORD dwProductType;
 426     const char* w = 0;
 427     const char* w64 = "32 bit";
 428     char *vptr;
 429     size_t vlen;
 430     unsigned major, minor;
 431     ULONGLONG major_equal, minor_equal;
 432     BOOL ws;
 433
 434     nWindowsVersion = WINDOWS_UNDEFINED;
 435     static_strcpy(WindowsVersionStr, "Windows Undefined");
 436
 437     // suppress the C4996 warning for GetVersionExA
 438     #pragma warning(push)
 439     #pragma warning(disable:4996)
 440
 441     memset(&vi, 0, sizeof(vi));
 442     vi.dwOSVersionInfoSize = sizeof(vi);
 443     if (!GetVersionExA((OSVERSIONINFOA *)&vi)) {
 444         memset(&vi, 0, sizeof(vi));
 445         vi.dwOSVersionInfoSize = sizeof(OSVERSIONINFOA);
 446         if (!GetVersionExA((OSVERSIONINFOA *)&vi))
 447             return;
 448     }
 449
 450     #pragma warning(pop)
 451
 452     if (vi.dwPlatformId == VER_PLATFORM_WIN32_NT) {
 453
 454         if (vi.dwMajorVersion > 6 || (vi.dwMajorVersion == 6 && vi.dwMinorVersion >= 2)) {
 455             // Starting with Windows 8.1 Preview, GetVersionEx() does no longer report the actual OS version
 456             // See: http://msdn.microsoft.com/en-us/library/windows/desktop/dn302074.aspx
 457             // And starting with Windows 10 Preview 2, Windows enforces the use of the application/supportedOS
 458             // manifest in order for VerSetConditionMask() to report the ACTUAL OS major and minor...
 459
 460             major_equal = VerSetConditionMask(0, VER_MAJORVERSION, VER_EQUAL);
 461             for (major = vi.dwMajorVersion; major <= 9; major++) {
 462                 memset(&vi2, 0, sizeof(vi2));
 463                 vi2.dwOSVersionInfoSize = sizeof(vi2); vi2.dwMajorVersion = major;
 464                 if (!VerifyVersionInfoA(&vi2, VER_MAJORVERSION, major_equal))
 465                     continue;
 466                 if (vi.dwMajorVersion < major) {
 467                     vi.dwMajorVersion = major; vi.dwMinorVersion = 0;
 468                 }
 469
 470                 minor_equal = VerSetConditionMask(0, VER_MINORVERSION, VER_EQUAL);
 471                 for (minor = vi.dwMinorVersion; minor <= 9; minor++) {
 472                     memset(&vi2, 0, sizeof(vi2)); vi2.dwOSVersionInfoSize = sizeof(vi2);
 473                     vi2.dwMinorVersion = minor;
 474                     if (!VerifyVersionInfoA(&vi2, VER_MINORVERSION, minor_equal))
 475                         continue;
 476                     vi.dwMinorVersion = minor;
 477                     break;
 478                 }
 479
 480                 break;
 481             }
 482         }
 483
 484         if (vi.dwMajorVersion <= 0xf && vi.dwMinorVersion <= 0xf) {
 485             ws = (vi.wProductType <= VER_NT_WORKSTATION);
 486             nWindowsVersion = vi.dwMajorVersion << 4 | vi.dwMinorVersion;
 487             switch (nWindowsVersion) {
 488             case WINDOWS_XP: w = "XP";
 489                 break;
 490             case WINDOWS_2003: w = (ws ? "XP_64" : (!GetSystemMetrics(89) ? "Server 2003" : "Server 2003_R2"));
 491                 break;
 492             case WINDOWS_VISTA: w = (ws ? "Vista" : "Server 2008");
 493                 break;
 494             case WINDOWS_7: w = (ws ? "7" : "Server 2008_R2");
 495                 break;
 496             case WINDOWS_8: w = (ws ? "8" : "Server 2012");
 497                 break;
 498             case WINDOWS_8_1: w = (ws ? "8.1" : "Server 2012_R2");
 499                 break;
 500             case WINDOWS_10_PREVIEW1: w = (ws ? "10 (Preview 1)" : "Server 10 (Preview 1)");
 501                 break;
 502                 // Starting with Windows 10 Preview 2, the major is the same as the public-facing version
 503             case WINDOWS_10:
 504                 if (vi.dwBuildNumber < 20000) {
 505                     w = (ws ? "10" : ((vi.dwBuildNumber < 17763) ? "Server 2016" : "Server 2019"));
 506                     break;
 507                 }
 508                 nWindowsVersion = WINDOWS_11;
 509                 // Fall through
 510             case WINDOWS_11: w = (ws ? "11" : "Server 2022");
 511                 break;
 512             default:
 513                 if (nWindowsVersion < WINDOWS_XP)
 514                     nWindowsVersion = WINDOWS_UNSUPPORTED;
 515                 else
 516                     w = "12 or later";
 517                 break;
 518             }
 519         }
 520     }
 521
 522     if (is_x64())
 523         w64 = "64-bit";
 524
 525     GetProductInfo(vi.dwMajorVersion, vi.dwMinorVersion, vi.wServicePackMajor, vi.wServicePackMinor, &dwProductType);
 526     vptr = WindowsVersionStr;
 527     vlen = sizeof(WindowsVersionStr) - 1;
 528
 529     if (!w)
 530         sprintf_s(vptr, vlen, "%s %u.%u %s", (vi.dwPlatformId == VER_PLATFORM_WIN32_NT ? "NT" : "??"),
 531         (unsigned)vi.dwMajorVersion, (unsigned)vi.dwMinorVersion, w64);
 532     else if (vi.wServicePackMinor)
 533         sprintf_s(vptr, vlen, "%s SP%u.%u %s", w, vi.wServicePackMajor, vi.wServicePackMinor, w64);
 534     else if (vi.wServicePackMajor)
 535         sprintf_s(vptr, vlen, "%s SP%u %s", w, vi.wServicePackMajor, w64);
 536     else
 537         sprintf_s(vptr, vlen, "%s%s%s, %s",
 538         w, (dwProductType != PRODUCT_UNDEFINED) ? " " : "", GetEdition(dwProductType), w64);
 539
 540     // Add the build number (including UBR if available) for Windows 8.0 and later
 541     nWindowsBuildNumber = vi.dwBuildNumber;
 542     if (nWindowsVersion >= 0x62) {
 543         int nUbr = ReadRegistryKey32(REGKEY_HKLM, "Software\\Microsoft\\Windows NT\\CurrentVersion\\UBR");
 544         vptr = WindowsVersionStr + strlen(WindowsVersionStr);
 545         vlen = sizeof(WindowsVersionStr) - strlen(WindowsVersionStr) - 1;
 546         if (nUbr > 0)
 547             sprintf_s(vptr, vlen, " (Build %d.%d)", nWindowsBuildNumber, nUbr);
 548         else
 549             sprintf_s(vptr, vlen, " (Build %d)", nWindowsBuildNumber);
 550     }
 551 }
 552
 553
 554
 555 void DumpWindowsVersion(void)
 556 {
 557     GetWindowsVersion();
 558     Log("Windows Version: <<Windows %s>>", WindowsVersionStr);
 559     return;
 560 }
 561
 562 BOOL IsVentoyLogicalDrive(CHAR DriveLetter)
 563 {
 564     int i;
 565     CONST CHAR *Files[] =
 566     {
 567         "EFI\\BOOT\\BOOTX64.EFI",
 568         "grub\\themes\\ventoy\\theme.txt",
 569         "ventoy\\ventoy.cpio",
 570     };
 571
 572     for (i = 0; i < sizeof(Files) / sizeof(Files[0]); i++)
 573     {
 574         if (!IsFileExist("%C:\\%s", DriveLetter, Files[i]))
 575         {
 576             return FALSE;
 577         }
 578     }
 579
 580     return TRUE;
 581 }
 582
 583
 584 int VentoyFillMBRLocation(UINT64 DiskSizeInBytes, UINT32 StartSectorId, UINT32 SectorCount, PART_TABLE *Table)
 585 {
 586     BYTE Head;
 587     BYTE Sector;
 588     BYTE nSector = 63;
 589     BYTE nHead = 8;
 590     UINT32 Cylinder;
 591     UINT32 EndSectorId;
 592
 593     while (nHead != 0 && (DiskSizeInBytes / 512 / nSector / nHead) > 1024)
 594     {
 595         nHead = (BYTE)nHead * 2;
 596     }
 597
 598     if (nHead == 0)
 599     {
 600         nHead = 255;
 601     }
 602
 603     Cylinder = StartSectorId / nSector / nHead;
 604     Head = StartSectorId / nSector % nHead;
 605     Sector = StartSectorId % nSector + 1;
 606
 607     Table->StartHead = Head;
 608     Table->StartSector = Sector;
 609     Table->StartCylinder = Cylinder;
 610
 611     EndSectorId = StartSectorId + SectorCount - 1;
 612     Cylinder = EndSectorId / nSector / nHead;
 613     Head = EndSectorId / nSector % nHead;
 614     Sector = EndSectorId % nSector + 1;
 615
 616     Table->EndHead = Head;
 617     Table->EndSector = Sector;
 618     Table->EndCylinder = Cylinder;
 619
 620     Table->StartSectorId = StartSectorId;
 621     Table->SectorCount = SectorCount;
 622
 623     return 0;
 624 }
 625
 626 int VentoyFillMBR(UINT64 DiskSizeBytes, MBR_HEAD *pMBR, int PartStyle)
 627 {
 628     GUID Guid;
 629         int ReservedValue;
 630     UINT32 DiskSignature;
 631     UINT32 DiskSectorCount;
 632     UINT32 PartSectorCount;
 633     UINT32 PartStartSector;
 634         UINT32 ReservedSector;
 635
 636     VentoyGetLocalBootImg(pMBR);
 637
 638     CoCreateGuid(&Guid);
 639
 640     memcpy(&DiskSignature, &Guid, sizeof(UINT32));
 641
 642     Log("Disk signature: 0x%08x", DiskSignature);
 643
 644     *((UINT32 *)(pMBR->BootCode + 0x1B8)) = DiskSignature;
 645
 646     if (DiskSizeBytes / 512 > 0xFFFFFFFF)
 647     {
 648         DiskSectorCount = 0xFFFFFFFF;
 649     }
 650     else
 651     {
 652         DiskSectorCount = (UINT32)(DiskSizeBytes / 512);
 653     }
 654
 655         ReservedValue = GetReservedSpaceInMB();
 656         if (ReservedValue <= 0)
 657         {
 658                 ReservedSector = 0;
 659         }
 660         else
 661         {
 662                 ReservedSector = (UINT32)(ReservedValue * 2048);
 663         }
 664
 665     if (PartStyle)
 666     {
 667         ReservedSector += 33; // backup GPT part table
 668     }
 669
 670     // check aligned with 4KB
 671     if (IsPartNeed4KBAlign())
 672     {
 673         UINT64 sectors = DiskSizeBytes / 512;
 674         if (sectors % 8)
 675         {
 676             Log("Disk need to align with 4KB %u", (UINT32)(sectors % 8));
 677             ReservedSector += (UINT32)(sectors % 8);
 678         }
 679     }
 680
 681         Log("ReservedSector: %u", ReservedSector);
 682
 683     //Part1
 684     PartStartSector = VENTOY_PART1_START_SECTOR;
 685         PartSectorCount = DiskSectorCount - ReservedSector - VENTOY_EFI_PART_SIZE / 512 - PartStartSector;
 686     VentoyFillMBRLocation(DiskSizeBytes, PartStartSector, PartSectorCount, pMBR->PartTbl);
 687
 688     pMBR->PartTbl[0].Active = 0x80; // bootable
 689     pMBR->PartTbl[0].FsFlag = 0x07; // exFAT/NTFS/HPFS
 690
 691     //Part2
 692     PartStartSector += PartSectorCount;
 693     PartSectorCount = VENTOY_EFI_PART_SIZE / 512;
 694     VentoyFillMBRLocation(DiskSizeBytes, PartStartSector, PartSectorCount, pMBR->PartTbl + 1);
 695
 696     pMBR->PartTbl[1].Active = 0x00;
 697     pMBR->PartTbl[1].FsFlag = 0xEF; // EFI System Partition
 698
 699     pMBR->Byte55 = 0x55;
 700     pMBR->ByteAA = 0xAA;
 701
 702     return 0;
 703 }
 704
 705
 706 static int VentoyFillProtectMBR(UINT64 DiskSizeBytes, MBR_HEAD *pMBR)
 707 {
 708     GUID Guid;
 709     UINT32 DiskSignature;
 710     UINT64 DiskSectorCount;
 711
 712     VentoyGetLocalBootImg(pMBR);
 713
 714     CoCreateGuid(&Guid);
 715
 716     memcpy(&DiskSignature, &Guid, sizeof(UINT32));
 717
 718     Log("Disk signature: 0x%08x", DiskSignature);
 719
 720     *((UINT32 *)(pMBR->BootCode + 0x1B8)) = DiskSignature;
 721
 722     DiskSectorCount = DiskSizeBytes / 512 - 1;
 723     if (DiskSectorCount > 0xFFFFFFFF)
 724     {
 725         DiskSectorCount = 0xFFFFFFFF;
 726     }
 727
 728     memset(pMBR->PartTbl, 0, sizeof(pMBR->PartTbl));
 729
 730     pMBR->PartTbl[0].Active = 0x00;
 731     pMBR->PartTbl[0].FsFlag = 0xee; // EE
 732
 733     pMBR->PartTbl[0].StartHead = 0;
 734     pMBR->PartTbl[0].StartSector = 1;
 735     pMBR->PartTbl[0].StartCylinder = 0;
 736     pMBR->PartTbl[0].EndHead = 254;
 737     pMBR->PartTbl[0].EndSector = 63;
 738     pMBR->PartTbl[0].EndCylinder = 1023;
 739
 740     pMBR->PartTbl[0].StartSectorId = 1;
 741     pMBR->PartTbl[0].SectorCount = (UINT32)DiskSectorCount;
 742
 743     pMBR->Byte55 = 0x55;
 744     pMBR->ByteAA = 0xAA;
 745
 746     pMBR->BootCode[92] = 0x22;
 747
 748     return 0;
 749 }
 750
 751 int VentoyFillWholeGpt(UINT64 DiskSizeBytes, VTOY_GPT_INFO *pInfo)
 752 {
 753     UINT64 Part1SectorCount = 0;
 754     UINT64 DiskSectorCount = DiskSizeBytes / 512;
 755     VTOY_GPT_HDR *Head = &pInfo->Head;
 756     VTOY_GPT_PART_TBL *Table = pInfo->PartTbl;
 757     static GUID WindowsDataPartType = { 0xebd0a0a2, 0xb9e5, 0x4433, { 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7 } };
 758
 759     VentoyFillProtectMBR(DiskSizeBytes, &pInfo->MBR);
 760
 761     Part1SectorCount = DiskSectorCount - 33 - 2048;
 762
 763     memcpy(Head->Signature, "EFI PART", 8);
 764     Head->Version[2] = 0x01;
 765     Head->Length = 92;
 766     Head->Crc = 0;
 767     Head->EfiStartLBA = 1;
 768     Head->EfiBackupLBA = DiskSectorCount - 1;
 769     Head->PartAreaStartLBA = 34;
 770     Head->PartAreaEndLBA = DiskSectorCount - 34;
 771     CoCreateGuid(&Head->DiskGuid);
 772     Head->PartTblStartLBA = 2;
 773     Head->PartTblTotNum = 128;
 774     Head->PartTblEntryLen = 128;
 775
 776
 777     memcpy(&(Table[0].PartType), &WindowsDataPartType, sizeof(GUID));
 778     CoCreateGuid(&(Table[0].PartGuid));
 779     Table[0].StartLBA = 2048;
 780     Table[0].LastLBA = 2048 + Part1SectorCount - 1;
 781     Table[0].Attr = 0;
 782     memcpy(Table[0].Name, L"Data", 4 * 2);
 783
 784     //Update CRC
 785     Head->PartTblCrc = VentoyCrc32(Table, sizeof(pInfo->PartTbl));
 786     Head->Crc = VentoyCrc32(Head, Head->Length);
 787
 788     return 0;
 789 }
 790
 791 int VentoyFillGpt(UINT64 DiskSizeBytes, VTOY_GPT_INFO *pInfo)
 792 {
 793     INT64 ReservedValue = 0;
 794     UINT64 ModSectorCount = 0;
 795     UINT64 ReservedSector = 33;
 796     UINT64 Part1SectorCount = 0;
 797     UINT64 DiskSectorCount = DiskSizeBytes / 512;
 798     VTOY_GPT_HDR *Head = &pInfo->Head;
 799     VTOY_GPT_PART_TBL *Table = pInfo->PartTbl;
 800     static GUID WindowsDataPartType = { 0xebd0a0a2, 0xb9e5, 0x4433, { 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7 } };
 801     static GUID EspPartType = { 0xc12a7328, 0xf81f, 0x11d2, { 0xba, 0x4b, 0x00, 0xa0, 0xc9, 0x3e, 0xc9, 0x3b } };
 802         static GUID BiosGrubPartType = { 0x21686148, 0x6449, 0x6e6f, { 0x74, 0x4e, 0x65, 0x65, 0x64, 0x45, 0x46, 0x49 } };
 803
 804     VentoyFillProtectMBR(DiskSizeBytes, &pInfo->MBR);
 805
 806     ReservedValue = GetReservedSpaceInMB();
 807     if (ReservedValue > 0)
 808     {
 809         ReservedSector += ReservedValue * 2048;
 810     }
 811
 812     Part1SectorCount = DiskSectorCount - ReservedSector - (VENTOY_EFI_PART_SIZE / 512) - 2048;
 813
 814     ModSectorCount = (Part1SectorCount % 8);
 815     if (ModSectorCount)
 816     {
 817         Log("Part1SectorCount:%llu is not aligned by 4KB (%llu)", (ULONGLONG)Part1SectorCount, (ULONGLONG)ModSectorCount);
 818     }
 819
 820     // check aligned with 4KB
 821     if (IsPartNeed4KBAlign())
 822     {
 823         if (ModSectorCount)
 824         {
 825             Log("Disk need to align with 4KB %u", (UINT32)ModSectorCount);
 826             Part1SectorCount -= ModSectorCount;
 827         }
 828         else
 829         {
 830             Log("no need to align with 4KB");
 831         }
 832     }
 833
 834     memcpy(Head->Signature, "EFI PART", 8);
 835     Head->Version[2] = 0x01;
 836     Head->Length = 92;
 837     Head->Crc = 0;
 838     Head->EfiStartLBA = 1;
 839     Head->EfiBackupLBA = DiskSectorCount - 1;
 840     Head->PartAreaStartLBA = 34;
 841     Head->PartAreaEndLBA = DiskSectorCount - 34;
 842     CoCreateGuid(&Head->DiskGuid);
 843     Head->PartTblStartLBA = 2;
 844     Head->PartTblTotNum = 128;
 845     Head->PartTblEntryLen = 128;
 846
 847
 848     memcpy(&(Table[0].PartType), &WindowsDataPartType, sizeof(GUID));
 849     CoCreateGuid(&(Table[0].PartGuid));
 850     Table[0].StartLBA = 2048;
 851     Table[0].LastLBA = 2048 + Part1SectorCount - 1;
 852     Table[0].Attr = 0;
 853     memcpy(Table[0].Name, L"Ventoy", 6 * 2);
 854
 855     // to fix windows issue
 856     //memcpy(&(Table[1].PartType), &EspPartType, sizeof(GUID));
 857     memcpy(&(Table[1].PartType), &WindowsDataPartType, sizeof(GUID));
 858     CoCreateGuid(&(Table[1].PartGuid));
 859     Table[1].StartLBA = Table[0].LastLBA + 1;
 860     Table[1].LastLBA = Table[1].StartLBA + VENTOY_EFI_PART_SIZE / 512 - 1;
 861     Table[1].Attr = 0xC000000000000001ULL;
 862     memcpy(Table[1].Name, L"VTOYEFI", 7 * 2);
 863
 864 #if 0
 865         memcpy(&(Table[2].PartType), &BiosGrubPartType, sizeof(GUID));
 866         CoCreateGuid(&(Table[2].PartGuid));
 867         Table[2].StartLBA = 34;
 868         Table[2].LastLBA = 2047;
 869         Table[2].Attr = 0;
 870 #endif
 871
 872     //Update CRC
 873     Head->PartTblCrc = VentoyCrc32(Table, sizeof(pInfo->PartTbl));
 874     Head->Crc = VentoyCrc32(Head, Head->Length);
 875
 876     return 0;
 877 }
 878
 879 int VentoyFillBackupGptHead(VTOY_GPT_INFO *pInfo, VTOY_GPT_HDR *pHead)
 880 {
 881     UINT64 LBA;
 882     UINT64 BackupLBA;
 883
 884     memcpy(pHead, &pInfo->Head, sizeof(VTOY_GPT_HDR));
 885
 886     LBA = pHead->EfiStartLBA;
 887     BackupLBA = pHead->EfiBackupLBA;
 888
 889     pHead->EfiStartLBA = BackupLBA;
 890     pHead->EfiBackupLBA = LBA;
 891     pHead->PartTblStartLBA = BackupLBA + 1 - 33;
 892
 893     pHead->Crc = 0;
 894     pHead->Crc = VentoyCrc32(pHead, pHead->Length);
 895
 896     return 0;
 897 }
 898
 899 CHAR GetFirstUnusedDriveLetter(void)
 900 {
 901     CHAR Letter = 'D';
 902     DWORD Drives = GetLogicalDrives();
 903
 904     Drives >>= 3;
 905     while (Drives & 0x1)
 906     {
 907         Letter++;
 908         Drives >>= 1;
 909     }
 910
 911     return Letter;
 912 }
 913
 914 const CHAR * GetBusTypeString(STORAGE_BUS_TYPE Type)
 915 {
 916     switch (Type)
 917     {
 918         case BusTypeUnknown: return "unknown";
 919         case BusTypeScsi: return "SCSI";
 920         case BusTypeAtapi: return "Atapi";
 921         case BusTypeAta: return "ATA";
 922         case BusType1394: return "1394";
 923         case BusTypeSsa: return "SSA";
 924         case BusTypeFibre: return "Fibre";
 925         case BusTypeUsb: return "USB";
 926         case BusTypeRAID: return "RAID";
 927         case BusTypeiScsi: return "iSCSI";
 928         case BusTypeSas: return "SAS";
 929         case BusTypeSata: return "SATA";
 930         case BusTypeSd: return "SD";
 931         case BusTypeMmc: return "MMC";
 932         case BusTypeVirtual: return "Virtual";
 933         case BusTypeFileBackedVirtual: return "FileBackedVirtual";
 934         case BusTypeSpaces: return "Spaces";
 935         case BusTypeNvme: return "Nvme";
 936     }
 937     return "unknown";
 938 }
 939
 940 int VentoyGetLocalBootImg(MBR_HEAD *pMBR)
 941 {
 942     int Len = 0;
 943     BYTE *ImgBuf = NULL;
 944     static int Loaded = 0;
 945     static MBR_HEAD MBR;
 946
 947     if (Loaded)
 948     {
 949         memcpy(pMBR, &MBR, 512);
 950         return 0;
 951     }
 952
 953     if (0 == ReadWholeFileToBuf(VENTOY_FILE_BOOT_IMG, 0, (void **)&ImgBuf, &Len))
 954     {
 955         Log("Copy boot img success");
 956         memcpy(pMBR, ImgBuf, 512);
 957         free(ImgBuf);
 958
 959         CoCreateGuid((GUID *)(pMBR->BootCode + 0x180));
 960
 961         memcpy(&MBR, pMBR, 512);
 962         Loaded = 1;
 963
 964         return 0;
 965     }
 966     else
 967     {
 968         Log("Copy boot img failed");
 969         return 1;
 970     }
 971 }
 972
 973 int GetHumanReadableGBSize(UINT64 SizeBytes)
 974 {
 975     int i;
 976     int Pow2 = 1;
 977     double Delta;
 978     double GB = SizeBytes * 1.0 / 1000 / 1000 / 1000;
 979
 980     if ((SizeBytes % 1073741824) == 0)
 981     {
 982         return (int)(SizeBytes / 1073741824);
 983     }
 984
 985     for (i = 0; i < 12; i++)
 986     {
 987         if (Pow2 > GB)
 988         {
 989             Delta = (Pow2 - GB) / Pow2;
 990         }
 991         else
 992         {
 993             Delta = (GB - Pow2) / Pow2;
 994         }
 995
 996         if (Delta < 0.05)
 997         {
 998             return Pow2;
 999         }
1000
1001         Pow2 <<= 1;
1002     }
1003
1004     return (int)GB;
1005 }
1006
1007 void TrimString(CHAR *String)
1008 {
1009     CHAR *Pos1 = String;
1010     CHAR *Pos2 = String;
1011     size_t Len = strlen(String);
1012
1013     while (Len > 0)
1014     {
1015         if (String[Len - 1] != ' ' && String[Len - 1] != '\t')
1016         {
1017             break;
1018         }
1019         String[Len - 1] = 0;
1020         Len--;
1021     }
1022
1023     while (*Pos1 == ' ' || *Pos1 == '\t')
1024     {
1025         Pos1++;
1026     }
1027
1028     while (*Pos1)
1029     {
1030         *Pos2++ = *Pos1++;
1031     }
1032     *Pos2++ = 0;
1033
1034     return;
1035 }
1036
1037 int GetRegDwordValue(HKEY Key, LPCSTR SubKey, LPCSTR ValueName, DWORD *pValue)
1038 {
1039     HKEY hKey;
1040     DWORD Type;
1041     DWORD Size;
1042     LSTATUS lRet;
1043     DWORD Value;
1044
1045     lRet = RegOpenKeyExA(Key, SubKey, 0, KEY_QUERY_VALUE, &hKey);
1046     Log("RegOpenKeyExA <%s> Ret:%ld", SubKey, lRet);
1047
1048     if (ERROR_SUCCESS == lRet)
1049     {
1050         Size = sizeof(Value);
1051         lRet = RegQueryValueExA(hKey, ValueName, NULL, &Type, (LPBYTE)&Value, &Size);
1052         Log("RegQueryValueExA <%s> ret:%u  Size:%u Value:%u", ValueName, lRet, Size, Value);
1053
1054         *pValue = Value;
1055         RegCloseKey(hKey);
1056
1057         return 0;
1058     }
1059     else
1060     {
1061         return 1;
1062     }
1063 }
1064
1065 int GetPhysicalDriveCount(void)
1066 {
1067     DWORD Value;
1068     int Count = 0;
1069
1070     if (GetRegDwordValue(HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\disk\\Enum", "Count", &Value) == 0)
1071     {
1072         Count = (int)Value;
1073     }
1074
1075     Log("GetPhysicalDriveCount: %d", Count);
1076     return Count;
1077 }
1078
1079
1080