Newer
Older
AMI-Aptio-BIOS-Reversed / RegAccessSMM / RegAccessSMM.c
@Ajax Dong Ajax Dong 2 days ago 24 KB Init
/** @file
  RegAccessSMM.c - Register Access SMM Driver

  This SMM driver provides register access services for the Purley/Coffee Lake
  platform, integrating:
  - Unified System Register Access (USRA) filter library
  - S3 Boot Script save/restore library
  - SMM LockBox library
  - PCIe MMIO address translation

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

#include "RegAccessSMM.h"

// ---------------------------------------------------------------------------
// Global Variables (.data segment: 0x6380-0x67E0)
// ---------------------------------------------------------------------------

// SMM System State
EFI_HANDLE                 gImageHandle         = NULL;
EFI_SYSTEM_TABLE          *gSystemTable         = NULL;
EFI_BOOT_SERVICES         *gBootServices        = NULL;  // 0x6508
EFI_RUNTIME_SERVICES      *gRuntimeServices     = NULL;  // 0x6510
EFI_SMM_SYSTEM_TABLE2     *gSmst                = NULL;  // 0x6520
VOID                       *gDebugProtocol      = NULL;  // 0x6528
PCD_PROTOCOL               *gPcdProtocol        = NULL;  // 0x6530
UINT64                      gPciExpressBaseAddr = 0;     // 0x6538
UINT64                      gMmCfgCount         = 0;     // 0x6540
UINT64                      gTscFrequency       = 0;     // 0x6548
BOOLEAN                     gLockBoxCtxInstalled= FALSE; // 0x6550 byte
volatile UINT64             gRegAccessSpinLock  = 0;     // 0x6558

// USRA / IIO Topology
VOID                       *gUsraProtocol       = NULL;  // 0x6560
BOOLEAN                     gCsrViaMsr          = FALSE; // 0x6568 byte
VOID                       *gUsraBatchCtx       = NULL;  // 0x65A8
BOOLEAN                     gUsraBatchActive    = FALSE; // 0x65B0 byte
UINT64                      gUsraSocketData     = 0;     // 0x65C0
UINT64                      gModuleStatus       = 0;     // 0x66C8
VOID                       *gIioProtocol        = NULL;  // 0x66D0
UINT8                      *gIioProtocolData    = NULL;  // 0x66D8
UINT8                       gSocketTopology[24];        // 0x66E0
UINT8                       gSocketTopoData[24];        // 0x66E4
UINT8                       gIioPresence        = 0;     // 0x670C
UINT8                       gSocketCnt          = 0;     // 0x670D
UINT32                      gSockEnMask         = 0;     // 0x670E
UINT32                      gSockPrMask         = 0;     // 0x6712
UINT32                      gSysMemSize         = 0;     // 0x6716
UINT8                       gCpuStepping        = 0;     // 0x671A
UINT8                       gCpuType            = 0;     // 0x671B
UINT8                       gMaxPciePorts       = 0;     // 0x671C
UINT8                       gSockTopoActive[24];        // 0x6740
UINT8                       gIioPresenceActive  = 0;     // 0x676C
UINT32                      gSockEnMaskActive   = 0;     // 0x676E
UINT32                      gSockPrMaskActive   = 0;     // 0x6772
UINT8                       gCpuSteppingActive  = 0;     // 0x677A
UINT8                       gMaxPciePortsActive = 0;     // 0x677C

// S3 Boot Script State
VOID                       *gS3ReadyToBootNotify= NULL;  // 0x6570
BOOLEAN                     gS3ReadyToLockFlag  = FALSE; // 0x6578 byte
EVENT                       gS3DxeRdyLockEvent  = NULL;  // 0x6580
EFI_SMM_S3_PROTOCOL        *gSmmS3Protocol      = NULL;  // 0x6588
VOID                       *gS3LockBoxReg       = NULL;  // 0x6590
VOID                       *gS3RdyBootReg       = NULL;  // 0x6598
BOOLEAN                     gS3MemAttrAlloc     = FALSE; // 0x65A0 byte
UINT8                      *gS3BootScriptData   = NULL;  // 0x67A0
UINT8                      *gS3BootScriptData2  = NULL;  // 0x67A8

// SmmLockBox State
GUID                        gLockBoxCommGuid;            // 0x67B0 16 bytes
VOID                       *gLockBoxDataAddr    = NULL;  // 0x67B8
UINT64                      gSmramRangeCnt      = 0;     // 0x67C0
EFI_SMRAM_DESCRIPTOR       *gSmramRanges        = NULL;  // 0x67C8

// ---------------------------------------------------------------------------
// Module Entry Point (AutoGen)
// ---------------------------------------------------------------------------

/**
  Main entry point called by UEFI SMM core.
  Instantiates all services: SMM tables, LockBox, S3 Boot Script, USRA.
**/
EFI_STATUS
EFIAPI
ModuleEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS  Status;
  UINT64      Ret;

  // 1. Full constructor: init SMM services, protocols, allocate state
  Status = RegAccessSmmConstructor (ImageHandle, SystemTable);
  gModuleStatus = EFI_SUCCESS;

  // 2. Set up error recovery via SetJmp/LongJmp
  if (SetJump (&gModuleStatus) == 0) {
    // 3. Post-entry: locate SMM S3, init LockBox context, save script data ptr
    Ret = RegAccessSmmEntryPoint ();
    if (Ret >= 0 || gModuleStatus < 0) {
      gModuleStatus = Ret;
    }

    // 4. Validate SetJmp buffer alignment (debug assertion checks)
    ValidateSetJumpBuffer (&gModuleStatus);
    LongJump (&gModuleStatus, -1);

    // These are debug ASSERT_NOT_REACHED markers from AutoGen
    DEBUG_ASSERT (!EFI_ERROR (Status), L"ModuleEntryPoint");
    DEBUG_ASSERT (!EFI_ERROR (Status), L"ModuleEntryPoint: Exit");
  }

  // 5. Cleanup on error: invoke destructor
  if (EFI_ERROR (gModuleStatus)) {
    RegAccessSmmDestructor ();
  }

  return gModuleStatus;
}

