Newer
Older
AMI-Aptio-BIOS-Reversed / SataController / SataController.c
@Ajax Dong Ajax Dong 2 days ago 35 KB Init
 * EFI_IDE_CONTROLLER_INIT_PROTOCOL implementation for SATA Controller driver.
 *
 * Provides channel management, device identification, and mode negotiation
 * for SATA controllers on Intel Purley PCH (Lewisburg).
 *
 * SataGetDeviceCount      - Returns number of channels and devices per channel
 * SataGetDevice            - Validates channel index
 * SataIdentifyDevice       - Copies 512-byte ATA identify data for a channel
 * SataModifyDevice         - Copies 44-byte channel info structure
 * SataSubmitAsyncCommand   - Submits async command and queries PIO/DMA modes
 * SataCalculateBestPioMode - Calculates best PIO mode from identify data
 * SataStopDevice           - No-op stub
 */

#include "SataController.h"

//
// ======================================================================
// EFI_IDE_CONTROLLER_INIT_PROTOCOL implementation
// ======================================================================
//

/**
 * Returns the number of SATA channels and maximum devices per channel.
 *
 * @param[in]  This                     Pointer to the protocol instance
 * @param[out] NumberOfChannels         Total SATA channels available
 * @param[out] MaximumDevicesPerChannel Max devices per channel
 *
 * @return EFI_SUCCESS or EFI_INVALID_PARAMETER
 */
