/** @file
  Produce Load File Protocol for UEFI Applications in Firmware Volumes

  Copyright (c) 2011 - 2016, Intel Corporation. All rights reserved.<BR>
  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <PiDxe.h>

#include <Guid/LzmaDecompress.h>
#include <Protocol/LoadFile.h>
#include <Protocol/DevicePath.h>
#include <Protocol/FirmwareVolume2.h>
#include <Protocol/FirmwareVolumeBlock.h>

#include <Library/DebugLib.h>
#include <Library/UefiLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/UefiDriverEntryPoint.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/DevicePathLib.h>

#define LOAD_FILE_ON_FV2_PRIVATE_DATA_SIGNATURE  SIGNATURE_32 ('l', 'f', 'f', 'v')

typedef struct {
  UINTN                            Signature;
  EFI_LOAD_FILE_PROTOCOL           LoadFile;
  EFI_DEVICE_PATH_PROTOCOL         *DevicePath;
  EFI_FIRMWARE_VOLUME2_PROTOCOL    *Fv;
  EFI_GUID                         NameGuid;
  LIST_ENTRY                       Link;
} LOAD_FILE_ON_FV2_PRIVATE_DATA;

#define LOAD_FILE_ON_FV2_PRIVATE_DATA_FROM_THIS(a)  CR (a, LOAD_FILE_ON_FV2_PRIVATE_DATA, LoadFile, LOAD_FILE_ON_FV2_PRIVATE_DATA_SIGNATURE)
#define LOAD_FILE_ON_FV2_PRIVATE_DATA_FROM_LINK(a)  CR (a, LOAD_FILE_ON_FV2_PRIVATE_DATA, Link, LOAD_FILE_ON_FV2_PRIVATE_DATA_SIGNATURE)

VOID        *mFvRegistration;
LIST_ENTRY  mPrivateDataList;

/**
  Causes the driver to load a specified file from firmware volume.

  @param[in]      This                Protocol instance pointer.
  @param[in]      FilePath            The device specific path of the file to load.
  @param[in]      BootPolicy          If TRUE, indicates that the request originates from the
                                      boot manager is attempting to load FilePath as a boot
                                      selection. If FALSE, then FilePath must match an exact file
                                      to be loaded.
  @param[in, out] BufferSize          On input the size of Buffer in bytes. On output with a return
                                      code of EFI_SUCCESS, the amount of data transferred to
                                      Buffer. On output with a return code of EFI_BUFFER_TOO_SMALL,
                                      the size of Buffer required to retrieve the requested file.
  @param[in]      Buffer              The memory buffer to transfer the file to. IF Buffer is NULL,
                                      then no the size of the requested file is returned in
                                      BufferSize.

  @retval EFI_SUCCESS                 The file was loaded.
  @retval EFI_UNSUPPORTED             The device does not support the provided BootPolicy.
  @retval EFI_INVALID_PARAMETER       FilePath is not a valid device path, or
                                      BufferSize is NULL.
  @retval EFI_DEVICE_ERROR            The file was not loaded due to a device error.
  @retval EFI_NOT_FOUND               The file was not found.
  @retval EFI_OUT_OF_RESOURCES        An allocation failure occurred.
  @retval EFI_ACCESS_DENIED           The firmware volume is configured to
                                      disallow reads.
**/
EFI_STATUS
EFIAPI
LoadFileOnFv2LoadFile (
  IN     EFI_LOAD_FILE_PROTOCOL    *This,
  IN     EFI_DEVICE_PATH_PROTOCOL  *FilePath,
  IN     BOOLEAN                   BootPolicy,
  IN OUT UINTN                     *BufferSize,
  IN     VOID                      *Buffer       OPTIONAL
  )
{
  EFI_STATUS                     Status;
  LOAD_FILE_ON_FV2_PRIVATE_DATA  *Private;
  VOID                           *Pe32Buffer;
  UINTN                          Pe32BufferSize;
  UINT32                         AuthenticationStatus;

  if ((This == NULL) || (BufferSize == NULL)) {
    return EFI_INVALID_PARAMETER;
  }

  //
  // Only support BootPolicy
  //
  if (!BootPolicy) {
    return EFI_UNSUPPORTED;
  }

  //
  // Get private context data
  //
  Private = LOAD_FILE_ON_FV2_PRIVATE_DATA_FROM_THIS (This);

  //
  // Determine the size of the PE32 section
  //
  Pe32Buffer     = NULL;
  Pe32BufferSize = 0;
  Status         = Private->Fv->ReadSection (
                                  Private->Fv,
                                  &Private->NameGuid,
                                  EFI_SECTION_PE32,
                                  0,
                                  &Pe32Buffer,
                                  &Pe32BufferSize,
                                  &AuthenticationStatus
                                  );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // If the buffer passed in is not large enough, return the size of the required
  // buffer in BufferSize and return EFI_BUFFER_TOO_SMALL
  //
  if ((*BufferSize < Pe32BufferSize) || (Buffer == NULL)) {
    *BufferSize = Pe32BufferSize;
    return EFI_BUFFER_TOO_SMALL;
  }

  //
  // The buffer passed in is large enough, so read the PE32 section directly into
  // the buffer, update BufferSize with the actual size read, and return the status
  // from ReadSection()
  //
  return Private->Fv->ReadSection (
                        Private->Fv,
                        &Private->NameGuid,
                        EFI_SECTION_PE32,
                        0,
                        &Buffer,
                        BufferSize,
                        &AuthenticationStatus
                        );
}

