/** @file
  SlotDataUpdateDxeLightningRidgeEXRP - Platform Slot Table (PSLT) update driver
  for LightningRidge EXRP platforms.

  This driver is part of the Lenovo UBA (Universal BIOS Architecture) framework.
  It provides platform-specific PCIe slot configuration data for the EXRP variant
  of the LightningRidge platform by registering PSLT entries through the UBA
  protocol.

  Build environment: VS2015, X64, DEBUG
  Source tree:      PurleyRpPkg\Uba\UbaMain\Dxe\TypeLightningRidgeEXRP\SlotDataUpdateDxe
  PDB:              SlotDataUpdateDxeLightningRidgeEXRP.pdb

  Copyright (C) 2026, Lenovo. All rights reserved.
  SPDX-License-Identifier: BSD-2-Clause-Patent
**/

#include "SlotDataUpdateDxeLightningRidgeEXRP.h"

//
// =============================================================================
// GUID Definitions (stored in .rdata section)
// =============================================================================
//

EFI_GUID gUbaProtocolGuid          = UBA_PROTOCOL_GUID;
EFI_GUID gUbaSlotDataPsl1Guid      = UBA_SLOT_DATA_PSLT1_GUID;
EFI_GUID gUbaSlotDataPsl2Guid      = UBA_SLOT_DATA_PSLT2_GUID;
EFI_GUID gEfiDxeServicesTableGuid  = DXE_SERVICES_TABLE_GUID;

//
// =============================================================================
// Global Variables
// =============================================================================
//

EFI_HANDLE            gImageHandle      = NULL;  ///< 0xBF8: Image handle from entry point
EFI_SYSTEM_TABLE     *gSystemTable      = NULL;  ///< 0xC00: System table from entry point
EFI_BOOT_SERVICES    *gBootServices     = NULL;  ///< 0xC08: Boot services table pointer
EFI_RUNTIME_SERVICES *gRuntimeServices  = NULL;  ///< 0xC10: Runtime services table pointer
VOID                 *gUbaProtocol      = NULL;  ///< 0xC18: Cached UBA protocol interface
VOID                 *gHobList          = NULL;  ///< 0xC20: Cached HOB list pointer
UINT8                 gSlotPlatformId   = 0;     ///< 0xC28: Cached platform slot type

//
// =============================================================================
// Platform Slot Data (embedded in .data section)
//
// Two PSLT entries are defined statically and registered with the UBA protocol
// during initialization.
//
// PSLT Entry 1 (at 0xBB0, 40 bytes):
//   - Full PSLT_HEADER with both function pointers
//   - GetSlotCount function reference (lazy-relocated via IMAGE_REL_BASED_DIR64)
//   - GetSlotData function reference (lazy-relocated via IMAGE_REL_BASED_DIR64)
//   - Flags field = 1 (one slot, or slot count)
//
// PSLT Entry 2 (at 0xBD8, 32 bytes):
//   - Compact PSLT_HEADER with single function pointer
//   - GetSlotCount function reference (lazy-relocated via IMAGE_REL_BASED_DIR64)
//   - No GetSlotData function pointer
//   - Flags field = 1 (one slot, or slot count)
//
// The function pointers at offsets 16 and 32 within Entry 1, and offset 16
// within Entry 2, are IMAGE_REL_BASED_DIR64 fixups applied by the PE loader
// at runtime to relocate from RVA to absolute address.
//
// =============================================================================
//

// PSLT Entry 1 (at 0xBB0, size 40): Primary PSLT with GetSlotCount + GetSlotData
STATIC CONST UINT8 mPslEntry1[40] = {
  'P', 'S', 'L', 'T',             // [0x00] Signature = "PSLT" (0x544C5350)
  0x01, 0x00, 0x00, 0x00,         // [0x04] Version = 1
  0x00, 0x00, 0x00, 0x00,         // [0x08] Reserved (low 32 bits)
  0x00, 0x00, 0x00, 0x00,         // [0x0C] Reserved (high 32 bits)
  0xB0, 0x04, 0x00, 0x00,         // [0x10] BoardDataFunc = RVA 0x4B0 -> GetSlotCount
  0x00, 0x00, 0x00, 0x00,         // [0x14] BoardDataFunc (high 32 bits, reloc target)
  0x01, 0x00, 0x00, 0x00,         // [0x18] Flags = 1 (slot count or config flag)
  0x00, 0x00, 0x00, 0x00,         // [0x1C] Flags (high 32 bits)
  0xB4, 0x04, 0x00, 0x00,         // [0x20] GetSlotDataFunc = RVA 0x4B4 -> GetSlotData
  0x00, 0x00, 0x00, 0x00          // [0x24] GetSlotDataFunc (high 32 bits, reloc target)
};