// ---------------------------------------------------------------------------
// Constructor: RegAccessSmmConstructor
// ---------------------------------------------------------------------------

/**
  Full module constructor. Initializes SMM services table, SMRAM allocation,
  PCIe MMIO, TSC frequency calibration, LockBox, USRA protocol detection,
  S3 boot script, and register spinlock.
**/
EFI_STATUS
RegAccessSmmConstructor (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE *SystemTable
  )
{
  EFI_STATUS                  Status;
  EFI_SMM_BASE2_PROTOCOL     *SmmBase2;
  EFI_SMM_CPU_PROTOCOL       *SmmCpu;
  UINTN                       SmramRangeCount;
  EFI_SMRAM_DESCRIPTOR       *SmramRanges;
  UINTN                       BufferSize;
  EFI_SMM_S3_PROTOCOL        *SmmS3;
  UINT64                      Ranges;
  UINT64                      TscCount;
  UINT64                      TscStart;
  UINT64                      TscEnd;
  UINTN                       Count;
  UINT16                      Flags;
  BOOLEAN                     InterruptsEnabled;
  UINT64                      i;

  // --- Boot Services Table Library init ---
  gImageHandle  = ImageHandle;
  DEBUG_ASSERT (gImageHandle != NULL, L"gImageHandle");
  gSystemTable  = SystemTable;
  DEBUG_ASSERT (gSystemTable != NULL, L"gST");
  gBootServices = SystemTable->BootServices;
  DEBUG_ASSERT (gBootServices != NULL, L"gBS");
  gRuntimeServices = SystemTable->RuntimeServices;
  DEBUG_ASSERT (gRuntimeServices != NULL, L"gRT");

  // --- SMM Services Table Library init ---
  Status = gBootServices->LocateProtocol (&gEfiSmmBase2ProtocolGuid, NULL, (VOID**)&SmmBase2);
  DEBUG_ASSERT (!EFI_ERROR (Status), L"Locate SmmBase2");
  SmmBase2->GetSmstLocation (SmmBase2, &gSmst);
  DEBUG_ASSERT (gSmst != NULL, L"gSmst");

  // --- SMM Memory Allocation Library init ---
  Status = gBootServices->LocateProtocol (&gEfiSmmCpuProtocolGuid, NULL, (VOID**)&SmmCpu);
  DEBUG_ASSERT (!EFI_ERROR (Status), L"Locate SmmCpu");

  // Query SMRAM range count
  Status = SmmCpu->GetSmramRangeCount (SmmCpu, &SmramRangeCount);
  DEBUG_ASSERT (Status == EFI_BUFFER_TOO_SMALL,
                L"Status == EFI_BUFFER_TOO_SMALL");

  BufferSize = SmramRangeCount * sizeof (EFI_SMRAM_DESCRIPTOR);
  gSmramRanges = (EFI_SMRAM_DESCRIPTOR*)SmramAllocatePool (SmramRangeCount);
  if (gSmramRanges == NULL) {
    DEBUG_ASSERT (FALSE, L"mSmramRanges != ((void *) 0)");
  }

  Status = SmmCpu->GetSmramRangeCount (SmmCpu, &SmramRangeCount);
  DEBUG_ASSERT (!EFI_ERROR (Status), L"GetSmramRangeCount");
  gSmramRangeCnt = SmramRangeCount >> 5;  // Count of UINT64 pairs

  // --- SMM PCI Express Library init ---
  gPciExpressBaseAddr = PcdGet64 (PcdPciExpressBaseAddress);
  SmmPciExpressInit ();

  // --- PCIe MMIO init (Purley: MMCFG at 0x80000000, size 1280) ---
  if ((INT8*)PciExpressBaseAddress (PcdPciExpressBaseRegister) >= 0) {
    IoWrite16 ((UINT16*)PciExpressBaseAddress (PcdPciExpressBaseRegister + 2),
               1280);
    *(UINT8*)PciExpressBaseAddress (PcdPciExpressBaseRegister + 4) |= 0x80;
  }

  // --- Timer / TSC calibration ---
  Flags         = (UINT16)GetFlags ();
  DisableInterrupts ();
  InterruptsEnabled = (Flags & 0x200) != 0;

  TscStart      = IoRead32 (RDTSC_PORT) & 0xFFFFFF;
  TscEnd        = ReadTsc ();
  while (((UINT32)(TscStart + 357 - (UINT32)IoRead32 (RDTSC_PORT)) & 0x800000) == 0) {
    CpuPause ();
  }
  TscEnd        = ReadTsc () - TscEnd;
  gTscFrequency = TscEnd * 10000;  // TSC ticks per 10 ms

  if (InterruptsEnabled) {
    EnableInterrupts ();
  } else {
    DisableInterrupts ();
  }

  // --- SmmLockBox Constructor ---
  DEBUG_PRINT (DEBUG_INFO, L"SmmLockBoxSmmLib SmmLockBoxSmmConstructor - Enter\n");

  if (SmmLockBoxFindContext () != NULL) {
    DEBUG_PRINT (DEBUG_INFO, L"SmmLockBoxSmmLib SmmLockBoxContext - already installed\n");
    DEBUG_PRINT (DEBUG_INFO, L"SmmLockBoxSmmLib SmmLockBoxSmmConstructor - Exit\n");
  } else {
    gLockBoxCommGuid.Data1 = '6' | ('_' << 8) | ('L' << 16) | ('O' << 24);
    gLockBoxCommGuid.Data2 = 'C';
    gLockBoxDataAddr       = &gLockBoxCommGuid;
    gLockBoxCtxInstalled   = TRUE;

    // Install LockBox communication configuration table
    Status = gSmst->SmmInstallConfigurationTable (
                      gSmst, &gEfiSmmLockBoxCommGuid,
                      &gLockBoxCommGuid, sizeof (GUID)
                      );
    if (EFI_ERROR (Status)) {
      DEBUG_ASSERT (!EFI_ERROR (Status), L"Install LockBox config table");
    }

    DEBUG_PRINT (DEBUG_INFO, L"SmmLockBoxSmmLib SmmLockBoxContext - %x\n", (UINTN)&gLockBoxCommGuid);
    DEBUG_PRINT (DEBUG_INFO, L"SmmLockBoxSmmLib LockBoxDataAddress - %x\n", (UINTN)&gLockBoxCommGuid);
    DEBUG_PRINT (DEBUG_INFO, L"SmmLockBoxSmmLib SmmLockBoxSmmConstructor - Exit\n");
  }

  // --- USRA protocol detection ---
  Status = gBootServices->LocateProtocol (&gEfiSmmUsraProtocolGuid, NULL, (VOID**)&gUsraProtocol);
  DEBUG_ASSERT (!EFI_ERROR (Status), L"Locate USRA protocol");

  // Detect IioProtocol and CSR/MSR routing
  Status = gBootServices->LocateProtocol (&gEfiSmmUsraProtocolGuid, NULL, (VOID**)&gIioProtocol);
  DEBUG_ASSERT (!EFI_ERROR (Status), L"Locate IIO protocol");

  if (gIioProtocol != NULL &&
      ((UINT8*)gIioProtocol)[1777] >= 9 &&
      *(UINT64*)((UINT8*)gIioProtocol + 1739) != 0 &&
      !*((UINT8*)gIioProtocol + 1747)) {
    gCsrViaMsr = TRUE;
    DEBUG_PRINT (DEBUG_INFO, L"CSR writes routed via MSR.\n");
  }

  // --- Determine cache line size from CPUID ---
  {
    UINT32 eax, ebx, ecx, edx;
    UINT64 CacheLineSize;

    Cpuid (1, &eax, &ebx, &ecx, &edx);
    CacheLineSize = 8 * (UINT8)(ebx >> 4);  // CLFLUSH line size in QWORDs

    if ((eax & 0xF00) == 0xF00) {
      // Intel hyper-threading: adjust for logical processors per core
      UINT8 Count = (UINT8)(eax & 0xF0) | ((eax >> 8) & 0xF00);
      Count >>= 4;
      if (Count <= 4 || Count == 6) {
        CacheLineSize *= 2;
      }
    }

    if (CacheLineSize < 32) {
      CacheLineSize = 32;
    }

    // Allocate cache-aligned region + spinlock at end
    if (CacheLineSize * 2 > 0) {
      gRegAccessSpinLock = (UINT64)SmramAllocatePool (0, 6, CacheLineSize * 2 / 4096 + 1, &gRegAccessSpinLock);
    }
    *(volatile UINT64*)gRegAccessSpinLock = 1;
  }

  // --- S3 Boot Script Library Initialize ---
  Status = S3BootScriptLibInitialize ();
  if (EFI_ERROR (Status)) {
    DEBUG_ASSERT (!EFI_ERROR (Status), L"S3BootScriptLibInitialize");
  }

  return Status;
}

