[Patch] MdeModulePkg/PiSmmCore: Control S3 related functionality with flag.

[Eric Dong <eric.dong@intel.com>]

https://bugzilla.tianocore.org/show_bug.cgi?id=1590

Use PcdAcpiS3Enable to control whether need to enable S3 related functionality
in Pi SMM Core.

Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Eric Dong <eric.dong@intel.com>
---
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.c   | 70 ++++++++++++++++++++++---------
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf |  1 +
 2 files changed, 51 insertions(+), 20 deletions(-)

diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
index d0bc65b564..bd19803c37 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
@@ -77,6 +77,12 @@ BOOLEAN  mInLegacyBoot = FALSE;
 //
 BOOLEAN  mDuringS3Resume = FALSE;
 
+//
+// Flag to determine if platform enabled S3.
+// Get the value from PcdAcpiS3Enable.
+//
+BOOLEAN  mAcpiS3Enable = FALSE;
+
 //
 // Table of SMI Handlers that are registered by the SMM Core when it is initialized
 //
@@ -87,6 +93,13 @@ SMM_CORE_SMI_HANDLERS  mSmmCoreSmiHandlers[] = {
   { SmmExitBootServicesHandler, &gEfiEventExitBootServicesGuid,      NULL, FALSE },
   { SmmReadyToBootHandler,      &gEfiEventReadyToBootGuid,           NULL, FALSE },
   { SmmEndOfDxeHandler,         &gEfiEndOfDxeEventGroupGuid,         NULL, TRUE },
+  { NULL,                       NULL,                                NULL, FALSE }
+};
+
+//
+// Table of SMI Handlers that are registered by the SMM Core when it is initialized
+//
+SMM_CORE_SMI_HANDLERS  mSmmCoreS3SmiHandlers[] = {
   { SmmS3SmmInitDoneHandler,    &gEdkiiS3SmmInitDoneGuid,            NULL, FALSE },
   { SmmEndOfS3ResumeHandler,    &gEdkiiEndOfS3ResumeGuid,            NULL, FALSE },
   { NULL,                       NULL,                                NULL, FALSE }
@@ -445,28 +458,30 @@ SmmEndOfDxeHandler (
              NULL
              );
 
-  //
-  // Locate SmmSxDispatch2 protocol.
-  //
-  Status = SmmLocateProtocol (
-             &gEfiSmmSxDispatch2ProtocolGuid,
-             NULL,
-             (VOID **)&SxDispatch
-             );
-  if (!EFI_ERROR (Status) && (SxDispatch != NULL)) {
+  if (mAcpiS3Enable) {
     //
-    // Register a S3 entry callback function to
-    // determine if it will be during S3 resume.
+    // Locate SmmSxDispatch2 protocol.
     //
-    EntryRegisterContext.Type  = SxS3;
-    EntryRegisterContext.Phase = SxEntry;
-    Status = SxDispatch->Register (
-                           SxDispatch,
-                           SmmS3EntryCallBack,
-                           &EntryRegisterContext,
-                           &S3EntryHandle
-                           );
-    ASSERT_EFI_ERROR (Status);
+    Status = SmmLocateProtocol (
+               &gEfiSmmSxDispatch2ProtocolGuid,
+               NULL,
+               (VOID **)&SxDispatch
+               );
+    if (!EFI_ERROR (Status) && (SxDispatch != NULL)) {
+      //
+      // Register a S3 entry callback function to
+      // determine if it will be during S3 resume.
+      //
+      EntryRegisterContext.Type  = SxS3;
+      EntryRegisterContext.Phase = SxEntry;
+      Status = SxDispatch->Register (
+                             SxDispatch,
+                             SmmS3EntryCallBack,
+                             &EntryRegisterContext,
+                             &S3EntryHandle
+                             );
+      ASSERT_EFI_ERROR (Status);
+    }
   }
 
   return EFI_SUCCESS;
@@ -883,6 +898,21 @@ SmmMain (
     ASSERT_EFI_ERROR (Status);
   }
 
+  mAcpiS3Enable = PcdGetBool (PcdAcpiS3Enable);
+  if (mAcpiS3Enable) {
+    //
+    // Register all S3 related SMI Handlers required by the SMM Core
+    //
+    for (Index = 0; mSmmCoreS3SmiHandlers[Index].HandlerType != NULL; Index++) {
+      Status = SmiHandlerRegister (
+                 mSmmCoreS3SmiHandlers[Index].Handler,
+                 mSmmCoreS3SmiHandlers[Index].HandlerType,
+                 &mSmmCoreS3SmiHandlers[Index].DispatchHandle
+                 );
+      ASSERT_EFI_ERROR (Status);
+    }
+  }
+
   RegisterSmramProfileHandler ();
   SmramProfileInstallProtocol ();
 
diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
index f3ece22121..9a31cb79d6 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
@@ -101,6 +101,7 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPageType                   ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPoolType                   ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPropertyMask               ## CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdAcpiS3Enable                        ## CONSUMES
 
 [Guids]
   gAprioriGuid                                  ## SOMETIMES_CONSUMES   ## File
-- 
2.15.0.windows.1