EFI_STATUS
EFIAPI
SataGetDeviceCount (
  IN  EFI_IDE_CONTROLLER_INIT_PROTOCOL  *This,
  OUT UINT8                             *NumberOfChannels,
  OUT UINT8                             *MaximumDevicesPerChannel
  )
{
  SATA_PRIVATE_DATA  *SataPrivateData;

  SataPrivateData = SATA_PRIVATE_DATA_FROM_THIS (This);
  if (SataPrivateData == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (Channel >= SataPrivateData->IdeRegistersBase) {
    *NumberOfChannels = 0;
    return EFI_INVALID_PARAMETER;
  }

  *NumberOfChannels = 1;
  *MaximumDevicesPerChannel = 1;
  return EFI_SUCCESS;
}

/**
 * Validates that a device exists on the given channel.
 *
 * @param[in]  This     Pointer to the protocol instance
 * @param[in]  Channel  Channel number
 * @param[out] Device   Device presence flags
 *
 * @return EFI_SUCCESS if channel valid, EFI_INVALID_PARAMETER otherwise
 */
EFI_STATUS
EFIAPI
SataGetDevice (
  IN  EFI_IDE_CONTROLLER_INIT_PROTOCOL  *This,
  IN  UINT8                             Channel,
  OUT UINT8                             *Device
  )
{
  SATA_PRIVATE_DATA  *SataPrivateData;

  SataPrivateData = SATA_PRIVATE_DATA_FROM_THIS (This);
  if (SataPrivateData == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (Channel >= SataPrivateData->IdeRegistersBase) {
    return EFI_INVALID_PARAMETER;
  }

  if (*Device > 6) {
    return EFI_UNSUPPORTED;
  }

  return EFI_SUCCESS;
}

/**
 * Copies ATA identify data (512 bytes) for the specified channel/device.
 *
 * The 512-byte identify data block is stored at:
 *   SataPrivateData + 424 + (Channel << 9)
 *
 * @param[in]  This         Pointer to the protocol instance
 * @param[in]  Channel      Channel number
 * @param[in]  Device       Device number (must be 0)
 * @param[in]  IdentifyData Buffer to receive 512-byte identify data
 *
 * @return EFI_SUCCESS or error code
 */
EFI_STATUS
EFIAPI
SataIdentifyDevice (
  IN  EFI_IDE_CONTROLLER_INIT_PROTOCOL  *This,
  IN  UINT8                             Channel,
  IN  UINT8                             Device,
  IN  EFI_IDE_IDENTIFY_DATA             *IdentifyData
  )
{
  SATA_PRIVATE_DATA  *SataPrivateData;

  SataPrivateData = SATA_PRIVATE_DATA_FROM_THIS (This);
  if (SataPrivateData == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (Channel >= SataPrivateData->IdeRegistersBase || Device != 0) {
    return EFI_INVALID_PARAMETER;
  }

  if (IdentifyData != NULL) {
    CopyMem (
      IdentifyData,
      SataPrivateData->DeviceIdentifyData[Channel],
      SATA_IDENTIFY_DATA_SIZE
      );
    SataPrivateData->DeviceValid[Channel] = 1;
  } else {
    SataPrivateData->DeviceValid[Channel] = 0;
  }

  return EFI_SUCCESS;
}

/**
 * Copies channel info (44 bytes) for the specified channel/device.
 *
 * Channel info is stored at SataPrivateData + 72 + (Channel * 44).
 *
 * @param[in]  This         Pointer to the protocol instance
 * @param[in]  Channel      Channel number
 * @param[in]  Device       Device number (must be 0)
 * @param[in]  ChannelInfo  Buffer to receive 44-byte channel info
 *
 * @return EFI_SUCCESS or error code
 */
EFI_STATUS
EFIAPI
SataModifyDevice (
  IN  EFI_IDE_CONTROLLER_INIT_PROTOCOL  *This,
  IN  UINT8                             Channel,
  IN  UINT8                             Device,
  IN  EFI_IDE_CHANNEL_INFO              *ChannelInfo
  )
{
  SATA_PRIVATE_DATA  *SataPrivateData;

  SataPrivateData = SATA_PRIVATE_DATA_FROM_THIS (This);
  if (SataPrivateData == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (Channel >= SataPrivateData->IdeRegistersBase ||
      Device != 0 ||
      ChannelInfo == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  CopyMem (&SataPrivateData->ChannelInfo[Channel], ChannelInfo, sizeof (EFI_IDE_CHANNEL_INFO));
  return EFI_SUCCESS;
}

/**
 * Submits an async command for the device on the given channel.
 *
 * Checks that the port has a valid device, allocates a result buffer,
 * queries the best PIO mode from identify data, and returns
 * negotiated DMA transfer mode.
 *
 * @param[in]  This     Pointer to the protocol instance
 * @param[in]  Channel  Channel number
 * @param[in]  Device   Device number (must be 0)
 *
 * @return EFI_SUCCESS or error code
 */
EFI_STATUS
EFIAPI
SataSubmitAsyncCommand (
  IN  EFI_IDE_CONTROLLER_INIT_PROTOCOL  *This,
  IN  UINT8                             Channel,
  IN  UINT8                             Device
  )
{
  SATA_PRIVATE_DATA  *SataPrivateData;
  EFI_STATUS         Status;
  UINT16             BestPioMode;
  UINT16             MaxMode;
  UINT8              *ResultBuffer;

  SataPrivateData = SATA_PRIVATE_DATA_FROM_THIS (This);
  if (SataPrivateData == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (Channel >= SataPrivateData->IdeRegistersBase ||
      Device != 0) {
    return EFI_INVALID_PARAMETER;
  }

  if (!SataPrivateData->DeviceValid[Channel]) {
    return EFI_NOT_READY;
  }

  //
  // Allocate result buffer for caller
  //
  ResultBuffer = AllocatePool (sizeof (EFI_IDE_CHANNEL_INFO));
  if (ResultBuffer == NULL) {
    ASSERT ((BOOLEAN)(0==1));
    return EFI_OUT_OF_RESOURCES;
  }
  ZeroMem (ResultBuffer, sizeof (EFI_IDE_CHANNEL_INFO));

  //
  // Query best PIO mode
  //
  MaxMode = 0;
  if (SataPrivateData->ChannelInfo[Channel].ChannelType != 0) {
    MaxMode = *(UINT16 *)((UINTN)&SataPrivateData->ChannelInfo[Channel] + 4);
  }

  Status = SataCalculateBestPioMode (
             (EFI_IDE_IDENTIFY_DATA *)SataPrivateData->DeviceIdentifyData[Channel],
             MaxMode != 0 ? &MaxMode : NULL,
             &BestPioMode
             );
  if (!EFI_ERROR (Status)) {
    ResultBuffer[0] = 1;                 // Mode valid
    *(UINT32 *)(ResultBuffer + 4) = BestPioMode;
  } else {
    ResultBuffer[0] = 0;                 // Mode invalid
  }

  //
  // Check DMA capability
  //
  if (SataPrivateData->DeviceIdentifyData[Channel][106] & 4) {
    //
    // Device supports DMA - calculate best DMA mode
    //
    UINT16 DmaMode = 0;
    UINT16 MaxDmaMode = (*(UINT16 *)&SataPrivateData->DeviceIdentifyData[Channel][600] >> 1) & 0x1F;

    while (MaxDmaMode) {
      DmaMode++;
      MaxDmaMode >>= 1;
    }

    if (MaxMode != 0 && MaxMode > 0) {
      if (DmaMode >= MaxMode) {
        DmaMode = MaxMode - 1;
      }
    }

    ResultBuffer[24] = 1;                // DMA valid
    *(UINT32 *)(ResultBuffer + 28) = (UINT32)DmaMode;
  } else {
    ResultBuffer[24] = 0;                // No DMA
  }

  *(UINTN *)((EFI_IDE_CONTROLLER_INIT_PROTOCOL *)This) = (UINTN)ResultBuffer;
  return EFI_SUCCESS;
}

/**
 * No-op stop device function.
 *
 * @return EFI_SUCCESS always
 */
EFI_STATUS
EFIAPI
SataStopDevice (
  IN  EFI_IDE_CONTROLLER_INIT_PROTOCOL  *This,
  IN  UINT8                             Channel,
  IN  UINT8                             Device
  )
{
  return EFI_SUCCESS;
}

//
// ======================================================================
// PIO Mode Calculation
// ======================================================================
//

/**
 * Calculates the best PIO mode for a device based on identify data.
 *
 * Reads the PIO mode timing from identify data fields 103, 128,
 * and 136 to determine the maximum supported PIO mode.
 *
 * @param[in]  IdentifyData    Pointer to 512-byte ATA identify data
 * @param[in]  ChannelInfoMaxMode Optional upper limit from channel info
 * @param[out] BestMode        Calculated best PIO mode
 *
 * @return EFI_SUCCESS or EFI_UNSUPPORTED
 */
EFI_STATUS
SataCalculateBestPioMode (
  IN  EFI_IDE_IDENTIFY_DATA   *IdentifyData,
  IN  UINT16                  *ChannelInfoMaxMode OPTIONAL,
  OUT UINT16                  *BestMode
  )
{
  UINT8   Field103;
  UINT16  Field128;
  UINT16  Field136;
  UINT8   BestPio;

  DEBUG ((DEBUG_ERROR, "CalculateBestPioMode() End\n"));

  //
  // Field 103: IORDY support and mode info
  //
  Field103 = IdentifyData->Field103;
  Field128 = *(UINT16 *)&IdentifyData->Data[128];

  if ((Field103 & 2) == 0) {
    //
    // No IORDY - use PIO mode 0-2 only
    //
    BestPio = Field103;
    if (ChannelInfoMaxMode != NULL) {
      if (*ChannelInfoMaxMode < 2) {
        return EFI_UNSUPPORTED;
      }
      if (BestPio >= *ChannelInfoMaxMode) {
        BestPio = (UINT8)(*ChannelInfoMaxMode - 1);
      }
    }
    if (BestPio >= 2) {
      *BestMode = 2;
      return EFI_SUCCESS;
    }
    *BestMode = 1;
    return EFI_SUCCESS;
  }

  //
  // IORDY supported - determine PIO mode 0-4
  //
  {
    UINT16  Field136;
    UINT16  PioModeBits;

    Field136 = *(UINT16 *)&IdentifyData->Data[136];
    PioModeBits = Field128 >> 8;

    //
    // Determine PIO mode from timing
    //
    if (PioModeBits <= 0x78) {
      BestPio = 4;
    } else if (PioModeBits <= 0xB4) {
      BestPio = 3;
    } else if (PioModeBits <= 0xF0) {
      BestPio = 2;
    } else {
      BestPio = 0;
    }
  }

  if (ChannelInfoMaxMode != NULL) {
    if (*ChannelInfoMaxMode < 2) {
      return EFI_UNSUPPORTED;
    }
    if (BestPio >= *ChannelInfoMaxMode) {
      BestPio = (UINT8)(*ChannelInfoMaxMode - 1);
    }
  }

  if (BestPio < 2) {
    *BestMode = 1;
  } else {
    *BestMode = BestPio;
  }

  return EFI_SUCCESS;
}

//
// ======================================================================
// PCH Series Detection
// ======================================================================
//

/**
 * Determines the PCH series from the LPC device ID.
 *
 * Checks cached value first, reads LPC device ID via PCI config if not set.
 * Supported: SPT (1), CNP (2), Unknown (3).
 *
 * @return PCH series identifier
 */
UINTN
PchGetPchSeries (
  VOID
  )
{
  if (gPchSeries == 3) {
    UINT64  PciBase;
    UINT16  LpcDeviceId;
    UINTN   Series = 1;

    PciBase = PcieGetExpressBase ();
    LpcDeviceId = IoRead16 ((UINT16 *)(PciBase + 2));

    //
    // Check if device ID falls in SPT range (0x9D00-0x9D3F with mask 0xFF70)
    //
    if (((LpcDeviceId + 24128) & 0xFF70) != 0) {
      //
      // Check CNP range
      //
      if ((UINT16)(LpcDeviceId + 25280) <= 8) {
        int Mask = 335;
        if ((Mask >> (LpcDeviceId + 25280)) & 1) {
          Series = 2;  // Cannon Point
        } else {
          DEBUG ((DEBUG_ERROR, "Unsupported PCH SKU, LpcDeviceId: 0x%04x!\n", LpcDeviceId));
          ASSERT ((BOOLEAN)(0==1));
        }
      }
    } else {
      Series = 1;  // Sunrise Point
    }

    gPchSeries = Series;
    return Series;
  }

  return gPchSeries;
}

//
// ======================================================================
// S3 Boot Script Library
// ======================================================================
//

/**
 * Initializes the S3 boot script save context.
 *
 * Allocates EfiReservedMemory for the boot script buffer,
 * registers DxeSmmReadyToLock event callback (S3BootScriptEventNotify),
 * locates SMM protocol, registers S3ReadyToLockNotify callback,
 * and sets up boot script close events.
 *
 * @return EFI_SUCCESS on success
 */
EFI_STATUS
S3BootScriptLibInit (
  VOID
  )
{
  EFI_STATUS  Status;

  //
  // Allocate EfiReservedMemory for boot script context
  //
  Status = GetPcdDbProtocol ();
  gBootScriptContext = (BOOT_SCRIPT_SAVE_CONTEXT *)
    ((UINT64 (*)(UINT32))GetPcdDbProtocol ()->GetPtr) (137);
  if (gBootScriptContext == NULL) {
    EFI_PHYSICAL_ADDRESS  PagesAddr = (EFI_PHYSICAL_ADDRESS)(UINTN)-1;
    Status = gBS->AllocatePages (AllocateAnyPages, EfiReservedMemoryType, 1, &PagesAddr);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "ASSERT_EFI_ERROR (Status = %r)\n", Status));
      ASSERT (!EFI_ERROR (Status));
    }
    gBootScriptContext = (BOOT_SCRIPT_SAVE_CONTEXT *)(UINTN)PagesAddr;
    gBootScriptAllocated = TRUE;

    Status = ((UINT64 (*)(UINT32, UINT64))GetPcdDbProtocol ()->SetPtr) (137, (UINT64)PagesAddr);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "ASSERT_EFI_ERROR (Status = %r)\n", Status));
      ASSERT (!EFI_ERROR (Status));
    }

    ZeroMem (gBootScriptContext, 32);
    gDxeSmmReadyToLockEvent = EfiLibCreateNotifyEvent (
                                &gEfiDxeSmmReadyToLockProtocolGuid,
                                S3BootScriptEventNotify,
                                &gDxeSmmReadyToLockEvent
                                );
    if (gDxeSmmReadyToLockEvent == NULL) {
      ASSERT (gDxeSmmReadyToLockEvent != ((VOID *)0));
    }
  }

  gBootScriptContext->Buffer = (UINT8 *)gBootScriptContext;
  gBootScriptContext->EntryCount = 0;

  //
  // Locate SMM protocol
  //
  Status = gBS->LocateProtocol (&gEfiSmmBase2ProtocolGuid, NULL, &gSmst);
  if (!EFI_ERROR (Status)) {
    VOID  *Buffer;
    Status = ((SMM_BASE2_PROTOCOL *)gSmst)->InSmm (gSmst);
    if (!Status) {
      Status = ((SMM_BASE2_PROTOCOL *)gSmst)->GetSmstLocation (gSmst, &gSmst);
    }

    //
    // Register S3 boot script close notification
    //
    Status = gBS->CreateEvent (
                    EVT_NOTIFY_SIGNAL,
                    TPL_CALLBACK,
                    S3BootScriptEventNotify,
                    NULL,
                    &Buffer
                    );
    if (!EFI_ERROR (Status)) {
      gEventExitBootSvc = Buffer;
      Status = gBS->RegisterProtocolNotify (
                      &gEfiExitBootServicesProtocolGuid,
                      (EFI_EVENT)Buffer,
                      &gEventExitBootSvc
                      );
      ASSERT (!EFI_ERROR (Status));
    }

    if (!EFI_ERROR (Status)) {
      Status = gBS->RegisterProtocolNotify (
                      &gEfiSetVirtualAddressMapProtocolGuid,
                      (EFI_EVENT)Buffer,
                      &gEventSetVirtualAddrMap
                      );
      ASSERT (!EFI_ERROR (Status));
    }
  }

  //
  // Register S3ReadyToLockNotify
  //
  Status = gBS->RegisterProtocolNotify (
                  &gEfiS3ReadyToLockProtocolGuid,
                  (EFI_EVENT)S3ReadyToLockNotify,
                  &gEventReadyToLock
                  );
  ASSERT (!EFI_ERROR (Status));

  return EFI_SUCCESS;
}

