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 VOID EFIAPI VDiskDebug(IN CONST CHAR8  *Format, ...)
  76 {
  77     VA_LIST  Marker;
  78     CHAR16   Buffer[512];
  79
  80     VA_START (Marker, Format);
  81     UnicodeVSPrintAsciiFormat(Buffer, sizeof(Buffer), Format, Marker);
  82     VA_END (Marker);
  83
  84     gST->ConOut->OutputString(gST->ConOut, Buffer);
  85 }
  86
  87 VOID EFIAPI vdisk_clear_input(VOID)
  88 {
  89     EFI_INPUT_KEY Key;
  90
  91     gST->ConIn->Reset(gST->ConIn, FALSE);
  92     while (EFI_SUCCESS == gST->ConIn->ReadKeyStroke(gST->ConIn, &Key))
  93     {
  94         ;
  95     }
  96     gST->ConIn->Reset(gST->ConIn, FALSE);
  97 }
  98
  99 STATIC EFI_STATUS EFIAPI vdisk_load_image
 100 (
 101     IN EFI_HANDLE ImageHandle,
 102     IN EFI_DEVICE_PATH_PROTOCOL *pDevicePath,
 103     IN CONST CHAR16 *FileName,
 104     IN UINTN FileNameLen,
 105     OUT EFI_HANDLE *Image
 106 )
 107 {
 108     EFI_STATUS Status = EFI_SUCCESS;
 109     CHAR16 TmpBuf[256] = {0};
 110     FILEPATH_DEVICE_PATH *pFilePath = NULL;
 111     EFI_DEVICE_PATH_PROTOCOL *pImgPath = NULL;
 112
 113     pFilePath = (FILEPATH_DEVICE_PATH *)TmpBuf;
 114     pFilePath->Header.Type = MEDIA_DEVICE_PATH;
 115     pFilePath->Header.SubType = MEDIA_FILEPATH_DP;
 116     pFilePath->Header.Length[0] = FileNameLen + sizeof(EFI_DEVICE_PATH_PROTOCOL);
 117     pFilePath->Header.Length[1] = 0;
 118     CopyMem(pFilePath->PathName, FileName, FileNameLen);
 119
 120     pImgPath = AppendDevicePathNode(pDevicePath, (EFI_DEVICE_PATH_PROTOCOL *)pFilePath);
 121     if (!pImgPath)
 122     {
 123         return EFI_NOT_FOUND;
 124     }
 125
 126     Status = gBS->LoadImage(FALSE, ImageHandle, pImgPath, NULL, 0, Image);
 127
 128     debug("Load Image File %r DP: <%s>", Status, ConvertDevicePathToText(pImgPath, FALSE, FALSE));
 129
 130     FreePool(pImgPath);
 131
 132     return Status;
 133 }
 134
 135 STATIC EFI_STATUS EFIAPI vdisk_decompress_vdisk(IN EFI_LOADED_IMAGE_PROTOCOL *pImageInfo)
 136 {
 137     UINT32 Size;
 138     UINT32 DestinationSize;
 139     UINT32 ScratchSize;
 140     UINT8 *buf;
 141     VOID  *ScratchBuf;
 142     EFI_STATUS Status = EFI_SUCCESS;
 143
 144     (VOID)pImageInfo;
 145
 146     vdisk_get_vdisk_raw(&buf, &Size);
 147     UefiDecompressGetInfo(buf + VDISK_MAGIC_LEN, Size - VDISK_MAGIC_LEN, &DestinationSize, &ScratchSize);
 148     debug("vdisk: size:%u realsize:%u", Size, DestinationSize);
 149
 150     g_disk_buf_size = DestinationSize;
 151     g_disk_buf_addr = AllocatePool(DestinationSize);
 152     ScratchBuf = AllocatePool(ScratchSize);
 153
 154     Status = UefiDecompress(buf + VDISK_MAGIC_LEN, g_disk_buf_addr, ScratchBuf);
 155     FreePool(ScratchBuf);
 156
 157     debug("Status:%r %p %u", Status, g_disk_buf_addr, (UINT32)g_disk_buf_size);
 158
 159     return EFI_SUCCESS;
 160 }
 161
 162 STATIC EFI_STATUS vdisk_patch_vdisk_path(CHAR16 *pos)
 163 {
 164     UINTN i;
 165     UINTN j;
 166     CHAR16 *end;
 167     CHAR8 *buf = (char *)g_disk_buf_addr;
 168
 169     if (*pos == L'\"')
 170     {
 171         pos++;
 172     }
 173
 174     end = StrStr(pos, L".vtoy");
 175     end += 5;//string length
 176
 177     for (i = 0; i < g_disk_buf_size; i++)
 178     {
 179         if (*(UINT32 *)(buf + i) == 0x59595959)
 180         {
 181             for (j = 0; j < 300; j++)
 182             {
 183                 if (buf[i + j] != 'Y')
 184                 {
 185                     break;
 186                 }
 187             }
 188
 189             if (j >= 300)
 190             {
 191                 break;
 192             }
 193         }
 194     }
 195
 196     if (i >= g_disk_buf_size)
 197     {
 198         debug("No need to fill vdisk path");
 199         return 0;
 200     }
 201
 202     debug("Fill vdisk path at %d", i);
 203
 204     while (pos != end)
 205     {
 206         buf[i++] = (CHAR8)(*pos++);
 207     }
 208
 209     buf[i++] = '\"';
 210
 211     while (buf[i] == 'Y' || buf[i] == '\"')
 212     {
 213         buf[i] = ' ';
 214         i++;
 215     }
 216
 217     return 0;
 218 }
 219
 220 STATIC EFI_STATUS EFIAPI vdisk_parse_cmdline(IN EFI_HANDLE ImageHandle)
 221 {
 222     CHAR16 *Pos = NULL;
 223     CHAR16 *pCmdLine = NULL;
 224     EFI_STATUS Status = EFI_SUCCESS;
 225     EFI_LOADED_IMAGE_PROTOCOL *pImageInfo = NULL;
 226
 227     Status = gBS->HandleProtocol(ImageHandle, &gEfiLoadedImageProtocolGuid, (VOID **)&pImageInfo);
 228     if (EFI_ERROR(Status))
 229     {
 230         VDiskDebug("Failed to handle load image protocol %r\n", Status);
 231         return Status;
 232     }
 233
 234     pCmdLine = (CHAR16 *)AllocatePool(pImageInfo->LoadOptionsSize + 4);
 235     SetMem(pCmdLine, pImageInfo->LoadOptionsSize + 4, 0);
 236     CopyMem(pCmdLine, pImageInfo->LoadOptions, pImageInfo->LoadOptionsSize);
 237
 238     if (StrStr(pCmdLine, L"debug"))
 239     {
 240         gVDiskDebugPrint = TRUE;
 241     }
 242
 243     debug("cmdline:<%s>", pCmdLine);
 244     vdisk_debug_pause();
 245
 246     Pos = StrStr(pCmdLine, L"vdisk=");
 247     if (NULL == Pos || NULL == StrStr(pCmdLine, L".vtoy"))
 248     {
 249         VDiskDebug("vdisk parameter not found!\n");
 250         return EFI_NOT_FOUND;
 251     }
 252
 253     vdisk_decompress_vdisk(pImageInfo);
 254
 255     vdisk_patch_vdisk_path(Pos + 6);
 256
 257     FreePool(pCmdLine);
 258     return EFI_SUCCESS;
 259 }
 260
 261 EFI_STATUS EFIAPI vdisk_boot(IN EFI_HANDLE ImageHandle)
 262 {
 263     UINTN t = 0;
 264     UINTN i = 0;
 265     UINTN j = 0;
 266     UINTN Find = 0;
 267     UINTN Count = 0;
 268     EFI_HANDLE Image = NULL;
 269     EFI_HANDLE *Handles = NULL;
 270     EFI_STATUS Status = EFI_SUCCESS;
 271     EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *pFile = NULL;
 272     EFI_DEVICE_PATH_PROTOCOL *pDevPath = NULL;
 273
 274     for (t = 0; t < 3; t++)
 275     {
 276         Count = 0;
 277         Handles = NULL;
 278
 279         Status = gBS->LocateHandleBuffer(ByProtocol, &gEfiSimpleFileSystemProtocolGuid,
 280                                      NULL, &Count, &Handles);
 281         if (EFI_ERROR(Status))
 282         {
 283             return Status;
 284         }
 285
 286         debug("vdisk_boot fs count:%u", Count);
 287
 288         for (i = 0; i < Count; i++)
 289         {
 290             Status = gBS->HandleProtocol(Handles[i], &gEfiSimpleFileSystemProtocolGuid, (VOID **)&pFile);
 291             if (EFI_ERROR(Status))
 292             {
 293                 continue;
 294             }
 295
 296             debug("FS:%u Protocol:%p  OpenVolume:%p", i, pFile, pFile->OpenVolume);
 297
 298             Status = gBS->OpenProtocol(Handles[i], &gEfiDevicePathProtocolGuid,
 299                                        (VOID **)&pDevPath,
 300                                        ImageHandle,
 301                                        Handles[i],
 302                                        EFI_OPEN_PROTOCOL_GET_PROTOCOL);
 303             if (EFI_ERROR(Status))
 304             {
 305                 debug("Failed to open device path protocol %r", Status);
 306                 continue;
 307             }
 308
 309             debug("Handle:%p FS DP: <%s>", Handles[i], ConvertDevicePathToText(pDevPath, FALSE, FALSE));
 310             if (CompareMem(gVDiskBlockData.Path, pDevPath, gVDiskBlockData.DevicePathCompareLen))
 311             {
 312                 debug("Not ventoy disk file system");
 313                 continue;
 314             }
 315
 316             for (j = 1; j < ARRAY_SIZE(gEfiBootFileName); j++)
 317             {
 318                 Status = vdisk_load_image(ImageHandle, pDevPath, gEfiBootFileName[j],
 319                                            StrSize(gEfiBootFileName[j]), &Image);
 320                 if (EFI_SUCCESS == Status)
 321                 {
 322                     break;
 323                 }
 324                 debug("Failed to load image %r <%s>", Status, gEfiBootFileName[j]);
 325             }
 326
 327             if (j >= ARRAY_SIZE(gEfiBootFileName))
 328             {
 329                 continue;
 330             }
 331
 332             Find++;
 333             debug("Find boot file, now try to boot .....");
 334             vdisk_debug_pause();
 335
 336             if (gVDiskDebugPrint)
 337             {
 338                 gST->ConIn->Reset(gST->ConIn, FALSE);
 339             }
 340
 341             /* can't add debug print here */
 342             //ventoy_wrapper_system();
 343             Status = gBS->StartImage(Image, NULL, NULL);
 344             if (EFI_ERROR(Status))
 345             {
 346                 debug("Failed to start image %r", Status);
 347                 sleep(3);
 348                 gBS->UnloadImage(Image);
 349                 break;
 350             }
 351         }
 352
 353         FreePool(Handles);
 354
 355         if (Find == 0)
 356         {
 357             debug("Fs not found, now wait and retry...");
 358             sleep(2);
 359         }
 360     }
 361
 362     if (Find == 0)
 363     {
 364         return EFI_NOT_FOUND;
 365     }
 366
 367     return EFI_SUCCESS;
 368 }
 369
 370 EFI_STATUS EFIAPI VDiskChainEfiMain
 371 (
 372     IN EFI_HANDLE         ImageHandle,
 373     IN EFI_SYSTEM_TABLE  *SystemTable
 374 )
 375 {
 376     EFI_STATUS Status = EFI_SUCCESS;
 377
 378     gST->ConOut->ClearScreen(gST->ConOut);
 379     vdisk_clear_input();
 380
 381     Status = vdisk_parse_cmdline(ImageHandle);
 382     if (EFI_ERROR(Status))
 383     {
 384         return Status;
 385     }
 386
 387     vdisk_install_blockio(ImageHandle, g_disk_buf_size);
 388     vdisk_debug_pause();
 389
 390     Status = vdisk_boot(ImageHandle);
 391
 392     gBS->DisconnectController(gVDiskBlockData.Handle, NULL, NULL);
 393     gBS->UninstallMultipleProtocolInterfaces(gVDiskBlockData.Handle,
 394             &gEfiBlockIoProtocolGuid, &gVDiskBlockData.BlockIo,
 395             &gEfiDevicePathProtocolGuid, gVDiskBlockData.Path,
 396             NULL);
 397
 398     if (EFI_NOT_FOUND == Status)
 399     {
 400         gST->ConOut->OutputString(gST->ConOut, L"No bootfile found for UEFI!\r\n");
 401         gST->ConOut->OutputString(gST->ConOut, L"Maybe the image does not support " VENTOY_UEFI_DESC  L"!\r\n");
 402         sleep(30);
 403     }
 404
 405     vdisk_clear_input();
 406     gST->ConOut->ClearScreen(gST->ConOut);
 407
 408     return EFI_SUCCESS;
 409 }
 410