1 files changed, 67 insertions, 0 deletions
diff --git a/MdePkg/Library/UefiDriverEntryPoint/DriverEntryPoint.c b/MdePkg/Library/UefiDriverEntryPoint/DriverEntryPoint.c
index 60177bc9c5..cd096c59b5 100644
--- 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
 
 **/
+
+EFI_EVENT  _mDriverExitBootServicesNotifyEvent;
+
 EFI_STATUS
 EFIAPI
 _DriverUnloadHandler (
@@ -41,6 +44,14 @@ _DriverUnloadHandler (
   // unloaded, and the library destructors should not be called
   //
   if (!EFI_ERROR (Status)) {
+    //
+    // Close our ExitBootServices () notify function
+    //
+    if (_gDriverExitBootServicesEvent[0] != NULL) {
+      Status = gBS->CloseEvent (_mDriverExitBootServicesNotifyEvent);
+      ASSERT_EFI_ERROR (Status);
+    }
+
     ProcessLibraryDestructorList (ImageHandle, gST);
   }
 
@@ -50,6 +61,39 @@ _DriverUnloadHandler (
   return Status;
 }
 
+VOID
+EFIAPI
+_DriverExitBootServices (
+  IN EFI_EVENT        Event,
+  IN VOID             *Context
+  )
+/*++
+
+Routine Description:
+
+  Set AtRuntime flag as TRUE after ExitBootServices
+
+Arguments:
+
+  Event   - The Event that is being processed
+  
+  Context - Event Context
+
+Returns: 
+
+  None
+
+--*/
+{
+  EFI_EVENT_NOTIFY  ChildNotifyEventHandler;
+  UINTN             Index;
+
+  for (Index = 0; _gDriverExitBootServicesEvent[Index] != NULL; Index++) {
+    ChildNotifyEventHandler = _gDriverExitBootServicesEvent[Index];
+    ChildNotifyEventHandler (Event, NULL);
+  }
+}
+
 /**
   Enrty point to DXE Driver.
 
@@ -85,6 +129,21 @@ _ModuleEntryPoint (
   ProcessLibraryConstructorList (ImageHandle, SystemTable);
 
   //
+  // Register our ExitBootServices () notify function
+  //
+  if (_gDriverExitBootServicesEvent[0] != NULL) {
+    Status = gBS->CreateEvent (
+                    EFI_EVENT_SIGNAL_EXIT_BOOT_SERVICES,
+                    EFI_TPL_NOTIFY,
+                    _DriverExitBootServices,
+                    NULL,
+                    &_mDriverExitBootServicesNotifyEvent
+                    );
+
+    ASSERT_EFI_ERROR (Status);
+  }
+
+  //
   //  Install unload handler...
   //
   if (_gDriverUnloadImageCount != 0) {
@@ -106,6 +165,14 @@ _ModuleEntryPoint (
   // If all of the drivers returned errors, then invoke all of the library destructors
   //
   if (EFI_ERROR (Status)) {
+    //
+    // Close our ExitBootServices () notify function
+    //
+    if (_gDriverExitBootServicesEvent[0] != NULL) {
+      Status = gBS->CloseEvent (_mDriverExitBootServicesNotifyEvent);
+      ASSERT_EFI_ERROR (Status);
+    }
+
     ProcessLibraryDestructorList (ImageHandle, SystemTable);
   }