1 files changed, 132 insertions, 0 deletions
diff --git a/src/VBox/Devices/EFI/Firmware/MdeModulePkg/Bus/Ata/AhciPei/AhciPeiS3.c b/src/VBox/Devices/EFI/Firmware/MdeModulePkg/Bus/Ata/AhciPei/AhciPeiS3.c
new file mode 100644
index 00000000..f3b4bfbc
--- /dev/null
+++ b/src/VBox/Devices/EFI/Firmware/MdeModulePkg/Bus/Ata/AhciPei/AhciPeiS3.c
@@ -0,0 +1,132 @@
+/** @file
+  The AhciPei driver is used to manage ATA hard disk device working under AHCI
+  mode at PEI phase.
+
+  Copyright (c) 2019, Intel Corporation. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include "AhciPei.h"
+
+#include <Guid/S3StorageDeviceInitList.h>
+
+#include <Library/LockBoxLib.h>
+
+/**
+  Collect the ports that need to be enumerated on a controller for S3 phase.
+
+  @param[in]  HcDevicePath          Device path of the controller.
+  @param[in]  HcDevicePathLength    Length of the device path specified by
+                                    HcDevicePath.
+  @param[out] PortBitMap            Bitmap that indicates the ports that need
+                                    to be enumerated on the controller.
+
+  @retval    The number of ports that need to be enumerated.
+
+**/
+UINT8
+AhciS3GetEumeratePorts (
+  IN  EFI_DEVICE_PATH_PROTOCOL    *HcDevicePath,
+  IN  UINTN                       HcDevicePathLength,
+  OUT UINT32                      *PortBitMap
+  )
+{
+  EFI_STATUS                  Status;
+  UINT8                       DummyData;
+  UINTN                       S3InitDevicesLength;
+  EFI_DEVICE_PATH_PROTOCOL    *S3InitDevices;
+  EFI_DEVICE_PATH_PROTOCOL    *DevicePathInst;
+  UINTN                       DevicePathInstLength;
+  BOOLEAN                     EntireEnd;
+  SATA_DEVICE_PATH            *SataDeviceNode;
+
+  *PortBitMap = 0;
+
+  //
+  // From the LockBox, get the list of device paths for devices need to be
+  // initialized in S3.
+  //
+  S3InitDevices       = NULL;
+  S3InitDevicesLength = sizeof (DummyData);
+  EntireEnd           = FALSE;
+  Status = RestoreLockBox (&gS3StorageDeviceInitListGuid, &DummyData, &S3InitDevicesLength);
+  if (Status != EFI_BUFFER_TOO_SMALL) {
+    return 0;
+  } else {
+    S3InitDevices = AllocatePool (S3InitDevicesLength);
+    if (S3InitDevices == NULL) {
+      return 0;
+    }
+
+    Status = RestoreLockBox (&gS3StorageDeviceInitListGuid, S3InitDevices, &S3InitDevicesLength);
+    if (EFI_ERROR (Status)) {
+      return 0;
+    }
+  }
+
+  if (S3InitDevices == NULL) {
+    return 0;
+  }
+
+  //
+  // Only enumerate the ports that exist in the device list.
+  //
+  do {
+    //
+    // Fetch the size of current device path instance.
+    //
+    Status = GetDevicePathInstanceSize (
+               S3InitDevices,
+               &DevicePathInstLength,
+               &EntireEnd
+               );
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+
+    DevicePathInst = S3InitDevices;
+    S3InitDevices  = (EFI_DEVICE_PATH_PROTOCOL *)((UINTN) S3InitDevices + DevicePathInstLength);
+
+    if (HcDevicePathLength >= DevicePathInstLength) {
+      continue;
+    }
+
+    //
+    // Compare the device paths to determine if the device is managed by this
+    // controller.
+    //
+    if (CompareMem (
+          DevicePathInst,
+          HcDevicePath,
+          HcDevicePathLength - sizeof (EFI_DEVICE_PATH_PROTOCOL)
+          ) == 0) {
+      //
+      // Get the port number.
+      //
+      while (DevicePathInst->Type != END_DEVICE_PATH_TYPE) {
+        if ((DevicePathInst->Type == MESSAGING_DEVICE_PATH) &&
+            (DevicePathInst->SubType == MSG_SATA_DP)) {
+          SataDeviceNode = (SATA_DEVICE_PATH *) DevicePathInst;
+          //
+          // For now, the driver only support upto AHCI_MAX_PORTS ports and
+          // devices directly connected to a HBA.
+          //
+          if ((SataDeviceNode->HBAPortNumber >= AHCI_MAX_PORTS) ||
+              (SataDeviceNode->PortMultiplierPortNumber != 0xFFFF)) {
+            break;
+          }
+          *PortBitMap |= (UINT32)BIT0 << SataDeviceNode->HBAPortNumber;
+          break;
+        }
+        DevicePathInst = NextDevicePathNode (DevicePathInst);
+      }
+    }
+  } while (!EntireEnd);
+
+  //
+  // Return the number of ports need to be enumerated on this controller.
+  //
+  return AhciGetNumberOfPortsFromMap (*PortBitMap);
+}