// PSLT Entry 2 (at 0xBD8, size 32): Secondary PSLT with GetSlotCount only
STATIC CONST UINT8 mPslEntry2[32] = {
  'P', 'S', 'L', 'T',             // [0x00] Signature = "PSLT" (0x544C5350)
  0x01, 0x00, 0x00, 0x00,         // [0x04] Version = 1
  0x00, 0x00, 0x00, 0x00,         // [0x08] Reserved (low 32 bits)
  0x00, 0x00, 0x00, 0x00,         // [0x0C] Reserved (high 32 bits)
  0xB0, 0x04, 0x00, 0x00,         // [0x10] BoardDataFunc = RVA 0x4B0 -> GetSlotCount
  0x00, 0x00, 0x00, 0x00,         // [0x14] BoardDataFunc (high 32 bits, reloc target)
  0x01, 0x00, 0x00, 0x00,         // [0x18] Flags = 1 (slot count or config flag)
  0x00, 0x00, 0x00, 0x00          // [0x1C] Flags (high 32 bits)
};

//
// =============================================================================
// Function Implementations
// =============================================================================
//

/**
  Returns the number of PCIe slots on this platform.

  For LightningRidge EXRP, the platform has 2 PCIe slots.

  @return UINT8  2 (fixed slot count for this platform).
**/
UINT8
GetSlotCount (
  VOID
  )
{
  //
  // LightningRidge EXRP platform has exactly 2 slots.
  //
  return 2;
}

/**
  Returns slot configuration data based on slot number.

  This is a pass-through function that returns the default slot data
  unchanged. For LightningRidge EXRP, slot configuration is fully
  determined by the PSLT entry data itself.

  @param[in] SlotNumber   The slot number (0-indexed) to query.
  @param[in] DefaultData  Default slot data value to return.

  @return UINT8  The DefaultData value, passed through unchanged.
**/
UINT8
GetSlotData (
  IN UINT8 SlotNumber,
  IN UINT8 DefaultData
  )
{
  //
  // Pass through the default slot data unchanged.
  // No platform-specific slot data remapping is needed.
  //
  return DefaultData;
}

/**
  Locates and caches the UBA protocol interface.

  This function checks the current TPL level before attempting to call
  LocateProtocol. If the TPL is above TPL_NOTIFY (0x10), the operation
  is skipped because LocateProtocol cannot be safely called at elevated
  TPL. The protocol interface is cached in gUbaProtocol for subsequent
  calls to avoid repeated lookups.

  @return Pointer to the UBA protocol interface, or NULL if:
          - TPL is above TPL_NOTIFY (not safe to call LocateProtocol)
          - LocateProtocol failed to find the UBA protocol
**/
VOID *
UbaGetProtocolInterface (
  VOID
  )
{
  //
  // Return cached protocol if already located.
  //
  if (gUbaProtocol != NULL) {
    return gUbaProtocol;
  }

  //
  // Check current TPL level by raising to TPL_HIGH_LEVEL and restoring.
  // If the old TPL was above TPL_NOTIFY (0x10), we cannot safely call
  // LocateProtocol (which may block/allocate memory).
  //
  {
    EFI_TPL  OldTpl;

    OldTpl = gBootServices->RaiseTPL (TPL_HIGH_LEVEL);
    gBootServices->RestoreTPL (OldTpl);

    if (OldTpl > TPL_NOTIFY) {
      //
      // TPL too high for LocateProtocol - skip and return NULL.
      //
      return NULL;
    }
  }

  //
  // Locate the UBA protocol using its GUID.
  //
  {
    EFI_STATUS  Status;

    Status = gBootServices->LocateProtocol (
                              &gUbaProtocolGuid,   // GUID: 36232936-0E76-31C8-A13A-3AF2FC1C3932
                              NULL,                 // No registration key
                              &gUbaProtocol         // Output: located protocol interface
                              );

    if (EFI_ERROR (Status)) {
      gUbaProtocol = NULL;
    }
  }

  return gUbaProtocol;
}