// ---------------------------------------------------------------------------
// Destructor: RegAccessSmmDestructor
// ---------------------------------------------------------------------------

/**
  Module destructor. Invoked when the module is unloaded or on error.
  Uninstalls LockBox config table, closes S3 boot script, frees SMRAM.
**/
EFI_STATUS
RegAccessSmmDestructor (
  VOID
  )
{
  EFI_STATUS  Status;

  // Close S3 boot script
  Status = S3BootScriptCloseTable ();
  DEBUG_ASSERT (!EFI_ERROR (Status), L"Destructor: CloseTable");

  DEBUG_PRINT (DEBUG_INFO, L"SmmLockBoxSmmLib SmmLockBoxSmmDestructor in %a module\n", L"RegAccessSMM");

  // Uninstall LockBox communication config table if installed
  if (gLockBoxCtxInstalled) {
    Status = gSmst->SmmInstallConfigurationTable (
                      gSmst, &gEfiSmmLockBoxCommGuid,
                      NULL, 0
                      );
    DEBUG_ASSERT (!EFI_ERROR (Status), L"Uninstall LockBox config table");
    DEBUG_PRINT (DEBUG_INFO,
                 L"SmmLockBoxSmmLib uninstall SmmLockBoxCommunication configuration table\n");
  }

  // Free SMRAM ranges pool
  SmramFreePool (gSmramRanges);

  return Status;
}

