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

   1 /******************************************************************************
   2  * VtoyUtil.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 <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 <Protocol/LoadedImage.h>
  32 #include <Guid/FileInfo.h>
  33 #include <Guid/FileSystemInfo.h>
  34 #include <Protocol/BlockIo.h>
  35 #include <Protocol/RamDisk.h>
  36 #include <Protocol/SimpleFileSystem.h>
  37 #include <VtoyUtil.h>
  38
  39 BOOLEAN gVtoyDebugPrint = FALSE;
  40 STATIC CONST CHAR16 *gCurFeature= NULL;
  41 STATIC CHAR16 *gCmdLine = NULL;
  42 STATIC grub_env_printf_pf g_env_printf = NULL;
  43
  44 STATIC VtoyUtilFeature gFeatureList[] =
  45 {
  46     { L"fix_windows_mmap", FixWindowsMemhole },
  47     { L"show_efi_drivers", ShowEfiDrivers    },
  48 };
  49
  50 EFI_STATUS VtoyGetComponentName(IN UINTN Ver, IN VOID *Protocol, OUT CHAR16 **DriverName)
  51 {
  52     EFI_STATUS Status = EFI_SUCCESS;
  53     CHAR16 *DrvName = NULL;
  54     EFI_COMPONENT_NAME_PROTOCOL *NameProtocol = NULL;
  55     EFI_COMPONENT_NAME2_PROTOCOL *Name2Protocol = NULL;
  56
  57     if (1 == Ver)
  58     {
  59         NameProtocol = (EFI_COMPONENT_NAME_PROTOCOL *)Protocol;
  60         Status = NameProtocol->GetDriverName(Protocol, "en", &DrvName);
  61         if (EFI_ERROR(Status) || NULL == DrvName)
  62         {
  63             Status = NameProtocol->GetDriverName(Protocol, "eng", &DrvName);
  64         }
  65     }
  66     else
  67     {
  68         Name2Protocol = (EFI_COMPONENT_NAME2_PROTOCOL *)Protocol;
  69         Status = Name2Protocol->GetDriverName(Protocol, "en", &DrvName);
  70         if (EFI_ERROR(Status) || NULL == DrvName)
  71         {
  72             Status = Name2Protocol->GetDriverName(Protocol, "eng", &DrvName);
  73         }
  74     }
  75
  76     *DriverName = DrvName;
  77     return Status;
  78 }
  79
  80 VOID EFIAPI VtoyUtilDebug(IN CONST CHAR8  *Format, ...)
  81 {
  82     VA_LIST  Marker;
  83     CHAR8    Buffer[512];
  84
  85     VA_START (Marker, Format);
  86     AsciiVSPrint(Buffer, sizeof(Buffer), Format, Marker);
  87     VA_END (Marker);
  88
  89     if (g_env_printf)
  90     {
  91         g_env_printf("%s", Buffer);
  92     }
  93 }
  94
  95 STATIC EFI_STATUS ParseCmdline(IN EFI_HANDLE ImageHandle)
  96 {
  97     CHAR16 *pPos = NULL;
  98     CHAR16 *pCmdLine = NULL;
  99     EFI_STATUS Status = EFI_SUCCESS;
 100     ventoy_grub_param *pGrubParam = NULL;
 101     EFI_LOADED_IMAGE_PROTOCOL *pImageInfo = NULL;
 102
 103     Status = gBS->HandleProtocol(ImageHandle, &gEfiLoadedImageProtocolGuid, (VOID **)&pImageInfo);
 104     if (EFI_ERROR(Status))
 105     {
 106         return Status;
 107     }
 108
 109     pCmdLine = (CHAR16 *)AllocatePool(pImageInfo->LoadOptionsSize + 4);
 110     SetMem(pCmdLine, pImageInfo->LoadOptionsSize + 4, 0);
 111     CopyMem(pCmdLine, pImageInfo->LoadOptions, pImageInfo->LoadOptionsSize);
 112
 113     if (StrStr(pCmdLine, L"vtoyefitest"))
 114     {
 115         gST->ConOut->OutputString(gST->ConOut, L"\r\n##########################");
 116         gST->ConOut->OutputString(gST->ConOut, L"\r\n#########  VTOY  #########");
 117         gST->ConOut->OutputString(gST->ConOut, L"\r\n##########################");
 118         return EFI_SUCCESS;
 119     }
 120
 121     if (StrStr(pCmdLine, L"debug"))
 122     {
 123         gVtoyDebugPrint = TRUE;
 124     }
 125
 126     pPos = StrStr(pCmdLine, L"env_param=");
 127     if (!pPos)
 128     {
 129         return EFI_INVALID_PARAMETER;
 130     }
 131
 132     pGrubParam = (ventoy_grub_param *)StrHexToUintn(pPos + StrLen(L"env_param="));
 133     g_env_printf = pGrubParam->grub_env_printf;
 134
 135     pPos = StrStr(pCmdLine, L"feature=");
 136     if (!pPos)
 137     {
 138         return EFI_INVALID_PARAMETER;
 139     }
 140
 141     gCurFeature = pPos + StrLen(L"feature=");
 142
 143     gCmdLine = pCmdLine;
 144
 145     return EFI_SUCCESS;
 146 }
 147
 148 EFI_STATUS EFIAPI VtoyUtilEfiMain
 149 (
 150     IN EFI_HANDLE         ImageHandle,
 151     IN EFI_SYSTEM_TABLE  *SystemTable
 152 )
 153 {
 154     UINTN i;
 155     UINTN Len;
 156
 157     ParseCmdline(ImageHandle);
 158
 159     for (i = 0; gCurFeature && i < ARRAY_SIZE(gFeatureList); i++)
 160     {
 161         Len = StrLen(gFeatureList[i].Cmd);
 162         if (StrnCmp(gFeatureList[i].Cmd, gCurFeature, Len) == 0)
 163         {
 164             debug("Find main proc <%s>", gFeatureList[i].Cmd);
 165             gFeatureList[i].MainProc(ImageHandle, gCurFeature + Len);
 166             break;
 167         }
 168     }
 169
 170     if (gCmdLine)
 171     {
 172         FreePool(gCmdLine);
 173         gCmdLine = NULL;
 174     }
 175
 176     return EFI_SUCCESS;
 177 }
 178