SecurityPkg-Opal(2): Enhance AHCI Bar MMIO region check.

[mirror_edk2.git] / SecurityPkg / Tcg / Opal / OpalPasswordSmm / OpalPasswordSmm.c

diff --git a/SecurityPkg/Tcg/Opal/OpalPasswordSmm/OpalPasswordSmm.c b/SecurityPkg/Tcg/Opal/OpalPasswordSmm/OpalPasswordSmm.c
index 57d10a17b77a42d60c44be990371e58697ace3d7..b7e2d552d5c2f25dd0e3ccecdc450e2779925d5b 100644 (file)
--- a/SecurityPkg/Tcg/Opal/OpalPasswordSmm/OpalPasswordSmm.c
+++ b/SecurityPkg/Tcg/Opal/OpalPasswordSmm/OpalPasswordSmm.c
@@ -61,6 +61,9 @@ VOID                 *mBuffer = NULL; // DMA can not read/write Data to smram, s
 // NVME\r
 NVME_CONTEXT         mNvmeContext;\r
 \r
+EFI_GCD_MEMORY_SPACE_DESCRIPTOR   *mGcdMemSpace        = NULL;\r
+UINTN                             mNumberOfDescriptors = 0;\r
+\r
 /**\r
   Add new bridge node or nvme device info to the device list.\r
 \r
@@ -592,6 +595,44 @@ S3SleepEntryCallBack (
   return Status;\r
 }\r
 \r
+/**\r
+  OpalPassword Notification for SMM EndOfDxe protocol.\r
+\r
+  @param[in] Protocol   Points to the protocol's unique identifier.\r
+  @param[in] Interface  Points to the interface instance.\r
+  @param[in] Handle     The handle on which the interface was installed.\r
+\r
+  @retval EFI_SUCCESS   Notification runs successfully.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+OpalPasswordEndOfDxeNotification (\r
+  IN CONST EFI_GUID  *Protocol,\r
+  IN VOID            *Interface,\r
+  IN EFI_HANDLE      Handle\r
+  )\r
+{\r
+  UINTN                            NumberOfDescriptors;\r
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *MemSpaceMap;\r
+  EFI_STATUS                       Status;\r
+\r
+  Status = gDS->GetMemorySpaceMap (&NumberOfDescriptors, &MemSpaceMap);\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
+  mGcdMemSpace = AllocateCopyPool (NumberOfDescriptors * sizeof (EFI_GCD_MEMORY_SPACE_DESCRIPTOR), MemSpaceMap);\r
+  if (EFI_ERROR (Status)) {\r
+    gBS->FreePool (MemSpaceMap);\r
+    return Status;\r
+  }\r
+\r
+  mNumberOfDescriptors = NumberOfDescriptors;\r
+  gBS->FreePool (MemSpaceMap);\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
 /**\r
   Main entry for this driver.\r
 \r
@@ -618,6 +659,7 @@ OpalPasswordSmmInit (
   EFI_SMM_VARIABLE_PROTOCOL             *SmmVariable;\r
   OPAL_EXTRA_INFO_VAR                   OpalExtraInfo;\r
   UINTN                                 DataSize;\r
+  EFI_EVENT                             EndOfDxeEvent;\r
   EFI_PHYSICAL_ADDRESS                  Address;\r
 \r
   mBuffer            = NULL;\r
@@ -726,6 +768,15 @@ OpalPasswordSmmInit (
   //\r
   mSwSmiValue = (UINT8) Context.SwSmiInputValue;\r
 \r
+  //\r
+  // Create event to record GCD descriptors at end of dxe for judging AHCI/NVMe PCI Bar\r
+  // is in MMIO space to avoid attack.\r
+  //\r
+  Status = gSmst->SmmRegisterProtocolNotify (&gEfiSmmEndOfDxeProtocolGuid, OpalPasswordEndOfDxeNotification, &EndOfDxeEvent);\r
+  if (EFI_ERROR (Status)) {\r
+    DEBUG((DEBUG_ERROR, "OpalPasswordSmm: Register SmmEndOfDxe fail, Status: %r\n", Status));\r
+    goto EXIT;\r
+  }\r
   Status = gSmst->SmmLocateProtocol (&gEfiSmmVariableProtocolGuid, NULL, (VOID**)&SmmVariable);\r
   if (!EFI_ERROR (Status)) {\r
     DataSize = sizeof (OPAL_EXTRA_INFO_VAR);\r