/** @file

  This file contains implementation for additional PE/COFF functionality needed to use
  Platform Runtime Mechanism (PRM) modules.

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

**/

#include <IndustryStandard/PeImage.h>
#include <Library/BaseLib.h>
#include <Library/DebugLib.h>
#include <Library/PeCoffLib.h>

#include <PrmExportDescriptor.h>
#include <PrmModuleImageContext.h>

#define _DBGMSGID_  "[PRMPECOFFLIB]"

/**
  Gets a pointer to the export directory in a given PE/COFF image.

  @param[in]  ImageExportDirectory        A pointer to an export directory table in a PE/COFF image.
  @param[in]  PeCoffLoaderImageContext    A pointer to a PE_COFF_LOADER_IMAGE_CONTEXT structure that contains the
                                          PE/COFF image context for the Image containing the PRM Module Export
                                          Descriptor table.
  @param[out] ExportDescriptor            A pointer to a pointer to the PRM Module Export Descriptor table found
                                          in the ImageExportDirectory given.

  @retval EFI_SUCCESS                     The PRM Module Export Descriptor table was found successfully.
  @retval EFI_INVALID_PARAMETER           A required parameter is NULL.
  @retval EFI_NOT_FOUND                   The PRM Module Export Descriptor table was not found in the given
                                          ImageExportDirectory.

**/
EFI_STATUS
GetPrmModuleExportDescriptorTable (
  IN  EFI_IMAGE_EXPORT_DIRECTORY          *ImageExportDirectory,
  IN  PE_COFF_LOADER_IMAGE_CONTEXT        *PeCoffLoaderImageContext,
  OUT PRM_MODULE_EXPORT_DESCRIPTOR_STRUCT **ExportDescriptor
  )
{
  UINTN                                   Index;
  EFI_PHYSICAL_ADDRESS                    CurrentImageAddress;
  UINT16                                  PrmModuleExportDescriptorOrdinal;
  CONST CHAR8                             *CurrentExportName;
  UINT16                                  *OrdinalTable;
  UINT32                                  *ExportNamePointerTable;
  UINT32                                  *ExportAddressTable;
  PRM_MODULE_EXPORT_DESCRIPTOR_STRUCT     *TempExportDescriptor;

  DEBUG ((DEBUG_INFO, "%a %a - Entry.\n", _DBGMSGID_, __FUNCTION__));

  if (ImageExportDirectory == NULL ||
      PeCoffLoaderImageContext == NULL ||
      PeCoffLoaderImageContext->ImageAddress == 0 ||
      ExportDescriptor == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  *ExportDescriptor = NULL;

  DEBUG ((
    DEBUG_INFO,
    "  %a %a: %d exported names found in this image.\n",
    _DBGMSGID_,
    __FUNCTION__,
    ImageExportDirectory->NumberOfNames
    ));

  //
  // The export name pointer table and export ordinal table form two parallel arrays associated by index.
  //
  CurrentImageAddress = PeCoffLoaderImageContext->ImageAddress;
  ExportAddressTable = (UINT32 *) ((UINTN) CurrentImageAddress + ImageExportDirectory->AddressOfFunctions);
  ExportNamePointerTable = (UINT32 *) ((UINTN) CurrentImageAddress + ImageExportDirectory->AddressOfNames);
  OrdinalTable = (UINT16 *) ((UINTN) CurrentImageAddress + ImageExportDirectory->AddressOfNameOrdinals);

  for (Index = 0; Index < ImageExportDirectory->NumberOfNames; Index++) {
    CurrentExportName = (CONST CHAR8 *) ((UINTN) CurrentImageAddress + ExportNamePointerTable[Index]);
    DEBUG ((
      DEBUG_INFO,
      "  %a %a: Export Name[0x%x] - %a.\n",
      _DBGMSGID_,
      __FUNCTION__,
      Index,
      CurrentExportName
      ));
    if (
      AsciiStrnCmp (
        PRM_STRING(PRM_MODULE_EXPORT_DESCRIPTOR_NAME),
        CurrentExportName,
        AsciiStrLen (PRM_STRING(PRM_MODULE_EXPORT_DESCRIPTOR_NAME))
        ) == 0) {
      PrmModuleExportDescriptorOrdinal = OrdinalTable[Index];
      DEBUG ((
        DEBUG_INFO,
        "  %a %a: PRM Module Export Descriptor found. Ordinal = %d.\n",
        _DBGMSGID_,
        __FUNCTION__,
        PrmModuleExportDescriptorOrdinal
        ));
      if (PrmModuleExportDescriptorOrdinal >= ImageExportDirectory->NumberOfFunctions) {
        DEBUG ((DEBUG_ERROR, "%a %a: The PRM Module Export Descriptor ordinal value is invalid.\n", _DBGMSGID_, __FUNCTION__));
        return EFI_NOT_FOUND;
      }
      TempExportDescriptor = (PRM_MODULE_EXPORT_DESCRIPTOR_STRUCT *) ((UINTN) CurrentImageAddress + ExportAddressTable[PrmModuleExportDescriptorOrdinal]);
      if (TempExportDescriptor->Header.Signature == PRM_MODULE_EXPORT_DESCRIPTOR_SIGNATURE) {
        *ExportDescriptor = TempExportDescriptor;
        DEBUG ((DEBUG_INFO, "  %a %a: PRM Module Export Descriptor found at 0x%x.\n", _DBGMSGID_, __FUNCTION__, (UINTN) ExportDescriptor));
      } else {
        DEBUG ((
          DEBUG_INFO,
          "  %a %a: PRM Module Export Descriptor found at 0x%x but signature check failed.\n",
          _DBGMSGID_,
          __FUNCTION__,
          (UINTN) TempExportDescriptor
          ));
      }
      DEBUG ((DEBUG_INFO, "  %a %a: Exiting export iteration since export descriptor found.\n", _DBGMSGID_, __FUNCTION__));
      return EFI_SUCCESS;
    }
  }

  return EFI_NOT_FOUND;
}

/**
  Gets a pointer to the export directory in a given PE/COFF image.

  @param[in]  Image                       A pointer to a PE32/COFF image base address that is loaded into memory
                                          and already relocated to the memory base address. RVAs in the image given
                                          should be valid.
  @param[in]  PeCoffLoaderImageContext    A pointer to a PE_COFF_LOADER_IMAGE_CONTEXT structure that contains the
                                          PE/COFF image context for the Image given.
  @param[out] ImageExportDirectory        A pointer to a pointer to the export directory found in the Image given.

  @retval EFI_SUCCESS                     The export directory was found successfully.
  @retval EFI_INVALID_PARAMETER           A required parameter is NULL.
  @retval EFI_UNSUPPORTED                 The PE/COFF image given is not supported as a PRM Module.
  @retval EFI_NOT_FOUND                   The image export directory could not be found for this image.

**/
EFI_STATUS
GetExportDirectoryInPeCoffImage (
  IN  VOID                                *Image,
  IN  PE_COFF_LOADER_IMAGE_CONTEXT        *PeCoffLoaderImageContext,
  OUT EFI_IMAGE_EXPORT_DIRECTORY          **ImageExportDirectory
  )
{
  UINT16                                  Magic;
  UINT32                                  NumberOfRvaAndSizes;
  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION     OptionalHeaderPtrUnion;
  EFI_IMAGE_DATA_DIRECTORY                *DirectoryEntry;
  EFI_IMAGE_EXPORT_DIRECTORY              *ExportDirectory;

  if (Image == NULL || PeCoffLoaderImageContext == NULL || ImageExportDirectory == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  DirectoryEntry  = NULL;
  ExportDirectory = NULL;

  //
  // NOTE: For backward compatibility, use the Machine field to identify a PE32/PE32+
  //       image instead of using the Magic field. Some systems might generate a PE32+
  //       image with PE32 magic.
  //
  switch (PeCoffLoaderImageContext->Machine) {
  case EFI_IMAGE_MACHINE_IA32:
    //
    // Assume PE32 image with IA32 Machine field.
    //
    Magic = EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC;
    break;
  case EFI_IMAGE_MACHINE_X64:
  case EFI_IMAGE_MACHINE_AARCH64:
    //
    // Assume PE32+ image with X64 Machine field
    //
    Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;
    break;
  default:
    //
    // For unknown Machine field, use Magic in optional header
    //
    DEBUG ((
      DEBUG_WARN,
      "%a %a: The machine type for this image is not valid for a PRM module.\n",
      _DBGMSGID_,
      __FUNCTION__
      ));
    return EFI_UNSUPPORTED;
  }

  OptionalHeaderPtrUnion.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) (
                                  (UINTN) Image +
                                  PeCoffLoaderImageContext->PeCoffHeaderOffset
                                  );

  //
  // Check the PE/COFF Header Signature. Determine if the image is valid and/or a TE image.
  //
  if (OptionalHeaderPtrUnion.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {
    DEBUG ((DEBUG_ERROR, "%a %a: The PE signature is not valid for the current image.\n", _DBGMSGID_, __FUNCTION__));
    return EFI_UNSUPPORTED;
  }

  if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
    //
    // Use the PE32 offset to get the Export Directory Entry
    //
    NumberOfRvaAndSizes = OptionalHeaderPtrUnion.Pe32->OptionalHeader.NumberOfRvaAndSizes;
    DirectoryEntry = (EFI_IMAGE_DATA_DIRECTORY *) &(OptionalHeaderPtrUnion.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_EXPORT]);
  } else if (OptionalHeaderPtrUnion.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
    //
    // Use the PE32+ offset get the Export Directory Entry
    //
    NumberOfRvaAndSizes = OptionalHeaderPtrUnion.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes;
    DirectoryEntry = (EFI_IMAGE_DATA_DIRECTORY *) &(OptionalHeaderPtrUnion.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_EXPORT]);
  } else {
    return EFI_UNSUPPORTED;
  }

  if (NumberOfRvaAndSizes <= EFI_IMAGE_DIRECTORY_ENTRY_EXPORT || DirectoryEntry->VirtualAddress == 0) {
    //
    // The export directory is not present
    //
    return EFI_NOT_FOUND;
  } else if (((UINT32) (~0) - DirectoryEntry->VirtualAddress) < DirectoryEntry->Size) {
    //
    // The directory address overflows
    //
    DEBUG ((DEBUG_ERROR, "%a %a: The export directory entry in this image results in overflow.\n", _DBGMSGID_, __FUNCTION__));
    return EFI_UNSUPPORTED;
  } else {
    DEBUG ((DEBUG_INFO, "%a %a: Export Directory Entry found in the image at 0x%x.\n", _DBGMSGID_, __FUNCTION__, (UINTN) OptionalHeaderPtrUnion.Pe32));
    DEBUG ((DEBUG_INFO, "  %a %a: Directory Entry Virtual Address = 0x%x.\n", _DBGMSGID_, __FUNCTION__, DirectoryEntry->VirtualAddress));

    ExportDirectory = (EFI_IMAGE_EXPORT_DIRECTORY *) ((UINTN) Image + DirectoryEntry->VirtualAddress);
    DEBUG ((
      DEBUG_INFO,
      "  %a %a: Export Directory Table found successfully at 0x%x. Name address = 0x%x. Name = %a.\n",
      _DBGMSGID_,
      __FUNCTION__,
      (UINTN) ExportDirectory,
      ((UINTN) Image + ExportDirectory->Name),
      (CHAR8 *) ((UINTN) Image + ExportDirectory->Name)
      ));
  }
  *ImageExportDirectory = ExportDirectory;

  return EFI_SUCCESS;
}

