]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Universal/Variable/RuntimeDxe/VariableSmm.c
Add EDKII_VARIABLE_LOCK_PROTOCOL and the implementation in MdeModulePkg variable...
[mirror_edk2.git] / MdeModulePkg / Universal / Variable / RuntimeDxe / VariableSmm.c
index 75f5cad469abdd0d4cc8941252ff78eda72539cd..4009fcb1712eaded50b18d740c91a814e19c46f6 100644 (file)
   implements  an SMI handler to communicate with the DXE runtime driver \r
   to provide variable services.\r
 \r
-Copyright (c) 2010, Intel Corporation. All rights reserved.<BR>\r
-This program and the accompanying materials                          \r
-are licensed and made available under the terms and conditions of the BSD License         \r
-which accompanies this distribution.  The full text of the license may be found at        \r
-http://opensource.org/licenses/bsd-license.php                                            \r
+  Caution: This module requires additional review when modified.\r
+  This driver will have external input - variable data and communicate buffer in SMM mode.\r
+  This external input must be validated carefully to avoid security issue like\r
+  buffer overflow, integer overflow.\r
+\r
+  SmmVariableHandler() will receive untrusted input and do basic validation.\r
+\r
+  Each sub function VariableServiceGetVariable(), VariableServiceGetNextVariableName(), \r
+  VariableServiceSetVariable(), VariableServiceQueryVariableInfo(), ReclaimForOS(), \r
+  SmmVariableGetStatistics() should also do validation based on its own knowledge.\r
+\r
+Copyright (c) 2010 - 2013, Intel Corporation. All rights reserved.<BR>\r
+This program and the accompanying materials \r
+are licensed and made available under the terms and conditions of the BSD License \r
+which accompanies this distribution.  The 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
 \r
 **/\r
+#include <Protocol/SmmVariable.h>\r
+#include <Protocol/SmmFirmwareVolumeBlock.h>\r
 #include <Protocol/SmmFaultTolerantWrite.h>\r
+#include <Protocol/SmmAccess2.h>\r
+#include <Protocol/SmmEndOfDxe.h>\r
+\r
 #include <Library/SmmServicesTableLib.h>\r
 \r
+#include <Guid/VariableFormat.h>\r
+#include <Guid/SmmVariableCommon.h>\r
 #include "Variable.h"\r
-#include "VariableSmmCommon.h"\r
 \r
-extern SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY  *gVariableInfo;\r
+EFI_SMRAM_DESCRIPTOR                                 *mSmramRanges;\r
+UINTN                                                mSmramRangeCount;\r
+\r
+extern VARIABLE_INFO_ENTRY                           *gVariableInfo;\r
 EFI_HANDLE                                           mSmmVariableHandle      = NULL;\r
 EFI_HANDLE                                           mVariableHandle         = NULL;\r
 BOOLEAN                                              mAtRuntime              = FALSE;\r
 EFI_GUID                                             mZeroGuid               = {0, 0, 0, {0, 0, 0, 0, 0, 0, 0, 0}};\r