LOAD_FILE_ON_FV2_PRIVATE_DATA  mLoadFileOnFv2PrivateDataTemplate = {
  LOAD_FILE_ON_FV2_PRIVATE_DATA_SIGNATURE,
  {
    LoadFileOnFv2LoadFile
  }
};

/**
  Check if the FFS has been installed LoadFileProtocol for it.

  @param[in] NameGuid Point to FFS File GUID to be checked.

  @retval TRUE        The FFS's FileLoadProtocol is in list.
  @retval FALSE       The FFS's FileLoadProtocol is not in list.

**/
BOOLEAN
EFIAPI
IsInPrivateList (
  IN EFI_GUID  *NameGuid
  )
{
  LIST_ENTRY                     *Entry;
  LOAD_FILE_ON_FV2_PRIVATE_DATA  *PrivateData;

  if (IsListEmpty (&mPrivateDataList)) {
    return FALSE;
  }

  for (Entry = (&mPrivateDataList)->ForwardLink; Entry != (&mPrivateDataList); Entry = Entry->ForwardLink) {
    PrivateData = LOAD_FILE_ON_FV2_PRIVATE_DATA_FROM_LINK (Entry);
    if (CompareGuid (NameGuid, &PrivateData->NameGuid)) {
      DEBUG ((DEBUG_INFO, "LoadFileOnFv2:FileLoadProtocol has been installed in:%g\n", NameGuid));
      return TRUE;
    }
  }

  return FALSE;
}