//
// ======================================================================
// S3 Boot Script Event Notify Callbacks
// ======================================================================
//

/**
 * Event notify for DxeSmmReadyToLock.
 *
 * On first invocation, finalizes the boot script and saves it to LockBox.
 *
 * @param[in] Event   Event being signaled
 * @param[in] Context Not used
 */
VOID
EFIAPI
S3BootScriptEventNotify (
  IN EFI_EVENT  Event,
  IN VOID       *Context
  )
{
  EFI_STATUS  Status;

  Status = gBS->LocateProtocol (&gEfiSmmBase2ProtocolGuid, NULL, &Context);
  if (!EFI_ERROR (Status)) {
    {
      SMM_BASE2_PROTOCOL  *SmmBase = (SMM_BASE2_PROTOCOL *)Context;

      Status = SmmBase->InSmm (SmmBase);
      if (!Status) {
        //
        // Inside SMM - check if already saved
        //
        if (gBootScriptContext != NULL && !gBootScriptContext->FlagSaved) {
          S3BootScriptFinalize ();
          SaveBootScriptDataToLockBox ();
          gBootScriptContext->FlagSaved = TRUE;
        }
      }
    }
  }
}

/**
 * Finalizes the boot script by appending a terminator entry.
 * Terminator: 0xFF, 0x00, 0x03
 */
