From cad149cc14211c4aefd85ef35b8743eefc62c3e1 Mon Sep 17 00:00:00 2001
From: Jiewen Yao <jiewen.yao@intel.com>
Date: Fri, 11 Aug 2017 22:27:58 +0800
Subject: KabylakeSiliconPkg: Publish DMAR table at PciEnumDone.

As such, we can start DMA protection as early as possible.

Cc: Chasel Chiu <chasel.chiu@intel.com>
Cc: Amy Chan <amy.chan@intel.com>
Cc: Rangasai V Chaganty <rangasai.v.chaganty@intel.com>
Cc: Michael A Kubacki <michael.a.kubacki@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiewen Yao <jiewen.yao@intel.com>
Reviewed-by: Chasel Chiu <chasel.chiu@intel.com>
Reviewed-by: Rangasai V Chaganty <rangasai.v.chaganty@intel.com>
---
 .../SystemAgent/SaInit/Dxe/SaAcpi.c                |  2 -
 .../SystemAgent/SaInit/Dxe/SaInitDxe.c             | 44 ++++++++++++++++++++++
 .../SystemAgent/SaInit/Dxe/SaInitDxe.h             | 19 ++++++++++
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaAcpi.c b/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaAcpi.c
index 4c7fc08200..40aa155ead 100644
--- a/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaAcpi.c
+++ b/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaAcpi.c
@@ -209,8 +209,6 @@ SaAcpiEndOfDxeCallback (
 {
   EFI_STATUS          Status;
 
-  UpdateDmarEndOfDxe ();
-
   if (MmioRead16 (MmPciBase (SA_MC_BUS, 2, 0) + R_SA_IGD_VID) != 0xFFFF) {
     Status = GetVBiosVbtEndOfDxe ();
     if (EFI_SUCCESS != Status) {
diff --git a/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaInitDxe.c b/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaInitDxe.c
index 6900cf8a5b..0948a91a72 100644
--- a/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaInitDxe.c
+++ b/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaInitDxe.c
@@ -33,15 +33,59 @@ SaInitEntryPointDxe (
   )
 {
   EFI_STATUS                Status;
+  VOID                      *Registration;
 
   DEBUG ((DEBUG_INFO, "SaInitDxe Start\n"));
 
 
   Status = SaAcpiInit (ImageHandle);
 
+  ///
+  /// Create PCI Enumeration Completed callback for SA
+  ///
+  EfiCreateProtocolNotifyEvent (
+    &gEfiPciEnumerationCompleteProtocolGuid,
+    TPL_CALLBACK,
+    SaPciEnumCompleteCallback,
+    NULL,
+    &Registration
+    );
 
   DEBUG ((DEBUG_INFO, "SaInitDxe End\n"));
 
   return EFI_SUCCESS;
 }
 
+/**
+  This function gets registered as a callback to perform SA initialization before EndOfDxe
+
+  @param[in] Event     - A pointer to the Event that triggered the callback.
+  @param[in] Context   - A pointer to private data registered with the callback function.
+**/
+VOID
+EFIAPI
+SaPciEnumCompleteCallback (
+  IN EFI_EVENT    Event,
+  IN VOID         *Context
+  )
+{
+  EFI_STATUS          Status;
+  VOID                *ProtocolPointer;
+
+  DEBUG ((DEBUG_INFO, "SaPciEnumCompleteCallback Start\n"));
+  ///
+  /// Check if this is first time called by EfiCreateProtocolNotifyEvent() or not,
+  /// if it is, we will skip it until real event is triggered
+  ///
+  Status = gBS->LocateProtocol (&gEfiPciEnumerationCompleteProtocolGuid, NULL, (VOID **) &ProtocolPointer);
+  if (EFI_SUCCESS != Status) {
+    return;
+  }
+
+  gBS->CloseEvent (Event);
+
+  UpdateDmarEndOfDxe ();
+
+  return;
+}
+
diff --git a/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaInitDxe.h b/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaInitDxe.h
index 55a87fef4d..9381941b87 100644
--- a/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaInitDxe.h
+++ b/Silicon/Intel/KabylakeSiliconPkg/SystemAgent/SaInit/Dxe/SaInitDxe.h
@@ -32,6 +32,25 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 extern EFI_GUID gSaSsdtAcpiTableStorageGuid;
 
+///
+/// Function Prototype
+///
+/**
+  This function gets registered as a callback to perform SA initialization before ExitPmAuth
+
+  @param[in] Event     - A pointer to the Event that triggered the callback.
+  @param[in] Context   - A pointer to private data registered with the callback function.
+
+  @retval EFI_SUCCESS   - Always.
+
+**/
+VOID
+EFIAPI
+SaPciEnumCompleteCallback (
+  IN EFI_EVENT Event,
+  IN VOID      *Context
+  );
+
 /**
   <b>System Agent Initialization DXE Driver Entry Point</b>
   - <b>Introduction</b> \n
-- 
cgit v1.2.3