/**
  Create file device path based on FFS file GUID and UI name.

  @param Device    Handle to Firmware Volume.
  @param NameGuid  Point to FFS file GUID.
  @param FileName  Point to FFS UI section name.

  @return the combined device path
**/
EFI_DEVICE_PATH_PROTOCOL *
EFIAPI
CreateFileDevicePath (
  IN EFI_HANDLE    Device,
  IN EFI_GUID      *NameGuid,
  IN CONST CHAR16  *FileName
  )
{
  UINTN                              Size;
  FILEPATH_DEVICE_PATH               *FilePath;
  EFI_DEVICE_PATH_PROTOCOL           *DevicePath;
  EFI_DEVICE_PATH_PROTOCOL           *FileDevicePath;
  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  FileNode;

  EfiInitializeFwVolDevicepathNode (&FileNode, NameGuid);
  DevicePath = AppendDevicePathNode (
                 DevicePathFromHandle (Device),
                 (EFI_DEVICE_PATH_PROTOCOL *)&FileNode
                 );

  Size           = StrSize (FileName);
  FileDevicePath = AllocatePool (Size + SIZE_OF_FILEPATH_DEVICE_PATH + END_DEVICE_PATH_LENGTH);
  if (FileDevicePath != NULL) {
    FilePath                 = (FILEPATH_DEVICE_PATH *)FileDevicePath;
    FilePath->Header.Type    = MEDIA_DEVICE_PATH;
    FilePath->Header.SubType = MEDIA_FILEPATH_DP;
    CopyMem (&FilePath->PathName, FileName, Size);
    SetDevicePathNodeLength (&FilePath->Header, Size + SIZE_OF_FILEPATH_DEVICE_PATH);
    SetDevicePathEndNode (NextDevicePathNode (&FilePath->Header));

    DevicePath = AppendDevicePath (DevicePath, FileDevicePath);
    FreePool (FileDevicePath);
  }

  return DevicePath;
}

/**
  Install LoadFile Protocol for Application FFS.

  @param Handle          FV Handle.

**/
VOID
EFIAPI
InstallFileLoadProtocol (
  EFI_HANDLE  Handle
  )
{
  EFI_STATUS                          Status;
  LOAD_FILE_ON_FV2_PRIVATE_DATA       *Private;
  EFI_FIRMWARE_VOLUME2_PROTOCOL       *Fv;
  EFI_FIRMWARE_VOLUME_BLOCK_PROTOCOL  *Fvb;
  EFI_PHYSICAL_ADDRESS                Address;
  EFI_FV_FILETYPE                     FileType;
  UINTN                               Key;
  EFI_GUID                            NameGuid;
  EFI_FV_FILE_ATTRIBUTES              Attributes;
  UINTN                               Size;
  EFI_HANDLE                          LoadFileHandle;
  UINT32                              AuthenticationStatus;
  CHAR16                              *UiName;
  UINTN                               UiNameSize;

  DEBUG ((DEBUG_INFO, "LoadFileOnFv2:Find a FV!\n"));
  Status = gBS->HandleProtocol (Handle, &gEfiFirmwareVolume2ProtocolGuid, (VOID **)&Fv);
  ASSERT_EFI_ERROR (Status);
  Status = gBS->HandleProtocol (Handle, &gEfiFirmwareVolumeBlockProtocolGuid, (VOID **)&Fvb);
  Fvb->GetPhysicalAddress (Fvb, &Address);
  DEBUG ((DEBUG_INFO, "LoadFileOnFv2:Fvb->Address=%x \n", Address));

  //
  // Use Firmware Volume 2 Protocol to search for a FFS files of type
  // EFI_FV_FILETYPE_APPLICATION and produce a LoadFile protocol for
  // each one found.
  //
  FileType = EFI_FV_FILETYPE_APPLICATION;
  Key      = 0;
  while (TRUE) {
    Status = Fv->GetNextFile (Fv, &Key, &FileType, &NameGuid, &Attributes, &Size);
    if (EFI_ERROR (Status)) {
      break;
    }

    UiName = NULL;
    Status = Fv->ReadSection (
                   Fv,
                   &NameGuid,
                   EFI_SECTION_USER_INTERFACE,
                   0,
                   (VOID **)&UiName,
                   &UiNameSize,
                   &AuthenticationStatus
                   );
    if (EFI_ERROR (Status)) {
      continue;
    }

    if (!IsInPrivateList (&NameGuid)) {
      Private = (LOAD_FILE_ON_FV2_PRIVATE_DATA *)AllocateCopyPool (sizeof (mLoadFileOnFv2PrivateDataTemplate), &mLoadFileOnFv2PrivateDataTemplate);
      ASSERT (Private != NULL);
      Private->Fv         = Fv;
      Private->DevicePath = CreateFileDevicePath (Handle, &NameGuid, UiName);
      CopyGuid (&Private->NameGuid, &NameGuid);
      LoadFileHandle = NULL;
      DEBUG ((DEBUG_INFO, "Find a APPLICATION in this FV!\n"));
      Status = gBS->InstallMultipleProtocolInterfaces (
                      &LoadFileHandle,
                      &gEfiDevicePathProtocolGuid,
                      Private->DevicePath,
                      &gEfiLoadFileProtocolGuid,
                      &Private->LoadFile,
                      NULL
                      );
      if (!EFI_ERROR (Status)) {
        InsertTailList (&mPrivateDataList, &Private->Link);
      } else {
        DEBUG ((DEBUG_ERROR, "Application with the same name %s has been installed.!\n", UiName));
        FreePool (Private->DevicePath);
        FreePool (Private);
      }
    }
  }
}

