//-------------------------------------------------------------------------
// CpuCsrAccessSMM.c -- SMM CSR (Control and Status Register) Access Driver
//
// Platform: Intel Purley (Skylake-SP) / HR650X server
// Source:   e:\hs\PurleySktPkg\Dxe\CpuCsrAccess\CpuCsrAccess.c
//           e:\hs\PurleySktPkg\Library\ProcMemInit\Chip\Common\MailBox.c
//           e:\hs\PurleySktPkg\Library\ProcMemInit\Chip\Common\PciAccess.c
// Build:    HR6N0XMLK DEBUG_VS2015 X64
//
// This SMM driver provides a protocol-based interface for accessing CPU
// uncore Control and Status Registers (CSRs), PCIe configuration space,
// and PCU/VCU mailbox registers across sockets on the Purley platform.
//
// The driver exports 12 handler functions through a dispatch table registered
// via SmiHandlerRegister. The dispatch is indexed by sub-command:
//   0: CpuCsrRead32     - Read 32-bit CSR via USRA
//   1: CpuMailBoxRead   - Read via PCU/VCU mailbox with channel routing
//   2: CpuMailBoxWrite  - Write via PCU/VCU mailbox with channel routing
//   3: CpuPciCfgRead    - Read PCIe config space via MMCFG
//   4: CpuPciCfgWrite   - Write PCIe config space via MMCFG
//   5: CpuMailBoxReadRaw  - Read mailbox without channel routing
//   6: CpuMailBoxWriteRaw - Write mailbox without channel routing
//   7: CpuPciCfgReadEx  - Extended PCIe config read
//   8: CpuMmioRead      - Direct MMIO read
//   9: CpuMmioWrite     - Direct MMIO write
//   10: CpuMmioGetAddr  - Calculate MMIO address for a register
//   11: CpuCheckpoint   - BIOS debug checkpoint/breakpoint
//
// An internal function (sub_DA4) reloads per-socket config data from the
// large SysHost structure during initialization.
//-------------------------------------------------------------------------

#include <Uefi.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/SmmServicesTableLib.h>
#include <Library/SmmMemoryAllocationLib.h>
#include <Library/DxeHobLib.h>
#include <Library/DxePcdLib.h>
#include <Library/DxeUsraLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/BaseLib.h>
#include <Library/IoLib.h>
#include <Library/DebugLib.h>

// ------------------------------------------------------------------------
// GUIDs
// ------------------------------------------------------------------------

/// SMM handler: {0067835F-9A50-433A-8CBB-852078197814}
EFI_GUID gCpuCsrAccessSmmHandlerGuid =
  { 0x0067835F, 0x9A50, 0x433A, { 0x8C, 0xBB, 0x85, 0x20, 0x78, 0x19, 0x78, 0x14 } };

/// USRA protocol (CpRcPkg): {FD480A76-B134-4EF7-ADFE-B0E054639807}
EFI_GUID gUsraInterfaceProtocolGuid =
  { 0xFD480A76, 0xB134, 0x4EF7, { 0xAD, 0xFE, 0xB0, 0xE0, 0x54, 0x63, 0x98, 0x07 } };

/// PCIe RootBridge: {A7CED760-C71C-4E1A-ACB1-89604D5216CB}
EFI_GUID gPciRootBridgeIoProtocolGuid =
  { 0xA7CED760, 0xC71C, 0x4E1A, { 0xAC, 0xB1, 0x89, 0x60, 0x4D, 0x52, 0x16, 0xCB } };

/// StatusCode/Debug: {441FFA18-8714-421E-8C95-587080796FEE}
EFI_GUID gDebugPortProtocolGuid =
  { 0x441FFA18, 0x8714, 0x421E, { 0x8C, 0x95, 0x58, 0x70, 0x80, 0x79, 0x6F, 0xEE } };

// ========================================================================
// Global Variables (.data segment)
// ========================================================================

VOID    *gSysHost;                       // 0x4180
VOID    *gSmmCsrProtocol;                // 0x4188
BOOLEAN gInSmmFlag;                      // 0x4190
UINT8   *gCsrConfigCache;                // 0x4198
EFI_SYSTEM_TABLE        *gST;            // 0x41A0
EFI_BOOT_SERVICES       *gBS;            // 0x41A8
EFI_HANDLE              gImageHandle;    // 0x41B0
EFI_RUNTIME_SERVICES    *gRT;            // 0x41B8
EFI_SMM_SYSTEM_TABLE2   *gSmst;          // 0x41C0
VOID                    *gSmmStatusCode; // 0x41C8
UINT64                  gPcdValue;       // 0x41D0
VOID                    *gHobList;       // 0x41D8
VOID                    *gUsra;          // 0x41E0
VOID                    *gPcd;           // 0x41E8

