/** @file

Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
This program and the accompanying materials are licensed and made available under
the terms and conditions of the BSD License that accompanies this distribution.
The full text of the license may be found at
http://opensource.org/licenses/bsd-license.php.

THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.

**/

#include <Uefi.h>
#include <PiDxe.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/TestPointCheckLib.h>
#include <Library/TestPointLib.h>
#include <Library/DebugLib.h>
#include <Library/UefiLib.h>
#include <Library/PrintLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/PeCoffGetEntryPointLib.h>
#include <Guid/MemoryAttributesTable.h>
#include <Protocol/Runtime.h>

CHAR8 *
ShortNameOfMemoryType(
  IN UINT32 Type
  );

VOID
TestPointDumpMemoryAttributesTable (
  IN EFI_MEMORY_ATTRIBUTES_TABLE                     *MemoryAttributesTable
  )
{
  UINTN                 Index;
  EFI_MEMORY_DESCRIPTOR *Entry;
  UINT64                Pages[EfiMaxMemoryType];

  ZeroMem (Pages, sizeof(Pages));

  DEBUG ((DEBUG_INFO, "MemoryAttributesTable:"));
  DEBUG ((DEBUG_INFO, " Version=0x%x", MemoryAttributesTable->Version));
  DEBUG ((DEBUG_INFO, " Count=0x%x", MemoryAttributesTable->NumberOfEntries));
  DEBUG ((DEBUG_INFO, " DescriptorSize=0x%x\n", MemoryAttributesTable->DescriptorSize));
  Entry = (EFI_MEMORY_DESCRIPTOR *)(MemoryAttributesTable + 1);
  DEBUG ((DEBUG_INFO, "Type       Start            End              # Pages          Attributes\n"));
  for (Index = 0; Index < MemoryAttributesTable->NumberOfEntries; Index++) {
    DEBUG ((DEBUG_INFO, ShortNameOfMemoryType(Entry->Type)));
    DEBUG ((DEBUG_INFO, " %016LX-%016LX %016LX %016LX\n",
      Entry->PhysicalStart,
      Entry->PhysicalStart+MultU64x64 (SIZE_4KB,Entry->NumberOfPages) - 1,
      Entry->NumberOfPages,
      Entry->Attribute
      ));
    if (Entry->Type < EfiMaxMemoryType) {
      Pages[Entry->Type] += Entry->NumberOfPages;
    }
    Entry = NEXT_MEMORY_DESCRIPTOR (Entry, MemoryAttributesTable->DescriptorSize);
  }
  
  DEBUG ((DEBUG_INFO, "\n"));
  DEBUG ((DEBUG_INFO, "  RT_Code   : %14ld Pages (%ld Bytes)\n", Pages[EfiRuntimeServicesCode],     MultU64x64(SIZE_4KB, Pages[EfiRuntimeServicesCode])));
  DEBUG ((DEBUG_INFO, "  RT_Data   : %14ld Pages (%ld Bytes)\n", Pages[EfiRuntimeServicesData],     MultU64x64(SIZE_4KB, Pages[EfiRuntimeServicesData])));
  DEBUG ((DEBUG_INFO, "              -------------- \n"));
}

