Ventoy2Disk/Ventoy2Disk/Ventoy2Disk.c

   1 /******************************************************************************
   2  * Ventoy2Disk.c
   3  *
   4  * Copyright (c) 2020, longpanda <admin@ventoy.net>
   5  *
   6  * This program is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU General Public License as
   8  * published by the Free Software Foundation; either version 3 of the
   9  * License, or (at your option) any later version.
  10  *
  11  * This program is distributed in the hope that it will be useful, but
  12  * WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, see <http://www.gnu.org/licenses/>.
  18  *
  19  */
  20
  21 #include <Windows.h>
  22 #include "resource.h"
  23 #include "Language.h"
  24 #include "Ventoy2Disk.h"
  25
  26 PHY_DRIVE_INFO *g_PhyDriveList = NULL;
  27 DWORD g_PhyDriveCount = 0;
  28 static int g_FilterRemovable = 0;
  29 int g_FilterUSB = 1;
  30 int g_ForceOperation = 1;
  31
  32 int ParseCmdLineOption(LPSTR lpCmdLine)
  33 {
  34     int i;
  35     char cfgfile[MAX_PATH];
  36
  37     if (lpCmdLine && lpCmdLine[0])
  38     {
  39         Log("CmdLine:<%s>", lpCmdLine);
  40     }
  41
  42     for (i = 0; i < __argc; i++)
  43     {
  44         if (strncmp(__argv[i], "-U", 2) == 0 ||
  45                         strncmp(__argv[i], "-u", 2) == 0)
  46         {
  47             g_FilterUSB = 0;
  48         }
  49         else if (strncmp(__argv[i], "-F", 2) == 0)
  50         {
  51             g_ForceOperation = 1;
  52         }
  53     }
  54
  55     GetCurrentDirectoryA(sizeof(cfgfile), cfgfile);
  56     strcat_s(cfgfile, sizeof(cfgfile), "\\Ventoy2Disk.ini");
  57
  58     if (0 == GetPrivateProfileIntA("Filter", "USB", 1, cfgfile))
  59     {
  60         g_FilterUSB = 0;
  61     }
  62
  63     if (1 == GetPrivateProfileIntA("Operation", "Force", 0, cfgfile))
  64     {
  65         g_ForceOperation = 1;
  66     }
  67
  68     Log("Control Flag: %d %d %d", g_FilterRemovable, g_FilterUSB, g_ForceOperation);
  69
  70     return 0;
  71 }
  72
  73 static BOOL IsVentoyPhyDrive(int PhyDrive, UINT64 SizeBytes, MBR_HEAD *pMBR, UINT64 *Part2StartSector)
  74 {
  75     int i;
  76     BOOL bRet;
  77     DWORD dwSize;
  78     HANDLE hDrive;
  79     MBR_HEAD MBR;
  80     UINT32 PartStartSector;
  81     UINT32 PartSectorCount;
  82     CHAR PhyDrivePath[128];
  83         VTOY_GPT_INFO *pGpt = NULL;
  84
  85     safe_sprintf(PhyDrivePath, "\\\\.\\PhysicalDrive%d", PhyDrive);
  86     hDrive = CreateFileA(PhyDrivePath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
  87     Log("Create file Handle:%p %s status:%u", hDrive, PhyDrivePath, LASTERR);
  88
  89     if (hDrive == INVALID_HANDLE_VALUE)
  90     {
  91         return FALSE;
  92     }
  93
  94         bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
  95         Log("Read MBR Ret:%u Size:%u code:%u", bRet, dwSize, LASTERR);
  96
  97     if ((!bRet) || (dwSize != sizeof(MBR)))
  98     {
  99                 CHECK_CLOSE_HANDLE(hDrive);
 100         return FALSE;
 101     }
 102
 103     if (MBR.Byte55 != 0x55 || MBR.ByteAA != 0xAA)
 104     {
 105         Log("Byte55 ByteAA not match 0x%x 0x%x", MBR.Byte55, MBR.ByteAA);
 106                 CHECK_CLOSE_HANDLE(hDrive);
 107         return FALSE;
 108     }
 109
 110         for (i = 0; i < 4; i++)
 111         {
 112                 Log("=========== Partition Table %d ============", i + 1);
 113                 Log("PartTbl.Active = 0x%x", MBR.PartTbl[i].Active);
 114                 Log("PartTbl.FsFlag = 0x%x", MBR.PartTbl[i].FsFlag);
 115                 Log("PartTbl.StartSectorId = %u", MBR.PartTbl[i].StartSectorId);
 116                 Log("PartTbl.SectorCount = %u", MBR.PartTbl[i].SectorCount);
 117                 Log("PartTbl.StartHead = %u", MBR.PartTbl[i].StartHead);
 118                 Log("PartTbl.StartSector = %u", MBR.PartTbl[i].StartSector);
 119                 Log("PartTbl.StartCylinder = %u", MBR.PartTbl[i].StartCylinder);
 120                 Log("PartTbl.EndHead = %u", MBR.PartTbl[i].EndHead);
 121                 Log("PartTbl.EndSector = %u", MBR.PartTbl[i].EndSector);
 122                 Log("PartTbl.EndCylinder = %u", MBR.PartTbl[i].EndCylinder);
 123         }
 124
 125         if (MBR.PartTbl[0].FsFlag == 0xEE)
 126         {
 127                 pGpt = malloc(sizeof(VTOY_GPT_INFO));
 128                 if (!pGpt)
 129                 {
 130                         CHECK_CLOSE_HANDLE(hDrive);
 131                         return FALSE;
 132                 }
 133
 134                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
 135                 bRet = ReadFile(hDrive, pGpt, sizeof(VTOY_GPT_INFO), &dwSize, NULL);
 136                 CHECK_CLOSE_HANDLE(hDrive);
 137                 if ((!bRet) || (dwSize != sizeof(VTOY_GPT_INFO)))
 138                 {
 139                         Log("Failed to read gpt info %d %u %d", bRet, dwSize, LASTERR);
 140                         return FALSE;
 141                 }
 142
 143                 if (memcmp(pGpt->Head.Signature, "EFI PART", 8))
 144                 {
 145                         Log("Invalid GPT signature");
 146                         return FALSE;
 147                 }
 148
 149                 if (memcmp(pGpt->PartTbl[1].Name, L"VTOYEFI", 7 * 2))
 150                 {
 151                         Log("Invalid ventoy efi part name");
 152                         return FALSE;
 153                 }
 154
 155         if (pGpt->PartTbl[0].StartLBA != 2048)
 156         {
 157             Log("Part1 not match %llu", pGpt->PartTbl[0].StartLBA);
 158             return FALSE;
 159         }
 160
 161         PartSectorCount = VENTOY_EFI_PART_SIZE / 512;
 162
 163         if (pGpt->PartTbl[1].StartLBA != pGpt->PartTbl[0].LastLBA + 1 ||
 164             (UINT32)(pGpt->PartTbl[1].LastLBA + 1 - pGpt->PartTbl[1].StartLBA) != PartSectorCount)
 165         {
 166             Log("Part2 not match [%llu %llu] [%llu %llu]",
 167                 pGpt->PartTbl[0].StartLBA, pGpt->PartTbl[0].LastLBA,
 168                 pGpt->PartTbl[1].StartLBA, pGpt->PartTbl[1].LastLBA);
 169             return FALSE;
 170         }
 171
 172                 *Part2StartSector = pGpt->PartTbl[1].StartLBA;
 173
 174         memcpy(pMBR, &(pGpt->MBR), sizeof(MBR_HEAD));
 175         }
 176         else
 177         {
 178                 CHECK_CLOSE_HANDLE(hDrive);
 179
 180                 if (MBR.PartTbl[0].StartSectorId != 2048)
 181                 {
 182                         Log("Part1 not match %u", MBR.PartTbl[0].StartSectorId);
 183                         return FALSE;
 184                 }
 185
 186                 PartStartSector = MBR.PartTbl[0].StartSectorId + MBR.PartTbl[0].SectorCount;
 187                 PartSectorCount = VENTOY_EFI_PART_SIZE / 512;
 188
 189                 if (MBR.PartTbl[1].StartSectorId != PartStartSector ||
 190                         MBR.PartTbl[1].SectorCount != PartSectorCount)
 191                 {
 192                         Log("Part2 not match [0x%x 0x%x] [%u %u] [%u %u]",
 193                                 MBR.PartTbl[1].FsFlag, 0xEF,
 194                                 MBR.PartTbl[1].StartSectorId, PartStartSector,
 195                                 MBR.PartTbl[1].SectorCount, PartSectorCount);
 196                         return FALSE;
 197                 }
 198
 199                 if (MBR.PartTbl[0].Active != 0x80 && MBR.PartTbl[1].Active != 0x80)
 200                 {
 201                         Log("Part1 and Part2 are both NOT active 0x%x 0x%x", MBR.PartTbl[0].Active, MBR.PartTbl[1].Active);
 202             if (MBR.PartTbl[2].Active != 0x80 && MBR.PartTbl[3].Active != 0x80)
 203             {
 204                 Log("Part3 and Part4 are both NOT active 0x%x 0x%x", MBR.PartTbl[2].Active, MBR.PartTbl[3].Active);
 205                 //return FALSE;
 206             }
 207                 }
 208
 209                 *Part2StartSector = MBR.PartTbl[1].StartSectorId;
 210
 211         memcpy(pMBR, &MBR, sizeof(MBR_HEAD));
 212         }
 213
 214     Log("PhysicalDrive%d is ventoy disk", PhyDrive);
 215     return TRUE;
 216 }
 217
 218
 219 static int FilterPhysicalDrive(PHY_DRIVE_INFO *pDriveList, DWORD DriveCount)
 220 {
 221     DWORD i;
 222     DWORD LogDrive;
 223     int Count = 0;
 224     int Letter = 'A';
 225     int Id = 0;
 226     int LetterCount = 0;
 227         UINT64 Part2StartSector = 0;
 228     PHY_DRIVE_INFO *CurDrive;
 229         MBR_HEAD MBR;
 230     int LogLetter[VENTOY_MAX_PHY_DRIVE];
 231     int PhyDriveId[VENTOY_MAX_PHY_DRIVE];
 232
 233     for (LogDrive = GetLogicalDrives(); LogDrive > 0; LogDrive >>= 1)
 234     {
 235         if (LogDrive & 0x01)
 236         {
 237             LogLetter[LetterCount] = Letter;
 238             PhyDriveId[LetterCount] = GetPhyDriveByLogicalDrive(Letter);
 239
 240             Log("Logical Drive:%C  ===> PhyDrive:%d", LogLetter[LetterCount], PhyDriveId[LetterCount]);
 241             LetterCount++;
 242         }
 243
 244         Letter++;
 245     }
 246
 247     for (i = 0; i < DriveCount; i++)
 248     {
 249         CurDrive = pDriveList + i;
 250
 251         CurDrive->Id = -1;
 252         memset(CurDrive->DriveLetters, 0, sizeof(CurDrive->DriveLetters));
 253
 254         if (g_FilterRemovable && (!CurDrive->RemovableMedia))
 255         {
 256             Log("<%s %s> is filtered for not removable.", CurDrive->VendorId, CurDrive->ProductId);
 257             continue;
 258         }
 259
 260         if (g_FilterUSB && CurDrive->BusType != BusTypeUsb)
 261         {
 262             Log("<%s %s> is filtered for not USB type.", CurDrive->VendorId, CurDrive->ProductId);
 263             continue;
 264         }
 265
 266         CurDrive->Id = Id++;
 267
 268         for (Count = 0, Letter = 0; Letter < LetterCount; Letter++)
 269         {
 270             if (PhyDriveId[Letter] == CurDrive->PhyDrive)
 271             {
 272                 if (Count + 1 < sizeof(CurDrive->DriveLetters) / sizeof(CHAR))
 273                 {
 274                     CurDrive->DriveLetters[Count] = LogLetter[Letter];
 275                 }
 276                 Count++;
 277             }
 278         }
 279
 280                 if (IsVentoyPhyDrive(CurDrive->PhyDrive, CurDrive->SizeInBytes, &MBR, &Part2StartSector))
 281         {
 282             memcpy(&(CurDrive->MBR), &MBR, sizeof(MBR));
 283             CurDrive->PartStyle = (MBR.PartTbl[0].FsFlag == 0xEE) ? 1 : 0;
 284             GetVentoyVerInPhyDrive(CurDrive, Part2StartSector, CurDrive->VentoyVersion, sizeof(CurDrive->VentoyVersion), &(CurDrive->SecureBootSupport));
 285             Log("PhyDrive %d is Ventoy Disk ver:%s SecureBoot:%u", CurDrive->PhyDrive, CurDrive->VentoyVersion, CurDrive->SecureBootSupport);
 286         }
 287     }
 288
 289     // for safe
 290     for (i = 0; i < DriveCount; i++)
 291     {
 292         CurDrive = pDriveList + i;
 293         if (CurDrive->Id < 0)
 294         {
 295             CurDrive->PhyDrive = 0x00FFFFFF;
 296         }
 297     }
 298
 299     return Id;
 300 }
 301
 302 PHY_DRIVE_INFO * GetPhyDriveInfoById(int Id)
 303 {
 304     DWORD i;
 305     for (i = 0; i < g_PhyDriveCount; i++)
 306     {
 307         if (g_PhyDriveList[i].Id >= 0 && g_PhyDriveList[i].Id == Id)
 308         {
 309             return g_PhyDriveList + i;
 310         }
 311     }
 312
 313     return NULL;
 314 }
 315
 316 int SortPhysicalDrive(PHY_DRIVE_INFO *pDriveList, DWORD DriveCount)
 317 {
 318         DWORD i, j;
 319         BOOL flag;
 320         PHY_DRIVE_INFO TmpDriveInfo;
 321
 322         for (i = 0; i < DriveCount; i++)
 323         {
 324                 for (j = i + 1; j < DriveCount; j++)
 325                 {
 326                         flag = FALSE;
 327
 328                         if (pDriveList[i].BusType == BusTypeUsb && pDriveList[j].BusType == BusTypeUsb)
 329                         {
 330                                 if (pDriveList[i].RemovableMedia == FALSE && pDriveList[j].RemovableMedia == TRUE)
 331                                 {
 332                                         flag = TRUE;
 333                                 }
 334                         }
 335                         else if (pDriveList[j].BusType == BusTypeUsb)
 336                         {
 337                                 flag = TRUE;
 338                         }
 339                         else
 340                         {
 341                                 if (pDriveList[j].PhyDrive < pDriveList[i].PhyDrive)
 342                                 {
 343                                         flag = TRUE;
 344                                 }
 345                         }
 346
 347                         if (flag)
 348                         {
 349                                 memcpy(&TmpDriveInfo, pDriveList + i, sizeof(PHY_DRIVE_INFO));
 350                                 memcpy(pDriveList + i, pDriveList + j, sizeof(PHY_DRIVE_INFO));
 351                                 memcpy(pDriveList + j, &TmpDriveInfo, sizeof(PHY_DRIVE_INFO));
 352                         }
 353                 }
 354         }
 355
 356         return 0;
 357 }
 358
 359 int Ventoy2DiskInit(void)
 360 {
 361     Log("\n===================== Enum All PhyDrives =====================");
 362     g_PhyDriveList = (PHY_DRIVE_INFO *)malloc(sizeof(PHY_DRIVE_INFO)* VENTOY_MAX_PHY_DRIVE);
 363     if (NULL == g_PhyDriveList)
 364     {
 365         Log("Failed to alloc phy drive memory");
 366         return FALSE;
 367     }
 368     memset(g_PhyDriveList, 0, sizeof(PHY_DRIVE_INFO)* VENTOY_MAX_PHY_DRIVE);
 369
 370     GetAllPhysicalDriveInfo(g_PhyDriveList, &g_PhyDriveCount);
 371
 372         SortPhysicalDrive(g_PhyDriveList, g_PhyDriveCount);
 373
 374     FilterPhysicalDrive(g_PhyDriveList, g_PhyDriveCount);
 375
 376     return 0;
 377 }
 378
 379 int Ventoy2DiskDestroy(void)
 380 {
 381     free(g_PhyDriveList);
 382     return 0;
 383 }