VOID
S3BootScriptFinalize (
  VOID
  )
{
  BOOT_SCRIPT_SAVE_CONTEXT  *Context;

  Context = gBootScriptContext;
  if (Context == NULL || Context->Buffer == NULL) {
    return;
  }

  UINT8  Terminator[3] = { 0xFF, 0x00, 0x03 };
  CopyMem (
    Context->Buffer + Context->EntryCount,
    Terminator,
    3
    );
  *(UINT32 *)(Context->Buffer + 5) = Context->EntryCount + 3;
}

/**
 * Event notify for S3 ready to lock.
 * Copies boot script context to SMM context if different.
 */
VOID
EFIAPI
S3ReadyToLockNotify (
  IN EFI_EVENT  Event,
  IN VOID       *Context
  )
{
  S3BootScriptEventNotify (Event, Context);

  if (gBootScriptContext != gAltBootScriptContext) {
    if (gAltBootScriptContext != NULL &&
        gAltBootScriptContext->Buffer == NULL) {

      CopyMem (
        gAltBootScriptContext,
        gBootScriptContext,
        sizeof (BOOT_SCRIPT_SAVE_CONTEXT)
        );
      gAltBootScriptContext->FlagInProgress = 1;
    }
    gBootScriptContext = gAltBootScriptContext;
  }
}

/**
 * Closes the S3 boot script event.
 * Saves boot script lockbox entry with final data.
 */
EFI_STATUS
S3BootScriptCloseEvent (
  VOID
  )
{
  EFI_STATUS  Status;

  gBootScriptContext->FlagInProgress = 1;

  Status = SmmLockBoxSave (&gBootScriptCloseGuid,
                           gBootScriptContext,
                           sizeof (BOOT_SCRIPT_SAVE_CONTEXT));
  ASSERT (!EFI_ERROR (Status));

  Status = SmmLockBoxSetAttributes (&gBootScriptCloseGuid);
  ASSERT (!EFI_ERROR (Status));

  gBootScriptContext->FlagInProgress = 0;
  return EFI_SUCCESS;
}

/**
 * Event notify for boot script close.
 */
VOID
EFIAPI
S3BootScriptCloseNotify (
  IN EFI_EVENT  Event,
  IN VOID       *Context
  )
{
  EFI_STATUS  Status;

  if (gBootScriptContext != NULL && !gBootScriptContext->FlagEntriesClosed) {
    gBootScriptContext->FinalEntryOffset = gBootScriptContext->EntryCount + 3;

    Status = S3BootScriptCloseEvent ();
    if (!EFI_ERROR (Status)) {
      gBootScriptContext->FlagEntriesClosed = 1;
    }
  }
}

//
// ======================================================================
// S3 Boot Script Data save to LockBox
// ======================================================================
//