/**
  Returns the image major and image minor version in a given PE/COFF image.

  @param[in]  Image                       A pointer to a PE32/COFF image base address that is loaded into memory
                                          and already relocated to the memory base address. RVAs in the image given
                                          should be valid.
  @param[in]  PeCoffLoaderImageContext    A pointer to a PE_COFF_LOADER_IMAGE_CONTEXT structure that contains the
                                          PE/COFF image context for the Image given.
  @param[out] ImageMajorVersion           A pointer to a UINT16 buffer to hold the image major version.
  @param[out] ImageMinorVersion           A pointer to a UINT16 buffer to hold the image minor version.

  @retval EFI_SUCCESS                     The image version was read successfully.
  @retval EFI_INVALID_PARAMETER           A required parameter is NULL.
  @retval EFI_UNSUPPORTED                 The PE/COFF image given is not supported.

**/
EFI_STATUS
GetImageVersionInPeCoffImage (
  IN  VOID                                *Image,
  IN  PE_COFF_LOADER_IMAGE_CONTEXT        *PeCoffLoaderImageContext,
  OUT UINT16                              *ImageMajorVersion,
  OUT UINT16                              *ImageMinorVersion
  )
{
  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION     OptionalHeaderPtrUnion;
  UINT16                                  Magic;

  DEBUG ((DEBUG_INFO, "    %a %a - Entry.\n", _DBGMSGID_, __FUNCTION__));

  if (Image == NULL || PeCoffLoaderImageContext == NULL || ImageMajorVersion == NULL || ImageMinorVersion == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  //
  // NOTE: For backward compatibility, use the Machine field to identify a PE32/PE32+
  //       image instead of using the Magic field. Some systems might generate a PE32+
  //       image with PE32 magic.
  //
  switch (PeCoffLoaderImageContext->Machine) {
  case EFI_IMAGE_MACHINE_IA32:
    //
    // Assume PE32 image
    //
    Magic = EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC;
    break;
  case EFI_IMAGE_MACHINE_X64:
  case EFI_IMAGE_MACHINE_AARCH64:
    //
    // Assume PE32+ image
    //
    Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;
    break;
  default:
    //
    // For unknown Machine field, use Magic in optional header
    //
    DEBUG ((
      DEBUG_WARN,
      "%a %a: The machine type for this image is not valid for a PRM module.\n",
      _DBGMSGID_,
      __FUNCTION__
      ));
    return EFI_UNSUPPORTED;
  }

  OptionalHeaderPtrUnion.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) (
                                  (UINTN) Image +
                                  PeCoffLoaderImageContext->PeCoffHeaderOffset
                                  );
  //
  // Check the PE/COFF Header Signature. Determine if the image is valid and/or a TE image.
  //
  if (OptionalHeaderPtrUnion.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {
    DEBUG ((DEBUG_ERROR, "%a %a: The PE signature is not valid for the current image.\n", _DBGMSGID_, __FUNCTION__));
    return EFI_UNSUPPORTED;
  }

  if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
    //
    // Use the PE32 offset to get the Export Directory Entry
    //
    *ImageMajorVersion = OptionalHeaderPtrUnion.Pe32->OptionalHeader.MajorImageVersion;
    *ImageMinorVersion = OptionalHeaderPtrUnion.Pe32->OptionalHeader.MinorImageVersion;
  } else {
    //
    // Use the PE32+ offset to get the Export Directory Entry
    //
    *ImageMajorVersion = OptionalHeaderPtrUnion.Pe32Plus->OptionalHeader.MajorImageVersion;
    *ImageMinorVersion = OptionalHeaderPtrUnion.Pe32Plus->OptionalHeader.MinorImageVersion;
  }

  DEBUG ((DEBUG_INFO, "      %a %a - Image Major Version: 0x%02x.\n", _DBGMSGID_, __FUNCTION__, *ImageMajorVersion));
  DEBUG ((DEBUG_INFO, "      %a %a - Image Minor Version: 0x%02x.\n", _DBGMSGID_, __FUNCTION__, *ImageMinorVersion));

  return EFI_SUCCESS;
}