// ---------------------------------------------------------------------------
// Post-Entry: RegAccessSmmEntryPoint
// ---------------------------------------------------------------------------

/**
  Post-constructor entry. Locates the SMM S3 protocol, initializes the LockBox
  context, and saves the boot script data pointer into the SMM S3 protocol.
**/
UINT64
RegAccessSmmEntryPoint (
  VOID
  )
{
  EFI_STATUS  Status;
  VOID       *SmmS3;
  VOID       *LockBoxContext;
  UINT64      Ret;

  // 1. Locate SMM S3 protocol
  Status = gBootServices->LocateProtocol (&gSmmS3ProtocolGuid, NULL, &SmmS3);
  DEBUG_ASSERT (!EFI_ERROR (Status), L"Locate SmmS3");

  // 2. Locate LockBox context
  LockBoxContext = SmmLockBoxFindContext ();
  if (LockBoxContext != NULL) {
    // 3. Get the LockBox queue pointer
    VOID *Queue = SmmLockBoxGetQueue ();
    if (Queue != NULL) {
      // 4. Save the LockBox queue address to SMM S3 protocol
      //    (so S3 resume can locate boot script data)
      gS3BootScriptData = (UINT8*)*(UINT64*)((UINT8*)SmmS3 + 0);
      // ... rest of S3 protocol init
    }
  }

  return EFI_SUCCESS;
}