/**
 * Saves the boot script data and memory ranges to SMM LockBox.
 *
 * 1. Saves boot script context to lockbox with set attributes
 * 2. Gets system memory map
 * 3. Builds memory ranges structure (excluding MMIO ranges)
 * 4. Saves memory ranges to lockbox
 */
EFI_STATUS
SaveBootScriptDataToLockBox (
  VOID
  )
{
  EFI_STATUS                     Status;
  UINTN                          MapKey;
  UINTN                          MapDescriptorSize;
  UINT32                         MapDescriptorVersion;
  EFI_MEMORY_DESCRIPTOR          *MemoryMap = NULL;
  UINTN                          MemoryMapSize;
  UINTN                          NumberOfRanges;
  EFI_PHYSICAL_ADDRESS           PagesAddr;
  MEMORY_RANGES_STRUCTURE        *MemoryRanges;
  UINT8                          *MemRangesBuffer;
  UINTN                          BufferSize;
  UINTN                          Index;
  UINTN                          RangeIndex;

  //
  // Step 1: Save boot script context to lockbox
  //
  Status = SmmLockBoxSave (&gBootScriptFinalGuid,
                           gBootScriptContext,
                           sizeof (BOOT_SCRIPT_SAVE_CONTEXT));
  ASSERT (!EFI_ERROR (Status));

  Status = SmmLockBoxSetAttributes (&gBootScriptFinalGuid);
  ASSERT (!EFI_ERROR (Status));

  //
  // Step 2: Get memory map
  //
  MemoryMapSize = 0;
  Status = gBS->GetMemoryMap (&MemoryMapSize, MemoryMap, &MapKey,
                              &MapDescriptorSize, &MapDescriptorVersion);

  if (Status == EFI_BUFFER_TOO_SMALL) {
    MemoryMap = AllocatePool (MemoryMapSize);
    if (MemoryMap == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }

    Status = gBS->GetMemoryMap (&MemoryMapSize, MemoryMap, &MapKey,
                                &MapDescriptorSize, &MapDescriptorVersion);
  }

  ASSERT (Status == EFI_SUCCESS);

  //
  // Step 3: Allocate memory ranges structure
  //
  PagesAddr = (EFI_PHYSICAL_ADDRESS)(UINTN)-1;
  Status = gBS->AllocatePages (AllocateAnyPages, EfiReservedMemoryType,
                               EFI_SIZE_TO_PAGES (sizeof (MEMORY_RANGES_STRUCTURE)),
                               &PagesAddr);
  ASSERT (!EFI_ERROR (Status));

  MemoryRanges = (MEMORY_RANGES_STRUCTURE *)(UINTN)PagesAddr;
  MemoryRanges->NumberOfMemoryRanges = 0;

  //
  // Count ranges that are not MMIO
  //
  for (Index = 0; Index < MemoryMapSize / MapDescriptorSize; Index++) {
    EFI_MEMORY_DESCRIPTOR *Desc = (EFI_MEMORY_DESCRIPTOR *)
      ((UINT8 *)MemoryMap + Index * MapDescriptorSize);
    if (Desc->Type > EfiMemoryMappedIO || Desc->Type == EfiRuntimeServicesCode ||
        Desc->Type == EfiRuntimeServicesData || Desc->Type == EfiBootServicesCode ||
        Desc->Type == EfiBootServicesData || Desc->Type == EfiConventionalMemory ||
        Desc->Type == EfiLoaderCode || Desc->Type == EfiLoaderData ||
        Desc->Type == EfiACPIMemoryNVS || Desc->Type == EfiACPIReclaimMemory) {
      MemoryRanges->NumberOfMemoryRanges++;
    }
  }

  //
  // Allocate memory ranges buffer
  //
  BufferSize = MemoryRanges->NumberOfMemoryRanges * sizeof (MEMORY_RANGE_ENTRY);
  PagesAddr = (EFI_PHYSICAL_ADDRESS)(UINTN)-1;
  Status = gBS->AllocatePages (AllocateAnyPages, EfiReservedMemoryType,
                               EFI_SIZE_TO_PAGES (BufferSize), &PagesAddr);
  ASSERT (!EFI_ERROR (Status));

  MemoryRanges->MemRanges = (MEMORY_RANGE_ENTRY *)(UINTN)PagesAddr;
  if (MemoryRanges->MemRanges == NULL) {
    MemoryRanges->NumberOfMemoryRanges = 0;
    DEBUG ((DEBUG_ERROR, "ERROR: No memory to save MemMap !!!\n"));
    ASSERT (MemoryRanges->MemRanges != NULL);
  }

  //
  // Step 4: Fill memory ranges
  //
  RangeIndex = 0;
  for (Index = 0; Index < MemoryMapSize / MapDescriptorSize; Index++) {
    EFI_MEMORY_DESCRIPTOR *Desc = (EFI_MEMORY_DESCRIPTOR *)
      ((UINT8 *)MemoryMap + Index * MapDescriptorSize);
    if (Desc->Type > EfiMemoryMappedIO || Desc->Type == EfiRuntimeServicesCode ||
        Desc->Type == EfiRuntimeServicesData || Desc->Type == EfiBootServicesCode ||
        Desc->Type == EfiBootServicesData || Desc->Type == EfiConventionalMemory ||
        Desc->Type == EfiLoaderCode || Desc->Type == EfiLoaderData ||
        Desc->Type == EfiACPIMemoryNVS || Desc->Type == EfiACPIReclaimMemory) {
      MemoryRanges->MemRanges[RangeIndex].Base = Desc->PhysicalStart;
      MemoryRanges->MemRanges[RangeIndex].Length = Desc->NumberOfPages * EFI_PAGE_SIZE;
      RangeIndex++;
    }
  }

  ASSERT (RangeIndex == MemoryRanges->NumberOfMemoryRanges);

  gBS->FreePool (MemoryMap);

  //
  // Step 5: Save memory ranges to lockbox
  //
  Status = SmmLockBoxSave (&gBootScriptFinalGuid, MemoryRanges,
                           sizeof (MEMORY_RANGES_STRUCTURE));
  ASSERT (!EFI_ERROR (Status));

  Status = SmmLockBoxSetAttributes (&gBootScriptFinalGuid);
  ASSERT (!EFI_ERROR (Status));

  Status = SmmLockBoxSave (&gLockBoxGuid2, MemoryRanges->MemRanges,
                           MemoryRanges->NumberOfMemoryRanges * sizeof (MEMORY_RANGE_ENTRY));
  ASSERT (!EFI_ERROR (Status));

  Status = SmmLockBoxSetAttributes (&gLockBoxGuid2);
  ASSERT (!EFI_ERROR (Status));

  gBootScriptContext->MemoryRanges = MemoryRanges;
  return EFI_SUCCESS;
}

