MdeModulePkg PiSmmCore: Use SxDispatch in SmmEndOfDxeHandler()

As some implementation of SMM Child Dispatcher (including SxDispatch)
may deny the handler registration after SmmReadyToLock, using
SxDispatch in SmmReadyToBootHandler() will be too late.

This patch updates code to use SxDispatch in SmmEndOfDxeHandler()
instead of SmmReadyToBootHandler().

Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Star Zeng <star.zeng@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>

diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
index f136b5d5693bda22ca0f3b49f2cc14a798912e01..d2f0207b5c83b78e39fde30176787131fb24a2f6 100644 (file)
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
@@ -296,9 +296,6 @@ SmmReadyToBootHandler (
 {\r
   EFI_STATUS                        Status;\r
   EFI_HANDLE                        SmmHandle;\r
-  EFI_SMM_SX_DISPATCH2_PROTOCOL     *SxDispatch;\r
-  EFI_SMM_SX_REGISTER_CONTEXT       EntryRegisterContext;\r
-  EFI_HANDLE                        S3EntryHandle;\r
 \r
   //\r
   // Install SMM Ready To Boot protocol.\r
@@ -313,31 +310,7 @@ SmmReadyToBootHandler (
 \r
   SmiHandlerUnRegister (DispatchHandle);\r
 \r
-  //\r
-  // Locate SmmSxDispatch2 protocol.\r
-  //\r
-  Status = SmmLocateProtocol (\r
-             &gEfiSmmSxDispatch2ProtocolGuid,\r
-             NULL,\r
-             (VOID **)&SxDispatch\r
-             );\r
-  if (!EFI_ERROR (Status) && (SxDispatch != NULL)) {\r
-    //\r
-    // Register a S3 entry callback function to\r
-    // determine if it will be during S3 resume.\r
-    //\r
-    EntryRegisterContext.Type  = SxS3;\r
-    EntryRegisterContext.Phase = SxEntry;\r
-    Status = SxDispatch->Register (\r
-                           SxDispatch,\r
-                           SmmS3EntryCallBack,\r
-                           &EntryRegisterContext,\r
-                           &S3EntryHandle\r
-                           );\r
-    ASSERT_EFI_ERROR (Status);\r
-  }\r
-\r
-  return EFI_SUCCESS;\r
+  return Status;\r
 }\r
 \r
 /**\r
@@ -452,10 +425,14 @@ SmmEndOfDxeHandler (
   IN OUT UINTN       *CommBufferSize  OPTIONAL\r
   )\r
 {\r
-  EFI_STATUS  Status;\r
-  EFI_HANDLE  SmmHandle;\r
+  EFI_STATUS                        Status;\r
+  EFI_HANDLE                        SmmHandle;\r
+  EFI_SMM_SX_DISPATCH2_PROTOCOL     *SxDispatch;\r
+  EFI_SMM_SX_REGISTER_CONTEXT       EntryRegisterContext;\r
+  EFI_HANDLE                        S3EntryHandle;\r
 \r
   DEBUG ((EFI_D_INFO, "SmmEndOfDxeHandler\n"));\r
+\r
   //\r
   // Install SMM EndOfDxe protocol\r
   //\r
@@ -466,7 +443,32 @@ SmmEndOfDxeHandler (
              EFI_NATIVE_INTERFACE,\r
              NULL\r
              );\r
-  return Status;\r
+\r
+  //\r
+  // Locate SmmSxDispatch2 protocol.\r
+  //\r
+  Status = SmmLocateProtocol (\r
+             &gEfiSmmSxDispatch2ProtocolGuid,\r
+             NULL,\r
+             (VOID **)&SxDispatch\r
+             );\r
+  if (!EFI_ERROR (Status) && (SxDispatch != NULL)) {\r
+    //\r
+    // Register a S3 entry callback function to\r
+    // determine if it will be during S3 resume.\r
+    //\r
+    EntryRegisterContext.Type  = SxS3;\r
+    EntryRegisterContext.Phase = SxEntry;\r
+    Status = SxDispatch->Register (\r
+                           SxDispatch,\r
+                           SmmS3EntryCallBack,\r
+                           &EntryRegisterContext,\r
+                           &S3EntryHandle\r
+                           );\r
+    ASSERT_EFI_ERROR (Status);\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -526,7 +528,6 @@ SmmEndOfS3ResumeHandler (
   ASSERT_EFI_ERROR (Status);\r
 \r
   mDuringS3Resume = FALSE;\r
-\r
   return Status;\r
 }\r
 \r