/** @file
  Provides 'initrd' dynamic UEFI shell command to load a Linux initrd
  via its GUIDed vendor media path

  Copyright (c) 2020, Arm, Ltd. All rights reserved.<BR>

  SPDX-License-Identifier: BSD-2-Clause-Patent
**/

#include <Uefi.h>

#include <Library/DebugLib.h>
#include <Library/DevicePathLib.h>
#include <Library/HiiLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/ShellLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiHiiServicesLib.h>

#include <Guid/LinuxEfiInitrdMedia.h>

#include <Protocol/DevicePath.h>
#include <Protocol/HiiPackageList.h>
#include <Protocol/LoadFile2.h>
#include <Protocol/ShellDynamicCommand.h>

#pragma pack (1)
typedef struct {
  VENDOR_DEVICE_PATH          VenMediaNode;
  EFI_DEVICE_PATH_PROTOCOL    EndNode;
} SINGLE_NODE_VENDOR_MEDIA_DEVPATH;
#pragma pack ()

STATIC EFI_HII_HANDLE         mLinuxInitrdShellCommandHiiHandle;
STATIC EFI_PHYSICAL_ADDRESS   mInitrdFileAddress;
STATIC UINTN                  mInitrdFileSize;
STATIC EFI_HANDLE             mInitrdLoadFile2Handle;

STATIC CONST SHELL_PARAM_ITEM ParamList[] = {
  {L"-u", TypeFlag},
  {NULL, TypeMax}
  };

STATIC CONST SINGLE_NODE_VENDOR_MEDIA_DEVPATH mInitrdDevicePath = {
  {
    {
      MEDIA_DEVICE_PATH, MEDIA_VENDOR_DP, { sizeof (VENDOR_DEVICE_PATH) }
    },
    LINUX_EFI_INITRD_MEDIA_GUID
  }, {
    END_DEVICE_PATH_TYPE, END_ENTIRE_DEVICE_PATH_SUBTYPE,
    { sizeof (EFI_DEVICE_PATH_PROTOCOL) }
  }
};

STATIC
BOOLEAN
IsOtherInitrdDevicePathAlreadyInstalled (
  VOID
  )
{
  EFI_STATUS                  Status;
  EFI_DEVICE_PATH_PROTOCOL    *DevicePath;
  EFI_HANDLE                  Handle;

  DevicePath = (EFI_DEVICE_PATH_PROTOCOL *)&mInitrdDevicePath;
  Status = gBS->LocateDevicePath (&gEfiLoadFile2ProtocolGuid, &DevicePath,
                  &Handle);
  if (EFI_ERROR (Status)) {
    return FALSE;
  }

  //
  // Check whether the existing instance is one that we installed during
  // a previous invocation.
  //
  if (Handle == mInitrdLoadFile2Handle) {
    return FALSE;
  }
  return TRUE;
}

STATIC
EFI_STATUS
EFIAPI
InitrdLoadFile2 (
  IN      EFI_LOAD_FILE2_PROTOCOL       *This,
  IN      EFI_DEVICE_PATH_PROTOCOL      *FilePath,
  IN      BOOLEAN                       BootPolicy,
  IN  OUT UINTN                         *BufferSize,
  OUT     VOID                          *Buffer     OPTIONAL
  )
{
  if (BootPolicy) {
    return EFI_UNSUPPORTED;
  }

  if (BufferSize == NULL || !IsDevicePathValid (FilePath, 0)) {
    return EFI_INVALID_PARAMETER;
  }

  if (FilePath->Type != END_DEVICE_PATH_TYPE ||
      FilePath->SubType != END_ENTIRE_DEVICE_PATH_SUBTYPE ||
      mInitrdFileSize == 0) {
    return EFI_NOT_FOUND;
  }

  if (Buffer == NULL || *BufferSize < mInitrdFileSize) {
    *BufferSize = mInitrdFileSize;
    return EFI_BUFFER_TOO_SMALL;
  }

  ASSERT (mInitrdFileAddress != 0);

  gBS->CopyMem (Buffer, (VOID *)(UINTN)mInitrdFileAddress, mInitrdFileSize);
  *BufferSize = mInitrdFileSize;
  return EFI_SUCCESS;
}