/**
  Gets the address of an entry in an image export table by ASCII name.

  @param[in]  ExportName                  A pointer to an ASCII name string of the entry name.
  @param[in]  ImageBaseAddress            The base address of the PE/COFF image.
  @param[in]  ImageExportDirectory        A pointer to the export directory in the image.
  @param[out] ExportPhysicalAddress       A pointer that will be updated with the address of the address of the
                                          export entry if found.

  @retval EFI_SUCCESS                     The export entry was found successfully.
  @retval EFI_INVALID_PARAMETER           A required pointer argument is NULL.
  @retval EFI_NOT_FOUND                   An entry with the given ExportName was not found.

**/
EFI_STATUS
GetExportEntryAddress (
  IN  CONST CHAR8                         *ExportName,
  IN  EFI_PHYSICAL_ADDRESS                ImageBaseAddress,
  IN  EFI_IMAGE_EXPORT_DIRECTORY          *ImageExportDirectory,
  OUT EFI_PHYSICAL_ADDRESS                *ExportPhysicalAddress
  )
{
  UINTN                                   ExportNameIndex;
  UINT16                                  CurrentExportOrdinal;
  UINT32                                  *ExportAddressTable;
  UINT32                                  *ExportNamePointerTable;
  UINT16                                  *OrdinalTable;
  CONST CHAR8                             *ExportNameTablePointerName;

  if (ExportName == NULL || ImageBaseAddress == 0 || ImageExportDirectory == NULL || ExportPhysicalAddress == NULL) {
    return EFI_INVALID_PARAMETER;
  }
  *ExportPhysicalAddress = 0;

  ExportAddressTable = (UINT32 *) ((UINTN) ImageBaseAddress + ImageExportDirectory->AddressOfFunctions);
  ExportNamePointerTable = (UINT32 *) ((UINTN) ImageBaseAddress + ImageExportDirectory->AddressOfNames);
  OrdinalTable = (UINT16 *) ((UINTN) ImageBaseAddress + ImageExportDirectory->AddressOfNameOrdinals);

  for (ExportNameIndex = 0; ExportNameIndex < ImageExportDirectory->NumberOfNames; ExportNameIndex++) {
    ExportNameTablePointerName = (CONST CHAR8 *) ((UINTN) ImageBaseAddress + ExportNamePointerTable[ExportNameIndex]);

    if (AsciiStrnCmp (ExportName, ExportNameTablePointerName, PRM_HANDLER_NAME_MAXIMUM_LENGTH) == 0) {
      CurrentExportOrdinal = OrdinalTable[ExportNameIndex];

      ASSERT (CurrentExportOrdinal < ImageExportDirectory->NumberOfFunctions);
      if (CurrentExportOrdinal >= ImageExportDirectory->NumberOfFunctions) {
        DEBUG ((DEBUG_ERROR, "  %a %a: The export ordinal value is invalid.\n", _DBGMSGID_, __FUNCTION__));
        break;
      }

      *ExportPhysicalAddress = (EFI_PHYSICAL_ADDRESS) ((UINTN) ImageBaseAddress + ExportAddressTable[CurrentExportOrdinal]);
      return EFI_SUCCESS;
    }
  }

  return EFI_NOT_FOUND;
}