/**
  Sends a debug message through the UBA protocol.

  Reads the platform type from RTC CMOS register 0x4B (I/O ports 0x70/0x71)
  and determines the debug routing. The platform type byte (N3) maps to a
  debug route mask. If the specified DebugLevel matches the route mask, the
  message is forwarded to the UBA protocol's DebugPrint function.

  Platform type detection:
    - N3 = 0: Read hardware strap from fixed address 0xFDAF0490 bit 1, OR with 1
    - N3 = 1: Debug route = 0x80000004
    - N3 >= 2: Debug route = 0x80000006

  @param[in] DebugLevel  Debug level mask. Messages with matching bits in the
                          routing mask are output (typically 0x80000000).
  @param[in] Format      Format string for the debug message.
  @param[in] ...         Variable argument list for format substitution.
**/
VOID
EFIAPI
UbaDebugPrint (
  IN UINTN        DebugLevel,
  IN CONST CHAR8 *Format,
  ...
  )
{
  UINTN        DebugRoute;
  UINT8        PlatformType;
  UINT8        N3;
  UINT64     (*UbaDebugFunc)(UINT64, CHAR8 *, UINT64 *);
  VA_LIST      Va;

  DebugRoute = 0;

  if (UbaGetProtocolInterface () != NULL) {
    //
    // Read platform type from RTC CMOS register 0x4B.
    // I/O port 0x70 = CMOS address register (preserve NMI mask bit 7).
    // I/O port 0x71 = CMOS data register.
    //
    PlatformType  = IoRead8 (RTC_CMOS_ADDRESS_PORT);
    IoWrite8 (RTC_CMOS_ADDRESS_PORT, (PlatformType & 0x80) | RTC_CMOS_REG_4B);
    N3 = IoRead8 (RTC_CMOS_DATA_PORT);

    //
    // Determine debug routing based on platform type.
    //
    if (N3 > 3) {
      //
      // N3 is in the valid platform type range (4-255), use as-is.
      // No routing change needed.
      //
    } else if (N3 == 0) {
      //
      // Platform type 0: read hardware strap from fixed memory address.
      // Bit 1 of the strap byte indicates debug capability.
      //
      N3 = (*(volatile UINT8 *)PLATFORM_STRAP_ADDRESS & 0x02) | 0x01;
    }
    // N3 values 1, 2, 3: use directly (valid)

    //
    // Map the platform type to a debug route value.
    // N3 of 0 would have been remapped above, so 1+ is expected.
    //
    if ((N3 - 1) <= 0xFD) {
      //
      // Valid range: N3 is 1 to 255 (effectively all non-zero values).
      //
      if (N3 == 1) {
        DebugRoute = UBA_DEBUG_ROUTE_A;    // 0x80000004
      } else {
        DebugRoute = UBA_DEBUG_ROUTE_B;    // 0x80000006
      }
    }

    //
    // If the debug level matches the computed route mask, forward the message.
    //
    if ((DebugRoute & DebugLevel) != 0) {
      //
      // UBA protocol function index 1 (offset 0x08): DebugPrint
      // The function is called with (Severity, FormatString, VaList).
      //
      UbaDebugFunc = (UINT64 (*)(UINT64, CHAR8 *, UINT64 *))
                     (*((UINT64 *)((UINT8 *)gUbaProtocol + UBA_PROTOCOL_DEBUG_PRINT_OFFSET)));

      VA_START (Va, Format);
      UbaDebugFunc ((UINT64)DebugLevel, (CHAR8 *)Format, (UINT64 *)Va);
      VA_END (Va);
    }
  }
}

/**
  UBA-aware ASSERT replacement.

  Prints the assertion failure information through the UBA protocol's
  debug output function and then triggers a system halt or deadloop.

  @param[in] FileName     Source file name where the assertion occurred.
  @param[in] LineNumber   Line number where the assertion occurred.
  @param[in] AssertString The assertion expression that failed.
**/
VOID
EFIAPI
UbaDebugAssert (
  IN CONST CHAR8 *FileName,
  IN UINTN        LineNumber,
  IN CONST CHAR8 *AssertString
  )
{
  UINT64  (*UbaAssertFunc)(UINT64, UINT64, UINT64);

  //
  // Get the UBA protocol and call its assert handler.
  // The assert handler is at the same function index as DebugPrint
  // (index 1, offset 0x08) but called with (FileName, LineNumber, AssertString).
  //
  if (UbaGetProtocolInterface () != NULL) {
    UbaAssertFunc = (UINT64 (*)(UINT64, UINT64, UINT64))
                    (*((UINT64 *)((UINT8 *)gUbaProtocol + UBA_PROTOCOL_DEBUG_PRINT_OFFSET)));
    UbaAssertFunc ((UINT64)FileName, (UINT64)LineNumber, (UINT64)AssertString);
  }
}