// ---------------------------------------------------------------------------
// S3 Boot Script Library
// ---------------------------------------------------------------------------

EFI_STATUS
S3BootScriptLibInitialize (
  VOID
  )
{
  EFI_STATUS              Status;
  UINT64                  ScriptData;
  BOOLEAN                 MemoryAllocated;
  EFI_SMM_S3_PROTOCOL    *SmmS3;
  VOID                   *Registration;
  UINT64                  PcdValue;

  // First call from constructor
  PcdValue = PcdGet64 (PcdS3BootScriptTableSize);
  if (PcdValue == 0) {
    // Allocate default-sized boot script table
    Status = gBootServices->AllocatePool (
               EfiBootServicesData, 1,
               1, &ScriptData
               );
    if (EFI_ERROR (Status)) {
      DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:586");
    }
    gS3BootScriptData = (UINT8*)(UINTN)ScriptData;
    gS3ReadyToLockFlag = TRUE;

    PcdValue = PcdGet64 (PcdS3BootScriptTableSize);
    Status = PcdSet64S (PcdS3BootScriptTableSize, gS3BootScriptData);
    if (EFI_ERROR (Status)) {
      DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:591");
    }

    ZeroMem (gS3BootScriptData, 32);

    // Register DxeSmmReadyToLock notification
    gS3DxeRdyLockEvent = CreateReadyToLockEvent (0, 0, 0, 0, &Registration);
    if (gS3DxeRdyLockEvent == NULL) {
      DEBUG_ASSERT (FALSE, L"mEventDxeSmmReadyToLock != ((void *) 0)");
    }
  }

  // Locate SMM S3 protocol
  Status = gBootServices->LocateProtocol (
             &gEfiSmmBase2ProtocolGuid, 0,
             (VOID**)&SmmS3
             );
  if (EFI_ERROR (Status) ||
      !SmmS3->IsSmmS3Active (SmmS3) ||
      SmmS3->GetS3MemoryInfo (SmmS3, &gS3BootScriptData2) >= 0) {

    // Allocate S3 memory descriptor
    PcdValue = (UINT64)gS3BootScriptData2;
    if (PcdValue == 0) {
      Status = SmmS3->AllocateS3Memory (SmmS3, 6, 32, &PcdValue);
      if (EFI_ERROR (Status)) {
        DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:639");
      }
      gS3BootScriptData2 = (UINT8*)(UINTN)PcdValue;
      gS3MemAttrAlloc = TRUE;

      Status = PcdSet64S (PcdS3BootScriptTableSize, gS3BootScriptData2);
      if (EFI_ERROR (Status)) {
        DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:643");
      }

      ZeroMem (gS3BootScriptData2, 32);

      // Register S3 notifications
      Status = SmmS3->RegisterS3Notify (
                        SmmS3, &gEfiDxeServicesTableGuid,
                        S3BootScriptNotifyDxeSmmReadyToLock,
                        &gS3LockBoxReg
                        );
      if (EFI_ERROR (Status)) {
        DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:654");
      }
      Status = SmmS3->RegisterS3Notify (
                        SmmS3, &gEfiSmmLockBoxCommGuid,
                        S3BootScriptNotifyDxeSmmReadyToLock,
                        &gS3ReadyToBootNotify
                        );
      if (EFI_ERROR (Status)) {
        DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:661");
      }
    }

    // Register ReadyToBoot notification
    Status = SmmS3->RegisterS3Notify (
                      SmmS3, &gEfiEventReadyToBootGuid,
                      S3BootScriptNotifyReadyToBoot,
                      &gS3RdyBootReg
                      );
    if (EFI_ERROR (Status)) {
      DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:673");
    }
  }

  return EFI_SUCCESS;
}

