/** @file
  OemReadyToBootDxe - OEM ReadyToBoot Event Handler

  This module registers a ReadyToBoot event callback that scans IIO NVMe
  ports on both CPU sockets and clears the power fault status.

  Copyright (c) Inspur Corporation. All rights reserved.
  SPDX-License-Identifier: BSD-2-Clause-Patent

  Binary: HR650X BIOS / 0102_OemReadyToBootDxe_fcf4696c02b0
**/

#include "OemReadyToBootDxe.h"

//
// Global protocol pointers
//
EFI_BOOT_SERVICES   *gBS = NULL;
EFI_RUNTIME_SERVICES *gRT = NULL;
EFI_SYSTEM_TABLE    *gST = NULL;
EFI_HANDLE           gImageHandle = NULL;

UINT64  gPciExpressBaseAddress = 0;
VOID   *gHobList = NULL;
VOID   *gIioUdsProtocol = NULL;
VOID   *gPcdProtocol = NULL;
VOID   *gCpuCsrAccess = NULL;
UINT8   gDebugPortValue = 0;

//
// GUID definitions (referenced in .rdata)
//
extern EFI_GUID  gEfiIioUdsProtocolGuid;
extern EFI_GUID  gEfiCpuCsrAccessGuid;
extern EFI_GUID  gEfiPcdProtocolGuid;
extern EFI_GUID  gEfiHobListGuid;

/**
  Read an unaligned 64-bit value from a buffer.

  @param[in]  Buffer  Pointer to the buffer to read from.

  @return The 64-bit value read from the buffer.
**/
UINT64
ReadUnaligned64 (
  IN CONST VOID  *Buffer
  )
{
  if (Buffer == NULL) {
    DebugAssert (
      __FILE__, 192,
      "Buffer != ((void *) 0)"
      );
  }
  return *(UINT64 *)Buffer;
}

/**
  Check whether two GUIDs match by comparing their first and second
  64-bit halves.

  @param[in]  Guid1  Pointer to the first GUID.
  @param[in]  Guid2  Pointer to the second GUID.

  @retval TRUE   The GUIDs match.
  @retval FALSE  The GUIDs do not match.
**/
BOOLEAN
IsHobGuidMatching (
  IN EFI_GUID  *Guid1,
  IN EFI_GUID  *Guid2
  )
{
  UINT64  Data1;
  UINT64  Data2;

  if ((Guid1 == NULL) || (Guid2 == NULL)) {
    return FALSE;
  }

  Data1 = ReadUnaligned64 (Guid1);
  Data2 = ReadUnaligned64 (Guid2);

  return (BOOLEAN)(Data1 == Data2);
}

/**
  Retrieve the HOB list pointer.

  Locates the HOB list from the system table configuration table and
  caches it in gHobList.

  @return Pointer to the HOB list, or NULL if not found.
**/
EFI_HOB_HANDOFF_INFO_TABLE *
GetHobList (
  VOID
  )
{
  EFI_STATUS                 Status;
  UINTN                      Index;
  EFI_CONFIGURATION_TABLE    *ConfigTable;
  EFI_GUID                   mHobListGuid = { 0x7739F24C, 0x1E93, 0x4D66, { 0xB4, 0x90, 0x9B, 0xDB, 0xEB, 0xE1, 0xEE, 0x55 } };

  if (gHobList != NULL) {
    return gHobList;
  }

  gHobList = NULL;

  if (gST->NumberOfTableEntries != 0) {
    ConfigTable = (EFI_CONFIGURATION_TABLE *)(gST->ConfigurationTable);
    for (Index = 0; Index < gST->NumberOfTableEntries; Index++) {
      if (IsHobGuidMatching (&mHobListGuid, &ConfigTable->VendorGuid)) {
        gHobList = (VOID *)(UINTN)ConfigTable->VendorTable;
        break;
      }
      ConfigTable++;
    }
  }

  if (gHobList == NULL) {
    DebugPrint (0x80000000, "\nASSERT_EFI_ERROR (Status = %r)\n", EFI_NOT_FOUND);
    DebugAssert (__FILE__, 54, "!EFI_ERROR (Status)");
  }

  if (gHobList == NULL) {
    DebugAssert (__FILE__, 55, "mHobList != ((void *) 0)");
  }

  return gHobList;
}