/**
  Reads a UINT64 value from a potentially unaligned memory address.

  Wraps the BaseLib ReadUnaligned64 function with a NULL pointer check.
  If Buffer is NULL, an assertion is triggered via UbaDebugAssert.

  @param[in] Buffer  Pointer to the memory to read. Must not be NULL.

  @return The UINT64 value at the given address.
**/
UINT64
ReadUnaligned64 (
  IN CONST VOID *Buffer
  )
{
  if (Buffer == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\BaseLib\\Unaligned.c",
      192,
      "Buffer != ((void *) 0)"
      );
  }

  return *(const UINT64 *)Buffer;
}

/**
  Checks whether two GUIDs are equal by comparing their 64-bit halves.

  This performs a double 64-bit compare rather than a byte-by-byte or
  128-bit compare. Each GUID is read as two unaligned 64-bit values.

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

  @return TRUE if the GUIDs are equal, FALSE otherwise.
**/
BOOLEAN
IsHobGuidMatch (
  IN CONST EFI_GUID *Guid1,
  IN CONST EFI_GUID *Guid2
  )
{
  UINT64  Guid1Low;
  UINT64  Guid1High;
  UINT64  Guid2Low;
  UINT64  Guid2High;

  //
  // Read each GUID as two unaligned 64-bit values and compare.
  //
  Guid1Low  = ReadUnaligned64 (Guid1);
  Guid1High = ReadUnaligned64 ((CONST UINT8 *)Guid1 + 8);
  Guid2Low  = ReadUnaligned64 (Guid2);
  Guid2High = ReadUnaligned64 ((CONST UINT8 *)Guid2 + 8);

  return (BOOLEAN)(Guid1Low == Guid2Low && Guid1High == Guid2High);
}

/**
  Retrieves the HOB (Hand-Off Block) list pointer by locating the
  DXE Services HOB in the system table's configuration table.

  Walks the system table's HOB list entries to find the DXE Services HOB
  identified by gEfiDxeServicesTableGuid (7739F24C-93D7-11D4-9A3A-0090273FC14D).

  SystemTable layout (x64):
    +0x68 (104): NumberOfTableEntries (UINTN) - number of config table entries
    +0x70 (112): ConfigurationTable (EFI_CONFIGURATION_TABLE *) - array of entries
    Each entry (24 bytes):
      +0x00: VendorGuid (EFI_GUID, 16 bytes)
      +0x10: VendorTable (VOID *, 8 bytes) - pointer to the actual table

  The located HOB list pointer is cached globally in gHobList.

  @return Pointer to the start of the HOB list, or NULL if:
          - The system table has no HOB entries (count = 0)
          - No HOB matches the DXE Services GUID
**/
VOID *
GetHobList (
  VOID
  )
{
  UINTN    HobCount;
  UINT8   *HobEntries;
  UINTN    Index;

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

  //
  // SystemTable + 0x68 (104) = NumberOfTableEntries
  // SystemTable + 0x70 (112) = ConfigurationTable pointer
  //
  HobCount   = *(UINTN *)((UINT8 *)gSystemTable + 104);
  HobEntries = *(UINT8 **)((UINT8 *)gSystemTable + 112);

  gHobList = NULL;

  if (HobCount > 0) {
    //
    // Walk the configuration table to find the DXE Services HOB entry.
    //
    for (Index = 0; Index < HobCount; Index++) {
      if (IsHobGuidMatch (
            (EFI_GUID *)(HobEntries + Index * sizeof(EFI_CONFIGURATION_TABLE)),
            &gEfiDxeServicesTableGuid
            )) {
        //
        // Found the DXE Services HOB - cache its VendorTable pointer.
        //
        gHobList = *(VOID **)(HobEntries + Index * sizeof(EFI_CONFIGURATION_TABLE) + sizeof(EFI_GUID));
        break;
      }
    }

    if (gHobList == NULL) {
      //
      // DXE Services HOB not found - report error.
      //
      UbaDebugPrint (0x80000000LL, "\nASSERT_EFI_ERROR (Status = %r)\n");
      UbaDebugAssert (
        "e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c",
        54,
        "!EFI_ERROR (Status)"
        );
    }
  }

  if (gHobList == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\DxeHobLib\\HobLib.c",
      55,
      "mHobList != ((void *) 0)"
      );
  }

  return gHobList;
}