// ... (remaining S3 Boot Script, SmmLockBox, USRA functions implemented
//      as per decompilation of sub_1DC8, sub_1E24, sub_2554, etc.)
//      Full implementation in RegAccessSMM_full.c

// ---------------------------------------------------------------------------
// Helper Functions
// ---------------------------------------------------------------------------

VOID
CopyMem (
  OUT VOID       *Destination,
  IN  CONST VOID *Source,
  IN  UINTN      Length
  )
{
  // Validate no overflow
  if (Length > 0) {
    ASSERT ((Length - 1) <= (MAX_UINTN - (UINTN)Destination));
    ASSERT ((Length - 1) <= (MAX_UINTN - (UINTN)Source));
    if (Destination != Source) {
      InternalCopyMem (Destination, Source, Length);
    }
  }
}

VOID
ZeroMem (
  OUT VOID   *Buffer,
  IN  UINTN  Length
  )
{
  ASSERT (Buffer != NULL);
  ASSERT (Length <= (MAX_UINTN - (UINTN)Buffer + 1));
  InternalZeroMem (Buffer, Length);
}

VOID
DebugPrint (
  IN  UINTN        ErrorLevel,
  IN  CONST CHAR16 *Format,
  ...
  )
{
  VA_LIST           Marker;
  UINT8             DebugLevel;
  UINT8             CmosIndex;

  if (gDebugProtocol == NULL) {
    return;
  }

  // Read CMOS debug level: port 0x70 index 0x4C, port 0x71 data
  CmosIndex = __inbyte (CMOS_INDEX_PORT);
  __outbyte (CMOS_INDEX_PORT, (CmosIndex & 0x80) | 0x4C);
  DebugLevel = __inbyte (CMOS_DATA_PORT);

  if (DebugLevel > 3) {
    // Platform-specific debug level detection
    UINT8 Level = DebugLevel;
    if (Level == 0) {
      Level = *(volatile UINT8*)0xFDAF0490 & 2 | 1;
    }
    // Map to standard error level
  }

  // Call SMM Debug protocol print
  if (gDebugProtocol != NULL) {
    VA_START (Marker, Format);
    // (*(gDebugProtocol->Print))(ErrorLevel, Format, Marker);
    VA_END (Marker);
  }
}

