/**@file

Copyright (c) 2006 - 2009, Intel Corporation. All rights reserved.<BR>
Portions copyright (c) 2008 - 2010, Apple Inc. All rights reserved.<BR>
Portions copyright (c) 2011 - 2018, ARM Ltd. All rights reserved.<BR>

SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <PiDxe.h>

#include <Library/ArmMmuLib.h>
#include <Library/BaseLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>
#include <Library/PcdLib.h>
#include <Library/PeCoffLib.h>
#include <Library/PeCoffExtraActionLib.h>

typedef RETURN_STATUS (*REGION_PERMISSION_UPDATE_FUNC) (
  IN  EFI_PHYSICAL_ADDRESS      BaseAddress,
  IN  UINT64                    Length
  );

STATIC
RETURN_STATUS
UpdatePeCoffPermissions (
  IN  CONST PE_COFF_LOADER_IMAGE_CONTEXT      *ImageContext,
  IN  REGION_PERMISSION_UPDATE_FUNC           NoExecUpdater,
  IN  REGION_PERMISSION_UPDATE_FUNC           ReadOnlyUpdater
  )
{
  RETURN_STATUS                         Status;
  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION   Hdr;
  EFI_IMAGE_OPTIONAL_HEADER_UNION       HdrData;
  UINTN                                 Size;
  UINTN                                 ReadSize;
  UINT32                                SectionHeaderOffset;
  UINTN                                 NumberOfSections;
  UINTN                                 Index;
  EFI_IMAGE_SECTION_HEADER              SectionHeader;
  PE_COFF_LOADER_IMAGE_CONTEXT          TmpContext;
  EFI_PHYSICAL_ADDRESS                  Base;

  //
  // We need to copy ImageContext since PeCoffLoaderGetImageInfo ()
  // will mangle the ImageAddress field
  //
  CopyMem (&TmpContext, ImageContext, sizeof (TmpContext));

  if (TmpContext.PeCoffHeaderOffset == 0) {
    Status = PeCoffLoaderGetImageInfo (&TmpContext);
    if (RETURN_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR,
        "%a: PeCoffLoaderGetImageInfo () failed (Status = %r)\n",
        __FUNCTION__, Status));
      return Status;
    }
  }

  if (TmpContext.IsTeImage &&
      TmpContext.ImageAddress == ImageContext->ImageAddress) {
    DEBUG ((DEBUG_INFO, "%a: ignoring XIP TE image at 0x%lx\n", __FUNCTION__,
      ImageContext->ImageAddress));
    return RETURN_SUCCESS;
  }

  if (TmpContext.SectionAlignment < EFI_PAGE_SIZE) {
    //
    // The sections need to be at least 4 KB aligned, since that is the
    // granularity at which we can tighten permissions. So just clear the
    // noexec permissions on the entire region.
    //
    if (!TmpContext.IsTeImage) {
      DEBUG ((DEBUG_WARN,
        "%a: non-TE Image at 0x%lx has SectionAlignment < 4 KB (%lu)\n",
        __FUNCTION__, ImageContext->ImageAddress, TmpContext.SectionAlignment));
    }
    Base = ImageContext->ImageAddress & ~(EFI_PAGE_SIZE - 1);
    Size = ImageContext->ImageAddress - Base + ImageContext->ImageSize;
    return NoExecUpdater (Base, ALIGN_VALUE (Size, EFI_PAGE_SIZE));
  }

  //
  // Read the PE/COFF Header. For PE32 (32-bit) this will read in too much
  // data, but that should not hurt anything. Hdr.Pe32->OptionalHeader.Magic
  // determines if this is a PE32 or PE32+ image. The magic is in the same
  // location in both images.
  //
  Hdr.Union = &HdrData;
  Size = sizeof (EFI_IMAGE_OPTIONAL_HEADER_UNION);
  ReadSize = Size;
  Status = TmpContext.ImageRead (TmpContext.Handle,
                         TmpContext.PeCoffHeaderOffset, &Size, Hdr.Pe32);
  if (RETURN_ERROR (Status) || (Size != ReadSize)) {
    DEBUG ((DEBUG_ERROR,
      "%a: TmpContext.ImageRead () failed (Status = %r)\n",
      __FUNCTION__, Status));
    return Status;
  }

  ASSERT (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE);

  SectionHeaderOffset = TmpContext.PeCoffHeaderOffset + sizeof (UINT32) +
                        sizeof (EFI_IMAGE_FILE_HEADER);
  NumberOfSections    = (UINTN)(Hdr.Pe32->FileHeader.NumberOfSections);

  switch (Hdr.Pe32->OptionalHeader.Magic) {
    case EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC:
      SectionHeaderOffset += Hdr.Pe32->FileHeader.SizeOfOptionalHeader;
      break;
    case EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC:
      SectionHeaderOffset += Hdr.Pe32Plus->FileHeader.SizeOfOptionalHeader;
      break;
    default:
      ASSERT (FALSE);
  }

  //
  // Iterate over the sections
  //
  for (Index = 0; Index < NumberOfSections; Index++) {
    //
    // Read section header from file
    //
    Size = sizeof (EFI_IMAGE_SECTION_HEADER);
    ReadSize = Size;
    Status = TmpContext.ImageRead (TmpContext.Handle, SectionHeaderOffset,
                                   &Size, &SectionHeader);
    if (RETURN_ERROR (Status) || (Size != ReadSize)) {
      DEBUG ((DEBUG_ERROR,
        "%a: TmpContext.ImageRead () failed (Status = %r)\n",
        __FUNCTION__, Status));
      return Status;
    }

    Base = TmpContext.ImageAddress + SectionHeader.VirtualAddress;

    if ((SectionHeader.Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE) == 0) {

      if ((SectionHeader.Characteristics & EFI_IMAGE_SCN_MEM_WRITE) == 0) {

        DEBUG ((DEBUG_INFO,
          "%a: Mapping section %d of image at 0x%lx with RO-XN permissions and size 0x%x\n",
          __FUNCTION__, Index, Base, SectionHeader.Misc.VirtualSize));
        ReadOnlyUpdater (Base, SectionHeader.Misc.VirtualSize);
      } else {
        DEBUG ((DEBUG_WARN,
          "%a: Mapping section %d of image at 0x%lx with RW-XN permissions and size 0x%x\n",
          __FUNCTION__, Index, Base, SectionHeader.Misc.VirtualSize));
      }
    } else {
        DEBUG ((DEBUG_INFO,
          "%a: Mapping section %d of image at 0x%lx with RO-X permissions and size 0x%x\n",
          __FUNCTION__, Index, Base, SectionHeader.Misc.VirtualSize));
        ReadOnlyUpdater (Base, SectionHeader.Misc.VirtualSize);
        NoExecUpdater (Base, SectionHeader.Misc.VirtualSize);
    }

    SectionHeaderOffset += sizeof (EFI_IMAGE_SECTION_HEADER);
  }
  return RETURN_SUCCESS;
}

/**
  Performs additional actions after a PE/COFF image has been loaded and relocated.

  If ImageContext is NULL, then ASSERT().

  @param  ImageContext  Pointer to the image context structure that describes the
                        PE/COFF image that has already been loaded and relocated.

**/
VOID
EFIAPI
PeCoffLoaderRelocateImageExtraAction (
  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
  )
{
  UpdatePeCoffPermissions (
    ImageContext,
    ArmClearMemoryRegionNoExec,
    ArmSetMemoryRegionReadOnly
    );
}



/**
  Performs additional actions just before a PE/COFF image is unloaded.  Any resources
  that were allocated by PeCoffLoaderRelocateImageExtraAction() must be freed.

  If ImageContext is NULL, then ASSERT().

  @param  ImageContext  Pointer to the image context structure that describes the
                        PE/COFF image that is being unloaded.

**/
VOID
EFIAPI
PeCoffLoaderUnloadImageExtraAction (
  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
  )
{
  UpdatePeCoffPermissions (
    ImageContext,
    ArmSetMemoryRegionNoExec,
    ArmClearMemoryRegionReadOnly
    );
}
