1 files changed, 52 insertions, 5 deletions
diff --git a/SecurityPkg/VariableAuthenticated/RuntimeDxe/VariableDxe.c b/SecurityPkg/VariableAuthenticated/RuntimeDxe/VariableDxe.c
index f91cb5dc56..a2e764cf61 100644
--- a/SecurityPkg/VariableAuthenticated/RuntimeDxe/VariableDxe.c
+++ b/SecurityPkg/VariableAuthenticated/RuntimeDxe/VariableDxe.c
@@ -16,11 +16,13 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include "Variable.h"
 #include "AuthService.h"
 
-extern VARIABLE_STORE_HEADER        *mNvVariableCache;
-extern VARIABLE_INFO_ENTRY          *gVariableInfo;
-EFI_HANDLE                          mHandle                    = NULL;
-EFI_EVENT                           mVirtualAddressChangeEvent = NULL;
-EFI_EVENT                           mFtwRegistration           = NULL;
+extern VARIABLE_STORE_HEADER   *mNvVariableCache;
+extern VARIABLE_INFO_ENTRY     *gVariableInfo;
+EFI_HANDLE                     mHandle                    = NULL;
+EFI_EVENT                      mVirtualAddressChangeEvent = NULL;
+EFI_EVENT                      mFtwRegistration           = NULL;
+extern BOOLEAN                 mEndOfDxe;
+EDKII_VARIABLE_LOCK_PROTOCOL   mVariableLock              = { VariableLockRequestToLock };
 
 /**
   Return TRUE if ExitBootServices () has been called.
@@ -257,12 +259,34 @@ OnReadyToBoot (
   VOID                                    *Context
   )
 {
+  //
+  // Set the End Of DXE bit in case the EFI_END_OF_DXE_EVENT_GROUP_GUID event is not signaled.
+  //
+  mEndOfDxe = TRUE;
   ReclaimForOS ();
   if (FeaturePcdGet (PcdVariableCollectStatistics)) {
     gBS->InstallConfigurationTable (&gEfiAuthenticatedVariableGuid, gVariableInfo);
   }
 }
 
+/**
+  Notification function of EFI_END_OF_DXE_EVENT_GROUP_GUID event group.
+
+  This is a notification function registered on EFI_END_OF_DXE_EVENT_GROUP_GUID event group.
+
+  @param  Event        Event whose notification function is being invoked.
+  @param  Context      Pointer to the notification function's context.
+
+**/
+VOID
+EFIAPI
+OnEndOfDxe (
+  EFI_EVENT                               Event,
+  VOID                                    *Context
+  )
+{
+  mEndOfDxe = TRUE;
+}
 
 /**
   Fault Tolerant Write protocol notification event handler.
@@ -378,10 +402,19 @@ VariableServiceInitialize (
 {
   EFI_STATUS                            Status;
   EFI_EVENT                             ReadyToBootEvent;
+  EFI_EVENT                             EndOfDxeEvent;
 
   Status = VariableCommonInitialize ();
   ASSERT_EFI_ERROR (Status);
 
+  Status = gBS->InstallMultipleProtocolInterfaces (
+                  &mHandle,
+                  &gEdkiiVariableLockProtocolGuid,
+                  &mVariableLock,
+                  NULL
+                  );
+  ASSERT_EFI_ERROR (Status);
+
   SystemTable->RuntimeServices->GetVariable         = VariableServiceGetVariable;
   SystemTable->RuntimeServices->GetNextVariableName = VariableServiceGetNextVariableName;
   SystemTable->RuntimeServices->SetVariable         = VariableServiceSetVariable;
@@ -428,6 +461,20 @@ VariableServiceInitialize (
              NULL,
              &ReadyToBootEvent
              );
+  ASSERT_EFI_ERROR (Status);
+
+  //
+  // Register the event handling function to set the End Of DXE flag.
+  //
+  Status = gBS->CreateEventEx (
+                  EVT_NOTIFY_SIGNAL,
+                  TPL_NOTIFY,
+                  OnEndOfDxe,
+                  NULL,
+                  &gEfiEndOfDxeEventGroupGuid,
+                  &EndOfDxeEvent
+                  );
+  ASSERT_EFI_ERROR (Status);
 
   return EFI_SUCCESS;
 }