EFI_STATUS
TestPointCheckMemoryAttribute (
  IN EFI_MEMORY_ATTRIBUTES_TABLE     *MemoryAttributesTable,
  IN EFI_PHYSICAL_ADDRESS            Base,
  IN UINT64                          Size,
  IN BOOLEAN                         IsCode,
  IN BOOLEAN                         IsFromSmm
  )
{
  UINTN                 Index;
  EFI_MEMORY_DESCRIPTOR *Entry;
  
  DEBUG ((DEBUG_ERROR, "Attribute Checking 0x%lx - 0x%lx\n", Base, Size));
  Entry = (EFI_MEMORY_DESCRIPTOR *)(MemoryAttributesTable + 1);
  for (Index = 0; Index < MemoryAttributesTable->NumberOfEntries; Index++) {
    if (Base >= Entry->PhysicalStart && Base+Size <= Entry->PhysicalStart+MultU64x64 (SIZE_4KB,Entry->NumberOfPages)) {
      if (IsFromSmm) {
        if (IsCode) {
          if (Entry->Type != EfiRuntimeServicesCode) {
            DEBUG ((DEBUG_ERROR, "Invalid Entry->Type %d\n", Entry->Type));
            return EFI_INVALID_PARAMETER;
          }
          if ((Entry->Attribute & (EFI_MEMORY_RO | EFI_MEMORY_XP)) != EFI_MEMORY_RO) {
            DEBUG ((DEBUG_ERROR, "Invalid Code Entry->Attribute 0x%lx\n", Entry->Attribute));
            return EFI_INVALID_PARAMETER;
          }
        } else {
          if (Entry->Type != EfiRuntimeServicesData) {
            DEBUG ((DEBUG_ERROR, "Invalid Entry->Type %d\n", Entry->Type));
            return EFI_INVALID_PARAMETER;
          }
          if ((Entry->Attribute & (EFI_MEMORY_RO | EFI_MEMORY_XP)) != EFI_MEMORY_XP) {
            DEBUG ((DEBUG_ERROR, "Invalid Data Entry->Attribute 0x%lx\n", Entry->Attribute));
            return EFI_INVALID_PARAMETER;
          }
        }
      } else {
        if (Entry->Type != EfiRuntimeServicesCode) {
          DEBUG ((DEBUG_ERROR, "Invalid Entry->Type %d\n", Entry->Type));
          return EFI_INVALID_PARAMETER;
        }
        if (IsCode) {
          if ((Entry->Attribute & (EFI_MEMORY_RO | EFI_MEMORY_XP)) != EFI_MEMORY_RO) {
            DEBUG ((DEBUG_ERROR, "Invalid Code Entry->Attribute 0x%lx\n", Entry->Attribute));
            return EFI_INVALID_PARAMETER;
          }
        } else {
          if ((Entry->Attribute & (EFI_MEMORY_RO | EFI_MEMORY_XP)) != EFI_MEMORY_XP) {
            DEBUG ((DEBUG_ERROR, "Invalid Data Entry->Attribute 0x%lx\n", Entry->Attribute));
            return EFI_INVALID_PARAMETER;
          }
        }
      }
    }
  }

  return EFI_SUCCESS;
}