STATIC CONST EFI_LOAD_FILE2_PROTOCOL     mInitrdLoadFile2 = {
  InitrdLoadFile2,
};

STATIC
EFI_STATUS
UninstallLoadFile2Protocol (
  VOID
  )
{
  EFI_STATUS    Status;

  if (mInitrdLoadFile2Handle != NULL) {
    Status = gBS->UninstallMultipleProtocolInterfaces (mInitrdLoadFile2Handle,
                    &gEfiDevicePathProtocolGuid,    &mInitrdDevicePath,
                    &gEfiLoadFile2ProtocolGuid,     &mInitrdLoadFile2,
                    NULL);
    if (!EFI_ERROR (Status)) {
      mInitrdLoadFile2Handle = NULL;
    }
    return Status;
  }
  return EFI_SUCCESS;
}

STATIC
VOID
FreeInitrdFile (
  VOID
  )
{
  if (mInitrdFileSize != 0) {
    gBS->FreePages (mInitrdFileAddress, EFI_SIZE_TO_PAGES (mInitrdFileSize));
    mInitrdFileSize = 0;
  }
}

STATIC
EFI_STATUS
CacheInitrdFile (
  IN  SHELL_FILE_HANDLE   FileHandle
  )
{
  EFI_STATUS              Status;
  UINT64                  FileSize;
  UINTN                   ReadSize;

  Status = gEfiShellProtocol->GetFileSize (FileHandle, &FileSize);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  if (FileSize == 0 || FileSize > MAX_UINTN) {
    return EFI_UNSUPPORTED;
  }

  Status = gBS->AllocatePages (AllocateAnyPages, EfiLoaderData,
                  EFI_SIZE_TO_PAGES ((UINTN)FileSize), &mInitrdFileAddress);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  ReadSize = (UINTN)FileSize;
  Status = gEfiShellProtocol->ReadFile (FileHandle, &ReadSize,
                                (VOID *)(UINTN)mInitrdFileAddress);
  if (EFI_ERROR (Status) || ReadSize < FileSize) {
    DEBUG ((DEBUG_WARN, "%a: failed to read initrd file - %r 0x%lx 0x%lx\n",
      __FUNCTION__, Status, (UINT64)ReadSize, FileSize));
    goto FreeMemory;
  }

  if (mInitrdLoadFile2Handle == NULL) {
    Status = gBS->InstallMultipleProtocolInterfaces (&mInitrdLoadFile2Handle,
                    &gEfiDevicePathProtocolGuid,  &mInitrdDevicePath,
                    &gEfiLoadFile2ProtocolGuid,   &mInitrdLoadFile2,
                    NULL);
    ASSERT_EFI_ERROR (Status);
  }

  mInitrdFileSize = (UINTN)FileSize;
  return EFI_SUCCESS;

FreeMemory:
  gBS->FreePages (mInitrdFileAddress, EFI_SIZE_TO_PAGES ((UINTN)FileSize));
  return Status;
}