/**
  Installs PSLT Slot Data entries via the UBA protocol.

  This function performs the following actions:
    1. Prints a debug banner via UbaDebugPrint identifying this module.
    2. Locates the UBA protocol via gBS->LocateProtocol().
    3. Calls the UBA protocol's SetData function (index 2, offset 0x10)
       twice with different PSLT configuration entries.

  The PSLT entries registered:
    - PSLT1 (GUID B93613E1-48F0-4B32-B3A8-4FEDFC7C1365): 40 bytes,
      includes GetSlotCount (returns 2) and GetSlotData (pass-through).
    - PSLT2 (GUID 226763AE-972C-4E3C-80D1-73B25E8CBBA3): 32 bytes,
      includes GetSlotCount only.

  Both entries signal through the Flags field that the configuration
  applies to slot count 1.

  @return EFI_STATUS  Result of the installation.
                       Returns the first error encountered, or EFI_SUCCESS
                       if both PSLT entries are registered successfully.
**/
EFI_STATUS
InstallSlotData (
  VOID
  )
{
  EFI_STATUS  Status;
  VOID       *UbaProtocol;
  UINT64    (*UbaSetData)(VOID *, EFI_GUID *, VOID *, UINT64);

  //
  // Print debug banner identifying this module.
  //
  UbaDebugPrint (0x80000000LL, "UBA:SlotDataUpdate-TypeLightningRidgeEXRP\n");

  //
  // Locate the UBA protocol interface.
  //
  Status = gBootServices->LocateProtocol (
                            &gUbaProtocolGuid,     // UBA protocol GUID
                            NULL,                   // No registration key
                            &UbaProtocol            // Output: located protocol
                            );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Retrieve the UBA SetData function at protocol index 2 (offset 0x10).
  // Signature: EFI_STATUS (*)(VOID *This, EFI_GUID *DataGuid,
  //                           VOID *Data, UINTN DataSize)
  //
  UbaSetData = (UINT64 (*)(VOID *, EFI_GUID *, VOID *, UINT64))
               (*((UINT64 *)((UINT8 *)UbaProtocol + UBA_PROTOCOL_SET_DATA_OFFSET)));

  //
  // Register PSLT Entry 1: Primary slot table with GetSlotCount + GetSlotData.
  // GUID: B93613E1-48F0-4B32-B3A8-4FEDFC7C1365
  // Size: 40 bytes (full PSLT_HEADER)
  //
  Status = UbaSetData (UbaProtocol, &gUbaSlotDataPsl1Guid, (VOID *)mPslEntry1, 40);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register PSLT Entry 2: Secondary slot table with GetSlotCount only.
  // GUID: 226763AE-972C-4E3C-80D1-73B25E8CBBA3
  // Size: 32 bytes (compact PSLT_HEADER_COMPACT)
  //
  Status = UbaSetData (UbaProtocol, &gUbaSlotDataPsl2Guid, (VOID *)mPslEntry2, 32);

  return Status;
}

/**
  Entry point of the SlotDataUpdateDxeLightningRidgeEXRP driver.

  This is the standard UEFI DXE driver entry point. It:
    1. Validates and caches ImageHandle, SystemTable, BootServices,
       and RuntimeServices pointers.
    2. Initializes the HOB list via GetHobList().
    3. Installs platform slot data via InstallSlotData().

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

  @retval EFI_SUCCESS           Slot data successfully installed.
  @retval EFI_NOT_FOUND         UBA protocol or required HOBs not found.
  @retval EFI_INVALID_PARAMETER One or more parameters are invalid.
  @retval Others                Error from InstallSlotData or protocol operations.
**/
EFI_STATUS
EFIAPI
SlotDataUpdateDxeEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS  Status;

  //
  // Cache UEFI global variables.
  //
  gImageHandle = ImageHandle;
  if (gImageHandle == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiBootServicesTableLib\\UefiBootServicesTableLib.c",
      51,
      "gImageHandle != ((void *) 0)"
      );
  }

  gSystemTable = SystemTable;
  if (gSystemTable == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiBootServicesTableLib\\UefiBootServicesTableLib.c",
      57,
      "gST != ((void *) 0)"
      );
  }

  gBootServices = SystemTable->BootServices;
  if (gBootServices == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiBootServicesTableLib\\UefiBootServicesTableLib.c",
      63,
      "gBS != ((void *) 0)"
      );
  }

  gRuntimeServices = SystemTable->RuntimeServices;
  if (gRuntimeServices == NULL) {
    UbaDebugAssert (
      "e:\\hs\\MdePkg\\Library\\UefiRuntimeServicesTableLib\\UefiRuntimeServicesTableLib.c",
      47,
      "gRT != ((void *) 0)"
      );
  }

  //
  // Retrieve the HOB list (cached globally in gHobList).
  //
  GetHobList ();

  //
  // Install platform-specific slot data via UBA protocol.
  //
  Status = InstallSlotData ();

  return Status;
}