-EFI_GUID                                             mSmmVariableWriteGuid   = EFI_SMM_VARIABLE_WRITE_GUID;\r
-  \r
+UINT8                                                *mVariableBufferPayload = NULL;\r
+UINTN                                                mVariableBufferPayloadSize;\r
+extern BOOLEAN                                       mEndOfDxe;\r
+extern BOOLEAN                                       mEnableLocking;\r
+\r
+/**\r
+\r
+  This code sets variable in storage blocks (Volatile or Non-Volatile).\r
+\r
+  @param VariableName                     Name of Variable to be found.\r
+  @param VendorGuid                       Variable vendor GUID.\r
+  @param Attributes                       Attribute value of the variable found\r
+  @param DataSize                         Size of Data found. If size is less than the\r
+                                          data, this value contains the required size.\r
+  @param Data                             Data pointer.\r
+\r
+  @return EFI_INVALID_PARAMETER           Invalid parameter.\r
+  @return EFI_SUCCESS                     Set successfully.\r
+  @return EFI_OUT_OF_RESOURCES            Resource not enough to set variable.\r
+  @return EFI_NOT_FOUND                   Not found.\r
+  @return EFI_WRITE_PROTECTED             Variable is read-only.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmmVariableSetVariable (\r
+  IN CHAR16                  *VariableName,\r
+  IN EFI_GUID                *VendorGuid,\r
+  IN UINT32                  Attributes,\r
+  IN UINTN                   DataSize,\r
+  IN VOID                    *Data\r
+  )\r
+{\r
+  EFI_STATUS                 Status;\r
+\r
+  //\r
+  // Disable write protection when the calling SetVariable() through EFI_SMM_VARIABLE_PROTOCOL.\r
+  //\r
+  mEnableLocking = FALSE;\r
+  Status         = VariableServiceSetVariable (\r
+                     VariableName,\r
+                     VendorGuid,\r
+                     Attributes,\r
+                     DataSize,\r
+                     Data\r
+                     );\r
+  mEnableLocking = TRUE;\r
+  return Status;\r
+}\r
+\r
 EFI_SMM_VARIABLE_PROTOCOL      gSmmVariable = {\r
   VariableServiceGetVariable,\r
   VariableServiceGetNextVariableName,\r
-  VariableServiceSetVariable,\r
+  SmmVariableSetVariable,\r
   VariableServiceQueryVariableInfo\r
 };\r
 \r
-\r
 /**\r
   Return TRUE if ExitBootServices () has been called.\r
   \r
@@ -48,6 +115,60 @@ AtRuntime (
   return mAtRuntime;\r
 }\r
 \r
+/**\r
+  This function check if the address is in SMRAM.\r
+\r
+  @param Buffer  the buffer address to be checked.\r
+  @param Length  the buffer length to be checked.\r
+\r
+  @retval TRUE  this address is in SMRAM.\r
+  @retval FALSE this address is NOT in SMRAM.\r
+**/\r
+BOOLEAN\r
+InternalIsAddressInSmram (\r
+  IN EFI_PHYSICAL_ADDRESS  Buffer,\r
+  IN UINT64                Length\r
+  )\r
+{\r
+  UINTN  Index;\r
+\r
+  for (Index = 0; Index < mSmramRangeCount; Index ++) {\r
+    if (((Buffer >= mSmramRanges[Index].CpuStart) && (Buffer < mSmramRanges[Index].CpuStart + mSmramRanges[Index].PhysicalSize)) ||\r
+        ((mSmramRanges[Index].CpuStart >= Buffer) && (mSmramRanges[Index].CpuStart < Buffer + Length))) {\r
+      return TRUE;\r
+    }\r
+  }\r
+\r
+  return FALSE;\r
+}\r
+\r
+/**\r
+  This function check if the address refered by Buffer and Length is valid.\r
+\r
+  @param Buffer  the buffer address to be checked.\r
+  @param Length  the buffer length to be checked.\r
+\r
+  @retval TRUE  this address is valid.\r
+  @retval FALSE this address is NOT valid.\r
+**/\r
+BOOLEAN\r
+InternalIsAddressValid (\r
+  IN UINTN                 Buffer,\r
+  IN UINTN                 Length\r
+  )\r
+{\r
+  if (Buffer > (MAX_ADDRESS - Length)) {\r
+    //\r
+    // Overflow happen\r
+    //\r
+    return FALSE;\r
+  }\r
+  if (InternalIsAddressInSmram ((EFI_PHYSICAL_ADDRESS)Buffer, (UINT64)Length)) {\r
+    return FALSE;\r
+  }\r
+  return TRUE;\r
+}\r
+\r
 /**\r
   Initializes a basic mutual exclusion lock.\r
 \r
@@ -230,6 +351,8 @@ GetFvbCountAndBuffer (
   *NumberHandles = BufferSize / sizeof(EFI_HANDLE);\r
   if (EFI_ERROR(Status)) {\r
     *NumberHandles = 0;\r
+    FreePool (*Buffer);\r
+    *Buffer = NULL;\r
   }\r
 \r
   return Status;\r
@@ -239,12 +362,15 @@ GetFvbCountAndBuffer (
 /**\r
   Get the variable statistics information from the information buffer pointed by gVariableInfo.\r
 \r
-  @param[in, out]  InfoEntry   A pointer to the buffer of variable information entry.\r
-                               On input, point to the variable information returned last time. if \r
-                               InfoEntry->VendorGuid is zero, return the first information.\r
-                               On output, point to the next variable information.\r
-  @param[in, out]  InfoSize    On input, the size of the variable information buffer.\r
-                               On output, the returned variable information size.\r
+  Caution: This function may be invoked at SMM runtime.\r
+  InfoEntry and InfoSize are external input. Care must be taken to make sure not security issue at runtime.\r
+\r
+  @param[in, out]  InfoEntry    A pointer to the buffer of variable information entry.\r
+                                On input, point to the variable information returned last time. if \r
+                                InfoEntry->VendorGuid is zero, return the first information.\r
+                                On output, point to the next variable information.\r
+  @param[in, out]  InfoSize     On input, the size of the variable information buffer.\r
+                                On output, the returned variable information size.\r
 \r
   @retval EFI_SUCCESS          The variable information is found and returned successfully.\r
   @retval EFI_UNSUPPORTED      No variable inoformation exists in variable driver. The \r
@@ -254,33 +380,36 @@ GetFvbCountAndBuffer (
 **/\r
 EFI_STATUS\r
 SmmVariableGetStatistics (\r
-  IN OUT SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY  *InfoEntry,\r
+  IN OUT VARIABLE_INFO_ENTRY                           *InfoEntry,\r
   IN OUT UINTN                                         *InfoSize\r
   )\r
 {\r
-  SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY         *VariableInfo;\r
+  VARIABLE_INFO_ENTRY                                  *VariableInfo;\r
   UINTN                                                NameLength;\r
   UINTN                                                StatisticsInfoSize;\r
   CHAR16                                               *InfoName;\r
\r
+  EFI_GUID                                             VendorGuid;\r
+\r
   ASSERT (InfoEntry != NULL);\r
   VariableInfo = gVariableInfo; \r
   if (VariableInfo == NULL) {\r
     return EFI_UNSUPPORTED;\r
   }\r
 \r
-  StatisticsInfoSize = sizeof (SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY) + StrSize (VariableInfo->Name);\r
-  if (*InfoSize < sizeof (SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY)) {\r
+  StatisticsInfoSize = sizeof (VARIABLE_INFO_ENTRY) + StrSize (VariableInfo->Name);\r
+  if (*InfoSize < StatisticsInfoSize) {\r
     *InfoSize = StatisticsInfoSize;\r
     return EFI_BUFFER_TOO_SMALL;\r
   }\r
   InfoName = (CHAR16 *)(InfoEntry + 1);\r
 \r
-  if (CompareGuid (&InfoEntry->VendorGuid, &mZeroGuid)) {\r
+  CopyGuid (&VendorGuid, &InfoEntry->VendorGuid);\r
+\r
+  if (CompareGuid (&VendorGuid, &mZeroGuid)) {\r
     //\r
     // Return the first variable info\r
     //\r
-    CopyMem (InfoEntry, VariableInfo, sizeof (SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY));\r
+    CopyMem (InfoEntry, VariableInfo, sizeof (VARIABLE_INFO_ENTRY));\r
     CopyMem (InfoName, VariableInfo->Name, StrSize (VariableInfo->Name));\r
     *InfoSize = StatisticsInfoSize;\r
     return EFI_SUCCESS;\r
@@ -290,7 +419,7 @@ SmmVariableGetStatistics (
   // Get the next variable info\r
   //\r
   while (VariableInfo != NULL) {\r
-    if (CompareGuid (&VariableInfo->VendorGuid, &InfoEntry->VendorGuid)) {\r
+    if (CompareGuid (&VariableInfo->VendorGuid, &VendorGuid)) {\r
       NameLength = StrSize (VariableInfo->Name);\r
       if (NameLength == StrSize (InfoName)) {\r
         if (CompareMem (VariableInfo->Name, InfoName, NameLength) == 0) {\r
@@ -313,13 +442,13 @@ SmmVariableGetStatistics (
   //\r
   // Output the new variable info\r
   //\r
-  StatisticsInfoSize = sizeof (SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY) + StrSize (VariableInfo->Name);\r
+  StatisticsInfoSize = sizeof (VARIABLE_INFO_ENTRY) + StrSize (VariableInfo->Name);\r
   if (*InfoSize < StatisticsInfoSize) {\r
     *InfoSize = StatisticsInfoSize;\r
     return EFI_BUFFER_TOO_SMALL;\r
   }\r
 \r
-  CopyMem (InfoEntry, VariableInfo, sizeof (SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY));\r
+  CopyMem (InfoEntry, VariableInfo, sizeof (VARIABLE_INFO_ENTRY));\r
   CopyMem (InfoName, VariableInfo->Name, StrSize (VariableInfo->Name));\r
   *InfoSize = StatisticsInfoSize;\r
   \r
@@ -332,6 +461,12 @@ SmmVariableGetStatistics (
 \r
   This SMI handler provides services for the variable wrapper driver.\r
 \r
+  Caution: This function may receive untrusted input.\r
+  This variable data and communicate buffer are external input, so this function will do basic validation.\r
+  Each sub function VariableServiceGetVariable(), VariableServiceGetNextVariableName(), \r
+  VariableServiceSetVariable(), VariableServiceQueryVariableInfo(), ReclaimForOS(), \r
+  SmmVariableGetStatistics() should also do validation based on its own knowledge.\r
+\r
   @param[in]     DispatchHandle  The unique handle assigned to this handler by SmiHandlerRegister().\r
   @param[in]     RegisterContext Points to an optional handler context which was specified when the\r
                                  handler was registered.\r
@@ -361,15 +496,74 @@ SmmVariableHandler (
   SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE         *SmmVariableHeader;\r
   SMM_VARIABLE_COMMUNICATE_GET_NEXT_VARIABLE_NAME  *GetNextVariableName;\r
   SMM_VARIABLE_COMMUNICATE_QUERY_VARIABLE_INFO     *QueryVariableInfo;\r
-  SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY     *VariableInfo;\r
+  VARIABLE_INFO_ENTRY                              *VariableInfo;\r
+  SMM_VARIABLE_COMMUNICATE_LOCK_VARIABLE           *VariableToLock;\r
   UINTN                                            InfoSize;\r
+  UINTN                                            NameBufferSize;\r
+  UINTN                                            CommBufferPayloadSize;\r
 \r
-  ASSERT (CommBuffer != NULL);\r
+  //\r
+  // If input is invalid, stop processing this SMI\r
+  //\r
+  if (CommBuffer == NULL || CommBufferSize == NULL) {\r
+    return EFI_SUCCESS;\r
+  }\r
+\r
+  if (*CommBufferSize < SMM_VARIABLE_COMMUNICATE_HEADER_SIZE) {\r
+    DEBUG ((EFI_D_ERROR, "SmmVariableHandler: SMM communication buffer size invalid!\n"));\r
+    return EFI_SUCCESS;\r
+  }\r
+  CommBufferPayloadSize = *CommBufferSize - SMM_VARIABLE_COMMUNICATE_HEADER_SIZE;\r
+  if (CommBufferPayloadSize > mVariableBufferPayloadSize) {\r
+    DEBUG ((EFI_D_ERROR, "SmmVariableHandler: SMM communication buffer payload size invalid!\n"));\r
+    return EFI_SUCCESS;\r
+  }\r
+\r
+  if (!InternalIsAddressValid ((UINTN)CommBuffer, *CommBufferSize)) {\r
+    DEBUG ((EFI_D_ERROR, "SmmVariableHandler: SMM communication buffer in SMRAM or overflow!\n"));\r
+    return EFI_SUCCESS;\r
+  }\r
 \r
   SmmVariableFunctionHeader = (SMM_VARIABLE_COMMUNICATE_HEADER *)CommBuffer;\r
   switch (SmmVariableFunctionHeader->Function) {\r
     case SMM_VARIABLE_FUNCTION_GET_VARIABLE:\r
-      SmmVariableHeader = (SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE *) SmmVariableFunctionHeader->Data;     \r
+      if (CommBufferPayloadSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name)) {\r
+        DEBUG ((EFI_D_ERROR, "GetVariable: SMM communication buffer size invalid!\n"));\r
+        return EFI_SUCCESS;\r
+      }\r
+      //\r
+      // Copy the input communicate buffer payload to pre-allocated SMM variable buffer payload.\r
+      //\r
+      CopyMem (mVariableBufferPayload, SmmVariableFunctionHeader->Data, CommBufferPayloadSize);\r
+      SmmVariableHeader = (SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE *) mVariableBufferPayload;\r
+      if (((UINTN)(~0) - SmmVariableHeader->DataSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name)) ||\r
+         ((UINTN)(~0) - SmmVariableHeader->NameSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name) + SmmVariableHeader->DataSize)) {\r
+        //\r
+        // Prevent InfoSize overflow happen\r
+        //\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+      InfoSize = OFFSET_OF(SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name) \r
+                 + SmmVariableHeader->DataSize + SmmVariableHeader->NameSize;\r
+\r
+      //\r
+      // SMRAM range check already covered before\r
+      //\r
+      if (InfoSize > CommBufferPayloadSize) {\r
+        DEBUG ((EFI_D_ERROR, "GetVariable: Data size exceed communication buffer size limit!\n"));\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+\r
+      if (SmmVariableHeader->NameSize < sizeof (CHAR16) || SmmVariableHeader->Name[SmmVariableHeader->NameSize/sizeof (CHAR16) - 1] != L'\0') {\r
+        //\r
+        // Make sure VariableName is A Null-terminated string.\r
+        //\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+\r
       Status = VariableServiceGetVariable (\r
                  SmmVariableHeader->Name,\r
                  &SmmVariableHeader->Guid,\r
@@ -377,19 +571,93 @@ SmmVariableHandler (
                  &SmmVariableHeader->DataSize,\r
                  (UINT8 *)SmmVariableHeader->Name + SmmVariableHeader->NameSize\r
                  );\r
+      CopyMem (SmmVariableFunctionHeader->Data, mVariableBufferPayload, CommBufferPayloadSize);\r
       break;\r
       \r
     case SMM_VARIABLE_FUNCTION_GET_NEXT_VARIABLE_NAME:\r
-      GetNextVariableName = (SMM_VARIABLE_COMMUNICATE_GET_NEXT_VARIABLE_NAME *) SmmVariableFunctionHeader->Data;\r
+      if (CommBufferPayloadSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_GET_NEXT_VARIABLE_NAME, Name)) {\r
+        DEBUG ((EFI_D_ERROR, "GetNextVariableName: SMM communication buffer size invalid!\n"));\r
+        return EFI_SUCCESS;\r
+      }\r
+      //\r
+      // Copy the input communicate buffer payload to pre-allocated SMM variable buffer payload.\r
+      //\r
+      CopyMem (mVariableBufferPayload, SmmVariableFunctionHeader->Data, CommBufferPayloadSize);\r
+      GetNextVariableName = (SMM_VARIABLE_COMMUNICATE_GET_NEXT_VARIABLE_NAME *) mVariableBufferPayload;\r
+      if ((UINTN)(~0) - GetNextVariableName->NameSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_GET_NEXT_VARIABLE_NAME, Name)) {\r
+        //\r
+        // Prevent InfoSize overflow happen\r
+        //\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+      InfoSize = OFFSET_OF(SMM_VARIABLE_COMMUNICATE_GET_NEXT_VARIABLE_NAME, Name) + GetNextVariableName->NameSize;\r
+\r
+      //\r
+      // SMRAM range check already covered before\r
+      //\r
+      if (InfoSize > CommBufferPayloadSize) {\r
+        DEBUG ((EFI_D_ERROR, "GetNextVariableName: Data size exceed communication buffer size limit!\n"));\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+\r
+      NameBufferSize = CommBufferPayloadSize - OFFSET_OF(SMM_VARIABLE_COMMUNICATE_GET_NEXT_VARIABLE_NAME, Name);\r
+      if (NameBufferSize < sizeof (CHAR16) || GetNextVariableName->Name[NameBufferSize/sizeof (CHAR16) - 1] != L'\0') {\r
+        //\r
+        // Make sure input VariableName is A Null-terminated string.\r
+        //\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+\r
       Status = VariableServiceGetNextVariableName (\r
                  &GetNextVariableName->NameSize,\r
                  GetNextVariableName->Name,\r
                  &GetNextVariableName->Guid\r
                  );\r
+      CopyMem (SmmVariableFunctionHeader->Data, mVariableBufferPayload, CommBufferPayloadSize);\r
       break;\r
       \r
     case SMM_VARIABLE_FUNCTION_SET_VARIABLE:\r
-      SmmVariableHeader = (SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE *) SmmVariableFunctionHeader->Data;\r
+      if (CommBufferPayloadSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name)) {\r
+        DEBUG ((EFI_D_ERROR, "SetVariable: SMM communication buffer size invalid!\n"));\r
+        return EFI_SUCCESS;\r
+      }\r
+      //\r
+      // Copy the input communicate buffer payload to pre-allocated SMM variable buffer payload.\r
+      //\r
+      CopyMem (mVariableBufferPayload, SmmVariableFunctionHeader->Data, CommBufferPayloadSize);\r
+      SmmVariableHeader = (SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE *) mVariableBufferPayload;\r
+      if (((UINTN)(~0) - SmmVariableHeader->DataSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name)) ||\r
+         ((UINTN)(~0) - SmmVariableHeader->NameSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name) + SmmVariableHeader->DataSize)) {\r
+        //\r
+        // Prevent InfoSize overflow happen\r
+        //\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+      InfoSize = OFFSET_OF(SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name)\r
+                 + SmmVariableHeader->DataSize + SmmVariableHeader->NameSize;\r
+\r
+      //\r
+      // SMRAM range check already covered before\r
+      // Data buffer should not contain SMM range\r
+      //\r
+      if (InfoSize > CommBufferPayloadSize) {\r
+        DEBUG ((EFI_D_ERROR, "SetVariable: Data size exceed communication buffer size limit!\n"));\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+\r
+      if (SmmVariableHeader->NameSize < sizeof (CHAR16) || SmmVariableHeader->Name[SmmVariableHeader->NameSize/sizeof (CHAR16) - 1] != L'\0') {\r
+        //\r
+        // Make sure VariableName is A Null-terminated string.\r
+        //\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+\r
       Status = VariableServiceSetVariable (\r
                  SmmVariableHeader->Name,\r
                  &SmmVariableHeader->Guid,\r
@@ -400,7 +668,12 @@ SmmVariableHandler (
       break;\r
       \r
     case SMM_VARIABLE_FUNCTION_QUERY_VARIABLE_INFO:\r
+      if (CommBufferPayloadSize < sizeof (SMM_VARIABLE_COMMUNICATE_QUERY_VARIABLE_INFO)) {\r
+        DEBUG ((EFI_D_ERROR, "QueryVariableInfo: SMM communication buffer size invalid!\n"));\r
+        return EFI_SUCCESS;\r
+      }\r
       QueryVariableInfo = (SMM_VARIABLE_COMMUNICATE_QUERY_VARIABLE_INFO *) SmmVariableFunctionHeader->Data;\r
+\r
       Status = VariableServiceQueryVariableInfo (\r
                  QueryVariableInfo->Attributes,\r
                  &QueryVariableInfo->MaximumVariableStorageSize,\r
@@ -410,6 +683,11 @@ SmmVariableHandler (
       break;\r
 \r
     case SMM_VARIABLE_FUNCTION_READY_TO_BOOT:\r
+      mEndOfDxe = TRUE;\r
+      if (AtRuntime()) {\r
+        Status = EFI_UNSUPPORTED;\r
+        break;\r
+      }\r
       ReclaimForOS ();\r
       Status = EFI_SUCCESS;\r
       break;\r
@@ -420,22 +698,102 @@ SmmVariableHandler (
       break;\r
 \r
     case SMM_VARIABLE_FUNCTION_GET_STATISTICS:\r
-      VariableInfo = (SMM_VARIABLE_COMMUNICATE_VARIABLE_INFO_ENTRY *) SmmVariableFunctionHeader->Data;\r
-      InfoSize = *CommBufferSize - OFFSET_OF (SMM_VARIABLE_COMMUNICATE_HEADER, Data);\r
+      VariableInfo = (VARIABLE_INFO_ENTRY *) SmmVariableFunctionHeader->Data;\r
+      InfoSize = *CommBufferSize - SMM_VARIABLE_COMMUNICATE_HEADER_SIZE;\r
+\r
+      //\r
+      // Do not need to check SmmVariableFunctionHeader->Data in SMRAM here. \r
+      // It is covered by previous CommBuffer check \r
+      //\r
+     \r
+      if (InternalIsAddressInSmram ((EFI_PHYSICAL_ADDRESS)(UINTN)CommBufferSize, sizeof(UINTN))) {\r
+        DEBUG ((EFI_D_ERROR, "GetStatistics: SMM communication buffer in SMRAM!\n"));\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }  \r
+\r
       Status = SmmVariableGetStatistics (VariableInfo, &InfoSize);\r
-      *CommBufferSize = InfoSize + OFFSET_OF (SMM_VARIABLE_COMMUNICATE_HEADER, Data);\r
+      *CommBufferSize = InfoSize + SMM_VARIABLE_COMMUNICATE_HEADER_SIZE;\r
+      break;\r
+\r
+    case SMM_VARIABLE_FUNCTION_LOCK_VARIABLE:\r
+      if (CommBufferPayloadSize < OFFSET_OF(SMM_VARIABLE_COMMUNICATE_LOCK_VARIABLE, Name)) {\r
+        DEBUG ((EFI_D_ERROR, "RequestToLock: SMM communication buffer size invalid!\n"));\r
+        return EFI_SUCCESS;\r
+      }\r
+      //\r
+      // Copy the input communicate buffer payload to pre-allocated SMM variable buffer payload.\r
+      //\r
+      CopyMem (mVariableBufferPayload, SmmVariableFunctionHeader->Data, CommBufferPayloadSize);\r
+      VariableToLock = (SMM_VARIABLE_COMMUNICATE_LOCK_VARIABLE *) mVariableBufferPayload;\r
+\r
+      if (VariableToLock->NameSize > MAX_ADDRESS - OFFSET_OF (SMM_VARIABLE_COMMUNICATE_LOCK_VARIABLE, Name)) {\r
+        //\r
+        // Prevent InfoSize overflow happen\r
+        //\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+\r
+      if (VariableToLock->NameSize < sizeof (CHAR16) || VariableToLock->Name[VariableToLock->NameSize/sizeof (CHAR16) - 1] != L'\0') {\r
+        //\r
+        // Make sure VariableName is A Null-terminated string.\r
+        //\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+      \r
+      InfoSize = OFFSET_OF (SMM_VARIABLE_COMMUNICATE_LOCK_VARIABLE, Name) + VariableToLock->NameSize;\r
+      \r
+      //\r
+      // SMRAM range check already covered before\r
+      //\r
+      if (InfoSize > CommBufferPayloadSize) {\r
+        DEBUG ((EFI_D_ERROR, "Data size exceed communication buffer size limit!\n"));\r
+        Status = EFI_ACCESS_DENIED;\r
+        goto EXIT;\r
+      }\r
+\r
+      Status = VariableLockRequestToLock (\r
+                 NULL,\r
+                 VariableToLock->Name,\r
+                 &VariableToLock->Guid\r
+                 );\r
       break;\r
 \r
     default:\r
-      ASSERT (FALSE);\r
       Status = EFI_UNSUPPORTED;\r
   }\r
 \r
+EXIT:\r
+\r
   SmmVariableFunctionHeader->ReturnStatus = Status;\r
 \r
   return EFI_SUCCESS;\r
 }\r
 \r
+/**\r
+  SMM END_OF_DXE protocol notification event handler.\r
+\r
+  @param  Protocol   Points to the protocol's unique identifier\r
+  @param  Interface  Points to the interface instance\r
+  @param  Handle     The handle on which the interface was installed\r
+\r
+  @retval EFI_SUCCESS   SmmEndOfDxeCallback runs successfully\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmmEndOfDxeCallback (\r
+  IN CONST EFI_GUID                       *Protocol,\r
+  IN VOID                                 *Interface,\r
+  IN EFI_HANDLE                           Handle\r
+  )\r
+{\r
+  DEBUG ((EFI_D_INFO, "[Variable]END_OF_DXE is signaled\n"));\r
+  mEndOfDxe = TRUE;\r
+  return EFI_SUCCESS;\r
+}\r
 \r
 /**\r
   SMM Fault Tolerant Write protocol notification event handler.\r
@@ -498,7 +856,7 @@ SmmFtwNotificationEvent (
   //\r
   Status = gBS->InstallProtocolInterface (\r
                   &mSmmVariableHandle,\r
-                  &mSmmVariableWriteGuid,\r
+                  &gSmmVariableWriteGuid,\r
                   EFI_NATIVE_INTERFACE,\r
                   NULL\r
                   );\r
@@ -511,7 +869,7 @@ SmmFtwNotificationEvent (
 /**\r
   Variable Driver main entry point. The Variable driver places the 4 EFI\r
   runtime services in the EFI System Table and installs arch protocols \r
-  for variable read and write services being availible. It also registers\r
+  for variable read and write services being available. It also registers\r
   a notification function for an EVT_SIGNAL_VIRTUAL_ADDRESS_CHANGE event.\r
 \r
   @param[in] ImageHandle    The firmware allocated handle for the EFI image.  \r
@@ -530,7 +888,10 @@ VariableServiceInitialize (
   EFI_STATUS                              Status;\r
   EFI_HANDLE                              VariableHandle;\r
   VOID                                    *SmmFtwRegistration;\r
-  \r
+  EFI_SMM_ACCESS2_PROTOCOL                *SmmAccess;\r
+  UINTN                                   Size;\r
+  VOID                                    *SmmEndOfDxeRegistration;\r
+\r
   //\r
   // Variable initialize.\r
   //\r
@@ -549,6 +910,38 @@ VariableServiceInitialize (
                     );\r
   ASSERT_EFI_ERROR (Status);\r
 \r
+  //\r
+  // Get SMRAM information\r
+  //\r
+  Status = gBS->LocateProtocol (&gEfiSmmAccess2ProtocolGuid, NULL, (VOID **)&SmmAccess);\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  Size = 0;\r
+  Status = SmmAccess->GetCapabilities (SmmAccess, &Size, NULL);\r
+  ASSERT (Status == EFI_BUFFER_TOO_SMALL);\r
+\r
+  Status = gSmst->SmmAllocatePool (\r
+                    EfiRuntimeServicesData,\r
+                    Size,\r
+                    (VOID **)&mSmramRanges\r
+                    );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  Status = SmmAccess->GetCapabilities (SmmAccess, &Size, mSmramRanges);\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  mSmramRangeCount = Size / sizeof (EFI_SMRAM_DESCRIPTOR);\r
+\r
+  mVariableBufferPayloadSize = MAX (PcdGet32 (PcdMaxVariableSize), PcdGet32 (PcdMaxHardwareErrorVariableSize)) +\r
+                               OFFSET_OF (SMM_VARIABLE_COMMUNICATE_ACCESS_VARIABLE, Name) - sizeof (VARIABLE_HEADER);\r
+\r
+  Status = gSmst->SmmAllocatePool (\r
+                    EfiRuntimeServicesData,\r
+                    mVariableBufferPayloadSize,\r
+                    (VOID **)&mVariableBufferPayload\r
+                    );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
   ///\r
   /// Register SMM variable SMI handler\r
   ///\r
@@ -567,6 +960,16 @@ VariableServiceInitialize (
                                         );\r
   ASSERT_EFI_ERROR (Status);\r
  \r
+  //\r
+  // Register EFI_SMM_END_OF_DXE_PROTOCOL_GUID notify function.\r
+  //\r
+  Status = gSmst->SmmRegisterProtocolNotify (\r
+                    &gEfiSmmEndOfDxeProtocolGuid,\r
+                    SmmEndOfDxeCallback,\r
+                    &SmmEndOfDxeRegistration\r
+                    );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
   //\r
   // Register FtwNotificationEvent () notify function.\r
   // \r