REF: https://bugzilla.tianocore.org/show_bug.cgi?id=3648

Register SmmExitBootServices and SmmLegacyBoot callback function to unregister this handler.

Signed-off-by: Hao Shi <hao.shi@intel.com>
Cc: Dandan Bi <dandan.bi@intel.com>
Cc: Liming Gao <gaoliming@byosoft.com.cn>
---
.../UserAuthenticationSmm.c                   | 34 +++++++++++++++++++
.../UserAuthenticationSmm.inf                 |  2 ++
2 files changed, 36 insertions(+)

diff --git a/Features/Intel/UserInterface/UserAuthFeaturePkg/UserAuthenticationDxeSmm/UserAuthenticationSmm.c b/Features/Intel/UserInterface/UserAuthFeaturePkg/UserAuthenticationDxeSmm/UserAuthenticationSmm.c
index 07e834eb..30f889dd 100644
--- a/Features/Intel/UserInterface/UserAuthFeaturePkg/UserAuthenticationDxeSmm/UserAuthenticationSmm.c
+++ b/Features/Intel/UserInterface/UserAuthFeaturePkg/UserAuthenticationDxeSmm/UserAuthenticationSmm.c
@@ -13,6 +13,7 @@ UINTN                           mAdminPasswordTryCount = 0;

BOOLEAN                         mNeedReVerify = TRUE;
BOOLEAN                         mPasswordVerified = FALSE;
+EFI_HANDLE                      mSmmHandle = NULL;

/**
Verify if the password is correct.
@@ -612,6 +613,30 @@ EXIT:
return EFI_SUCCESS;
}

+/**
+  Performs Exit Boot Services UserAuthentication actions
+
+  @param[in] Protocol   Points to the protocol's unique identifier.
+  @param[in] Interface  Points to the interface instance.
+  @param[in] Handle     The handle on which the interface was installed.
+
+  @retval EFI_SUCCESS   Notification runs successfully.
+**/
+EFI_STATUS
+EFIAPI
+UaExitBootServices (
+  IN CONST EFI_GUID     *Protocol,
+  IN VOID               *Interface,
+  IN EFI_HANDLE         Handle
+  )
+{
+  DEBUG ((DEBUG_INFO, "Unregister User Authentication Smi\n"));
+
+  gSmst->SmiHandlerUnRegister(mSmmHandle);
+
+  return EFI_SUCCESS;
+}
+
/**
Main entry for this driver.

@@ -633,6 +658,7 @@ PasswordSmmInit (
EDKII_VARIABLE_LOCK_PROTOCOL          *VariableLock;
CHAR16                                PasswordHistoryName[sizeof(USER_AUTHENTICATION_VAR_NAME)/sizeof(CHAR16) + 5];
UINTN                                 Index;
+  EFI_EVENT                           ExitBootServicesEvent;

ASSERT (PASSWORD_HASH_SIZE == SHA256_DIGEST_SIZE);
ASSERT (PASSWORD_HISTORY_CHECK_COUNT < 0xFFFF);
@@ -663,6 +689,14 @@ PasswordSmmInit (
if (EFI_ERROR (Status)) {
return Status;
}
+  mSmmHandle = SmmHandle;
+  //
+  // Register for SmmExitBootServices and SmmLegacyBoot notification.
+  //
+  Status = gSmst->SmmRegisterProtocolNotify (&gEdkiiSmmExitBootServicesProtocolGuid, UaExitBootServices, &ExitBootServicesEvent);
+  ASSERT_EFI_ERROR (Status);
+  Status = gSmst->SmmRegisterProtocolNotify (&gEdkiiSmmLegacyBootProtocolGuid, UaExitBootServices, &ExitBootServicesEvent);
+  ASSERT_EFI_ERROR (Status);

if (IsPasswordCleared()) {
DEBUG ((DEBUG_INFO, "IsPasswordCleared\n"));
diff --git a/Features/Intel/UserInterface/UserAuthFeaturePkg/UserAuthenticationDxeSmm/UserAuthenticationSmm.inf b/Features/Intel/UserInterface/UserAuthFeaturePkg/UserAuthenticationDxeSmm/UserAuthenticationSmm.inf
index 0b33b194..d73a2fe2 100644
--- a/Features/Intel/UserInterface/UserAuthFeaturePkg/UserAuthenticationDxeSmm/UserAuthenticationSmm.inf
+++ b/Features/Intel/UserInterface/UserAuthFeaturePkg/UserAuthenticationDxeSmm/UserAuthenticationSmm.inf
@@ -48,6 +48,8 @@
[Protocols]
gEdkiiVariableLockProtocolGuid                ## CONSUMES
gEfiSmmVariableProtocolGuid                   ## CONSUMES
+  gEdkiiSmmExitBootServicesProtocolGuid         ## CONSUMES
+  gEdkiiSmmLegacyBootProtocolGuid               ## CONSUMES

[Depex]
gEfiSmmVariableProtocolGuid AND gEfiVariableWriteArchProtocolGuid
--
2.26.2.windows.1