/** @file
  This library retrieve the EFI_BOOT_SERVICES pointer from EFI system table in
  library's constructor and allow overriding its properties.

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

**/

#include <Uefi.h>

#include <Library/OcBootServicesTableLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>
#include <Library/DevicePathLib.h>

#include <Protocol/DevicePath.h>
#include <Protocol/LoadedImage.h>

typedef struct OC_REGISTERED_PROTOCOL_ {
  EFI_GUID   *ProtocolGuid;
  VOID       *ProtocolInstance;
  BOOLEAN    Override;
} OC_REGISTERED_PROTOCOL;

EFI_HANDLE                      gImageHandle       = NULL;
EFI_SYSTEM_TABLE                *gST               = NULL;
EFI_BOOT_SERVICES               *gBS               = NULL;

STATIC EFI_CONNECT_CONTROLLER    mConnectController = NULL;
STATIC EFI_OPEN_PROTOCOL         mOpenProtocol = NULL;
STATIC EFI_LOCATE_HANDLE_BUFFER  mLocateHandleBuffer = NULL;
STATIC EFI_LOCATE_PROTOCOL       mLocateProtocol    = NULL;
STATIC OC_REGISTERED_PROTOCOL    mRegisteredProtocols[16];
STATIC UINTN                     mRegisteredProtocolCount = 0;

