[Patch v3 3/3] MdeModulePkg/PiSmmCore: Install Protocol when S3 resume finished.

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

Install EdkiiSmmEndOfS3ResumeProtocol when S3 resume finished.
S3ResumePei will send S3 resume finished event to SmmCore through
communication buffer.

V2 change:
 None.

V3 change:
1. Uninstall the protocol right after install it to avoid run out of memory.

Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Eric Dong <eric.dong@intel.com>
---
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.c   | 69 +++++++++++++++++++++++++++----
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.h   | 24 +++++++++++
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf |  1 +
 3 files changed, 87 insertions(+), 7 deletions(-)

diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
index 9e4390e..b833763 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
@@ -75,13 +75,14 @@ BOOLEAN  mInLegacyBoot = FALSE;
 // Table of SMI Handlers that are registered by the SMM Core when it is initialized
 //
 SMM_CORE_SMI_HANDLERS  mSmmCoreSmiHandlers[] = {
-  { SmmDriverDispatchHandler,   &gEfiEventDxeDispatchGuid,          NULL, TRUE  },
-  { SmmReadyToLockHandler,      &gEfiDxeSmmReadyToLockProtocolGuid, NULL, TRUE }, 
-  { SmmLegacyBootHandler,       &gEfiEventLegacyBootGuid,           NULL, FALSE },
-  { SmmExitBootServicesHandler, &gEfiEventExitBootServicesGuid,     NULL, FALSE },
-  { SmmReadyToBootHandler,      &gEfiEventReadyToBootGuid,          NULL, FALSE },
-  { SmmEndOfDxeHandler,         &gEfiEndOfDxeEventGroupGuid,        NULL, TRUE },
-  { NULL,                       NULL,                               NULL, FALSE }
+  { SmmDriverDispatchHandler,   &gEfiEventDxeDispatchGuid,           NULL, TRUE  },
+  { SmmReadyToLockHandler,      &gEfiDxeSmmReadyToLockProtocolGuid,  NULL, TRUE }, 
+  { SmmLegacyBootHandler,       &gEfiEventLegacyBootGuid,            NULL, FALSE },
+  { SmmExitBootServicesHandler, &gEfiEventExitBootServicesGuid,      NULL, FALSE },
+  { SmmReadyToBootHandler,      &gEfiEventReadyToBootGuid,           NULL, FALSE },
+  { SmmEndOfDxeHandler,         &gEfiEndOfDxeEventGroupGuid,         NULL, TRUE },
+  { SmmEndOfS3ResumeHandler,    &gEdkiiSmmEndOfS3ResumeProtocolGuid, NULL, FALSE },
+  { NULL,                       NULL,                                NULL, FALSE }
 };
 
 UINTN                           mFullSmramRangeCount;
@@ -383,6 +384,60 @@ SmmEndOfDxeHandler (
 }
 
 /**
+  Software SMI handler that is called when the EndOfS3Resume event is trigged.
+  This function installs the SMM EndOfS3Resume Protocol so SMM Drivers are informed that
+  S3 resume has finished.
+
+  @param  DispatchHandle  The unique handle assigned to this handler by SmiHandlerRegister().
+  @param  Context         Points to an optional handler context which was specified when the handler was registered.
+  @param  CommBuffer      A pointer to a collection of data in memory that will
+                          be conveyed from a non-SMM environment into an SMM environment.
+  @param  CommBufferSize  The size of the CommBuffer.
+
+  @return Status Code
+
+**/
+EFI_STATUS
+EFIAPI
+SmmEndOfS3ResumeHandler (
+  IN     EFI_HANDLE  DispatchHandle,
+  IN     CONST VOID  *Context,        OPTIONAL
+  IN OUT VOID        *CommBuffer,     OPTIONAL
+  IN OUT UINTN       *CommBufferSize  OPTIONAL
+  )
+{
+  EFI_STATUS  Status;
+  EFI_HANDLE  SmmHandle;
+
+  DEBUG ((EFI_D_INFO, "SmmEndOfS3ResumeHandler\n"));
+
+  //
+  // Install SMM EndOfS3Resume protocol
+  //
+  SmmHandle = NULL;
+  Status = SmmInstallProtocolInterface (
+             &SmmHandle,
+             &gEdkiiSmmEndOfS3ResumeProtocolGuid,
+             EFI_NATIVE_INTERFACE,
+             NULL
+             );
+  ASSERT_EFI_ERROR (Status);
+
+  //
+  // Uninstall the protocol here because the comsume just hook the
+  // installation event.
+  //
+  Status = SmmUninstallProtocolInterface (
+           SmmHandle,
+           &gEdkiiSmmEndOfS3ResumeProtocolGuid,
+           NULL
+           );
+  ASSERT_EFI_ERROR (Status);
+
+  return Status;
+}
+
+/**
   Determine if two buffers overlap in memory.
 
   @param[in] Buff1  Pointer to first buffer
diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.h b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.h
index b6f815c..6cc824b 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.h
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.h
@@ -32,6 +32,7 @@
 #include <Protocol/SmmExitBootServices.h>
 #include <Protocol/SmmLegacyBoot.h>
 #include <Protocol/SmmReadyToBoot.h>
+#include <Protocol/SmmEndOfS3Resume.h>
 
 #include <Guid/Apriori.h>
 #include <Guid/EventGroup.h>
@@ -802,6 +803,29 @@ SmmReadyToBootHandler (
   );
 
 /**
+  Software SMI handler that is called when the EndOfS3Resume event is trigged.
+  This function installs the SMM EndOfS3Resume Protocol so SMM Drivers are informed that
+  S3 resume has finished.
+
+  @param  DispatchHandle  The unique handle assigned to this handler by SmiHandlerRegister().
+  @param  Context         Points to an optional handler context which was specified when the handler was registered.
+  @param  CommBuffer      A pointer to a collection of data in memory that will
+                          be conveyed from a non-SMM environment into an SMM environment.
+  @param  CommBufferSize  The size of the CommBuffer.
+
+  @return Status Code
+
+**/
+EFI_STATUS
+EFIAPI
+SmmEndOfS3ResumeHandler (
+  IN     EFI_HANDLE  DispatchHandle,
+  IN     CONST VOID  *Context,        OPTIONAL
+  IN OUT VOID        *CommBuffer,     OPTIONAL
+  IN OUT UINTN       *CommBufferSize  OPTIONAL
+  );
+
+/**
   Place holder function until all the SMM System Table Service are available.
 
   @param  Arg1                   Undefined
diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
index 95e34bd..a724189 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
@@ -79,6 +79,7 @@
   gEdkiiSmmExitBootServicesProtocolGuid         ## SOMETIMES_PRODUCES
   gEdkiiSmmLegacyBootProtocolGuid               ## SOMETIMES_PRODUCES
   gEdkiiSmmReadyToBootProtocolGuid              ## PRODUCES
+  gEdkiiSmmEndOfS3ResumeProtocolGuid            ## SOMETIMES_PRODUCES
 
   gEfiSmmSwDispatch2ProtocolGuid                ## SOMETIMES_CONSUMES
   gEfiSmmSxDispatch2ProtocolGuid                ## SOMETIMES_CONSUMES
-- 
2.7.0.windows.1