/**
  Locate the PCD protocol.

  Locates gEfiPcdProtocolGuid and caches it in gPcdProtocol.

  @return Pointer to the PCD protocol interface, or NULL.
**/
VOID *
GetPcdProtocol (
  VOID
  )
{
  EFI_STATUS  Status;

  if (gPcdProtocol != NULL) {
    return gPcdProtocol;
  }

  Status = gBS->LocateProtocol (
                  &gEfiPcdProtocolGuid,
                  NULL,
                  &gPcdProtocol
                  );
  if (EFI_ERROR (Status)) {
    DebugPrint (0x80000000, "\nASSERT_EFI_ERROR (Status = %r)\n", Status);
    DebugAssert (__FILE__, 78, "!EFI_ERROR (Status)");
  }

  if (gPcdProtocol == NULL) {
    DebugAssert (__FILE__, 79, "mPcd != ((void *) 0)");
  }

  return gPcdProtocol;
}

/**
  Translate a PCI Express register address to a memory-mapped address.

  Validates that the address fits within the PCIe MMIO window
  (upper 4 bits of the 32-bit address must be clear), then adds the
  cached PCIe base address.

  @param[in]  Address  PCI Express register address (bus/device/function/offset).

  @return The full memory-mapped PCIe address.
**/
UINT64
PciExpressRead (
  IN UINT64  Address
  )
{
  if ((Address & 0xFFFFFFFFF0000000ULL) != 0) {
    DebugAssert (
      __FILE__, 118,
      "((Address) & ~0xfffffff) == 0"
      );
  }

  return gPciExpressBaseAddress + Address;
}

/**
  Scan IIO NVMe ports on a given socket and clear the power fault status.

  For each of the 4 NVMe ports on the specified socket, reads the
  port status register and, if the power fault bit (bit 1) is set,
  writes it back to clear the fault.

  @param[in]  Socket  The IIO socket number (byte index into CPUBUSNO).
**/
VOID
IioScanNvmePort (
  IN UINT8  Socket
  )
{
  UINT8   Port;
  UINT32  NvmePortValue;
  UINT32  BaseBus;

  Port = 0;
  BaseBus = Socket;

  do {
    //
    // Calculate the PCIe address: offset 0xAA in the NVMe port's
    // PCI configuration space.
    //
    NvmePortValue = *(volatile UINT32 *)PciExpressRead (
                                          ((UINT32)(Socket << 5) | (Port & 0x1F)) << 15 | 0xAA
                                          );
    DebugPrint (64, "IioScanNvmePort:%x,%x, %x\n", BaseBus, Port, NvmePortValue);

    if ((NvmePortValue & 0x2) != 0) {
      //
      // Clear the power fault status (bit 1 = write-1-to-clear)
      //
      *(volatile UINT32 *)PciExpressRead (
                            ((UINT32)(Socket << 5) | (Port & 0x1F)) << 15 | 0xAA
                            ) = NvmePortValue | 0x2;
    }

    Port++;
  } while (Port < 4);
}