/**
  This notification function is invoked when an instance of the
  LzmaCustomDecompressGuid is produced. It installs another instance of the
  EFI_FIRMWARE_VOLUME_PROTOCOL on the handle of the FFS. This notification function
  also handles the situation when LZMA decoder driver loaded later than FirmwareVolume driver.

  @param  Event                 The event that occurred
  @param  Context               Context of event. Not used in this nofication function.

**/
VOID
EFIAPI
FvNotificationEvent (
  IN  EFI_EVENT  Event,
  IN  VOID       *Context
  )
{
  EFI_STATUS  Status;
  UINTN       BufferSize;
  EFI_HANDLE  *Handle;
  UINTN       Index;
  EFI_HANDLE  *CurHandle;

  Handle     = NULL;
  Index      = 0;
  BufferSize = sizeof (EFI_HANDLE);
  Handle     = AllocateZeroPool (BufferSize);
  if (Handle == NULL) {
    return;
  }

  Status = gBS->LocateHandle (
                  ByProtocol,
                  &gEfiFirmwareVolume2ProtocolGuid,
                  NULL,
                  &BufferSize,
                  Handle
                  );
  if (EFI_BUFFER_TOO_SMALL == Status) {
    FreePool (Handle);
    Handle = AllocateZeroPool (BufferSize);
    if (Handle == NULL) {
      return;
    }

    Status = gBS->LocateHandle (
                    ByProtocol,
                    &gEfiFirmwareVolume2ProtocolGuid,
                    NULL,
                    &BufferSize,
                    Handle
                    );
    if (EFI_ERROR (Status)) {
      return;
    }
  } else if (EFI_ERROR (Status)) {
    return;
  }

  CurHandle = Handle;
  for (Index = 0; Index < BufferSize/sizeof (EFI_HANDLE); Index++) {
    CurHandle = Handle + Index;
    //
    // Install LoadFile Protocol
    //
    InstallFileLoadProtocol (*CurHandle);
  }

  if (Handle != NULL) {
    FreePool (Handle);
  }
}

/**
  Entry point function initializes global variables and installs notifications.

  @param[in] ImageHandle    The firmware allocated handle for the EFI image.
  @param[in] SystemTable    A pointer to the EFI System Table.

  @retval EFI_SUCCESS       The entry point is executed successfully.
  @retval other             Some error occurs when executing this entry point.
**/
EFI_STATUS
EFIAPI
LoadFileOnFv2Intialize (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  InitializeListHead (&mPrivateDataList);

  EfiCreateProtocolNotifyEvent (
    &gEfiFirmwareVolume2ProtocolGuid,
    TPL_CALLBACK,
    FvNotificationEvent,
    NULL,
    &mFvRegistration
    );

  EfiCreateProtocolNotifyEvent (
    &gLzmaCustomDecompressGuid,
    TPL_CALLBACK,
    FvNotificationEvent,
    NULL,
    &mFvRegistration
    );

  return EFI_SUCCESS;
}