//
// ======================================================================
// SMM LockBox Wrappers
// ======================================================================
//

/**
 * Gets the SMM LockBox protocol interface.
 *
 * @return SMM_LOCK_BOX_PROTOCOL pointer or NULL
 */
SMM_LOCK_BOX_PROTOCOL *
GetSmmLockBoxProtocol (
  VOID
  )
{
  EFI_STATUS  Status;

  if (gSmmLockBoxProtocol == NULL) {
    Status = gBS->LocateProtocol (&gEfiSmmLockBoxProtocolGuid, NULL,
                                  &gSmmLockBoxProtocol);
    if (EFI_ERROR (Status)) {
      gSmmLockBoxProtocol = NULL;
    }
  }
  return gSmmLockBoxProtocol;
}

/**
 * Gets SMM communication buffer for LockBox operations.
 * Searches the PiSmmCommunicationRegionTable for a suitable buffer.
 *
 * @return Communication buffer address or NULL
 */
VOID *
SmmLockBoxGetCommBuffer (
  VOID
  )
{
  EFI_STATUS  Status;
  VOID        *Table;
  UINT32      *Entry;
  UINT32      Index;

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

  Status = EfiGetSystemConfigurationTable (&gEfiSmmCommunicationRegionTableGuid, &Table);
  if (EFI_ERROR (Status)) {
    return NULL;
  }

  ASSERT (Table != NULL);

  Entry = (UINT32 *)((UINT8 *)Table + 16);
  for (Index = 0; Index < *(UINT32 *)((UINT8 *)Table + 4); Index++) {
    if (*Entry == 7 && (*(UINT64 *)((UINT8 *)Entry + 24) << 12) >= 0x50) {
      gLockBoxCommBuffer = *(VOID **)((UINT8 *)Entry + 8);
      break;
    }
    Entry = (UINT32 *)((UINT8 *)Entry + *(UINT32 *)((UINT8 *)Table + 8));
  }

  return gLockBoxCommBuffer;
}

/**
 * Saves a data buffer to SMM LockBox.
 *
 * @param[in] Guid       LockBox entry GUID
 * @param[in] Buffer     Data to save
 * @param[in] Length     Length of data in bytes
 *
 * @return EFI_SUCCESS or error code
 */
EFI_STATUS
SmmLockBoxSave (
  IN EFI_GUID  *Guid,
  IN VOID      *Buffer,
  IN UINTN     Length
  )
{
  SMM_LOCK_BOX_PROTOCOL  *Protocol;
  EFI_STATUS             Status;
  UINT8                  StackBuffer[72];
  UINT8                  *CommBuffer;
  UINTN                  CommSize;

  DEBUG ((DEBUG_INFO, "SmmLockBoxDxeLib SaveLockBox - Enter\n"));

  if (Guid == NULL || Buffer == NULL || Length == 0) {
    return EFI_INVALID_PARAMETER;
  }

  Protocol = GetSmmLockBoxProtocol ();
  if (Protocol == NULL) {
    return EFI_NOT_STARTED;
  }

  CommBuffer = SmmLockBoxGetCommBuffer ();
  if (CommBuffer == NULL) {
    CommBuffer = StackBuffer;
  }

  CopyMem (CommBuffer, &gSmmLockBoxProtocolGuid, 16);
  *(UINT64 *)(CommBuffer + 16) = 48;
  *(UINT64 *)(CommBuffer + 32) = (UINT64)-1;  // ReturnStatus
  *(UINT32 *)(CommBuffer + 24) = 1;            // Function: Save
  *(UINT32 *)(CommBuffer + 28) = 48;           // HeaderSize
  CopyMem (CommBuffer + 40, Guid, 16);
  *(UINT64 *)(CommBuffer + 56) = (UINT64)Buffer;
  *(UINT64 *)(CommBuffer + 64) = Length;

  CommSize = 72;
  Status = Protocol->SmmLockBox (Protocol, (VOID *)CommBuffer, &CommSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "ASSERT_EFI_ERROR (Status = %r)\n", Status));
    ASSERT (!EFI_ERROR (Status));
  }

  DEBUG ((DEBUG_INFO, "SmmLockBoxDxeLib SaveLockBox - Exit (%r)\n",
          *(UINT64 *)(CommBuffer + 32)));

  return *(UINT64 *)(CommBuffer + 32);
}

/**
 * Sets LockBox entry attributes.
 *
 * @param[in] Guid  LockBox entry GUID
 *
 * @return EFI_SUCCESS or error code
 */