EFI_STATUS
TestPointCheckImageMemoryAttribute (
  IN EFI_MEMORY_ATTRIBUTES_TABLE     *MemoryAttributesTable,
  IN EFI_PHYSICAL_ADDRESS            ImageBase,
  IN UINT64                          ImageSize,
  IN BOOLEAN                         IsFromSmm
  )
{
  EFI_STATUS                           Status;
  EFI_STATUS                           ReturnStatus;
  VOID                                 *ImageAddress;
  EFI_IMAGE_DOS_HEADER                 *DosHdr;
  UINT32                               PeCoffHeaderOffset;
  UINT32                               SectionAlignment;
  UINT16                               ImageType;
  EFI_IMAGE_SECTION_HEADER             *Section;
  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
  UINT16                               Magic;
  UINT8                                *Name;
  UINTN                                Index;
  CHAR8                                *PdbPointer;

  ReturnStatus = EFI_SUCCESS;
  //
  // Check whole region
  //
  ImageAddress = (VOID *)(UINTN)ImageBase;

  PdbPointer = PeCoffLoaderGetPdbPointer (ImageAddress);
  if (PdbPointer != NULL) {
    DEBUG ((EFI_D_INFO, "  Image - %a\n", PdbPointer));
  }

  //
  // Check PE/COFF image
  //
  DosHdr = (EFI_IMAGE_DOS_HEADER *) (UINTN) ImageAddress;
  PeCoffHeaderOffset = 0;
  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
    PeCoffHeaderOffset = DosHdr->e_lfanew;
  }

  Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *) (UINTN) ImageAddress + PeCoffHeaderOffset);
  if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {
    DEBUG ((EFI_D_INFO, "Hdr.Pe32->Signature invalid - 0x%x\n", Hdr.Pe32->Signature));
    return EFI_INVALID_PARAMETER;
  }
  
  //
  // Measuring PE/COFF Image Header;
  // But CheckSum field and SECURITY data directory (certificate) are excluded
  //
  if (Hdr.Pe32->FileHeader.Machine == IMAGE_FILE_MACHINE_IA64 && Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
    //
    // NOTE: Some versions of Linux ELILO for Itanium have an incorrect magic value 
    //       in the PE/COFF Header. If the MachineType is Itanium(IA64) and the 
    //       Magic value in the OptionalHeader is EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC
    //       then override the magic value to EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC
    //
    Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;
    ImageType         = Hdr.Pe32->OptionalHeader.Subsystem;
    SectionAlignment  = Hdr.Pe32->OptionalHeader.SectionAlignment;
  } else {
    //
    // Get the magic value from the PE/COFF Optional Header
    //
    Magic = Hdr.Pe32->OptionalHeader.Magic;
    ImageType         = Hdr.Pe32Plus->OptionalHeader.Subsystem;
    SectionAlignment  = Hdr.Pe32Plus->OptionalHeader.SectionAlignment;
  }

  if ((SectionAlignment & (RUNTIME_PAGE_ALLOCATION_GRANULARITY - 1)) != 0) {
    DEBUG ((EFI_D_INFO, "!!!!!!!!  RecordImageMemoryMap - Section Alignment(0x%x) is not %dK  !!!!!!!!\n", SectionAlignment, RUNTIME_PAGE_ALLOCATION_GRANULARITY >> 10));
    PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);
    if (PdbPointer != NULL) {
      DEBUG ((EFI_D_INFO, "!!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
    }
    return EFI_INVALID_PARAMETER;
  }

  Section = (EFI_IMAGE_SECTION_HEADER *) (
               (UINT8 *) (UINTN) ImageAddress +
               PeCoffHeaderOffset +
               sizeof(UINT32) +
               sizeof(EFI_IMAGE_FILE_HEADER) +
               Hdr.Pe32->FileHeader.SizeOfOptionalHeader
               );

  Status = TestPointCheckMemoryAttribute (
             MemoryAttributesTable,
             (UINTN)ImageAddress,
             (UINTN)&Section[Hdr.Pe32->FileHeader.NumberOfSections] - (UINTN)ImageAddress,
             FALSE,
             IsFromSmm
             );
  if (EFI_ERROR(Status)) {
    ReturnStatus = Status;
  }

  for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {
    Name = Section[Index].Name;
    DEBUG ((
      EFI_D_INFO,
      "  Section - '%c%c%c%c%c%c%c%c'\n",
      Name[0],
      Name[1],
      Name[2],
      Name[3],
      Name[4],
      Name[5],
      Name[6],
      Name[7]
      ));
      
    DEBUG ((EFI_D_INFO, "    VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));
    DEBUG ((EFI_D_INFO, "    VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));
    DEBUG ((EFI_D_INFO, "    SizeOfRawData        - 0x%08x\n", Section[Index].SizeOfRawData));
    DEBUG ((EFI_D_INFO, "    PointerToRawData     - 0x%08x\n", Section[Index].PointerToRawData));
    DEBUG ((EFI_D_INFO, "    PointerToRelocations - 0x%08x\n", Section[Index].PointerToRelocations));
    DEBUG ((EFI_D_INFO, "    PointerToLinenumbers - 0x%08x\n", Section[Index].PointerToLinenumbers));
    DEBUG ((EFI_D_INFO, "    NumberOfRelocations  - 0x%08x\n", Section[Index].NumberOfRelocations));
    DEBUG ((EFI_D_INFO, "    NumberOfLinenumbers  - 0x%08x\n", Section[Index].NumberOfLinenumbers));
    DEBUG ((EFI_D_INFO, "    Characteristics      - 0x%08x\n", Section[Index].Characteristics));
    if ((Section[Index].Characteristics & EFI_IMAGE_SCN_CNT_CODE) != 0) {
      //
      // Check code section
      //
      Status = TestPointCheckMemoryAttribute (
                 MemoryAttributesTable,
                 (UINTN)ImageAddress + Section[Index].VirtualAddress,
                 Section[Index].SizeOfRawData,
                 TRUE,
                 IsFromSmm
                 );
    } else {
      //
      // Check data section
      //
      Status = TestPointCheckMemoryAttribute (
                 MemoryAttributesTable,
                 (UINTN)ImageAddress + Section[Index].VirtualAddress,
                 Section[Index].SizeOfRawData,
                 FALSE,
                 IsFromSmm
                 );
    }
    if (EFI_ERROR(Status)) {
      ReturnStatus = Status;
    }
  }

  return ReturnStatus;
}

EFI_STATUS
TestPointCheckUefiMemoryAttributesTable (
  IN EFI_MEMORY_ATTRIBUTES_TABLE                     *MemoryAttributesTable
  )
{
  EFI_STATUS                 Status;
  EFI_STATUS                 ReturnStatus;
  EFI_RUNTIME_ARCH_PROTOCOL  *RuntimeArch;
  LIST_ENTRY                 *Link;
  EFI_RUNTIME_IMAGE_ENTRY    *RuntimeImage;

  Status = gBS->LocateProtocol (
                  &gEfiRuntimeArchProtocolGuid,
                  NULL,
                  &RuntimeArch
                  );
  if (EFI_ERROR (Status)) {
    return EFI_SUCCESS;
  }

  ReturnStatus = EFI_SUCCESS;
  for (Link = RuntimeArch->ImageHead.ForwardLink; Link != &(RuntimeArch->ImageHead); Link = Link->ForwardLink) {
    RuntimeImage = BASE_CR (Link, EFI_RUNTIME_IMAGE_ENTRY, Link);
    Status = TestPointCheckImageMemoryAttribute (
               MemoryAttributesTable,
               (EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase,
               RuntimeImage->ImageSize,
               FALSE
               );
    if (EFI_ERROR(Status)) {
      ReturnStatus = Status;
    }
  }

  return ReturnStatus;
}

EFI_STATUS
TestPointCheckUefiMemAttribute (
  VOID
  )
{
  EFI_STATUS  Status;
  VOID        *MemoryAttributesTable;
  
  DEBUG ((DEBUG_INFO, "==== TestPointCheckUefiMemAttribute - Enter\n"));
  Status = EfiGetSystemConfigurationTable (&gEfiMemoryAttributesTableGuid, &MemoryAttributesTable);
  if (!EFI_ERROR (Status)) {
    TestPointDumpMemoryAttributesTable(MemoryAttributesTable);
    Status = TestPointCheckUefiMemoryAttributesTable(MemoryAttributesTable);
  }

  if (EFI_ERROR (Status)) {
    TestPointLibAppendErrorString (
      PLATFORM_TEST_POINT_ROLE_PLATFORM_IBV,
      NULL,
      TEST_POINT_BYTE4_READY_TO_BOOT_UEFI_MEMORY_ATTRIBUTE_TABLE_FUNCTIONAL_ERROR_CODE \
        TEST_POINT_READY_TO_BOOT \
        TEST_POINT_BYTE4_READY_TO_BOOT_UEFI_MEMORY_ATTRIBUTE_TABLE_FUNCTIONAL_ERROR_STRING
      );
  }
  DEBUG ((DEBUG_INFO, "==== TestPointCheckUefiMemAttribute - Exit\n"));

  return Status;
}