/**
  Function for 'initrd' command.

  @param[in] ImageHandle  Handle to the Image (NULL if Internal).
  @param[in] SystemTable  Pointer to the System Table (NULL if Internal).
**/
STATIC
SHELL_STATUS
EFIAPI
RunInitrd (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS            Status;
  LIST_ENTRY            *Package;
  CHAR16                *ProblemParam;
  CONST CHAR16          *Param;
  CHAR16                *Filename;
  SHELL_STATUS          ShellStatus;
  SHELL_FILE_HANDLE     FileHandle;

  ProblemParam        = NULL;
  ShellStatus         = SHELL_SUCCESS;

  Status = ShellInitialize ();
  ASSERT_EFI_ERROR (Status);

  //
  // parse the command line
  //
  Status = ShellCommandLineParse (ParamList, &Package, &ProblemParam, TRUE);
  if (EFI_ERROR (Status)) {
    if (Status == EFI_VOLUME_CORRUPTED && ProblemParam != NULL) {
      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM),
        mLinuxInitrdShellCommandHiiHandle, L"initrd", ProblemParam);
      FreePool (ProblemParam);
      ShellStatus = SHELL_INVALID_PARAMETER;
    } else {
      ASSERT(FALSE);
    }
  } else if (IsOtherInitrdDevicePathAlreadyInstalled ()) {
    ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_ALREADY_INSTALLED),
      mLinuxInitrdShellCommandHiiHandle, L"initrd");
    ShellStatus = SHELL_UNSUPPORTED;
  } else {
    if (ShellCommandLineGetCount (Package) > 2) {
      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_TOO_MANY),
        mLinuxInitrdShellCommandHiiHandle, L"initrd");
      ShellStatus = SHELL_INVALID_PARAMETER;
    } else if (ShellCommandLineGetCount (Package) < 2) {
      if (ShellCommandLineGetFlag (Package, L"-u")) {
        FreeInitrdFile ();
        UninstallLoadFile2Protocol ();
      } else {
        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_TOO_FEW),
          mLinuxInitrdShellCommandHiiHandle, L"initrd");
        ShellStatus = SHELL_INVALID_PARAMETER;
      }
    } else {
      Param = ShellCommandLineGetRawValue (Package, 1);
      ASSERT (Param != NULL);

      Filename = ShellFindFilePath (Param);
      if (Filename == NULL) {
        ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_FIND_FAIL),
          mLinuxInitrdShellCommandHiiHandle, L"initrd", Param);
        ShellStatus = SHELL_NOT_FOUND;
      } else {
        Status = ShellOpenFileByName (Filename, &FileHandle,
                   EFI_FILE_MODE_READ, 0);
        if (!EFI_ERROR (Status)) {
          FreeInitrdFile ();
          Status = CacheInitrdFile (FileHandle);
          ShellCloseFile (&FileHandle);
        }
        if (EFI_ERROR (Status)) {
          ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_FILE_OPEN_FAIL),
            mLinuxInitrdShellCommandHiiHandle, L"initrd", Param);
          ShellStatus = SHELL_NOT_FOUND;
        }
        FreePool (Filename);
      }
    }
  }
  return ShellStatus;
}


/**
  This is the shell command handler function pointer callback type.  This
  function handles the command when it is invoked in the shell.

  @param[in] This                   The instance of the
                                    EFI_SHELL_DYNAMIC_COMMAND_PROTOCOL.
  @param[in] SystemTable            The pointer to the system table.
  @param[in] ShellParameters        The parameters associated with the command.
  @param[in] Shell                  The instance of the shell protocol used in
                                    the context of processing this command.

  @return EFI_SUCCESS               the operation was successful
  @return other                     the operation failed.
**/
SHELL_STATUS
EFIAPI
LinuxInitrdCommandHandler (
  IN EFI_SHELL_DYNAMIC_COMMAND_PROTOCOL    *This,
  IN EFI_SYSTEM_TABLE                      *SystemTable,
  IN EFI_SHELL_PARAMETERS_PROTOCOL         *ShellParameters,
  IN EFI_SHELL_PROTOCOL                    *Shell
  )
{
  gEfiShellParametersProtocol = ShellParameters;
  gEfiShellProtocol           = Shell;

  return RunInitrd (gImageHandle, SystemTable);
}

/**
  This is the command help handler function pointer callback type.  This
  function is responsible for displaying help information for the associated
  command.

  @param[in] This                   The instance of the
                                    EFI_SHELL_DYNAMIC_COMMAND_PROTOCOL.
  @param[in] Language               The pointer to the language string to use.

  @return string                    Pool allocated help string, must be freed
                                    by caller
**/
STATIC
CHAR16 *
EFIAPI
LinuxInitrdGetHelp (
  IN EFI_SHELL_DYNAMIC_COMMAND_PROTOCOL    *This,
  IN CONST CHAR8                           *Language
  )
{
  return HiiGetString (mLinuxInitrdShellCommandHiiHandle,
           STRING_TOKEN (STR_GET_HELP_INITRD), Language);
}