EFI_STATUS
SmmLockBoxSetAttributes (
  IN EFI_GUID  *Guid
  )
{
  SMM_LOCK_BOX_PROTOCOL  *Protocol;
  EFI_STATUS             Status;
  UINT8                  StackBuffer[64];
  UINT8                  *CommBuffer;
  UINTN                  CommSize;

  DEBUG ((DEBUG_INFO, "SmmLockBoxDxeLib SetLockBoxAttributes - Enter\n"));

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

  Protocol = GetSmmLockBoxProtocol ();
  if (Protocol == NULL) {
    return EFI_NOT_STARTED;
  }

  CommBuffer = SmmLockBoxGetCommBuffer ();
  if (CommBuffer == NULL) {
    CommBuffer = StackBuffer;
  }

  CopyMem (CommBuffer, &gSmmLockBoxProtocolGuid, 16);
  *(UINT64 *)(CommBuffer + 16) = 40;
  *(UINT64 *)(CommBuffer + 32) = (UINT64)-1;
  *(UINT32 *)(CommBuffer + 24) = 4;            // Function: SetAttributes
  *(UINT32 *)(CommBuffer + 28) = 40;
  CopyMem (CommBuffer + 40, Guid, 16);
  *(UINT64 *)(CommBuffer + 56) = 1;             // Attributes: LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY

  CommSize = 64;
  Status = Protocol->SmmLockBox (Protocol, (VOID *)CommBuffer, &CommSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "ASSERT_EFI_ERROR (Status = %r)\n", Status));
    ASSERT (!EFI_ERROR (Status));
  }

  DEBUG ((DEBUG_INFO, "SmmLockBoxDxeLib SetLockBoxAttributes - Exit (%r)\n",
          *(UINT64 *)(CommBuffer + 32)));

  return *(UINT64 *)(CommBuffer + 32);
}

/**
 * Restores a data buffer from SMM LockBox.
 *
 * @param[in] Guid  LockBox entry GUID
 *
 * @return EFI_SUCCESS or error code
 */
EFI_STATUS
SmmLockBoxRestore (
  IN EFI_GUID  *Guid
  )
{
  SMM_LOCK_BOX_PROTOCOL  *Protocol;
  EFI_STATUS             Status;
  UINT8                  StackBuffer[72];
  UINT8                  *CommBuffer;
  UINTN                  CommSize;

  DEBUG ((DEBUG_INFO, "SmmLockBoxDxeLib RestoreLockBox - Enter\n"));

  Protocol = GetSmmLockBoxProtocol ();
  if (Protocol == NULL) {
    return EFI_NOT_STARTED;
  }

  CommBuffer = SmmLockBoxGetCommBuffer ();
  if (CommBuffer == NULL) {
    CommBuffer = StackBuffer;
  }

  CopyMem (CommBuffer, &gSmmLockBoxProtocolGuid, 16);
  *(UINT64 *)(CommBuffer + 16) = 48;
  *(UINT64 *)(CommBuffer + 32) = (UINT64)-1;
  *(UINT32 *)(CommBuffer + 24) = 3;            // Function: Restore
  *(UINT32 *)(CommBuffer + 28) = 48;
  CopyMem (CommBuffer + 40, &gBootScriptFinalGuid, 16);
  *(UINT64 *)(CommBuffer + 56) = 0;
  *(UINT64 *)(CommBuffer + 64) = 0;

  CommSize = 72;
  Status = Protocol->SmmLockBox (Protocol, (VOID *)CommBuffer, &CommSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "ASSERT_EFI_ERROR (Status = %r)\n", Status));
    ASSERT (!EFI_ERROR (Status));
  }

  DEBUG ((DEBUG_INFO, "SmmLockBoxDxeLib RestoreLockBox - Exit (%r)\n",
          *(UINT64 *)(CommBuffer + 32)));

  return *(UINT64 *)(CommBuffer + 32);
}

//
// ======================================================================
// Tear-down / Deinitialization
// ======================================================================
//

/**
 * Tears down S3 boot script resources when the driver unloads.
 *
 * Closes events, frees reserved memory, and unregisters protocol notifies.
 */
EFI_STATUS
S3BootScriptLibDeinit (
  VOID
  )
{
  DEBUG ((DEBUG_INFO, "%a() in %a module\n", __FUNCTION__, "BootScriptSave"));

  //
  // Close DxeSmmReadyToLock event
  //
  if (gDxeSmmReadyToLockEvent != NULL) {
    EFI_STATUS  Status;
    Status = gBS->CloseEvent (gDxeSmmReadyToLockEvent);
    ASSERT (!EFI_ERROR (Status));
  }

  //
  // Unregister notifies
  //
  if (gSmst != NULL) {
    if (gEventExitBootSvc != NULL) {
      gBS->RegisterProtocolNotify (&gEfiExitBootServicesProtocolGuid, NULL, &gEventExitBootSvc);
    }
    if (gEventSetVirtualAddrMap != NULL) {
      gBS->RegisterProtocolNotify (&gEfiSetVirtualAddressMapProtocolGuid, NULL, &gEventSetVirtualAddrMap);
    }
    if (gEventReadyToLock != NULL) {
      gBS->RegisterProtocolNotify (&gEfiS3ReadyToLockProtocolGuid, NULL, &gEventReadyToLock);
    }
  }

  //
  // Free boot script buffer
  //
  if (gBootScriptAllocated && gBootScriptContext != NULL) {
    gBS->FreePages ((EFI_PHYSICAL_ADDRESS)(UINTN)gBootScriptContext, 1);
    ((UINT64 (*)(UINT32, UINT64))GetPcdDbProtocol ()->SetPtr) (137, 0);
  }

  //
  // Free SMM allocated buffer
  //
  if (gSmstAllocated && gSmst != NULL) {
    ((SMM_SYSTEM_TABLE2 *)gSmst)->SmmAllocatePool (gSmst, 0);
  }

  return EFI_SUCCESS;
}