STATIC
EFI_STATUS
EFIAPI
OcConnectController (
  IN  EFI_HANDLE                    ControllerHandle,
  IN  EFI_HANDLE                    *DriverImageHandle   OPTIONAL,
  IN  EFI_DEVICE_PATH_PROTOCOL      *RemainingDevicePath OPTIONAL,
  IN  BOOLEAN                       Recursive
  )
{
  EFI_STATUS  Status;
  VOID        *DevicePath;

  if (ControllerHandle == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  //
  // Do not connect controllers without device paths.
  //
  // REF: https://bugzilla.tianocore.org/show_bug.cgi?id=2460
  //
  // Doing this reduces the amount of needless work during device
  // connection and resolves issues with firmwares that freeze
  // when connecting handles without device paths.
  //
  Status = gBS->HandleProtocol (
    ControllerHandle,
    &gEfiDevicePathProtocolGuid,
    &DevicePath
    );
  if (EFI_ERROR (Status)) {
    return EFI_NOT_FOUND;
  }

  Status = mConnectController (
    ControllerHandle,
    DriverImageHandle,
    RemainingDevicePath,
    Recursive
    );

  return Status;
}

STATIC
EFI_STATUS
EFIAPI
OcOpenProtocol (
  IN  EFI_HANDLE                Handle,
  IN  EFI_GUID                  *Protocol,
  OUT VOID                      **Interface OPTIONAL,
  IN  EFI_HANDLE                AgentHandle,
  IN  EFI_HANDLE                ControllerHandle,
  IN  UINT32                    Attributes
  )
{
  EFI_STATUS                 Status;
  UINTN                      Index;
  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;

  Status = mOpenProtocol (
    Handle,
    Protocol,
    Interface,
    AgentHandle,
    ControllerHandle,
    Attributes
    );

  //
  // On Apple EFI some paths may not be valid for our DevicePath library.
  // As a result we can end up in an infinite loop in GetImageNameFromHandle
  // calling "drivers" in the Shell. Workaround by discarding unsupported paths.
  //
  if (!EFI_ERROR (Status)
    && CompareGuid (Protocol, &gEfiLoadedImageProtocolGuid)
    && Interface != NULL) {
    LoadedImage = *Interface;

    if (LoadedImage->FilePath != NULL && !IsDevicePathValid (LoadedImage->FilePath, 0)) {
      LoadedImage->FilePath = NULL;
    }
  }

  if (Status != EFI_UNSUPPORTED || Handle != gImageHandle) {
    return Status;
  }

  if (Interface == NULL && Attributes != EFI_OPEN_PROTOCOL_TEST_PROTOCOL) {
    return EFI_INVALID_PARAMETER;
  }

  //
  // Process all protocols as overrides are not specific to handle.
  //
  for (Index = 0; Index < mRegisteredProtocolCount; ++Index) {
    if (CompareGuid (mRegisteredProtocols[Index].ProtocolGuid, Protocol)) {
      if (Interface != NULL) {
        *Interface = mRegisteredProtocols[Index].ProtocolInstance;
      }
      return EFI_SUCCESS;
    }
  }

  return Status;
}

STATIC
EFI_STATUS
EFIAPI
OcLocateHandleBuffer (
  IN     EFI_LOCATE_SEARCH_TYPE       SearchType,
  IN     EFI_GUID                     *Protocol      OPTIONAL,
  IN     VOID                         *SearchKey     OPTIONAL,
  IN OUT UINTN                        *NoHandles,
  OUT    EFI_HANDLE                   **Buffer
  )
{
  EFI_STATUS  Status;
  UINTN       Index;

  Status = mLocateHandleBuffer (
    SearchType,
    Protocol,
    SearchKey,
    NoHandles,
    Buffer
    );

  if (Status == EFI_NOT_FOUND && SearchType == ByProtocol) {
    //
    // Process all protocols as overrides are not specific to handle.
    //
    for (Index = 0; Index < mRegisteredProtocolCount; ++Index) {
      if (CompareGuid (mRegisteredProtocols[Index].ProtocolGuid, Protocol)) {
        Status = gBS->AllocatePool (
          EfiBootServicesData,
          sizeof (**Buffer),
          (VOID **) Buffer
          );
        if (!EFI_ERROR (Status)) {
          *NoHandles = 1;
          **Buffer   = gImageHandle;
          return EFI_SUCCESS;
        }

        return EFI_NOT_FOUND;
      }
    }
  }

  return Status;
}

STATIC
EFI_STATUS
EFIAPI
OcLocateProtocol (
  IN  EFI_GUID  *Protocol,
  IN  VOID      *Registration OPTIONAL,
  OUT VOID      **Interface
  )
{
  EFI_STATUS  Status;
  UINTN       Index;

  if (Protocol == NULL || Interface == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  //
  // Process overriding protocols first.
  //
  for (Index = 0; Index < mRegisteredProtocolCount; ++Index) {
    if (mRegisteredProtocols[Index].Override
      && CompareGuid (mRegisteredProtocols[Index].ProtocolGuid, Protocol)) {
      *Interface = mRegisteredProtocols[Index].ProtocolInstance;
      return EFI_SUCCESS;
    }
  }

  //
  // Call the original function second.
  //
  Status = mLocateProtocol (
    Protocol,
    Registration,
    Interface
    );

  //
  // Process fallback protocols third.
  //
  if (EFI_ERROR (Status)) {
    for (Index = 0; Index < mRegisteredProtocolCount; ++Index) {
      if (!mRegisteredProtocols[Index].Override
        && CompareGuid (mRegisteredProtocols[Index].ProtocolGuid, Protocol)) {
        *Interface = mRegisteredProtocols[Index].ProtocolInstance;
        return EFI_SUCCESS;
      }
    }
  }

  return Status;
}

EFI_STATUS
OcRegisterBootServicesProtocol (
  IN EFI_GUID   *ProtocolGuid,
  IN VOID       *ProtocolInstance,
  IN BOOLEAN    Override
  )
{
  if (mRegisteredProtocolCount == ARRAY_SIZE (mRegisteredProtocols)) {
    return EFI_OUT_OF_RESOURCES;
  }

  mRegisteredProtocols[mRegisteredProtocolCount].ProtocolGuid     = ProtocolGuid;
  mRegisteredProtocols[mRegisteredProtocolCount].ProtocolInstance = ProtocolInstance;
  mRegisteredProtocols[mRegisteredProtocolCount].Override         = Override;

  ++mRegisteredProtocolCount;

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
OcBootServicesTableLibConstructor (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS  Status;

  //
  // Assert on invalid startup arguments.
  //
  ASSERT (ImageHandle != NULL);
  ASSERT (SystemTable != NULL);
  ASSERT (SystemTable->BootServices != NULL);

  //
  // Cache the Image Handle.
  //
  gImageHandle = ImageHandle;

  //
  // Cache EFI System Table.
  // Note, we cannot override it as it may be used for ConOut
  // spoofing to redirect output for tools.
  //
  gST = SystemTable;

  //
  // Allocate a copy of EFI Boot Services Table.
  //
  Status = SystemTable->BootServices->AllocatePool (
    EfiBootServicesData,
    SystemTable->BootServices->Hdr.HeaderSize,
    (VOID **) &gBS
    );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Copy the original EFI Boot Services Table into our custom table.
  //
  CopyMem (
    gBS,
    SystemTable->BootServices,
    SystemTable->BootServices->Hdr.HeaderSize
    );

  //
  // Override boot services functions and rehash.
  //
  mConnectController      = gBS->ConnectController;
  mOpenProtocol           = gBS->OpenProtocol;
  mLocateHandleBuffer     = gBS->LocateHandleBuffer;
  mLocateProtocol         = gBS->LocateProtocol;
  gBS->ConnectController  = OcConnectController;
  gBS->OpenProtocol       = OcOpenProtocol;
  gBS->LocateHandleBuffer = OcLocateHandleBuffer;
  gBS->LocateProtocol     = OcLocateProtocol;
  gBS->Hdr.CRC32          = 0;
  SystemTable->BootServices->CalculateCrc32 (
    gBS,
    gBS->Hdr.HeaderSize,
    &gBS->Hdr.CRC32
    );

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
OcBootServicesTableLibDestructor (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  //
  // Free memory for table copies.
  //
  SystemTable->BootServices->FreePool (gBS);
  return EFI_SUCCESS;
}