EDK2/edk2_mod/edk2-edk2-stable201911/MdeModulePkg/Application/VDiskChain/VDiskChain.c

   1 /******************************************************************************
   2  * VDiskChain.c
   3  *
   4  * Copyright (c) 2021, 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 <Uefi.h>
  22 #include <Library/DebugLib.h>
  23 #include <Library/PrintLib.h>
  24 #include <Library/UefiLib.h>
  25 #include <Library/BaseMemoryLib.h>
  26 #include <Library/DevicePathLib.h>
  27 #include <Library/MemoryAllocationLib.h>
  28 #include <Library/UefiBootServicesTableLib.h>
  29 #include <Library/UefiRuntimeServicesTableLib.h>
  30 #include <Library/UefiApplicationEntryPoint.h>
  31 #include <Library/UefiDecompressLib.h>
  32 #include <Protocol/LoadedImage.h>
  33 #include <Guid/FileInfo.h>
  34 #include <Guid/FileSystemInfo.h>
  35 #include <Protocol/BlockIo.h>
  36 #include <Protocol/RamDisk.h>
  37 #include <Protocol/SimpleFileSystem.h>
  38 #include <VDiskChain.h>
  39
  40 BOOLEAN gVDiskDebugPrint = FALSE;
  41 vdisk_block_data gVDiskBlockData;
  42
  43 /* Boot filename */
  44 CONST CHAR16 *gEfiBootFileName[] =
  45 {
  46     L"@",
  47     EFI_REMOVABLE_MEDIA_FILE_NAME,
  48 #if   defined (MDE_CPU_IA32)
  49     L"\\EFI\\BOOT\\GRUBIA32.EFI",
  50     L"\\EFI\\BOOT\\BOOTia32.EFI",
  51     L"\\EFI\\BOOT\\bootia32.efi",
  52     L"\\efi\\boot\\bootia32.efi",
  53 #elif defined (MDE_CPU_X64)
  54     L"\\EFI\\BOOT\\GRUBX64.EFI",
  55     L"\\EFI\\BOOT\\BOOTx64.EFI",
  56     L"\\EFI\\BOOT\\bootx64.efi",
  57     L"\\efi\\boot\\bootx64.efi",
  58 #elif defined (MDE_CPU_ARM)
  59     L"\\EFI\\BOOT\\GRUBARM.EFI",
  60     L"\\EFI\\BOOT\\BOOTarm.EFI",
  61     L"\\EFI\\BOOT\\bootarm.efi",
  62     L"\\efi\\boot\\bootarm.efi",
  63 #elif defined (MDE_CPU_AARCH64)
  64     L"\\EFI\\BOOT\\GRUBAA64.EFI",
  65     L"\\EFI\\BOOT\\BOOTaa64.EFI",
  66     L"\\EFI\\BOOT\\bootaa64.efi",
  67     L"\\efi\\boot\\bootaa64.efi",
  68 #endif
  69
  70 };
  71
  72 UINT8 *g_disk_buf_addr = NULL;
  73 UINT64 g_disk_buf_size = 0;
  74
  75 STATIC EFI_GET_VARIABLE g_org_get_variable = NULL;
  76 STATIC EFI_EXIT_BOOT_SERVICES g_org_exit_boot_service = NULL;
  77
  78 VOID EFIAPI VDiskDebug(IN CONST CHAR8  *Format, ...)
  79 {
  80     VA_LIST  Marker;
  81     CHAR16   Buffer[512];
  82
  83     VA_START (Marker, Format);
  84     UnicodeVSPrintAsciiFormat(Buffer, sizeof(Buffer), Format, Marker);
  85     VA_END (Marker);
  86
  87     gST->ConOut->OutputString(gST->ConOut, Buffer);
  88 }
  89
  90 VOID EFIAPI vdisk_clear_input(VOID)
  91 {
  92     EFI_INPUT_KEY Key;
  93
  94     gST->ConIn->Reset(gST->ConIn, FALSE);
  95     while (EFI_SUCCESS == gST->ConIn->ReadKeyStroke(gST->ConIn, &Key))
  96     {
  97         ;
  98     }
  99     gST->ConIn->Reset(gST->ConIn, FALSE);
 100 }
 101
 102 STATIC EFI_STATUS EFIAPI vdisk_load_image
 103 (
 104     IN EFI_HANDLE ImageHandle,
 105     IN EFI_DEVICE_PATH_PROTOCOL *pDevicePath,
 106     IN CONST CHAR16 *FileName,
 107     IN UINTN FileNameLen,
 108     OUT EFI_HANDLE *Image
 109 )
 110 {
 111     EFI_STATUS Status = EFI_SUCCESS;
 112     CHAR16 TmpBuf[256] = {0};
 113     FILEPATH_DEVICE_PATH *pFilePath = NULL;
 114     EFI_DEVICE_PATH_PROTOCOL *pImgPath = NULL;
 115
 116     pFilePath = (FILEPATH_DEVICE_PATH *)TmpBuf;
 117     pFilePath->Header.Type = MEDIA_DEVICE_PATH;
 118     pFilePath->Header.SubType = MEDIA_FILEPATH_DP;
 119     pFilePath->Header.Length[0] = FileNameLen + sizeof(EFI_DEVICE_PATH_PROTOCOL);
 120     pFilePath->Header.Length[1] = 0;
 121     CopyMem(pFilePath->PathName, FileName, FileNameLen);
 122
 123     pImgPath = AppendDevicePathNode(pDevicePath, (EFI_DEVICE_PATH_PROTOCOL *)pFilePath);
 124     if (!pImgPath)
 125     {
 126         return EFI_NOT_FOUND;
 127     }
 128
 129     Status = gBS->LoadImage(FALSE, ImageHandle, pImgPath, NULL, 0, Image);
 130
 131     debug("Load Image File %r DP: <%s>", Status, ConvertDevicePathToText(pImgPath, FALSE, FALSE));
 132
 133     FreePool(pImgPath);
 134
 135     return Status;
 136 }
 137
 138 STATIC EFI_STATUS EFIAPI vdisk_decompress_vdisk(IN EFI_LOADED_IMAGE_PROTOCOL *pImageInfo)
 139 {
 140     UINT32 Size;
 141     UINT32 DestinationSize;
 142     UINT32 ScratchSize;
 143     UINT8 *buf;
 144     VOID  *ScratchBuf;
 145     EFI_STATUS Status = EFI_SUCCESS;
 146
 147     (VOID)pImageInfo;
 148
 149     vdisk_get_vdisk_raw(&buf, &Size);
 150     UefiDecompressGetInfo(buf + VDISK_MAGIC_LEN, Size - VDISK_MAGIC_LEN, &DestinationSize, &ScratchSize);
 151     debug("vdisk: size:%u realsize:%u", Size, DestinationSize);
 152
 153     g_disk_buf_size = DestinationSize;
 154     g_disk_buf_addr = AllocatePool(DestinationSize);
 155     ScratchBuf = AllocatePool(ScratchSize);
 156
 157     Status = UefiDecompress(buf + VDISK_MAGIC_LEN, g_disk_buf_addr, ScratchBuf);
 158     FreePool(ScratchBuf);
 159
 160     debug("Status:%r %p %u", Status, g_disk_buf_addr, (UINT32)g_disk_buf_size);
 161
 162     return EFI_SUCCESS;
 163 }
 164
 165 STATIC EFI_STATUS vdisk_patch_vdisk_path(CHAR16 *pos)
 166 {
 167     UINTN i;
 168     UINTN j;
 169     CHAR16 *end;
 170     CHAR8 *buf = (char *)g_disk_buf_addr;
 171
 172     if (*pos == L'\"')
 173     {
 174         pos++;
 175     }
 176
 177     end = StrStr(pos, L".vtoy");
 178     end += 5;//string length
 179
 180     for (i = 0; i < g_disk_buf_size; i++)
 181     {
 182         if (*(UINT32 *)(buf + i) == 0x59595959)
 183         {
 184             for (j = 0; j < 300; j++)
 185             {
 186                 if (buf[i + j] != 'Y')
 187                 {
 188                     break;
 189                 }
 190             }
 191
 192             if (j >= 300)
 193             {
 194                 break;
 195             }
 196         }
 197     }
 198
 199     if (i >= g_disk_buf_size)
 200     {
 201         debug("No need to fill vdisk path");
 202         return 0;
 203     }
 204
 205     debug("Fill vdisk path at %d", i);
 206
 207     while (pos != end)
 208     {
 209         buf[i++] = (CHAR8)(*pos++);
 210     }
 211
 212     buf[i++] = '\"';
 213
 214     while (buf[i] == 'Y' || buf[i] == '\"')
 215     {
 216         buf[i] = ' ';
 217         i++;
 218     }
 219
 220     return 0;
 221 }
 222
 223 EFI_STATUS EFIAPI vdisk_get_variable_wrapper
 224 (
 225     IN     CHAR16                      *VariableName,
 226     IN     EFI_GUID                    *VendorGuid,
 227     OUT    UINT32                      *Attributes,    OPTIONAL
 228     IN OUT UINTN                       *DataSize,
 229     OUT    VOID                        *Data           OPTIONAL
 230 )
 231 {
 232     EFI_STATUS Status = EFI_SUCCESS;
 233
 234     Status = g_org_get_variable(VariableName, VendorGuid, Attributes, DataSize, Data);
 235     if (StrCmp(VariableName, L"SecureBoot") == 0)
 236     {
 237         if ((*DataSize == 1) && Data)
 238         {
 239             *(UINT8 *)Data = 0;
 240         }
 241     }
 242
 243     return Status;
 244 }
 245
 246 EFI_STATUS EFIAPI vdisk_exit_boot_service_wrapper
 247 (
 248     IN  EFI_HANDLE                   ImageHandle,
 249     IN  UINTN                        MapKey
 250 )
 251 {
 252     if (g_org_get_variable)
 253     {
 254         gRT->GetVariable = g_org_get_variable;
 255         g_org_get_variable = NULL;
 256     }
 257
 258     return g_org_exit_boot_service(ImageHandle, MapKey);
 259 }
 260
 261 STATIC EFI_STATUS EFIAPI vdisk_disable_secure_boot(IN EFI_HANDLE ImageHandle)
 262 {
 263     /* step1: wrapper security protocol. */
 264     /* Do we still need it since we have been loaded ? */
 265
 266
 267     /* step2: fake SecureBoot variable */
 268     g_org_exit_boot_service = gBS->ExitBootServices;
 269     gBS->ExitBootServices = vdisk_exit_boot_service_wrapper;
 270
 271     g_org_get_variable = gRT->GetVariable;
 272     gRT->GetVariable = vdisk_get_variable_wrapper;
 273
 274     return EFI_SUCCESS;
 275 }
 276
 277 STATIC EFI_STATUS EFIAPI vdisk_parse_cmdline(IN EFI_HANDLE ImageHandle)
 278 {
 279     CHAR16 *Pos = NULL;
 280     CHAR16 *pCmdLine = NULL;
 281     EFI_STATUS Status = EFI_SUCCESS;
 282     EFI_LOADED_IMAGE_PROTOCOL *pImageInfo = NULL;
 283
 284     Status = gBS->HandleProtocol(ImageHandle, &gEfiLoadedImageProtocolGuid, (VOID **)&pImageInfo);
 285     if (EFI_ERROR(Status))
 286     {
 287         VDiskDebug("Failed to handle load image protocol %r\n", Status);
 288         return Status;
 289     }
 290
 291     pCmdLine = (CHAR16 *)AllocatePool(pImageInfo->LoadOptionsSize + 4);
 292     SetMem(pCmdLine, pImageInfo->LoadOptionsSize + 4, 0);
 293     CopyMem(pCmdLine, pImageInfo->LoadOptions, pImageInfo->LoadOptionsSize);
 294
 295     if (StrStr(pCmdLine, L"debug"))
 296     {
 297         gVDiskDebugPrint = TRUE;
 298     }
 299
 300     debug("cmdline:<%s>", pCmdLine);
 301     vdisk_debug_pause();
 302
 303     Pos = StrStr(pCmdLine, L"vdisk=");
 304     if (NULL == Pos || NULL == StrStr(pCmdLine, L".vtoy"))
 305     {
 306         VDiskDebug("vdisk parameter not found!\n");
 307         return EFI_NOT_FOUND;
 308     }
 309
 310     vdisk_decompress_vdisk(pImageInfo);
 311
 312     vdisk_patch_vdisk_path(Pos + 6);
 313
 314     if (StrStr(pCmdLine, L"secureboot=off"))
 315     {
 316         vdisk_disable_secure_boot(ImageHandle);
 317     }
 318
 319     FreePool(pCmdLine);
 320     return EFI_SUCCESS;
 321 }
 322
 323 EFI_STATUS EFIAPI vdisk_boot(IN EFI_HANDLE ImageHandle)
 324 {
 325     UINTN t = 0;
 326     UINTN i = 0;
 327     UINTN j = 0;
 328     UINTN Find = 0;
 329     UINTN Count = 0;
 330     EFI_HANDLE Image = NULL;
 331     EFI_HANDLE *Handles = NULL;
 332     EFI_STATUS Status = EFI_SUCCESS;
 333     EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *pFile = NULL;
 334     EFI_DEVICE_PATH_PROTOCOL *pDevPath = NULL;
 335
 336     for (t = 0; t < 3; t++)
 337     {
 338         Count = 0;
 339         Handles = NULL;
 340
 341         Status = gBS->LocateHandleBuffer(ByProtocol, &gEfiSimpleFileSystemProtocolGuid,
 342                                      NULL, &Count, &Handles);
 343         if (EFI_ERROR(Status))
 344         {
 345             return Status;
 346         }
 347
 348         debug("vdisk_boot fs count:%u", Count);
 349
 350         for (i = 0; i < Count; i++)
 351         {
 352             Status = gBS->HandleProtocol(Handles[i], &gEfiSimpleFileSystemProtocolGuid, (VOID **)&pFile);
 353             if (EFI_ERROR(Status))
 354             {
 355                 continue;
 356             }
 357
 358             debug("FS:%u Protocol:%p  OpenVolume:%p", i, pFile, pFile->OpenVolume);
 359
 360             Status = gBS->OpenProtocol(Handles[i], &gEfiDevicePathProtocolGuid,
 361                                        (VOID **)&pDevPath,
 362                                        ImageHandle,
 363                                        Handles[i],
 364                                        EFI_OPEN_PROTOCOL_GET_PROTOCOL);
 365             if (EFI_ERROR(Status))
 366             {
 367                 debug("Failed to open device path protocol %r", Status);
 368                 continue;
 369             }
 370
 371             debug("Handle:%p FS DP: <%s>", Handles[i], ConvertDevicePathToText(pDevPath, FALSE, FALSE));
 372             if (CompareMem(gVDiskBlockData.Path, pDevPath, gVDiskBlockData.DevicePathCompareLen))
 373             {
 374                 debug("Not ventoy disk file system");
 375                 continue;
 376             }
 377
 378             for (j = 1; j < ARRAY_SIZE(gEfiBootFileName); j++)
 379             {
 380                 Status = vdisk_load_image(ImageHandle, pDevPath, gEfiBootFileName[j],
 381                                            StrSize(gEfiBootFileName[j]), &Image);
 382                 if (EFI_SUCCESS == Status)
 383                 {
 384                     break;
 385                 }
 386                 debug("Failed to load image %r <%s>", Status, gEfiBootFileName[j]);
 387             }
 388
 389             if (j >= ARRAY_SIZE(gEfiBootFileName))
 390             {
 391                 continue;
 392             }
 393
 394             Find++;
 395             debug("Find boot file, now try to boot .....");
 396             vdisk_debug_pause();
 397
 398             if (gVDiskDebugPrint)
 399             {
 400                 gST->ConIn->Reset(gST->ConIn, FALSE);
 401             }
 402
 403             /* can't add debug print here */
 404             //ventoy_wrapper_system();
 405             Status = gBS->StartImage(Image, NULL, NULL);
 406             if (EFI_ERROR(Status))
 407             {
 408                 debug("Failed to start image %r", Status);
 409                 sleep(3);
 410                 gBS->UnloadImage(Image);
 411                 break;
 412             }
 413         }
 414
 415         FreePool(Handles);
 416
 417         if (Find == 0)
 418         {
 419             debug("Fs not found, now wait and retry...");
 420             sleep(2);
 421         }
 422     }
 423
 424     if (Find == 0)
 425     {
 426         return EFI_NOT_FOUND;
 427     }
 428
 429     return EFI_SUCCESS;
 430 }
 431
 432 EFI_STATUS EFIAPI VDiskChainEfiMain
 433 (
 434     IN EFI_HANDLE         ImageHandle,
 435     IN EFI_SYSTEM_TABLE  *SystemTable
 436 )
 437 {
 438     EFI_STATUS Status = EFI_SUCCESS;
 439
 440     gST->ConOut->ClearScreen(gST->ConOut);
 441     vdisk_clear_input();
 442
 443     Status = vdisk_parse_cmdline(ImageHandle);
 444     if (EFI_ERROR(Status))
 445     {
 446         return Status;
 447     }
 448
 449     vdisk_install_blockio(ImageHandle, g_disk_buf_size);
 450     vdisk_debug_pause();
 451
 452     Status = vdisk_boot(ImageHandle);
 453
 454     gBS->DisconnectController(gVDiskBlockData.Handle, NULL, NULL);
 455     gBS->UninstallMultipleProtocolInterfaces(gVDiskBlockData.Handle,
 456             &gEfiBlockIoProtocolGuid, &gVDiskBlockData.BlockIo,
 457             &gEfiDevicePathProtocolGuid, gVDiskBlockData.Path,
 458             NULL);
 459
 460     if (EFI_NOT_FOUND == Status)
 461     {
 462         gST->ConOut->OutputString(gST->ConOut, L"No bootfile found for UEFI!\r\n");
 463         gST->ConOut->OutputString(gST->ConOut, L"Maybe the image does not support " VENTOY_UEFI_DESC  L"!\r\n");
 464         sleep(30);
 465     }
 466
 467     vdisk_clear_input();
 468     gST->ConOut->ClearScreen(gST->ConOut);
 469
 470     return EFI_SUCCESS;
 471 }
 472