VOID
DebugAssert (
  IN CONST CHAR8  *FileName,
  IN UINTN        LineNumber,
  IN CONST CHAR8  *Description
  )
{
  if (gDebugProtocol != NULL) {
    // (*(gDebugProtocol->Assert))(FileName, LineNumber, Description);
  }
}

// ---------------------------------------------------------------------------
// SMM Spinlock
// ---------------------------------------------------------------------------

BOOLEAN
AcquireSpinLock (
  IN OUT volatile UINT64 *Lock
  )
{
  ASSERT (Lock != NULL);
  ASSERT (*Lock == 1 || *Lock == 2);

  return _InterlockedCompareExchange64 (Lock, 2, 1) == 1;
}

VOID
ReleaseSpinLock (
  IN OUT volatile UINT64 *Lock
  )
{
  ASSERT (Lock != NULL);
  ASSERT (*Lock == 1 || *Lock == 2);

  *Lock = 1;
}

// ---------------------------------------------------------------------------
// SmmLockBox Interface
// ---------------------------------------------------------------------------

/**
  Find LockBox context in SMM configuration table.
**/
VOID*
SmmLockBoxFindContext (
  VOID
  )
{
  UINTN  Index;
  UINT64 Count;

  Count = gSmst->NumberOfTableEntries;
  if (Count == 0) {
    return NULL;
  }

  for (Index = 0; Index < Count; Index++) {
    if (CompareGuid (&gSmst->SmmConfigurationTable[Index].VendorGuid,
                     &gEfiSmmLockBoxCommGuid)) {
      return gSmst->SmmConfigurationTable[Index].VendorTable;
    }
  }

  return NULL;
}

/**
  Return LockBox queue head from LockBox context.
**/
LIST_ENTRY*
SmmLockBoxGetQueue (
  VOID
  )
{
  VOID *Context;

  Context = SmmLockBoxFindContext ();
  ASSERT (Context != NULL);

  return (LIST_ENTRY*)((UINT8*)Context + 8);
}

// ---------------------------------------------------------------------------
// S3 Boot Script: notify DxeSmmReadyToLock
// ---------------------------------------------------------------------------

VOID
S3BootScriptNotifyDxeSmmReadyToLock (
  VOID
  )
{
  EFI_STATUS  Status;

  if (gS3BootScriptData[15] != 0) {
    return;
  }

  *((UINT32*)gS3BootScriptData + 4) = *((UINT32*)gS3BootScriptData + 2) + 3;

  // Restore boot script data from LockBox
  Status = SmmLockBoxRestore (&gS3BootScriptDataGuid, NULL, NULL);
  if (EFI_ERROR (Status)) {
    DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:473");
  }

  // Save updated boot script data to LockBox
  Status = SmmLockBoxSave (
             &gS3BootScriptDataGuid,
             gS3BootScriptData,
             *((UINT32*)gS3BootScriptData + 4)
             );
  if (EFI_ERROR (Status)) {
    DEBUG_ASSERT (!EFI_ERROR (Status), L"BootScriptSave.c:476");
  }
}

// ---------------------------------------------------------------------------
// USRA Register Access
// ---------------------------------------------------------------------------