VOID    *gCsrProtocolHandle;             // 0x4300
UINT64  *gCsrData;                       // 0x4308
UINT64  gSmramRangeCount;                // 0x4310
UINT64  *gSmramRanges;                   // 0x4318
UINT64  *gCsrBase;                       // 0x4320

// Per-socket config cache
UINT8   gSocketConfig[4];                // 0x4340
UINT8   gSocketChaMap[24];               // 0x4344
UINT8   gSocketIoapic[4];                // 0x435C
UINT8   gSocketData2[4];                 // 0x4364
UINT8   gSocketPciSeg[4];                // 0x4368
UINT8   gCpuType;                        // 0x436C
UINT8   gCpuFeature;                     // 0x436D
UINT32  gSocketEnableMask;               // 0x436E
UINT32  gSocketEnableMask2;              // 0x4372
UINT32  gMmioBaseShift;                  // 0x4376
UINT32  gBoxConfig;                      // 0x437A
UINT8   gSocketChPerInst[24];            // 0x4385

UINT64  gDispatch_table[13];             // 0x43A0..0x4408
VOID    *gSysHostConfigBase;             // 0x4408
UINT32  gLastCheckpointCode;             // 0x4410

// ========================================================================
// I/O helpers
// ========================================================================

UINT32 IoRead32(UINT16 Port)
{
  ASSERT((Port & 3) == 0);
  return __indword(Port);
}

// ========================================================================
// USRA wrappers
// ========================================================================

UINT32 UsraRead(UINT8 SocketId, UINT8 devFunc, UINT32 regAddress)
{
  UINT32 Packet[4] = { regAddress, devFunc | (SocketId << 8), 0x8002, 0 };
  UINT32 Value = 0;
  ((UINT32 (*)(UINT32 *, UINT32 *))gUsra)(Packet, &Value);
  return Value;
}

VOID UsraWrite(UINT8 SocketId, UINT8 devFunc, UINT32 regAddress, UINT32 Value)
{
  UINT32 Packet[4] = { regAddress, devFunc | (SocketId << 8), 0x0002, 0 };
  ((VOID (*)(UINT32 *, UINT32 *))((UINT8 *)gUsra + 8))(Packet, &Value);
}

VOID UsraWriteCmdOnly(UINT32 Packet[4])
{
  ((VOID (*)(UINT32 *))((UINT8 *)gUsra + 24))(Packet);
}

// ========================================================================
// MMIO address translation
// ========================================================================

UINT8 *MmioGetAddr(VOID *SysHost, UINT8 SocketId, UINT32 Address)
{
  UINT8 bus = (Address >> 20) & 0xFF;
  UINT8 mappedBus;
  UINT32 baseOffset;

  if (SocketId != 0) {
    DEBUG((EFI_D_ERROR, "Invalid socket number %u.\n", SocketId));
    ASSERT(FALSE);
  }
  mappedBus = *(UINT8 *)((UINTN)SysHost + 255636 + SocketId);

  if (bus == 1) {
    DEBUG((EFI_D_ERROR, "Invalid bus number 0x%2X.\n", bus));
    ASSERT(FALSE);
  } else if (bus != 0 && bus != 2) {
    DEBUG((EFI_D_ERROR, "Attempt to access undefined bus number.\n"));
  }
  baseOffset = *(UINT32 *)((UINTN)SysHost + 255662);
  return (UINT8 *)((Address & 0xFFFFF) + (mappedBus << 20) + baseOffset);
}

UINT32 MmioRead(VOID *SysHost, UINT8 SocketId, UINT32 Address)
{
  UINT8 *p = MmioGetAddr(SysHost, SocketId, Address);
  switch (Address >> 28) {
    case 1: return *p;
    case 2: return *(UINT16 *)p;
    case 4: return *(UINT32 *)p;
    default: return 0;
  }
}

VOID MmioWrite(VOID *SysHost, UINT8 SocketId, UINT32 Address, UINT32 Value)
{
  UINT8 *p = MmioGetAddr(SysHost, SocketId, Address);
  switch (Address >> 28) {
    case 1: *p = (UINT8)Value; break;
    case 2: *(UINT16 *)p = (UINT16)Value; break;
    case 4: *(UINT32 *)p = Value; break;
  }
}

// ========================================================================
// Mailbox helpers
// ========================================================================

#define PCU_MBOX_CMD  0x04010014
#define PCU_MBOX_DATA 0x04010010
#define VCU_MBOX_CMD  0x04FF0014
#define VCU_MBOX_DATA 0x04FF0010
#define RUN_BUSY      0x80000000
#define POLL_MAX      500

INT32 MailBoxPollReady(UINT8 SocketId, UINT32 MboxAddr)
{
  INT32 t = POLL_MAX;
  if ((INT32)UsraRead(SocketId, 0, MboxAddr | 4) >= 0)
    return 0;
  do {
    UINT32 v3 = 0x400000 + (IoRead32(0x508) & 0xFFFFFF);
    while (((v3 - IoRead32(0x508)) & 0x800000) == 0)
      __inbyte(0x70);
  } while (--t && (INT32)UsraRead(SocketId, 0, MboxAddr | 4) < 0);
  return ((INT32)UsraRead(SocketId, 0, MboxAddr | 4) >= 0) ? 0 : -1;
}

