]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Library/SmmLockBoxLib/SmmLockBoxSmmLib.c
MdeModulePkg: Replace BSD License with BSD+Patent License
[mirror_edk2.git] / MdeModulePkg / Library / SmmLockBoxLib / SmmLockBoxSmmLib.c
index a332634aa7d66a411a2204e4a0bc4b10aa2e5e55..4cc0e8b78e5bf27ce82aeb90b6abb403148d89e3 100644 (file)
@@ -1,15 +1,8 @@
 /** @file\r
 \r
-Copyright (c) 2010 ~ 2012, Intel Corporation. All rights reserved.<BR>\r
+Copyright (c) 2010 - 2019, Intel Corporation. All rights reserved.<BR>\r
 \r
-This program and the accompanying materials\r
-are licensed and made available under the terms and conditions\r
-of the BSD License which accompanies this distribution.  The\r
-full text of the license may be found at\r
-http://opensource.org/licenses/bsd-license.php\r
-\r
-THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
-WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
@@ -20,6 +13,10 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/LockBoxLib.h>\r
 #include <Library/DebugLib.h>\r
 #include <Guid/SmmLockBox.h>\r
+#include <Guid/EndOfS3Resume.h>\r
+#include <Protocol/SmmReadyToLock.h>\r
+#include <Protocol/SmmEndOfDxe.h>\r
+#include <Protocol/SmmSxDispatch2.h>\r
 \r
 #include "SmmLockBoxLibPrivate.h"\r
 \r
@@ -30,6 +27,13 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 SMM_LOCK_BOX_CONTEXT mSmmLockBoxContext;\r
 LIST_ENTRY           mLockBoxQueue = INITIALIZE_LIST_HEAD_VARIABLE (mLockBoxQueue);\r
 \r
+BOOLEAN              mSmmConfigurationTableInstalled = FALSE;\r
+VOID                 *mSmmLockBoxRegistrationSmmEndOfDxe = NULL;\r
+VOID                 *mSmmLockBoxRegistrationSmmReadyToLock = NULL;\r
+VOID                 *mSmmLockBoxRegistrationEndOfS3Resume = NULL;\r
+BOOLEAN              mSmmLockBoxSmmReadyToLock = FALSE;\r
+BOOLEAN              mSmmLockBoxDuringS3Resume = FALSE;\r
+\r
 /**\r
   This function return SmmLockBox context from SMST.\r
 \r
@@ -61,6 +65,128 @@ InternalGetSmmLockBoxContext (
   return NULL;\r
 }\r
 \r
+/**\r
+  Notification for SMM ReadyToLock protocol.\r
+\r
+  @param[in] Protocol   Points to the protocol's unique identifier.\r
+  @param[in] Interface  Points to the interface instance.\r
+  @param[in] Handle     The handle on which the interface was installed.\r
+\r
+  @retval EFI_SUCCESS   Notification runs successfully.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmmLockBoxSmmReadyToLockNotify (\r
+  IN CONST EFI_GUID  *Protocol,\r
+  IN VOID            *Interface,\r
+  IN EFI_HANDLE      Handle\r
+  )\r
+{\r
+  mSmmLockBoxSmmReadyToLock = TRUE;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Main entry point for an SMM handler dispatch or communicate-based callback.\r
+\r
+  @param[in]     DispatchHandle  The unique handle assigned to this handler by SmiHandlerRegister().\r
+  @param[in]     Context         Points to an optional handler context which was specified when the\r
+                                 handler was registered.\r
+  @param[in,out] CommBuffer      A pointer to a collection of data in memory that will\r
+                                 be conveyed from a non-SMM environment into an SMM environment.\r
+  @param[in,out] CommBufferSize  The size of the CommBuffer.\r
+\r
+  @retval EFI_SUCCESS                         The interrupt was handled and quiesced. No other handlers\r
+                                              should still be called.\r
+  @retval EFI_WARN_INTERRUPT_SOURCE_QUIESCED  The interrupt has been quiesced but other handlers should\r
+                                              still be called.\r
+  @retval EFI_WARN_INTERRUPT_SOURCE_PENDING   The interrupt is still pending and other handlers should still\r
+                                              be called.\r
+  @retval EFI_INTERRUPT_PENDING               The interrupt could not be quiesced.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmmLockBoxS3EntryCallBack (\r
+  IN           EFI_HANDLE           DispatchHandle,\r
+  IN     CONST VOID                 *Context         OPTIONAL,\r
+  IN OUT       VOID                 *CommBuffer      OPTIONAL,\r
+  IN OUT       UINTN                *CommBufferSize  OPTIONAL\r
+  )\r
+{\r
+  mSmmLockBoxDuringS3Resume = TRUE;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Notification for SMM EndOfDxe protocol.\r
+\r
+  @param[in] Protocol   Points to the protocol's unique identifier.\r
+  @param[in] Interface  Points to the interface instance.\r
+  @param[in] Handle     The handle on which the interface was installed.\r
+\r
+  @retval EFI_SUCCESS   Notification runs successfully.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmmLockBoxSmmEndOfDxeNotify (\r
+  IN CONST EFI_GUID  *Protocol,\r
+  IN VOID            *Interface,\r
+  IN EFI_HANDLE      Handle\r
+  )\r
+{\r
+  EFI_STATUS                        Status;\r
+  EFI_SMM_SX_DISPATCH2_PROTOCOL     *SxDispatch;\r
+  EFI_SMM_SX_REGISTER_CONTEXT       EntryRegisterContext;\r
+  EFI_HANDLE                        S3EntryHandle;\r
+\r
+  //\r
+  // Locate SmmSxDispatch2 protocol.\r
+  //\r
+  Status = gSmst->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
+                           SmmLockBoxS3EntryCallBack,\r
+                           &EntryRegisterContext,\r
+                           &S3EntryHandle\r
+                           );\r
+    ASSERT_EFI_ERROR (Status);\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Notification for SMM EndOfS3Resume protocol.\r
+\r
+  @param[in] Protocol   Points to the protocol's unique identifier.\r
+  @param[in] Interface  Points to the interface instance.\r
+  @param[in] Handle     The handle on which the interface was installed.\r
+\r
+  @retval EFI_SUCCESS   Notification runs successfully.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmmLockBoxEndOfS3ResumeNotify (\r
+  IN CONST EFI_GUID  *Protocol,\r
+  IN VOID            *Interface,\r
+  IN EFI_HANDLE      Handle\r
+  )\r
+{\r
+  mSmmLockBoxDuringS3Resume = FALSE;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
 /**\r
   Constructor for SmmLockBox library.\r
   This is used to set SmmLockBox context, which will be used in PEI phase in S3 boot path later.\r
@@ -68,12 +194,12 @@ InternalGetSmmLockBoxContext (
   @param[in] ImageHandle  Image handle of this driver.\r
   @param[in] SystemTable  A Pointer to the EFI System Table.\r
 \r
-  @retval EFI_SUCEESS     \r
+  @retval EFI_SUCEESS\r
   @return Others          Some error occurs.\r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
-SmmLockBoxSmmConstructuor (\r
+SmmLockBoxSmmConstructor (\r
   IN EFI_HANDLE        ImageHandle,\r
   IN EFI_SYSTEM_TABLE  *SystemTable\r
   )\r
@@ -81,7 +207,37 @@ SmmLockBoxSmmConstructuor (
   EFI_STATUS           Status;\r
   SMM_LOCK_BOX_CONTEXT *SmmLockBoxContext;\r
 \r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SmmLockBoxSmmConstructuor - Enter\n"));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SmmLockBoxSmmConstructor - Enter\n"));\r
+\r
+  //\r
+  // Register SmmReadyToLock notification.\r
+  //\r
+  Status = gSmst->SmmRegisterProtocolNotify (\r
+                    &gEfiSmmReadyToLockProtocolGuid,\r
+                    SmmLockBoxSmmReadyToLockNotify,\r
+                    &mSmmLockBoxRegistrationSmmReadyToLock\r
+                    );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  //\r
+  // Register SmmEndOfDxe notification.\r
+  //\r
+  Status = gSmst->SmmRegisterProtocolNotify (\r
+                    &gEfiSmmEndOfDxeProtocolGuid,\r
+                    SmmLockBoxSmmEndOfDxeNotify,\r
+                    &mSmmLockBoxRegistrationSmmEndOfDxe\r
+                    );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  //\r
+  // Register EndOfS3Resume notification.\r
+  //\r
+  Status = gSmst->SmmRegisterProtocolNotify (\r
+                    &gEdkiiEndOfS3ResumeGuid,\r
+                    SmmLockBoxEndOfS3ResumeNotify,\r
+                    &mSmmLockBoxRegistrationEndOfS3Resume\r
+                    );\r
+  ASSERT_EFI_ERROR (Status);\r
 \r
   //\r
   // Check if gEfiSmmLockBoxCommunicationGuid is installed by someone\r
@@ -92,8 +248,8 @@ SmmLockBoxSmmConstructuor (
     // Find it. That means some other library instance is already run.\r
     // No need to install again, just return.\r
     //\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SmmLockBoxContext - already installed\n"));\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SmmLockBoxSmmConstructuor - Exit\n"));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SmmLockBoxContext - already installed\n"));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SmmLockBoxSmmConstructor - Exit\n"));\r
     return EFI_SUCCESS;\r
   }\r
 \r
@@ -114,14 +270,85 @@ SmmLockBoxSmmConstructuor (
                     sizeof(mSmmLockBoxContext)\r
                     );\r
   ASSERT_EFI_ERROR (Status);\r
+  mSmmConfigurationTableInstalled = TRUE;\r
 \r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SmmLockBoxContext - %x\n", (UINTN)&mSmmLockBoxContext));\r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib LockBoxDataAddress - %x\n", (UINTN)&mLockBoxQueue));\r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SmmLockBoxSmmConstructuor - Exit\n"));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SmmLockBoxContext - %x\n", (UINTN)&mSmmLockBoxContext));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib LockBoxDataAddress - %x\n", (UINTN)&mLockBoxQueue));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SmmLockBoxSmmConstructor - Exit\n"));\r
 \r
   return Status;\r
 }\r
 \r
+/**\r
+  Destructor for SmmLockBox library.\r
+  This is used to uninstall SmmLockBoxCommunication configuration table\r
+  if it has been installed in Constructor.\r
+\r
+  @param[in] ImageHandle    Image handle of this driver.\r
+  @param[in] SystemTable    A Pointer to the EFI System Table.\r
+\r
+  @retval EFI_SUCEESS       The destructor always returns EFI_SUCCESS.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmmLockBoxSmmDestructor (\r
+  IN EFI_HANDLE         ImageHandle,\r
+  IN EFI_SYSTEM_TABLE   *SystemTable\r
+  )\r
+{\r
+  EFI_STATUS            Status;\r
+\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SmmLockBoxSmmDestructor in %a module\n", gEfiCallerBaseName));\r
+\r
+  if (mSmmConfigurationTableInstalled) {\r
+    Status = gSmst->SmmInstallConfigurationTable (\r
+                      gSmst,\r
+                      &gEfiSmmLockBoxCommunicationGuid,\r
+                      NULL,\r
+                      0\r
+                      );\r
+    ASSERT_EFI_ERROR (Status);\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib uninstall SmmLockBoxCommunication configuration table\n"));\r
+  }\r
+\r
+  if (mSmmLockBoxRegistrationSmmReadyToLock != NULL) {\r
+    //\r
+    // Unregister SmmReadyToLock notification.\r
+    //\r
+    Status = gSmst->SmmRegisterProtocolNotify (\r
+                      &gEfiSmmReadyToLockProtocolGuid,\r
+                      NULL,\r
+                      &mSmmLockBoxRegistrationSmmReadyToLock\r
+                      );\r
+    ASSERT_EFI_ERROR (Status);\r
+  }\r
+  if (mSmmLockBoxRegistrationSmmEndOfDxe != NULL) {\r
+    //\r
+    // Unregister SmmEndOfDxe notification.\r
+    //\r
+    Status = gSmst->SmmRegisterProtocolNotify (\r
+                      &gEfiSmmEndOfDxeProtocolGuid,\r
+                      NULL,\r
+                      &mSmmLockBoxRegistrationSmmEndOfDxe\r
+                      );\r
+    ASSERT_EFI_ERROR (Status);\r
+  }\r
+  if (mSmmLockBoxRegistrationEndOfS3Resume != NULL) {\r
+    //\r
+    // Unregister EndOfS3Resume notification.\r
+    //\r
+    Status = gSmst->SmmRegisterProtocolNotify (\r
+                      &gEdkiiEndOfS3ResumeGuid,\r
+                      NULL,\r
+                      &mSmmLockBoxRegistrationEndOfS3Resume\r
+                      );\r
+    ASSERT_EFI_ERROR (Status);\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
 /**\r
   This function return SmmLockBox queue address.\r
 \r
@@ -204,13 +431,13 @@ SaveLockBox (
   EFI_STATUS                  Status;\r
   LIST_ENTRY                  *LockBoxQueue;\r
 \r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SaveLockBox - Enter\n"));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SaveLockBox - Enter\n"));\r
 \r
   //\r
   // Basic check\r
   //\r
   if ((Guid == NULL) || (Buffer == NULL) || (Length == 0)) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
@@ -219,7 +446,7 @@ SaveLockBox (
   //\r
   LockBox = InternalFindLockBoxByGuid (Guid);\r
   if (LockBox != NULL) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_ALREADY_STARTED));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_ALREADY_STARTED));\r
     return EFI_ALREADY_STARTED;\r
   }\r
 \r
@@ -234,7 +461,7 @@ SaveLockBox (
                     );\r
   ASSERT_EFI_ERROR (Status);\r
   if (EFI_ERROR (Status)) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_OUT_OF_RESOURCES));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_OUT_OF_RESOURCES));\r
     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
@@ -249,7 +476,7 @@ SaveLockBox (
   ASSERT_EFI_ERROR (Status);\r
   if (EFI_ERROR (Status)) {\r
     gSmst->SmmFreePages (SmramBuffer, EFI_SIZE_TO_PAGES (Length));\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_OUT_OF_RESOURCES));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_OUT_OF_RESOURCES));\r
     return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
@@ -267,7 +494,15 @@ SaveLockBox (
   LockBox->Length      = (UINT64)Length;\r
   LockBox->Attributes  = 0;\r
   LockBox->SmramBuffer = SmramBuffer;\r
-  \r
+\r
+  DEBUG ((\r
+    DEBUG_INFO,\r
+    "LockBoxGuid - %g, SmramBuffer - 0x%lx, Length - 0x%lx\n",\r
+    &LockBox->Guid,\r
+    LockBox->SmramBuffer,\r
+    LockBox->Length\r
+    ));\r
+\r
   LockBoxQueue = InternalGetLockBoxQueue ();\r
   ASSERT (LockBoxQueue != NULL);\r
   InsertTailList (LockBoxQueue, &LockBox->Link);\r
@@ -275,7 +510,7 @@ SaveLockBox (
   //\r
   // Done\r
   //\r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_SUCCESS));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SaveLockBox - Exit (%r)\n", EFI_SUCCESS));\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -301,14 +536,22 @@ SetLockBoxAttributes (
 {\r
   SMM_LOCK_BOX_DATA              *LockBox;\r
 \r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Enter\n"));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Enter\n"));\r
 \r
   //\r
   // Basic check\r
   //\r
   if ((Guid == NULL) ||\r
-      ((Attributes & ~LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE) != 0)) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
+      ((Attributes & ~(LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE | LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY)) != 0)) {\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  if (((Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE) != 0) &&\r
+      ((Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY) != 0)) {\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
+    DEBUG ((DEBUG_INFO, "  LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE and LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY\n\n"));\r
+    DEBUG ((DEBUG_INFO, "  can not be set together\n"));\r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
@@ -317,10 +560,20 @@ SetLockBoxAttributes (
   //\r
   LockBox = InternalFindLockBoxByGuid (Guid);\r
   if (LockBox == NULL) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Exit (%r)\n", EFI_NOT_FOUND));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Exit (%r)\n", EFI_NOT_FOUND));\r
     return EFI_NOT_FOUND;\r
   }\r
 \r
+  if ((((Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE) != 0) &&\r
+      ((LockBox->Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY) != 0)) ||\r
+      (((LockBox->Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE) != 0) &&\r
+      ((Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY) != 0))) {\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes 0x%lx 0x%lx - Exit (%r)\n", LockBox->Attributes, Attributes, EFI_INVALID_PARAMETER));\r
+    DEBUG ((DEBUG_INFO, "  LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE and LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY\n\n"));\r
+    DEBUG ((DEBUG_INFO, "  can not be set together\n"));\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
   //\r
   // Update data\r
   //\r
@@ -329,7 +582,7 @@ SetLockBoxAttributes (
   //\r
   // Done\r
   //\r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Exit (%r)\n", EFI_SUCCESS));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib SetLockBoxAttributes - Exit (%r)\n", EFI_SUCCESS));\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -344,7 +597,10 @@ SetLockBoxAttributes (
   @retval RETURN_SUCCESS            the information is saved successfully.\r
   @retval RETURN_INVALID_PARAMETER  the Guid is NULL, or Buffer is NULL, or Length is 0.\r
   @retval RETURN_NOT_FOUND          the requested GUID not found.\r
-  @retval RETURN_BUFFER_TOO_SMALL   the original buffer to too small to hold new information.\r
+  @retval RETURN_BUFFER_TOO_SMALL   for lockbox without attribute LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY,\r
+                                    the original buffer to too small to hold new information.\r
+  @retval RETURN_OUT_OF_RESOURCES   for lockbox with attribute LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY,\r
+                                    no enough resource to save the information.\r
   @retval RETURN_ACCESS_DENIED      it is too late to invoke this interface\r
   @retval RETURN_NOT_STARTED        it is too early to invoke this interface\r
   @retval RETURN_UNSUPPORTED        the service is not supported by implementaion.\r
@@ -359,14 +615,17 @@ UpdateLockBox (
   )\r
 {\r
   SMM_LOCK_BOX_DATA             *LockBox;\r
+  EFI_PHYSICAL_ADDRESS          SmramBuffer;\r
+  EFI_STATUS                    Status;\r
 \r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib UpdateLockBox - Enter\n"));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib UpdateLockBox - Enter\n"));\r
 \r
   //\r
   // Basic check\r
   //\r
-  if ((Guid == NULL) || (Buffer == NULL) || (Length == 0)) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
+  if ((Guid == NULL) || (Buffer == NULL) || (Length == 0) ||\r
+      (Length > MAX_UINTN - Offset)) {\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
@@ -375,7 +634,7 @@ UpdateLockBox (
   //\r
   LockBox = InternalFindLockBoxByGuid (Guid);\r
   if (LockBox == NULL) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_NOT_FOUND));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_NOT_FOUND));\r
     return EFI_NOT_FOUND;\r
   }\r
 \r
@@ -383,15 +642,74 @@ UpdateLockBox (
   // Update data\r
   //\r
   if (LockBox->Length < Offset + Length) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_BUFFER_TOO_SMALL));\r
-    return EFI_BUFFER_TOO_SMALL;\r
+    if ((LockBox->Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY) != 0) {\r
+      //\r
+      // If 'LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY' attribute is set, enlarge the\r
+      // LockBox.\r
+      //\r
+      DEBUG ((\r
+        DEBUG_INFO,\r
+        "SmmLockBoxSmmLib UpdateLockBox - Origin LockBox too small, enlarge.\n"\r
+        ));\r
+\r
+      if (EFI_PAGES_TO_SIZE (EFI_SIZE_TO_PAGES ((UINTN)LockBox->Length)) < Offset + Length) {\r
+        //\r
+        // In SaveLockBox(), the SMRAM buffer allocated for LockBox is of page\r
+        // granularity. Here, if the required size is larger than the origin size\r
+        // of the pages, allocate new buffer from SMRAM to enlarge the LockBox.\r
+        //\r
+        DEBUG ((\r
+          DEBUG_INFO,\r
+          "SmmLockBoxSmmLib UpdateLockBox - Allocate new buffer to enlarge.\n"\r
+          ));\r
+        Status = gSmst->SmmAllocatePages (\r
+                          AllocateAnyPages,\r
+                          EfiRuntimeServicesData,\r
+                          EFI_SIZE_TO_PAGES (Offset + Length),\r
+                          &SmramBuffer\r
+                          );\r
+        if (EFI_ERROR (Status)) {\r
+          DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_OUT_OF_RESOURCES));\r
+          return EFI_OUT_OF_RESOURCES;\r
+        }\r
+\r
+        //\r
+        // Copy origin data to the new SMRAM buffer and wipe the content in the\r
+        // origin SMRAM buffer.\r
+        //\r
+        CopyMem ((VOID *)(UINTN)SmramBuffer, (VOID *)(UINTN)LockBox->SmramBuffer, (UINTN)LockBox->Length);\r
+        ZeroMem ((VOID *)(UINTN)LockBox->SmramBuffer, (UINTN)LockBox->Length);\r
+        gSmst->SmmFreePages (LockBox->SmramBuffer, EFI_SIZE_TO_PAGES ((UINTN)LockBox->Length));\r
+\r
+        LockBox->SmramBuffer = SmramBuffer;\r
+      }\r
+\r
+      //\r
+      // Handle uninitialized content in the LockBox.\r
+      //\r
+      if (Offset > LockBox->Length) {\r
+        ZeroMem (\r
+          (VOID *)((UINTN)LockBox->SmramBuffer + (UINTN)LockBox->Length),\r
+          Offset - (UINTN)LockBox->Length\r
+          );\r
+      }\r
+      LockBox->Length = Offset + Length;\r
+    } else {\r
+      //\r
+      // If 'LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY' attribute is NOT set, return\r
+      // EFI_BUFFER_TOO_SMALL directly.\r
+      //\r
+      DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_BUFFER_TOO_SMALL));\r
+      return EFI_BUFFER_TOO_SMALL;\r
+    }\r
   }\r
+  ASSERT ((UINTN)LockBox->SmramBuffer <= (MAX_ADDRESS - Offset));\r
   CopyMem ((VOID *)((UINTN)LockBox->SmramBuffer + Offset), Buffer, Length);\r
 \r
   //\r
   // Done\r
   //\r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_SUCCESS));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib UpdateLockBox - Exit (%r)\n", EFI_SUCCESS));\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -405,7 +723,7 @@ UpdateLockBox (
 \r
   @retval RETURN_SUCCESS            the information is restored successfully.\r
   @retval RETURN_INVALID_PARAMETER  the Guid is NULL, or one of Buffer and Length is NULL.\r
-  @retval RETURN_WRITE_PROTECTED    Buffer and Length are NULL, but the LockBox has no \r
+  @retval RETURN_WRITE_PROTECTED    Buffer and Length are NULL, but the LockBox has no\r
                                     LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE attribute.\r
   @retval RETURN_BUFFER_TOO_SMALL   the Length is too small to hold the confidential information.\r
   @retval RETURN_NOT_FOUND          the requested GUID not found.\r
@@ -424,7 +742,7 @@ RestoreLockBox (
   SMM_LOCK_BOX_DATA              *LockBox;\r
   VOID                           *RestoreBuffer;\r
 \r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib RestoreLockBox - Enter\n"));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib RestoreLockBox - Enter\n"));\r
 \r
   //\r
   // Restore this, Buffer and Length MUST be both NULL or both non-NULL\r
@@ -432,7 +750,7 @@ RestoreLockBox (
   if ((Guid == NULL) ||\r
       ((Buffer == NULL) && (Length != NULL)) ||\r
       ((Buffer != NULL) && (Length == NULL))) {\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_INVALID_PARAMETER));\r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
@@ -444,10 +762,20 @@ RestoreLockBox (
     //\r
     // Not found\r
     //\r
-    DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_NOT_FOUND));\r
+    DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_NOT_FOUND));\r
     return EFI_NOT_FOUND;\r
   }\r
 \r
+  if (((LockBox->Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY) != 0) &&\r
+      mSmmLockBoxSmmReadyToLock &&\r
+      !mSmmLockBoxDuringS3Resume) {\r
+    //\r
+    // With LOCK_BOX_ATTRIBUTE_RESTORE_IN_S3_ONLY,\r
+    // this LockBox can be restored in S3 resume only.\r
+    //\r
+    return EFI_ACCESS_DENIED;\r
+  }\r
+\r
   //\r
   // Set RestoreBuffer\r
   //\r
@@ -461,7 +789,7 @@ RestoreLockBox (
     // restore to original buffer\r
     //\r
     if ((LockBox->Attributes & LOCK_BOX_ATTRIBUTE_RESTORE_IN_PLACE) == 0) {\r
-      DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_WRITE_PROTECTED));\r
+      DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_WRITE_PROTECTED));\r
       return EFI_WRITE_PROTECTED;\r
     }\r
     RestoreBuffer = (VOID *)(UINTN)LockBox->Buffer;\r
@@ -476,7 +804,7 @@ RestoreLockBox (
       // Input buffer is too small to hold all data.\r
       //\r
       *Length = (UINTN)LockBox->Length;\r
-      DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_BUFFER_TOO_SMALL));\r
+      DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_BUFFER_TOO_SMALL));\r
       return EFI_BUFFER_TOO_SMALL;\r
     }\r
     *Length = (UINTN)LockBox->Length;\r
@@ -490,7 +818,7 @@ RestoreLockBox (
   //\r
   // Done\r
   //\r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_SUCCESS));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib RestoreLockBox - Exit (%r)\n", EFI_SUCCESS));\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -511,7 +839,7 @@ RestoreAllLockBoxInPlace (
   LIST_ENTRY                    *Link;\r
   LIST_ENTRY                    *LockBoxQueue;\r
 \r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib RestoreAllLockBoxInPlace - Enter\n"));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib RestoreAllLockBoxInPlace - Enter\n"));\r
 \r
   LockBoxQueue = InternalGetLockBoxQueue ();\r
   ASSERT (LockBoxQueue != NULL);\r
@@ -537,7 +865,7 @@ RestoreAllLockBoxInPlace (
   //\r
   // Done\r
   //\r
-  DEBUG ((EFI_D_INFO, "SmmLockBoxSmmLib RestoreAllLockBoxInPlace - Exit (%r)\n", EFI_SUCCESS));\r
+  DEBUG ((DEBUG_INFO, "SmmLockBoxSmmLib RestoreAllLockBoxInPlace - Exit (%r)\n", EFI_SUCCESS));\r
   return EFI_SUCCESS;\r
 }\r
 \r