/** @file
 *  PE/COFF emulator protocol implementation to start Linux kernel
 *  images from non-native firmware
 *
 *  Copyright (c) 2020, ARM Ltd. All rights reserved.<BR>
 *
 *  SPDX-License-Identifier: BSD-2-Clause-Patent
 *
 */

#include <PiDxe.h>

#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>
#include <Library/UefiBootServicesTableLib.h>

#include <Protocol/PeCoffImageEmulator.h>

#pragma pack (1)
typedef struct {
  UINT8     Type;
  UINT8     Size;
  UINT16    MachineType;
  UINT32    EntryPoint;
} PE_COMPAT_TYPE1;
#pragma pack ()

STATIC
BOOLEAN
EFIAPI
IsImageSupported (
  IN  EDKII_PECOFF_IMAGE_EMULATOR_PROTOCOL  *This,
  IN  UINT16                                ImageType,
  IN  EFI_DEVICE_PATH_PROTOCOL              *DevicePath   OPTIONAL
  )
{
  return ImageType == EFI_IMAGE_SUBSYSTEM_EFI_APPLICATION;
}

STATIC
EFI_IMAGE_ENTRY_POINT
EFIAPI
GetCompatEntryPoint (
  IN  EFI_PHYSICAL_ADDRESS  ImageBase
  )
{
  EFI_IMAGE_DOS_HEADER      *DosHdr;
  UINTN                     PeCoffHeaderOffset;
  EFI_IMAGE_NT_HEADERS32    *Pe32;
  EFI_IMAGE_SECTION_HEADER  *Section;
  UINTN                     NumberOfSections;
  PE_COMPAT_TYPE1           *PeCompat;
  UINTN                     PeCompatEnd;

  DosHdr = (EFI_IMAGE_DOS_HEADER *)(UINTN)ImageBase;
  if (DosHdr->e_magic != EFI_IMAGE_DOS_SIGNATURE) {
    return NULL;
  }

  PeCoffHeaderOffset = DosHdr->e_lfanew;
  Pe32               = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)ImageBase + PeCoffHeaderOffset);

  Section = (EFI_IMAGE_SECTION_HEADER *)((UINTN)&Pe32->OptionalHeader +
                                         Pe32->FileHeader.SizeOfOptionalHeader);
  NumberOfSections = (UINTN)Pe32->FileHeader.NumberOfSections;

  while (NumberOfSections--) {
    if (!CompareMem (Section->Name, ".compat", sizeof (Section->Name))) {
      //
      // Dereference the section contents to find the mixed mode entry point
      //
      PeCompat    = (PE_COMPAT_TYPE1 *)((UINTN)ImageBase + Section->VirtualAddress);
      PeCompatEnd = (UINTN)(VOID *)PeCompat + Section->Misc.VirtualSize;

      while (PeCompat->Type != 0 && (UINTN)(VOID *)PeCompat < PeCompatEnd) {
        if ((PeCompat->Type == 1) &&
            (PeCompat->Size >= sizeof (PE_COMPAT_TYPE1)) &&
            EFI_IMAGE_MACHINE_TYPE_SUPPORTED (PeCompat->MachineType))
        {
          return (EFI_IMAGE_ENTRY_POINT)((UINTN)ImageBase + PeCompat->EntryPoint);
        }

        PeCompat = (PE_COMPAT_TYPE1 *)((UINTN)PeCompat + PeCompat->Size);
        ASSERT ((UINTN)(VOID *)PeCompat < PeCompatEnd);
      }
    }

    Section++;
  }

  return NULL;
}

STATIC
EFI_STATUS
EFIAPI
RegisterImage (
  IN      EDKII_PECOFF_IMAGE_EMULATOR_PROTOCOL  *This,
  IN      EFI_PHYSICAL_ADDRESS                  ImageBase,
  IN      UINT64                                ImageSize,
  IN  OUT EFI_IMAGE_ENTRY_POINT                 *EntryPoint
  )
{
  EFI_IMAGE_ENTRY_POINT  CompatEntryPoint;

  CompatEntryPoint = GetCompatEntryPoint (ImageBase);
  if (CompatEntryPoint == NULL) {
    return EFI_UNSUPPORTED;
  }

  *EntryPoint = CompatEntryPoint;
  return EFI_SUCCESS;
}

STATIC
EFI_STATUS
EFIAPI
UnregisterImage (
  IN  EDKII_PECOFF_IMAGE_EMULATOR_PROTOCOL  *This,
  IN  EFI_PHYSICAL_ADDRESS                  ImageBase
  )
{
  return EFI_SUCCESS;
}

STATIC EDKII_PECOFF_IMAGE_EMULATOR_PROTOCOL  mCompatLoaderPeCoffEmuProtocol = {
  IsImageSupported,
  RegisterImage,
  UnregisterImage,
  EDKII_PECOFF_IMAGE_EMULATOR_VERSION,
  EFI_IMAGE_MACHINE_X64
};

EFI_STATUS
EFIAPI
CompatImageLoaderDxeEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  return gBS->InstallProtocolInterface (
                &ImageHandle,
                &gEdkiiPeCoffImageEmulatorProtocolGuid,
                EFI_NATIVE_INTERFACE,
                &mCompatLoaderPeCoffEmuProtocol
                );
}