//
// ======================================================================
// Component Name Protocol
// ======================================================================
//

/**
 * Gets the driver name string.
 */
EFI_STATUS
EFIAPI
SataComponentNameGetDriverName (
  IN  EFI_COMPONENT_NAME_PROTOCOL  *This,
  IN  CHAR8                        *Language,
  OUT CHAR16                       **DriverName
  )
{
  if (Language == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  return LookupStringTable (
           Language,
           (CHAR8 **)((EFI_COMPONENT_NAME2_PROTOCOL *)This)->SupportedLanguages,
           gDriverNameTable,
           DriverName,
           This == (EFI_COMPONENT_NAME_PROTOCOL *)&gSataControllerComponentName2
           );
}

/**
 * Gets the controller name string.
 */
EFI_STATUS
EFIAPI
SataComponentNameGetControllerName (
  IN  EFI_COMPONENT_NAME_PROTOCOL  *This,
  IN  EFI_HANDLE                   ControllerHandle,
  IN  EFI_HANDLE                   ChildHandle      OPTIONAL,
  IN  CHAR8                        *Language,
  OUT CHAR16                       **ControllerName
  )
{
  if (ChildHandle != NULL) {
    return EFI_UNSUPPORTED;
  }

  //
  // Verify the controller is managed by this driver
  //
  EFI_STATUS  Status;
  gControllerHandle = ControllerHandle;
  Status = gBS->OpenProtocol (
                  ControllerHandle,
                  &gEfiPciIoProtocolGuid,
                  NULL,
                  gImageHandle,
                  ControllerHandle,
                  EFI_OPEN_PROTOCOL_TEST_PROTOCOL
                  );
  if (EFI_ERROR (Status)) {
    //
    // Not PCI IO - return unsupported
    //
    if (Status != EFI_UNSUPPORTED) {
      return EFI_UNSUPPORTED;
    }
  } else {
    gBS->CloseProtocol (
           ControllerHandle,
           &gEfiPciIoProtocolGuid,
           gImageHandle,
           ControllerHandle
           );
    return EFI_UNSUPPORTED;
  }

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

  return LookupStringTable (
           Language,
           (CHAR8 **)This->SupportedLanguages,
           gControllerNameTable,
           ControllerName,
           This == (EFI_COMPONENT_NAME_PROTOCOL *)&gSataControllerComponentName2
           );
}

//
// ======================================================================
// String Table Lookup
// ======================================================================
//

/**
 * Looks up a string in the component name string table by language token.
 *
 * The string table is an array of { (CHAR8 *)token, (CHAR16 *)string } pairs,
 * terminated by a NULL token.
 *
 * @param[in]  Language       Language string (ISO 639-2)
 * @param[in]  SupportedLanguages Array of supported language strings
 * @param[in]  StringTable    Array of {Token, String} pairs
 * @param[out] FoundString    Resulting string pointer
 * @param[in]  IsComponentName2 TRUE for Component Name 2 (RFC 4646)
 *
 * @return EFI_SUCCESS if found, EFI_UNSUPPORTED if not
 */
EFI_STATUS
LookupStringTable (
  IN  CHAR8          *Language,
  IN  CHAR8          **SupportedLanguages,
  IN  VOID           *StringTable,
  OUT CHAR16         **FoundString,
  IN  BOOLEAN        IsComponentName2
  )
{
  CHAR8   *Lang;
  CHAR8   *Token;
  UINTN   TokenLen;
  UINTN   Index;

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

  if (Language == NULL || StringTable == NULL) {
    return EFI_UNSUPPORTED;
  }

  Lang = Language;
  while (*Lang != '\0') {
    if (IsComponentName2) {
      //
      // RFC 4646 language tags: compare 3-byte prefix
      //
      if (ReadUint24 ((UINT32 *)Lang) == ReadUint24 ((UINT32 *)SupportedLanguages)) {
        goto FOUND;
      }
    } else {
      //
      // ISO 639-2: split by ';', compare each
      //
      TokenLen = 0;
      while (Lang[TokenLen] != '\0' && Lang[TokenLen] != ';') {
        TokenLen++;
      }
      if (AsciiStrCmp (Lang, (CHAR8 *)SupportedLanguages, TokenLen) == 0) {
        goto FOUND;
      }
      //
      // Skip to next language token
      //
      Lang += TokenLen;
      while (*Lang == ';') {
        Lang++;
      }
      continue;
    }
    break;
  }

  return EFI_UNSUPPORTED;

FOUND:
  {
    CHAR8   **Entry;

    Entry = (CHAR8 **)StringTable;
    while (*Entry != NULL) {
      //
      // Compare entry token with language
      //
      CHAR8   *StrTok = *Entry;
      TokenLen = 0;
      while (StrTok[TokenLen] != '\0' && StrTok[TokenLen] != ';') {
        TokenLen++;
      }
      if (AsciiStrCmp (StrTok, Language, TokenLen) == 0) {
        *FoundString = (CHAR16 *)(*(UINTN *)(Entry + 1));
        return EFI_SUCCESS;
      }
      //
      // Skip semicolons
      //
      StrTok += TokenLen;
      while (*StrTok == ';') {
        StrTok++;
      }
      Entry += 2;
    }
  }

  return EFI_UNSUPPORTED;
}