/**
  Clear NVMe power fault status on all sockets.

  Locates the IIO UDS and CPU CSR access protocols, then reads the
  CPUBUSNO for each socket to determine the bus number.  Calls
  IioScanNvmePort for each socket that has NVMe ports.

  For Socket 0: always present.
  For Socket 1: only present if bit 1 of byte at (IioUds + 2067) is set.
**/
VOID
ClearNvmePowerFaultStatus (
  VOID
  )
{
  EFI_STATUS  Status;
  UINT32      CpuBusNo;
  UINT32      Socket1CpuBusNo;
  VOID        *CpuCsrAccess;
  VOID        *IioUds;

  DebugPrint (64, "ClearNvmePowerFaultStatus\n");

  //
  // Locate IIO UDS protocol
  //
  Status = gBS->LocateProtocol (
                  &gEfiIioUdsProtocolGuid,
                  NULL,
                  &IioUds
                  );
  if (EFI_ERROR (Status)) {
    DebugPrint (0x80000000, "Locate gEfiIioUdsProtocolGuid FAILED\n");
    return;
  }

  //
  // Locate CPU CSR access protocol
  //
  Status = gBS->LocateProtocol (
                  &gEfiCpuCsrAccessGuid,
                  NULL,
                  &CpuCsrAccess
                  );
  if (EFI_ERROR (Status)) {
    DebugPrint (0x80000000, "Locate EfiCpuCsrAccessGuid FAILED %r\n", Status);
    return;
  }

  //
  // Read CPUBUSNO for Socket 0 (SocketId=0, StackId=0, Register=0x13023A0C)
  //
  CpuBusNo = ((UINT32 (*)(UINT8, UINT8, UINT32))(CpuCsrAccess)) (0, 0, 318914764);
  DebugPrint (64, "Socket: 0, CPUBUSNO: %08x\n", CpuBusNo);

  IioScanNvmePort ((UINT8)(CpuBusNo >> 8));
  IioScanNvmePort ((UINT8)(CpuBusNo >> 16));
  IioScanNvmePort ((UINT8)(CpuBusNo >> 24));

  //
  // Check if Socket 1 is present (bit 1 in byte 2067 of IIO UDS)
  //
  if (((*(UINT8 *)(*(UINT64 *)IioUds + 2067)) & 0x2) != 0) {
    Socket1CpuBusNo = ((UINT32 (*)(UINT8, UINT8, UINT32))(CpuCsrAccess)) (1, 0, 318914764);
    DebugPrint (64, "Socket: 1, CPUBUSNO: %08x\n", Socket1CpuBusNo);

    IioScanNvmePort ((UINT8)(Socket1CpuBusNo >> 8));
    IioScanNvmePort ((UINT8)(Socket1CpuBusNo >> 16));
    IioScanNvmePort ((UINT8)(Socket1CpuBusNo >> 24));
  }
}

/**
  The ReadyToBoot callback.

  Called when the ReadyToBoot event is signaled, which occurs just before
  the BDS phase hands off to the OS boot loader.

  This function calls ClearNvmePowerFaultStatus to ensure all NVMe
  power fault indicators are cleared before OS boot.

  @param[in]  Event     The ReadyToBoot event.
  @param[in]  Context   Not used (must be NULL).

  @return EFI_STATUS code (always EFI_SUCCESS in practice).
**/
EFI_STATUS
EFIAPI
OemReadyToBootDxeEntry (
  VOID
  )
{
  EFI_STATUS  Status;

  DebugPrint (64, "OemReadyToBootDxeEntry...Start\n");

  Status = EFI_INVALID_PARAMETER;

  //
  // Register the NVMe power fault clearing routine with the
  // ReadyToBoot event group.
  //
  Status = gBS->CreateEventEx (
                  EVT_NOTIFY_SIGNAL,
                  TPL_CALLBACK,
                  ClearNvmePowerFaultStatus,
                  NULL,
                  &gEfiEventReadyToBootGuid,
                  &gReadyToBootEvent
                  );

  if (EFI_ERROR (Status)) {
    DebugPrint (0x80000000, "OemReadyToBootDxeEntry:%r...End\n", Status);
  } else {
    DebugPrint (64, "OemReadyToBootDxeEntry:%r...End\n", Status);
  }

  return Status;
}

/**
  User driver entry point.

  Standard UEFI DXE driver entry point.  Initializes global variables,
  caches protocols, gets the PciExpressBaseAddress from PCD, and calls
  OemReadyToBootDxeEntry().

  @param[in]  ImageHandle     The firmware-allocated handle for this image.
  @param[in]  SystemTable     A pointer to the EFI System Table.

  @retval EFI_SUCCESS         The entry point executed successfully.
  @retval Others              An error occurred.
**/
EFI_STATUS
EFIAPI
ModuleEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS  Status;
  VOID        *PcdProtocol;

  //
  // Initialize UEFI global variables (UefiBootServicesTableLib)
  //
  gImageHandle = ImageHandle;
  gST          = SystemTable;
  gBS          = SystemTable->BootServices;
  gRT          = SystemTable->RuntimeServices;

  //
  // Cache the HOB list
  //
  GetHobList ();

  //
  // Get the PCD protocol to read PciExpressBaseAddress
  //
  PcdProtocol = GetPcdProtocol ();
  gPciExpressBaseAddress = ((UINT64 (*)(UINTN))(*(UINT64 **)PcdProtocol)[4])(5);

  //
  // Register the ReadyToBoot callback
  //
  Status = OemReadyToBootDxeEntry ();
  return Status;
}