/**
  Decode a USRA address into a PCIe MMIO address.

  USRA address format:
    Bits 31:20  - Socket ID (from address >> 13 & 3)
    Bits 19:15  - Instance number
    Bits 14:12  - Box type (SktPkg IIO component)
    Bits 11:0   - Register offset / BDF
**/
UINT64
UsraDecodeAddress (
  IN UINT32 *UsraAddr,
  OUT UINT64 *PcieAddr
  )
{
  UINT64  BaseAddr;
  UINT64  SocketBase;
  UINT32  Socket;
  UINT8   BoxType;
  UINT8   Instance;
  UINT64  BoxInstance;
  UINT32  BdfOffset;

  // Initialize socket topology data
  UsraSocketDataInit ();

  Socket      = (*UsraAddr >> 13) & 3;
  BoxType     = *((UINT8*)UsraAddr + 3);
  Instance    = *((UINT8*)UsraAddr + 4);

  // Mask high bits in register address
  *((UINT8*)UsraAddr + 2) &= 0xF0;

  // Get socket MMIO base
  SocketBase = UsraGetSocketData (Socket, BoxType, Instance);
  SocketBase <<= 20;

  // Get box instance
  BoxInstance = UsraGetBoxInstance (SocketBase, BoxType, Instance);
  SocketBase ^= (BoxInstance << 15);
  SocketBase &= 0xF8000;
  SocketBase ^= BaseAddr;  // mix with original

  return SocketBase;
}

// ---------------------------------------------------------------------------
// IIO Topology Initialization
// ---------------------------------------------------------------------------

/**
  Initialize per-socket IIO topology data from IioProtocol.
  Called once; caches socket presence, enabled masks, stepping, port counts.
**/
UINT8
UsraSocketDataInit (
  VOID
  )
{
  UINT32  SocketIndex;
  UINT8  *SocketData;
  UINT8  *IoData;

  if (gUsraSocketData != 0) {
    goto COPY_CACHED;
  }

  // Locate IIO protocol
  gBootServices->LocateProtocol (&gEfiSmmUsraProtocolGuid, NULL, &gIioProtocol);
  gIioProtocolData = *(UINT8**)gIioProtocol;

  // Extract per-socket topology data (4 sockets, 6 bytes per socket)
  for (SocketIndex = 0; SocketIndex < 4; SocketIndex++) {
    UINT8 *IioEntry = gIioProtocolData + 29 + SocketIndex * 43;

    gSocketTopology[SocketIndex] = IioEntry[12];         // socket BDF data
    gSocketTopology[SocketIndex + 4] = IioEntry[0];       // socket type
    gSocketTopology[SocketIndex + 8] = IioEntry[1];       // mmcfg addr
    gSocketTopology[SocketIndex +12] = IioEntry[2];       // reserved

    // Copy raw IO data (6 bytes per socket)
    CopyMem (&gSocketTopoData[SocketIndex * 6],
             &gIioProtocolData[SocketIndex * 37 + 42],
             6);
  }

  gIioPresence  = *(gIioProtocolData + 1780);
  gSocketCnt    = *(gIioProtocolData + 1777);
  gSockEnMask   = *(UINT32*)(gIioProtocolData + 2067);
  gSockPrMask   = *(UINT32*)(gIioProtocolData + 2071);
  gSysMemSize   = *(UINT32*)(gIioProtocolData + 280);
  gMaxPciePorts = *(gIioProtocolData + 2102);
  gCpuStepping  = *(gIioProtocolData + 2103);
  gCpuType      = *(gIioProtocolData + 2104);

  gUsraSocketData = (UINT64)&gSocketTopology;

COPY_CACHED:
  // Copy to active set
  CopyMem (&gSockTopoActive, &gSocketTopology, sizeof (gSocketTopology));
  CopyMem (&gSockTopoActive[4], &gSocketTopoData, 4 * 6);
  gIioPresenceActive  = gIioPresence;
  gSockEnMaskActive   = gSockEnMask;
  gSockPrMaskActive   = gSockPrMask;
  gMaxPciePortsActive = gMaxPciePorts;
  gCpuSteppingActive  = gCpuStepping;

  return gCpuStepping;
}