STATIC EFI_SHELL_DYNAMIC_COMMAND_PROTOCOL mLinuxInitrdDynamicCommand = {
  L"initrd",
  LinuxInitrdCommandHandler,
  LinuxInitrdGetHelp
};

/**
  Retrieve HII package list from ImageHandle and publish to HII database.

  @param ImageHandle            The image handle of the process.

  @return HII handle.
**/
STATIC
EFI_HII_HANDLE
InitializeHiiPackage (
  EFI_HANDLE                  ImageHandle
  )
{
  EFI_STATUS                  Status;
  EFI_HII_PACKAGE_LIST_HEADER *PackageList;
  EFI_HII_HANDLE              HiiHandle;

  //
  // Retrieve HII package list from ImageHandle
  //
  Status = gBS->OpenProtocol (ImageHandle, &gEfiHiiPackageListProtocolGuid,
                  (VOID **)&PackageList, ImageHandle, NULL,
                  EFI_OPEN_PROTOCOL_GET_PROTOCOL);
  ASSERT_EFI_ERROR (Status);
  if (EFI_ERROR (Status)) {
    return NULL;
  }

  //
  // Publish HII package list to HII Database.
  //
  Status = gHiiDatabase->NewPackageList (gHiiDatabase, PackageList, NULL,
                           &HiiHandle);
  ASSERT_EFI_ERROR (Status);
  if (EFI_ERROR (Status)) {
    return NULL;
  }
  return HiiHandle;
}

/**
  Entry point of Linux Initrd dynamic UEFI Shell command.

  Produce the DynamicCommand protocol to handle "initrd" command.

  @param ImageHandle            The image handle of the process.
  @param SystemTable            The EFI System Table pointer.

  @retval EFI_SUCCESS           Initrd command is executed successfully.
  @retval EFI_ABORTED           HII package was failed to initialize.
  @retval others                Other errors when executing Initrd command.
**/
EFI_STATUS
EFIAPI
LinuxInitrdDynamicShellCommandEntryPoint (
  IN EFI_HANDLE               ImageHandle,
  IN EFI_SYSTEM_TABLE         *SystemTable
  )
{
  EFI_STATUS                  Status;

  mLinuxInitrdShellCommandHiiHandle = InitializeHiiPackage (ImageHandle);
  if (mLinuxInitrdShellCommandHiiHandle == NULL) {
    return EFI_ABORTED;
  }

  Status = gBS->InstallProtocolInterface (&ImageHandle,
                  &gEfiShellDynamicCommandProtocolGuid,
                  EFI_NATIVE_INTERFACE,
                  &mLinuxInitrdDynamicCommand);
  ASSERT_EFI_ERROR (Status);
  return Status;
}

/**
  Unload the dynamic UEFI Shell command.

  @param ImageHandle            The image handle of the process.

  @retval EFI_SUCCESS           The image is unloaded.
  @retval Others                Failed to unload the image.
**/
EFI_STATUS
EFIAPI
LinuxInitrdDynamicShellCommandUnload (
  IN EFI_HANDLE               ImageHandle
)
{
  EFI_STATUS                  Status;

  FreeInitrdFile ();

  Status = UninstallLoadFile2Protocol ();
  if (EFI_ERROR (Status)) {
    return Status;
  }

  Status = gBS->UninstallProtocolInterface (ImageHandle,
                  &gEfiShellDynamicCommandProtocolGuid,
                  &mLinuxInitrdDynamicCommand);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  HiiRemovePackages (mLinuxInitrdShellCommandHiiHandle);
  return EFI_SUCCESS;
}