Move registration and processing of ExitBootServices() events into UefiDriverEntryPoi...

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@2056 6f19259b-4bc3-4df7-8a09-765794883524

diff --git a/MdePkg/Library/UefiDriverEntryPoint/DriverEntryPoint.c b/MdePkg/Library/UefiDriverEntryPoint/DriverEntryPoint.c
index 60177bc9c5b223fba534dee3f5f1b52ef666fd2d..cd096c59b5edca493bd483f7f263ca11c919b18a 100644 (file)
--- a/MdePkg/Library/UefiDriverEntryPoint/DriverEntryPoint.c
+++ b/MdePkg/Library/UefiDriverEntryPoint/DriverEntryPoint.c
@@ -22,6 +22,9 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
   @retval EFI_SUCCESS\r
 \r
 **/\r
+\r
+EFI_EVENT  _mDriverExitBootServicesNotifyEvent;\r
+\r
 EFI_STATUS\r
 EFIAPI\r
 _DriverUnloadHandler (\r
@@ -41,6 +44,14 @@ _DriverUnloadHandler (
   // unloaded, and the library destructors should not be called\r
   //\r
   if (!EFI_ERROR (Status)) {\r
+    //\r
+    // Close our ExitBootServices () notify function\r
+    //\r
+    if (_gDriverExitBootServicesEvent[0] != NULL) {\r
+      Status = gBS->CloseEvent (_mDriverExitBootServicesNotifyEvent);\r
+      ASSERT_EFI_ERROR (Status);\r
+    }\r
+\r
     ProcessLibraryDestructorList (ImageHandle, gST);\r
   }\r
 \r
@@ -50,6 +61,39 @@ _DriverUnloadHandler (
   return Status;\r
 }\r
 \r
+VOID\r
+EFIAPI\r
+_DriverExitBootServices (\r
+  IN EFI_EVENT        Event,\r
+  IN VOID             *Context\r
+  )\r
+/*++\r
+\r
+Routine Description:\r
+\r
+  Set AtRuntime flag as TRUE after ExitBootServices\r
+\r
+Arguments:\r
+\r
+  Event   - The Event that is being processed\r
+  \r
+  Context - Event Context\r
+\r
+Returns: \r
+\r
+  None\r
+\r
+--*/\r
+{\r
+  EFI_EVENT_NOTIFY  ChildNotifyEventHandler;\r
+  UINTN             Index;\r
+\r
+  for (Index = 0; _gDriverExitBootServicesEvent[Index] != NULL; Index++) {\r
+    ChildNotifyEventHandler = _gDriverExitBootServicesEvent[Index];\r
+    ChildNotifyEventHandler (Event, NULL);\r
+  }\r
+}\r
+\r
 /**\r
   Enrty point to DXE Driver.\r
 \r
@@ -84,6 +128,21 @@ _ModuleEntryPoint (
   //\r
   ProcessLibraryConstructorList (ImageHandle, SystemTable);\r
 \r
+  //\r
+  // Register our ExitBootServices () notify function\r
+  //\r
+  if (_gDriverExitBootServicesEvent[0] != NULL) {\r
+    Status = gBS->CreateEvent (\r
+                    EFI_EVENT_SIGNAL_EXIT_BOOT_SERVICES,\r
+                    EFI_TPL_NOTIFY,\r
+                    _DriverExitBootServices,\r
+                    NULL,\r
+                    &_mDriverExitBootServicesNotifyEvent\r
+                    );\r
+\r
+    ASSERT_EFI_ERROR (Status);\r
+  }\r
+\r
   //\r
   //  Install unload handler...\r
   //\r
@@ -106,6 +165,14 @@ _ModuleEntryPoint (
   // If all of the drivers returned errors, then invoke all of the library destructors\r
   //\r
   if (EFI_ERROR (Status)) {\r
+    //\r
+    // Close our ExitBootServices () notify function\r
+    //\r
+    if (_gDriverExitBootServicesEvent[0] != NULL) {\r
+      Status = gBS->CloseEvent (_mDriverExitBootServicesNotifyEvent);\r
+      ASSERT_EFI_ERROR (Status);\r
+    }\r
+\r
     ProcessLibraryDestructorList (ImageHandle, SystemTable);\r
   }\r
 \r