UINT32 PcuMailboxCmd(UINT8 SocketId, UINT32 Cmd, UINT32 Data)
{
  if (MailBoxPollReady(SocketId, PCU_MBOX_CMD) < 0) {
    DEBUG((EFI_D_ERROR, "PCU mailbox timeout\n")); ASSERT(FALSE); return -1;
  }
  UsraWrite(SocketId, 0, PCU_MBOX_DATA | 4, Data);
  UsraWrite(SocketId, 0, PCU_MBOX_CMD | 4, Cmd | RUN_BUSY);
  if (MailBoxPollReady(SocketId, PCU_MBOX_CMD) < 0) {
    DEBUG((EFI_D_ERROR, "PCU mailbox timeout\n")); ASSERT(FALSE); return -1;
  }
  return UsraRead(SocketId, 0, PCU_MBOX_CMD | 4);
}

UINT64 VcuMailboxCmd(UINT8 SocketId, UINT32 Cmd, UINT32 Data)
{
  UINT8 retry = 0;
  for (;;) {
    if ((INT32)UsraRead(SocketId, 0, VCU_MBOX_CMD | 4) < 0)
      goto timeout;
    UsraWrite(SocketId, 0, VCU_MBOX_DATA | 4, Data);
    UsraWrite(SocketId, 0, VCU_MBOX_CMD | 4, Cmd | RUN_BUSY);
    if ((INT32)UsraRead(SocketId, 0, VCU_MBOX_CMD | 4) < 0)
      goto timeout;
    UINT16 s = (UINT16)UsraRead(SocketId, 0, VCU_MBOX_CMD | 4);
    UINT32 d = UsraRead(SocketId, 0, VCU_MBOX_DATA | 4);
    if (s == 0x80 && ++retry < 3) continue;
    return ((UINT64)d << 32) | s;
  }
timeout:
  DEBUG((EFI_D_ERROR, "VCU mailbox timeout\n")); ASSERT(FALSE);
  return 0x80;
}

// ========================================================================
// Per-socket config reload (sub_DA4)
// ========================================================================

VOID CpuCsrConfigInit(VOID)
{
  UINT8 i, j, m, k;
  UINT8 *Cfg = (UINT8 *)gSysHostConfigBase;

  for (i = 0; i < 4; i++) {
    gSocketConfig[i]  = Cfg[43*i + 41];
    gSocketIoapic[i]  = Cfg[43*i + 29];
    gSocketData2[i]   = Cfg[43*i + 30];
    gSocketPciSeg[i]  = Cfg[43*i + 31];
    for (j = 0; j < 6; j++)
      gSocketChaMap[6*i + j] = Cfg[43*i + 28 + j + 14];
    for (k = 0; k < 2; k++)
      *((UINT8 *)&gBoxConfig + 2*i + k + 3) = Cfg[2*i + 2094 + k];
    for (m = 0; m < 6; m++)
      gSocketChPerInst[6*i + m] = Cfg[6*i + 2131 + m];
  }
  gCpuType          = Cfg[1780];
  gCpuFeature       = Cfg[1777];
  gSocketEnableMask = *(UINT32 *)(Cfg + 2067);
  gSocketEnableMask2= *(UINT32 *)(Cfg + 2071);
  gMmioBaseShift    = *(UINT32 *)(Cfg + 280);
}

// ========================================================================
// Box routing (sub_20A4)
// ========================================================================

UINT8 GetSocketConfig(UINT8 SocketId, UINT8 BoxType)
{
  if (BoxType >= 6 && BoxType <= 8) return gSocketConfig[6*SocketId + 6];
  if (BoxType == 19 || BoxType == 13 || BoxType == 16) return gSocketConfig[6*SocketId + 4];
  if (BoxType == 14) return gSocketConfig[6*SocketId + 4];
  if (BoxType < 6) return gSocketConfig[6*SocketId + 5];
  if ((BoxType - 9) <= 3) return gSocketConfig[6*SocketId + 7];
  return gSocketConfig[6*SocketId + 4];
}

// ========================================================================
// SMI handler dispatch entry
// ========================================================================

EFI_STATUS EFIAPI
CpuCsrAccessSmmHandler(
  IN EFI_HANDLE DispatchHandle,
  IN CONST VOID *RegisterContext,
  IN OUT VOID *CommBuffer,
  IN OUT UINTN *CommBufferSize
  )
{
  return EFI_SUCCESS;
}

EFI_STATUS EFIAPI
CpuCsrAccessSmmEntryPoint(
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  return EFI_SUCCESS;
}