]> git.proxmox.com Git - mirror_edk2.git/blobdiff - EdkCompatibilityPkg/Compatibility/SmmBaseHelper/SmmBaseHelper.c
Add function comments and refine function parameter IN OUT.
[mirror_edk2.git] / EdkCompatibilityPkg / Compatibility / SmmBaseHelper / SmmBaseHelper.c
index 381f6ae2edbf9bda8e70b4641f7bb1a17465513f..08ac649b3b1857fb8630ede6d3b0a9a1659d6093 100644 (file)
@@ -4,7 +4,7 @@
   This driver is the counterpart of the SMM Base On SMM Base2 Thunk driver. It\r
   provides helping services in SMM to the SMM Base On SMM Base2 Thunk driver.\r
 \r
-  Copyright (c) 2009, Intel Corporation\r
+  Copyright (c) 2009 - 2010, Intel Corporation\r
   All rights reserved. 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
 \r
 **/\r
 \r
-#include "SmmBaseHelper.h"\r
+#include <PiSmm.h>\r
+#include <Library/DebugLib.h>\r
+#include <Library/UefiBootServicesTableLib.h>\r
+#include <Library/SmmServicesTableLib.h>\r
+#include <Library/BaseLib.h>\r
+#include <Library/BaseMemoryLib.h>\r
+#include <Library/PeCoffLib.h>\r
+#include <Library/DevicePathLib.h>\r
+#include <Library/CacheMaintenanceLib.h>\r
+#include <Library/MemoryAllocationLib.h>\r
+#include <Library/SynchronizationLib.h>\r
+#include <Library/CpuLib.h>\r
+#include <Guid/SmmBaseThunkCommunication.h>\r
+#include <Protocol/SmmBaseHelperReady.h>\r
+#include <Protocol/SmmCpu.h>\r
+#include <Protocol/LoadedImage.h>\r
+#include <Protocol/SmmCpuSaveState.h>\r
+#include <Protocol/MpService.h>\r
+#include <Protocol/LoadPe32Image.h>\r
+#include <Protocol/SmmReadyToLock.h>\r
+\r
+///\r
+/// Structure for tracking paired information of registered Framework SMI handler\r
+/// and correpsonding dispatch handle for SMI handler thunk.\r
+///\r
+typedef struct {\r
+  LIST_ENTRY                    Link;\r
+  EFI_HANDLE                    DispatchHandle;\r
+  EFI_HANDLE                    SmmImageHandle;\r
+  EFI_SMM_CALLBACK_ENTRY_POINT  CallbackAddress;\r
+  VOID                          *CommunicationBuffer;\r
+  UINTN                         *SourceSize;\r
+} CALLBACK_INFO;\r
+\r
+typedef struct {\r
+  ///\r
+  /// PI SMM CPU Save State register index\r
+  ///\r
+  EFI_SMM_SAVE_STATE_REGISTER   Register;\r
+  ///\r
+  /// Offset in Framework SMST\r
+  ///\r
+  UINTN                         Offset;\r
+} CPU_SAVE_STATE_CONVERSION;\r
+\r
+#define CPU_SAVE_STATE_GET_OFFSET(Field)  (UINTN)(&(((EFI_SMM_CPU_SAVE_STATE *) 0)->Ia32SaveState.Field))\r
+\r
 \r
 EFI_HANDLE                         mDispatchHandle;\r
 EFI_SMM_CPU_PROTOCOL               *mSmmCpu;\r
+EFI_PE32_IMAGE_PROTOCOL            *mLoadPe32Image;\r
 EFI_GUID                           mEfiSmmCpuIoGuid = EFI_SMM_CPU_IO_GUID;\r
 EFI_SMM_BASE_HELPER_READY_PROTOCOL *mSmmBaseHelperReady;\r
 EFI_SMM_SYSTEM_TABLE               *mFrameworkSmst;\r
+UINTN                              mNumberOfProcessors;\r
+BOOLEAN                            mLocked = FALSE;\r
+BOOLEAN                            mPageTableHookEnabled;\r
+BOOLEAN                            mHookInitialized;\r
+UINT64                             *mCpuStatePageTable;\r
+SPIN_LOCK                          mPFLock;\r
+UINT64                             mPhyMask;\r
+VOID                               *mOriginalHandler;\r
+EFI_SMM_CPU_SAVE_STATE             *mShadowSaveState;\r
 \r
 LIST_ENTRY mCallbackInfoListHead = INITIALIZE_LIST_HEAD_VARIABLE (mCallbackInfoListHead);\r
 \r
@@ -50,6 +106,354 @@ CPU_SAVE_STATE_CONVERSION mCpuSaveStateConvTable[] = {
   {EFI_SMM_SAVE_STATE_REGISTER_CR3      , CPU_SAVE_STATE_GET_OFFSET(CR3)}\r
 };\r
 \r
+VOID\r
+PageFaultHandlerHook (\r
+  VOID\r
+  );\r
+\r
+/**\r
+  Read CpuSaveStates from PI for Framework use.\r
+\r
+  The function reads PI style CpuSaveStates of CpuIndex-th CPU for Framework driver use. If\r
+  ToRead is specified, the CpuSaveStates will be copied to ToRead, otherwise copied to\r
+  mFrameworkSmst->CpuSaveState[CpuIndex].\r
+\r
+  @param[in]      CpuIndex        The zero-based CPU index.\r
+  @param[in, out] ToRead          If not NULL, CpuSaveStates will be copied to it.\r
+\r
+**/\r
+VOID\r
+ReadCpuSaveState (\r
+  IN     UINTN                   CpuIndex,\r
+  IN OUT EFI_SMM_CPU_SAVE_STATE  *ToRead\r
+  )\r
+{\r
+  EFI_STATUS Status;\r
+  UINTN Index;\r
+  EFI_SMM_CPU_STATE *State;\r
+  EFI_SMI_CPU_SAVE_STATE *SaveState;\r
+\r
+  State = (EFI_SMM_CPU_STATE *)gSmst->CpuSaveState[CpuIndex];\r
+  if (ToRead != NULL) {\r
+    SaveState = &ToRead->Ia32SaveState;\r
+  } else {\r
+    SaveState = &mFrameworkSmst->CpuSaveState[CpuIndex].Ia32SaveState;\r
+  }\r
+\r
+  if (State->x86.SMMRevId < EFI_SMM_MIN_REV_ID_x64) {\r
+    SaveState->SMBASE = State->x86.SMBASE;\r
+    SaveState->SMMRevId = State->x86.SMMRevId;\r
+    SaveState->IORestart = State->x86.IORestart;\r
+    SaveState->AutoHALTRestart = State->x86.AutoHALTRestart;\r
+  } else {\r
+    SaveState->SMBASE = State->x64.SMBASE;\r
+    SaveState->SMMRevId = State->x64.SMMRevId;\r
+    SaveState->IORestart = State->x64.IORestart;\r
+    SaveState->AutoHALTRestart = State->x64.AutoHALTRestart;\r
+  }\r
+\r
+  for (Index = 0; Index < sizeof (mCpuSaveStateConvTable) / sizeof (CPU_SAVE_STATE_CONVERSION); Index++) {\r
+    ///\r
+    /// Try to use SMM CPU Protocol to access CPU save states if possible\r
+    ///\r
+    Status = mSmmCpu->ReadSaveState (\r
+                        mSmmCpu,\r
+                        (UINTN)sizeof (UINT32),\r
+                        mCpuSaveStateConvTable[Index].Register,\r
+                        CpuIndex,\r
+                        ((UINT8 *)SaveState) + mCpuSaveStateConvTable[Index].Offset\r
+                        );\r
+    ASSERT_EFI_ERROR (Status);\r
+  }\r
+}\r
+\r
+/**\r
+  Write CpuSaveStates from Framework into PI.\r
+\r
+  The function writes back CpuSaveStates of CpuIndex-th CPU from PI to Framework. If\r
+  ToWrite is specified, it contains the CpuSaveStates to write from, otherwise CpuSaveStates\r
+  to write from mFrameworkSmst->CpuSaveState[CpuIndex].\r
+\r
+  @param[in] CpuIndex      The zero-based CPU index.\r
+  @param[in] ToWrite       If not NULL, CpuSaveStates to write from.\r
+\r
+**/\r
+VOID\r
+WriteCpuSaveState (\r
+  IN UINTN                   CpuIndex,\r
+  IN EFI_SMM_CPU_SAVE_STATE  *ToWrite\r
+  )\r
+{\r
+  EFI_STATUS Status;\r
+  UINTN      Index;\r
+  EFI_SMI_CPU_SAVE_STATE *SaveState;\r
+\r
+  if (ToWrite != NULL) {\r
+    SaveState = &ToWrite->Ia32SaveState;\r
+  } else {\r
+    SaveState = &mFrameworkSmst->CpuSaveState[CpuIndex].Ia32SaveState;\r
+  }\r
+  \r
+  for (Index = 0; Index < sizeof (mCpuSaveStateConvTable) / sizeof (CPU_SAVE_STATE_CONVERSION); Index++) {\r
+    Status = mSmmCpu->WriteSaveState (\r
+                        mSmmCpu,\r
+                        (UINTN)sizeof (UINT32),\r
+                        mCpuSaveStateConvTable[Index].Register,\r
+                        CpuIndex,\r
+                        ((UINT8 *)SaveState) + \r
+                        mCpuSaveStateConvTable[Index].Offset\r
+                        );\r
+  }\r
+}\r
+\r
+/**\r
+  Read or write a page that contains CpuSaveStates. Read is from PI to Framework.\r
+  Write is from Framework to PI.\r
+\r
+  This function reads or writes a page that contains CpuSaveStates. The page contains Framework\r
+  CpuSaveStates. On read, it reads PI style CpuSaveStates and fill the page up. On write, it\r
+  writes back from the page content to PI CpuSaveStates struct.\r
+  The first Framework CpuSaveStates (for CPU 0) is from mFrameworkSmst->CpuSaveState which is\r
+  page aligned. Because Framework CpuSaveStates are continuous, we can know which CPUs' SaveStates\r
+  are in the page start from PageAddress.\r
+\r
+  @param[in] PageAddress   The base address for a page.\r
+  @param[in] IsRead        TRUE for Read, FALSE for Write.\r
+\r
+**/\r
+VOID\r
+ReadWriteCpuStatePage (\r
+  IN UINT64  PageAddress,\r
+  IN BOOLEAN IsRead\r
+  )\r
+{\r
+  UINTN          FirstSSIndex;   // Index of first CpuSaveState in the page\r
+  UINTN          LastSSIndex;    // Index of last CpuSaveState in the page\r
+  BOOLEAN        FirstSSAligned; // Whether first CpuSaveState is page-aligned\r
+  BOOLEAN        LastSSAligned;  // Whether the end of last CpuSaveState is page-aligned\r
+  UINTN          ClippedSize;\r
+  UINTN          CpuIndex;\r
+\r
+  FirstSSIndex = ((UINTN)PageAddress - (UINTN)mFrameworkSmst->CpuSaveState) / sizeof (EFI_SMM_CPU_SAVE_STATE);\r
+  FirstSSAligned = TRUE;\r
+  if (((UINTN)PageAddress - (UINTN)mFrameworkSmst->CpuSaveState) % sizeof (EFI_SMM_CPU_SAVE_STATE) != 0) {\r
+    FirstSSIndex++;\r
+    FirstSSAligned = FALSE;\r
+  }\r
+  LastSSIndex = ((UINTN)PageAddress + SIZE_4KB - (UINTN)mFrameworkSmst->CpuSaveState - 1) / sizeof (EFI_SMM_CPU_SAVE_STATE);\r
+  LastSSAligned = TRUE;\r
+  if (((UINTN)PageAddress + SIZE_4KB - (UINTN)mFrameworkSmst->CpuSaveState) % sizeof (EFI_SMM_CPU_SAVE_STATE) != 0) {\r
+    LastSSIndex--;\r
+    LastSSAligned = FALSE;\r
+  }\r
+  for (CpuIndex = FirstSSIndex; CpuIndex <= LastSSIndex && CpuIndex < mNumberOfProcessors; CpuIndex++) {\r
+    if (IsRead) {\r
+      ReadCpuSaveState (CpuIndex, NULL);\r
+    } else {\r
+      WriteCpuSaveState (CpuIndex, NULL);\r
+    }\r
+  }\r
+  if (!FirstSSAligned) {\r
+    ReadCpuSaveState (FirstSSIndex - 1, mShadowSaveState);\r
+    ClippedSize = (UINTN)&mFrameworkSmst->CpuSaveState[FirstSSIndex] & (SIZE_4KB - 1);\r
+    if (IsRead) {\r
+      CopyMem ((VOID*)(UINTN)PageAddress, (VOID*)((UINTN)(mShadowSaveState + 1) - ClippedSize), ClippedSize);\r
+    } else {\r
+      CopyMem ((VOID*)((UINTN)(mShadowSaveState + 1) - ClippedSize), (VOID*)(UINTN)PageAddress, ClippedSize);\r
+      WriteCpuSaveState (FirstSSIndex - 1, mShadowSaveState);\r
+    }\r
+  }\r
+  if (!LastSSAligned && LastSSIndex + 1 < mNumberOfProcessors) {\r
+    ReadCpuSaveState (LastSSIndex + 1, mShadowSaveState);\r
+    ClippedSize = SIZE_4KB - ((UINTN)&mFrameworkSmst->CpuSaveState[LastSSIndex + 1] & (SIZE_4KB - 1));\r
+    if (IsRead) {\r
+      CopyMem (&mFrameworkSmst->CpuSaveState[LastSSIndex + 1], mShadowSaveState, ClippedSize);\r
+    } else {\r
+      CopyMem (mShadowSaveState, &mFrameworkSmst->CpuSaveState[LastSSIndex + 1], ClippedSize);\r
+      WriteCpuSaveState (LastSSIndex + 1, mShadowSaveState);\r
+    }\r
+  }\r
+}\r
+\r
+/**\r
+  The page fault handler that on-demand read PI CpuSaveStates for framework use. If the fault\r
+  is not targeted to mFrameworkSmst->CpuSaveState range, the function will return FALSE to let\r
+  PageFaultHandlerHook know it needs to pass the fault over to original page fault handler.\r
+  \r
+  @retval TRUE     The page fault is correctly handled.\r
+  @retval FALSE    The page fault is not handled and is passed through to original handler.\r
+\r
+**/\r
+BOOLEAN\r
+PageFaultHandler (\r
+  VOID\r
+  )\r
+{\r
+  BOOLEAN        IsHandled;\r
+  UINT64         *PageTable;\r
+  UINT64         PFAddress;\r
+  UINTN          NumCpuStatePages;\r
+  \r
+  ASSERT (mPageTableHookEnabled);\r
+  AcquireSpinLock (&mPFLock);\r
+\r
+  PageTable = (UINT64*)(UINTN)(AsmReadCr3 () & mPhyMask);\r
+  PFAddress = AsmReadCr2 ();\r
+  NumCpuStatePages = EFI_SIZE_TO_PAGES (mNumberOfProcessors * sizeof (EFI_SMM_CPU_SAVE_STATE));\r
+  IsHandled = FALSE;\r
+  if (((UINTN)mFrameworkSmst->CpuSaveState & ~(SIZE_2MB-1)) == (PFAddress & ~(SIZE_2MB-1))) {\r
+    if ((UINTN)mFrameworkSmst->CpuSaveState <= PFAddress &&\r
+        PFAddress < (UINTN)mFrameworkSmst->CpuSaveState + EFI_PAGES_TO_SIZE (NumCpuStatePages)\r
+        ) {\r
+      mCpuStatePageTable[BitFieldRead64 (PFAddress, 12, 20)] |= BIT0 | BIT1; // present and rw\r
+      CpuFlushTlb ();\r
+      ReadWriteCpuStatePage (PFAddress & ~(SIZE_4KB-1), TRUE);\r
+      IsHandled = TRUE;\r
+    } else {\r
+      ASSERT (FALSE);\r
+    }\r
+  }\r
+\r
+  ReleaseSpinLock (&mPFLock);\r
+  return IsHandled;\r
+}\r
+\r
+/**\r
+  Write back the dirty Framework CpuSaveStates to PI.\r
+  \r
+  The function scans the page table for dirty pages in mFrameworkSmst->CpuSaveState\r
+  to write back to PI CpuSaveStates. It is meant to be called on each SmmBaseHelper SMI\r
+  callback after Framework handler is called.\r
+\r
+**/\r
+VOID\r
+WriteBackDirtyPages (\r
+  VOID\r
+  )\r
+{\r
+  UINTN  NumCpuStatePages;\r
+  UINTN  PTIndex;\r
+  UINTN  PTStartIndex;\r
+  UINTN  PTEndIndex;\r
+\r
+  NumCpuStatePages = EFI_SIZE_TO_PAGES (mNumberOfProcessors * sizeof (EFI_SMM_CPU_SAVE_STATE));\r
+  PTStartIndex = (UINTN)BitFieldRead64 ((UINT64)mFrameworkSmst->CpuSaveState, 12, 20);\r
+  PTEndIndex   = (UINTN)BitFieldRead64 ((UINT64)mFrameworkSmst->CpuSaveState + EFI_PAGES_TO_SIZE(NumCpuStatePages) - 1, 12, 20);\r
+  for (PTIndex = PTStartIndex; PTIndex <= PTEndIndex; PTIndex++) {\r
+    if ((mCpuStatePageTable[PTIndex] & (BIT0|BIT6)) == (BIT0|BIT6)) { // present and dirty?\r
+      ReadWriteCpuStatePage (mCpuStatePageTable[PTIndex] & mPhyMask, FALSE);\r
+    }\r
+  }\r
+}\r
+\r
+/**\r
+  Hook IDT with our page fault handler so that the on-demand paging works on page fault.\r
+  \r
+  The function hooks the IDT with PageFaultHandlerHook to get on-demand paging work for\r
+  PI<->Framework CpuSaveStates marshalling. It also saves original handler for pass-through\r
+  purpose.\r
+\r
+**/\r
+VOID\r
+HookPageFaultHandler (\r
+  VOID\r
+  )\r
+{\r
+  IA32_DESCRIPTOR           Idtr;\r
+  IA32_IDT_GATE_DESCRIPTOR  *IdtGateDesc;\r
+  UINT32                    OffsetUpper;\r
+  \r
+  InitializeSpinLock (&mPFLock);\r
+  \r
+  AsmReadIdtr (&Idtr);\r
+  IdtGateDesc = (IA32_IDT_GATE_DESCRIPTOR *) Idtr.Base;\r
+  OffsetUpper = *(UINT32*)((UINT64*)IdtGateDesc + 1);\r
+  mOriginalHandler = (VOID *)(UINTN)(LShiftU64 (OffsetUpper, 32) + IdtGateDesc[14].Bits.OffsetLow + (IdtGateDesc[14].Bits.OffsetHigh << 16));\r
+  IdtGateDesc[14].Bits.OffsetLow = (UINT32)((UINTN)PageFaultHandlerHook & ((1 << 16) - 1));\r
+  IdtGateDesc[14].Bits.OffsetHigh = (UINT32)(((UINTN)PageFaultHandlerHook >> 16) & ((1 << 16) - 1));\r
+}\r
+\r
+/**\r
+  Initialize page table for pages contain HookData.\r
+  \r
+  The function initialize PDE for 2MB range that contains HookData. If the related PDE points\r
+  to a 2MB page, a page table will be allocated and initialized for 4KB pages. Otherwise we juse\r
+  use the original page table.\r
+\r
+  @param[in] HookData   Based on which to initialize page table.\r
+\r
+  @return    The pointer to a Page Table that points to 4KB pages which contain HookData.\r
+**/\r
+UINT64 *\r
+InitCpuStatePageTable (\r
+  IN VOID *HookData\r
+  )\r
+{\r
+  UINTN  Index;\r
+  UINT64 *PageTable;\r
+  UINT64 *PDPTE;\r
+  UINT64 HookAddress;\r
+  UINT64 PDE;\r
+  UINT64 Address;\r
+  \r
+  //\r
+  // Initialize physical address mask\r
+  // NOTE: Physical memory above virtual address limit is not supported !!!\r
+  //\r
+  AsmCpuid (0x80000008, (UINT32*)&Index, NULL, NULL, NULL);\r
+  mPhyMask = LShiftU64 (1, (UINT8)Index) - 1;\r
+  mPhyMask &= (1ull << 48) - EFI_PAGE_SIZE;\r
+  \r
+  HookAddress = (UINT64)(UINTN)HookData;\r
+  PageTable   = (UINT64 *)(UINTN)(AsmReadCr3 () & mPhyMask);\r
+  PageTable = (UINT64 *)(UINTN)(PageTable[BitFieldRead64 (HookAddress, 39, 47)] & mPhyMask);\r
+  PageTable = (UINT64 *)(UINTN)(PageTable[BitFieldRead64 (HookAddress, 30, 38)] & mPhyMask);\r
+  \r
+  PDPTE = (UINT64 *)(UINTN)PageTable;\r
+  PDE = PDPTE[BitFieldRead64 (HookAddress, 21, 29)];\r
+  ASSERT ((PDE & BIT0) != 0); // Present and 2M Page\r
+  \r
+  if ((PDE & BIT7) == 0) { // 4KB Page Directory\r
+    PageTable = (UINT64 *)(UINTN)(PDE & mPhyMask);\r
+  } else {\r
+    ASSERT ((PDE & mPhyMask) == (HookAddress & ~(SIZE_2MB-1))); // 2MB Page Point to HookAddress\r
+    PageTable = AllocatePages (1);\r
+    Address = HookAddress & ~(SIZE_2MB-1);\r
+    for (Index = 0; Index < 512; Index++) {\r
+      PageTable[Index] = Address | BIT0 | BIT1; // Present and RW\r
+      Address += SIZE_4KB;\r
+    }\r
+    PDPTE[BitFieldRead64 (HookAddress, 21, 29)] = (UINT64)(UINTN)PageTable | BIT0 | BIT1; // Present and RW\r
+  }\r
+  return PageTable;\r
+}\r
+\r
+/**\r
+  Mark all the CpuSaveStates as not present.\r
+  \r
+  The function marks all CpuSaveStates memory range as not present so that page fault can be triggered\r
+  on CpuSaveStates access. It is meant to be called on each SmmBaseHelper SMI callback before Framework\r
+  handler is called.\r
+\r
+  @param[in] CpuSaveState   The base of CpuSaveStates.\r
+\r
+**/\r
+VOID\r
+HookCpuStateMemory (\r
+  IN EFI_SMM_CPU_SAVE_STATE *CpuSaveState\r
+  )\r
+{\r
+  UINT64 Index;\r
+  UINT64 PTStartIndex;\r
+  UINT64 PTEndIndex;\r
+\r
+  PTStartIndex = BitFieldRead64 ((UINTN)CpuSaveState, 12, 20);\r
+  PTEndIndex = BitFieldRead64 ((UINTN)CpuSaveState + mNumberOfProcessors * sizeof (EFI_SMM_CPU_SAVE_STATE) - 1, 12, 20);\r
+  for (Index = PTStartIndex; Index <= PTEndIndex; Index++) {\r
+    mCpuStatePageTable[Index] &= ~(BIT0|BIT5|BIT6); // not present nor accessed nor dirty\r
+  }\r
+}  \r
+\r
 /**\r
   Framework SMST SmmInstallConfigurationTable() Thunk.\r
 \r
@@ -86,6 +490,69 @@ SmmInstallConfigurationTable (
   return Status;         \r
 }\r
 \r
+/**\r
+  Initialize all the stuff needed for on-demand paging hooks for PI<->Framework\r
+  CpuSaveStates marshalling.\r
+\r
+  @param[in] FrameworkSmst   Framework SMM system table pointer.\r
+\r
+**/\r
+VOID\r
+InitHook (\r
+  IN EFI_SMM_SYSTEM_TABLE  *FrameworkSmst\r
+  )\r
+{\r
+  UINTN                 NumCpuStatePages;\r
+  UINTN                 CpuStatePage;\r
+  UINTN                 Bottom2MPage;\r
+  UINTN                 Top2MPage;\r
+  \r
+  mPageTableHookEnabled = FALSE;\r
+  NumCpuStatePages = EFI_SIZE_TO_PAGES (mNumberOfProcessors * sizeof (EFI_SMM_CPU_SAVE_STATE));\r
+  //\r
+  // Only hook page table for X64 image and less than 2MB needed to hold all CPU Save States\r
+  //\r
+  if (EFI_IMAGE_MACHINE_TYPE_SUPPORTED(EFI_IMAGE_MACHINE_X64) && NumCpuStatePages <= EFI_SIZE_TO_PAGES (SIZE_2MB)) {\r
+    //\r
+    // Allocate double page size to make sure all CPU Save States are in one 2MB page.\r
+    //\r
+    CpuStatePage = (UINTN)AllocatePages (NumCpuStatePages * 2);\r
+    ASSERT (CpuStatePage != 0);\r
+    Bottom2MPage = CpuStatePage & ~(SIZE_2MB-1);\r
+    Top2MPage    = (CpuStatePage + EFI_PAGES_TO_SIZE (NumCpuStatePages * 2) - 1) & ~(SIZE_2MB-1);\r
+    if (Bottom2MPage == Top2MPage ||\r
+        CpuStatePage + EFI_PAGES_TO_SIZE (NumCpuStatePages * 2) - Top2MPage >= EFI_PAGES_TO_SIZE (NumCpuStatePages)\r
+        ) {\r
+      //\r
+      // If the allocated 4KB pages are within the same 2MB page or higher portion is larger, use higher portion pages.\r
+      //\r
+      FrameworkSmst->CpuSaveState = (EFI_SMM_CPU_SAVE_STATE *)(CpuStatePage + EFI_PAGES_TO_SIZE (NumCpuStatePages));\r
+      FreePages ((VOID*)CpuStatePage, NumCpuStatePages);\r
+    } else {\r
+      FrameworkSmst->CpuSaveState = (EFI_SMM_CPU_SAVE_STATE *)CpuStatePage;\r
+      FreePages ((VOID*)(CpuStatePage + EFI_PAGES_TO_SIZE (NumCpuStatePages)), NumCpuStatePages);\r
+    }\r
+    //\r
+    // Add temporary working buffer for hooking\r
+    //\r
+    mShadowSaveState = (EFI_SMM_CPU_SAVE_STATE*) AllocatePool (sizeof (EFI_SMM_CPU_SAVE_STATE));\r
+    ASSERT (mShadowSaveState != NULL);\r
+    //\r
+    // Allocate and initialize 4KB Page Table for hooking CpuSaveState.\r
+    // Replace the original 2MB PDE with new 4KB page table.\r
+    //\r
+    mCpuStatePageTable = InitCpuStatePageTable (FrameworkSmst->CpuSaveState);\r
+    //\r
+    // Mark PTE for CpuSaveState as non-exist.\r
+    //\r
+    HookCpuStateMemory (FrameworkSmst->CpuSaveState);\r
+    HookPageFaultHandler ();\r
+    CpuFlushTlb ();\r
+    mPageTableHookEnabled = TRUE;\r
+  }\r
+  mHookInitialized = TRUE;\r
+}\r
+\r
 /**\r
   Construct a Framework SMST based on the PI SMM SMST.\r
 \r
@@ -96,15 +563,10 @@ ConstructFrameworkSmst (
   VOID\r
   )\r
 {\r
-  EFI_STATUS            Status;\r
   EFI_SMM_SYSTEM_TABLE  *FrameworkSmst;\r
 \r
-  Status = gSmst->SmmAllocatePool (\r
-                    EfiRuntimeServicesData,\r
-                    sizeof (EFI_SMM_SYSTEM_TABLE),\r
-                    (VOID **)&FrameworkSmst\r
-                    );\r
-  ASSERT_EFI_ERROR (Status);\r
+  FrameworkSmst = (EFI_SMM_SYSTEM_TABLE  *)AllocatePool (sizeof (EFI_SMM_SYSTEM_TABLE));\r
+  ASSERT (FrameworkSmst != NULL);\r
 \r
   ///\r
   /// Copy same things from PI SMST to Framework SMST\r
@@ -122,13 +584,9 @@ ConstructFrameworkSmst (
   FrameworkSmst->Hdr.Revision = EFI_SMM_SYSTEM_TABLE_REVISION;\r
   CopyGuid (&FrameworkSmst->EfiSmmCpuIoGuid, &mEfiSmmCpuIoGuid);\r
 \r
-  Status = gSmst->SmmAllocatePool (\r
-                    EfiRuntimeServicesData,\r
-                    gSmst->NumberOfCpus * sizeof (EFI_SMI_CPU_SAVE_STATE),\r
-                    (VOID **)&FrameworkSmst->CpuSaveState\r
-                    );\r
-  ASSERT_EFI_ERROR (Status);\r
-  ZeroMem (FrameworkSmst->CpuSaveState, gSmst->NumberOfCpus * sizeof (EFI_SMI_CPU_SAVE_STATE));\r
+  mHookInitialized = FALSE;\r
+  FrameworkSmst->CpuSaveState = (EFI_SMM_CPU_SAVE_STATE *)AllocateZeroPool (mNumberOfProcessors * sizeof (EFI_SMM_CPU_SAVE_STATE));\r
+  ASSERT (FrameworkSmst->CpuSaveState != NULL);\r
 \r
   ///\r
   /// Do not support floating point state now\r
@@ -143,6 +601,7 @@ ConstructFrameworkSmst (
 /**\r
   Load a given Framework SMM driver into SMRAM and invoke its entry point.\r
 \r
+  @param[in]   ParentImageHandle     Parent Image Handle.\r
   @param[in]   FilePath              Location of the image to be installed as the handler.\r
   @param[in]   SourceBuffer          Optional source buffer in case the image file\r
                                      is in memory.\r
@@ -158,154 +617,77 @@ ConstructFrameworkSmst (
 **/\r
 EFI_STATUS\r
 LoadImage (\r
+  IN      EFI_HANDLE                ParentImageHandle,\r
   IN      EFI_DEVICE_PATH_PROTOCOL  *FilePath,\r
   IN      VOID                      *SourceBuffer,\r
   IN      UINTN                     SourceSize,\r
   OUT     EFI_HANDLE                *ImageHandle\r
   )\r
 {\r
-  EFI_STATUS                    Status;\r
-  UINTN                         PageCount;\r
-  EFI_PHYSICAL_ADDRESS          Buffer;\r
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;\r
-  EFI_HANDLE                    PesudoImageHandle;\r
-  UINTN                         NumHandles;\r
-  UINTN                         Index;\r
-  EFI_HANDLE                    *HandleBuffer;\r
-  EFI_LOADED_IMAGE_PROTOCOL     *LoadedImage;\r
-  EFI_DEVICE_PATH               *LoadedImageDevicePath;\r
-  UINTN                         DevicePathSize;\r
+  EFI_STATUS            Status;\r
+  UINTN                 PageCount;\r
+  UINTN                 OrgPageCount;\r
+  EFI_PHYSICAL_ADDRESS  DstBuffer;\r
 \r
   if (FilePath == NULL || ImageHandle == NULL) {    \r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
-  ///\r
-  /// Assume Framework SMM driver has an image copy in memory before registering itself into SMRAM.\r
-  /// Currently only supports load Framework SMM driver from existing image copy in memory.\r
-  /// Load PE32 Image Protocol can be used to support loading Framework SMM driver directly from FV.\r
-  ///\r
-  if (SourceBuffer == NULL) {\r
-    Status = gBS->LocateHandleBuffer (\r
-                    ByProtocol,\r
-                    &gEfiLoadedImageDevicePathProtocolGuid,\r
-                    NULL,\r
-                    &NumHandles,\r
-                    &HandleBuffer\r
-                    );\r
-    if (EFI_ERROR (Status)) {\r
-      return EFI_UNSUPPORTED;\r
-    }\r
-\r
-    DevicePathSize = GetDevicePathSize (FilePath);\r
-\r
-    for (Index = 0; Index < NumHandles; Index++) {\r
-      Status = gBS->HandleProtocol (\r
-                      HandleBuffer[Index],\r
-                      &gEfiLoadedImageDevicePathProtocolGuid,\r
-                      (VOID **)&LoadedImageDevicePath\r
+  PageCount = 1;\r
+  do {\r
+    OrgPageCount = PageCount;\r
+    DstBuffer = (UINTN)-1;\r
+    Status = gSmst->SmmAllocatePages (\r
+                      AllocateMaxAddress,\r
+                      EfiRuntimeServicesCode,\r
+                      PageCount,\r
+                      &DstBuffer\r
                       );\r
-      ASSERT_EFI_ERROR (Status);\r
-\r
-      if (GetDevicePathSize (LoadedImageDevicePath) == DevicePathSize &&\r
-          CompareMem (LoadedImageDevicePath, FilePath, DevicePathSize) == 0) {\r
-          break;\r
-      }     \r
+    if (EFI_ERROR (Status)) {\r
+      return Status;\r
     }\r
 \r
-    if (Index < NumHandles) {\r
-      Status = gBS->HandleProtocol (\r
-                      HandleBuffer[Index],\r
-                      &gEfiLoadedImageProtocolGuid,\r
-                      (VOID **)&LoadedImage\r
-                      );\r
-      ASSERT_EFI_ERROR (Status);\r
-      \r
-      SourceBuffer = LoadedImage->ImageBase;\r
-      gBS->FreePool (HandleBuffer);\r
-    } else {\r
-      gBS->FreePool (HandleBuffer);\r
-      return EFI_UNSUPPORTED;\r
+    Status = mLoadPe32Image->LoadPeImage (\r
+                               mLoadPe32Image,\r
+                               ParentImageHandle,\r
+                               FilePath,\r
+                               SourceBuffer,\r
+                               SourceSize,\r
+                               DstBuffer,\r
+                               &PageCount,\r
+                               ImageHandle,\r
+                               NULL,\r
+                               EFI_LOAD_PE_IMAGE_ATTRIBUTE_NONE\r
+                               );\r
+    if (EFI_ERROR (Status)) {\r
+      FreePages ((VOID *)(UINTN)DstBuffer, OrgPageCount);\r
     }\r
-  }\r
-\r
-  ImageContext.Handle = SourceBuffer;\r
-  ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;\r
-\r
-  ///\r
-  /// Get information about the image being loaded\r
-  ///\r
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
-\r
-  ///\r
-  /// Allocate buffer for loading image into SMRAM\r
-  ///\r
-  PageCount = (UINTN)EFI_SIZE_TO_PAGES (ImageContext.ImageSize + ImageContext.SectionAlignment);\r
-  Status = gSmst->SmmAllocatePages (AllocateAnyPages, EfiRuntimeServicesCode, PageCount, &Buffer);\r
-  if (EFI_ERROR (Status)) {\r
-    return Status;\r
-  }\r
-\r
-  ImageContext.ImageAddress = (PHYSICAL_ADDRESS)Buffer;\r
-\r
-  ///\r
-  /// Align buffer on section boundry\r
-  ///\r
-  ImageContext.ImageAddress += ImageContext.SectionAlignment - 1;\r
-  ImageContext.ImageAddress &= ~(ImageContext.SectionAlignment - 1);\r
-\r
-  ///\r
-  /// Load the image into SMRAM\r
-  ///\r
-  Status = PeCoffLoaderLoadImage (&ImageContext);\r
-  if (EFI_ERROR (Status)) {\r
-    goto Error;\r
-  }\r
-\r
-  ///\r
-  /// Relocate the image in our new buffer\r
-  ///\r
-  Status = PeCoffLoaderRelocateImage (&ImageContext);\r
-  if (EFI_ERROR (Status)) {\r
-    goto Error;\r
-  }\r
-\r
-  ///\r
-  /// Flush the instruction cache so the image data are written before we execute it\r
-  ///\r
-  InvalidateInstructionCacheRange ((VOID *)(UINTN) ImageContext.ImageAddress, (UINTN) ImageContext.ImageSize);\r
-\r
-  ///\r
-  /// Update MP state in Framework SMST before transferring control to Framework SMM driver entry point\r
-  /// in case it may invoke AP\r
-  ///\r
-  mFrameworkSmst->CurrentlyExecutingCpu = gSmst->CurrentlyExecutingCpu;\r
+  } while (Status == EFI_BUFFER_TOO_SMALL);\r
 \r
-  ///\r
-  /// For Framework SMM, ImageHandle does not have to be a UEFI image handle.  The only requirement is that the \r
-  /// ImageHandle is a unique value.  Use image base address as the unique value.\r
-  ///\r
-  PesudoImageHandle = (EFI_HANDLE)(UINTN)ImageContext.ImageAddress;\r
-\r
-  Status = ((EFI_IMAGE_ENTRY_POINT)(UINTN)ImageContext.EntryPoint) (PesudoImageHandle, gST);\r
   if (!EFI_ERROR (Status)) {\r
-    *ImageHandle = PesudoImageHandle;\r
-    return EFI_SUCCESS;\r
+    ///\r
+    /// Update MP state in Framework SMST before transferring control to Framework SMM driver entry point\r
+    /// in case it may invoke AP\r
+    ///\r
+    mFrameworkSmst->CurrentlyExecutingCpu = gSmst->CurrentlyExecutingCpu;\r
+\r
+    Status = gBS->StartImage (*ImageHandle, NULL, NULL);\r
+    if (EFI_ERROR (Status)) {\r
+      mLoadPe32Image->UnLoadPeImage (mLoadPe32Image, *ImageHandle);\r
+      *ImageHandle = NULL;\r
+      FreePages ((VOID *)(UINTN)DstBuffer, PageCount);\r
+    }\r
   }\r
 \r
-Error:\r
-  gSmst->SmmFreePages (Buffer, PageCount);\r
   return Status;\r
 }\r
 \r
+\r
 /** \r
   Thunk service of EFI_SMM_BASE_PROTOCOL.Register().\r
 \r
-  @param[in] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
-*/\r
+  @param[in, out] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
+**/\r
 VOID\r
 Register (\r
   IN OUT SMMBASE_FUNCTION_DATA *FunctionData\r
@@ -313,10 +695,11 @@ Register (
 {\r
   EFI_STATUS Status;\r
 \r
-  if (FunctionData->Args.Register.LegacyIA32Binary) {\r
+  if (mLocked || FunctionData->Args.Register.LegacyIA32Binary) {\r
     Status = EFI_UNSUPPORTED;\r
   } else {\r
     Status = LoadImage (\r
+               FunctionData->SmmBaseImageHandle,\r
                FunctionData->Args.Register.FilePath,\r
                FunctionData->Args.Register.SourceBuffer,\r
                FunctionData->Args.Register.SourceSize,\r
@@ -329,8 +712,8 @@ Register (
 /** \r
   Thunk service of EFI_SMM_BASE_PROTOCOL.UnRegister().\r
 \r
-  @param[in] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
-*/\r
+  @param[in, out] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
+**/\r
 VOID\r
 UnRegister (\r
   IN OUT SMMBASE_FUNCTION_DATA *FunctionData\r
@@ -347,7 +730,7 @@ UnRegister (
 \r
   @param[in] DispatchHandle  The unique handle assigned by SmiHandlerRegister().  \r
 \r
-  @return  Pointer to CALLBACK_INFO.\r
+  @return  Pointer to CALLBACK_INFO. If NULL, no callback info record is found.\r
 **/\r
 CALLBACK_INFO *\r
 GetCallbackInfo (\r
@@ -375,9 +758,9 @@ GetCallbackInfo (
   @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
+  @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
+  @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
@@ -398,47 +781,27 @@ CallbackThunk (
 {\r
   EFI_STATUS        Status;\r
   CALLBACK_INFO     *CallbackInfo;\r
-  UINTN             Index;\r
   UINTN             CpuIndex;\r
-  EFI_SMM_CPU_STATE *State;\r
-  EFI_SMI_CPU_SAVE_STATE *SaveState;\r
 \r
   ///\r
   /// Before transferring the control into the Framework SMI handler, update CPU Save States\r
   /// and MP states in the Framework SMST.\r
   ///\r
 \r
-  for (CpuIndex = 0; CpuIndex < gSmst->NumberOfCpus; CpuIndex++) {\r
-    State = (EFI_SMM_CPU_STATE *)gSmst->CpuSaveState[CpuIndex];\r
-    SaveState = &mFrameworkSmst->CpuSaveState[CpuIndex].Ia32SaveState;\r
-\r
-    if (State->x86.SMMRevId < EFI_SMM_MIN_REV_ID_x64) {\r
-      SaveState->SMBASE = State->x86.SMBASE;\r
-      SaveState->SMMRevId = State->x86.SMMRevId;\r
-      SaveState->IORestart = State->x86.IORestart;\r
-      SaveState->AutoHALTRestart = State->x86.AutoHALTRestart;\r
-    } else {\r
-      SaveState->SMBASE = State->x64.SMBASE;\r
-      SaveState->SMMRevId = State->x64.SMMRevId;\r
-      SaveState->IORestart = State->x64.IORestart;\r
-      SaveState->AutoHALTRestart = State->x64.AutoHALTRestart;\r
-    }\r
-\r
-    for (Index = 0; Index < sizeof (mCpuSaveStateConvTable) / sizeof (CPU_SAVE_STATE_CONVERSION); Index++) {\r
-      ///\r
-      /// Try to use SMM CPU Protocol to access CPU save states if possible\r
-      ///\r
-      Status = mSmmCpu->ReadSaveState (\r
-                          mSmmCpu,\r
-                          EFI_SMM_SAVE_STATE_IO_WIDTH_UINT32,\r
-                          mCpuSaveStateConvTable[Index].Register,\r
-                          CpuIndex,\r
-                          ((UINT8 *)SaveState) + mCpuSaveStateConvTable[Index].Offset\r
-                          );\r
-      ASSERT_EFI_ERROR (Status);\r
+  if (!mHookInitialized) {\r
+    InitHook (mFrameworkSmst);\r
+  }\r
+  if (mPageTableHookEnabled) {\r
+    HookCpuStateMemory (mFrameworkSmst->CpuSaveState);\r
+    CpuFlushTlb ();\r
+  } else {\r
+    for (CpuIndex = 0; CpuIndex < mNumberOfProcessors; CpuIndex++) {\r
+      ReadCpuSaveState (CpuIndex, NULL);\r
     }\r
   }\r
 \r
+  mFrameworkSmst->SmmStartupThisAp      = gSmst->SmmStartupThisAp;\r
+  mFrameworkSmst->NumberOfCpus          = mNumberOfProcessors;\r
   mFrameworkSmst->CurrentlyExecutingCpu = gSmst->CurrentlyExecutingCpu;\r
 \r
   ///\r
@@ -452,22 +815,17 @@ CallbackThunk (
   ///\r
   Status = (CallbackInfo->CallbackAddress) (\r
                             CallbackInfo->SmmImageHandle,\r
-                            CommBuffer,\r
-                            CommBufferSize\r
+                            CallbackInfo->CommunicationBuffer,\r
+                            CallbackInfo->SourceSize\r
                             );\r
   ///\r
   /// Save CPU Save States in case any of them was modified\r
   ///\r
-  for (CpuIndex = 0; CpuIndex < gSmst->NumberOfCpus; CpuIndex++) {\r
-    for (Index = 0; Index < sizeof (mCpuSaveStateConvTable) / sizeof (CPU_SAVE_STATE_CONVERSION); Index++) {\r
-      Status = mSmmCpu->WriteSaveState (\r
-                          mSmmCpu,\r
-                          EFI_SMM_SAVE_STATE_IO_WIDTH_UINT32,\r
-                          mCpuSaveStateConvTable[Index].Register,\r
-                          CpuIndex,\r
-                          ((UINT8 *)&mFrameworkSmst->CpuSaveState[CpuIndex].Ia32SaveState) + \r
-                          mCpuSaveStateConvTable[Index].Offset\r
-                          );\r
+  if (mPageTableHookEnabled) {\r
+    WriteBackDirtyPages ();\r
+  } else {\r
+    for (CpuIndex = 0; CpuIndex < mNumberOfProcessors; CpuIndex++) {\r
+      WriteCpuSaveState (CpuIndex, NULL);\r
     }\r
   }\r
 \r
@@ -492,16 +850,20 @@ CallbackThunk (
 /** \r
   Thunk service of EFI_SMM_BASE_PROTOCOL.RegisterCallback().\r
 \r
-  @param[in] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
-*/\r
+  @param[in, out] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
+**/\r
 VOID\r
 RegisterCallback (\r
-  IN OUT SMMBASE_FUNCTION_DATA *FunctionData\r
+  IN OUT SMMBASE_FUNCTION_DATA  *FunctionData\r
   )\r
 {\r
-  EFI_STATUS     Status;\r
   CALLBACK_INFO  *Buffer;\r
 \r
+  if (mLocked) {\r
+    FunctionData->Status = EFI_UNSUPPORTED;\r
+    return;\r
+  }\r
+\r
   ///\r
   /// Note that MakeLast and FloatingPointSave options are not supported in PI SMM\r
   ///\r
@@ -509,69 +871,113 @@ RegisterCallback (
   ///\r
   /// Allocate buffer for callback thunk information\r
   ///\r
-  Status = gSmst->SmmAllocatePool (\r
-                    EfiRuntimeServicesCode,\r
-                    sizeof (CALLBACK_INFO),\r
-                    (VOID **)&Buffer\r
-                    );\r
-  if (!EFI_ERROR (Status)) {\r
-    ///\r
-    /// Fill SmmImageHandle and CallbackAddress into the thunk\r
-    ///\r
-    Buffer->SmmImageHandle = FunctionData->Args.RegisterCallback.SmmImageHandle;\r
-    Buffer->CallbackAddress = FunctionData->Args.RegisterCallback.CallbackAddress;\r
+  Buffer = (CALLBACK_INFO *)AllocateZeroPool (sizeof (CALLBACK_INFO));\r
+  if (Buffer == NULL) {\r
+    FunctionData->Status = EFI_OUT_OF_RESOURCES;\r
+    return;\r
+  }\r
 \r
-    ///\r
-    /// Register the thunk code as a root SMI handler\r
-    ///\r
-    Status = gSmst->SmiHandlerRegister (\r
-                      CallbackThunk,\r
-                      NULL,\r
-                      &Buffer->DispatchHandle\r
-                      );\r
-    if (!EFI_ERROR (Status)) {\r
-      ///\r
-      /// Save this callback info\r
-      ///\r
-      InsertTailList (&mCallbackInfoListHead, &Buffer->Link);\r
-    } else {\r
-      gSmst->SmmFreePool (Buffer);\r
-    }\r
+  ///\r
+  /// Fill SmmImageHandle and CallbackAddress into the thunk\r
+  ///\r
+  Buffer->SmmImageHandle = FunctionData->Args.RegisterCallback.SmmImageHandle;\r
+  Buffer->CallbackAddress = FunctionData->Args.RegisterCallback.CallbackAddress;\r
+\r
+  ///\r
+  /// Register the thunk code as a root SMI handler\r
+  ///\r
+  FunctionData->Status = gSmst->SmiHandlerRegister (\r
+                                  CallbackThunk,\r
+                                  NULL,\r
+                                  &Buffer->DispatchHandle\r
+                                  );\r
+  if (EFI_ERROR (FunctionData->Status)) {\r
+    FreePool (Buffer);\r
+    return;\r
   }\r
-  FunctionData->Status = Status;\r
+\r
+  ///\r
+  /// Save this callback info\r
+  ///\r
+  InsertTailList (&mCallbackInfoListHead, &Buffer->Link);\r
 }\r
 \r
 \r
 /** \r
   Thunk service of EFI_SMM_BASE_PROTOCOL.SmmAllocatePool().\r
 \r
-  @param[in] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
-*/\r
+  @param[in, out] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
+**/\r
 VOID\r
 HelperAllocatePool (\r
   IN OUT SMMBASE_FUNCTION_DATA *FunctionData\r
   )\r
 {\r
-  FunctionData->Status = gSmst->SmmAllocatePool (\r
-                                  FunctionData->Args.AllocatePool.PoolType,\r
-                                  FunctionData->Args.AllocatePool.Size,\r
-                                  FunctionData->Args.AllocatePool.Buffer\r
-                                  );\r
+  if (mLocked) {\r
+    FunctionData->Status =  EFI_UNSUPPORTED;\r
+  } else {\r
+    FunctionData->Status = gSmst->SmmAllocatePool (\r
+                                    FunctionData->Args.AllocatePool.PoolType,\r
+                                    FunctionData->Args.AllocatePool.Size,\r
+                                    FunctionData->Args.AllocatePool.Buffer\r
+                                    );\r
+  }\r
 }\r
 \r
 /** \r
   Thunk service of EFI_SMM_BASE_PROTOCOL.SmmFreePool().\r
 \r
-  @param[in] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
-*/\r
+  @param[in, out] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
+**/\r
 VOID\r
 HelperFreePool (\r
   IN OUT SMMBASE_FUNCTION_DATA *FunctionData\r
   )\r
 {\r
-  FunctionData->Status = gSmst->SmmFreePool (\r
-                                  FunctionData->Args.FreePool.Buffer\r
-                                  );\r
+  if (mLocked) {\r
+    FunctionData->Status =  EFI_UNSUPPORTED;\r
+  } else {\r
+    FreePool (FunctionData->Args.FreePool.Buffer);\r
+    FunctionData->Status = EFI_SUCCESS;\r
+  }\r
+}\r
+\r
+/** \r
+  Thunk service of EFI_SMM_BASE_PROTOCOL.Communicate().\r
+\r
+  @param[in, out] FunctionData  Pointer to SMMBASE_FUNCTION_DATA.\r
+**/\r
+VOID\r
+HelperCommunicate (\r
+  IN OUT SMMBASE_FUNCTION_DATA *FunctionData\r
+  )\r
+{\r
+  LIST_ENTRY     *Node;\r
+  CALLBACK_INFO  *CallbackInfo;\r
+\r
+  if (FunctionData->Args.Communicate.CommunicationBuffer == NULL) {\r
+    FunctionData->Status = EFI_INVALID_PARAMETER;\r
+    return;\r
+  }\r
+\r
+  Node = GetFirstNode (&mCallbackInfoListHead);\r
+  while (!IsNull (&mCallbackInfoListHead, Node)) {\r
+    CallbackInfo = (CALLBACK_INFO *)Node;\r
+\r
+    if (FunctionData->Args.Communicate.ImageHandle == CallbackInfo->SmmImageHandle) {\r
+      CallbackInfo->CommunicationBuffer = FunctionData->Args.Communicate.CommunicationBuffer;\r
+      CallbackInfo->SourceSize          = FunctionData->Args.Communicate.SourceSize;\r
+\r
+      ///\r
+      /// The message was successfully posted.\r
+      ///\r
+      FunctionData->Status = EFI_SUCCESS;\r
+      return;\r
+    }\r
+    Node = GetNextNode (&mCallbackInfoListHead, Node);\r
+  }\r
+\r
+  FunctionData->Status = EFI_INVALID_PARAMETER;\r
 }\r
 \r
 /**\r
@@ -580,11 +986,11 @@ HelperFreePool (
   This SMI handler provides services for the SMM Base Thunk driver.\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
+  @param[in]     RegisterContext 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
+  @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
+  @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
@@ -611,21 +1017,24 @@ SmmHandlerEntry (
   FunctionData = (SMMBASE_FUNCTION_DATA *)CommBuffer;\r
 \r
   switch (FunctionData->Function) {\r
-    case SMMBASE_REGISTER:\r
+    case SmmBaseFunctionRegister:\r
       Register (FunctionData);\r
       break;\r
-    case SMMBASE_UNREGISTER:\r
+    case SmmBaseFunctionUnregister:\r
       UnRegister (FunctionData);\r
       break;\r
-    case SMMBASE_REGISTER_CALLBACK:\r
+    case SmmBaseFunctionRegisterCallback:\r
       RegisterCallback (FunctionData);\r
       break;\r
-    case SMMBASE_ALLOCATE_POOL:\r
+    case SmmBaseFunctionAllocatePool:\r
       HelperAllocatePool (FunctionData);\r
       break;\r
-    case SMMBASE_FREE_POOL:\r
+    case SmmBaseFunctionFreePool:\r
       HelperFreePool (FunctionData);\r
       break;\r
+    case SmmBaseFunctionCommunicate:\r
+      HelperCommunicate (FunctionData);\r
+      break;\r
     default:\r
       ASSERT (FALSE);\r
       FunctionData->Status = EFI_UNSUPPORTED;\r
@@ -633,6 +1042,29 @@ SmmHandlerEntry (
   return EFI_SUCCESS;\r
 }\r
 \r
+/**\r
+  Smm Ready To Lock event notification handler.\r
+\r
+  It sets a flag indicating that SMRAM has been locked.\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 handler runs successfully.\r
+ **/\r
+EFI_STATUS\r
+EFIAPI\r
+SmmReadyToLockEventNotify (\r
+  IN CONST EFI_GUID  *Protocol,\r
+  IN VOID            *Interface,\r
+  IN EFI_HANDLE      Handle\r
+  )\r
+{\r
+  mLocked = TRUE;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
 /**\r
   Entry point function of the SMM Base Helper SMM driver.\r
 \r
@@ -650,14 +1082,36 @@ SmmBaseHelperMain (
   )\r
 {\r
   EFI_STATUS  Status;\r
-  EFI_HANDLE  Handle = NULL;\r
-\r
+  EFI_MP_SERVICES_PROTOCOL   *MpServices;\r
+  EFI_HANDLE                 Handle;\r
+  UINTN                      NumberOfEnabledProcessors;\r
+  VOID                       *Registration;\r
+  \r
+  Handle = NULL;\r
   ///\r
-  /// Locate SMM CPU Protocol which is used later to update CPU Save States\r
+  /// Locate SMM CPU Protocol which is used later to retrieve/update CPU Save States\r
   ///\r
   Status = gSmst->SmmLocateProtocol (&gEfiSmmCpuProtocolGuid, NULL, (VOID **) &mSmmCpu);\r
   ASSERT_EFI_ERROR (Status);\r
 \r
+  ///\r
+  /// Locate PE32 Image Protocol which is used later to load Framework SMM driver\r
+  ///\r
+  Status = SystemTable->BootServices->LocateProtocol (&gEfiLoadPeImageProtocolGuid, NULL, (VOID **) &mLoadPe32Image);\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  //\r
+  // Get MP Services Protocol\r
+  //\r
+  Status = SystemTable->BootServices->LocateProtocol (&gEfiMpServiceProtocolGuid, NULL, (VOID **)&MpServices);\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
+  //\r
+  // Use MP Services Protocol to retrieve the number of processors and number of enabled processors\r
+  //\r
+  Status = MpServices->GetNumberOfProcessors (MpServices, &mNumberOfProcessors, &NumberOfEnabledProcessors);\r
+  ASSERT_EFI_ERROR (Status);\r
+  \r
   ///\r
   /// Interface structure of SMM BASE Helper Ready Protocol is allocated from UEFI pool\r
   /// instead of SMM pool so that SMM Base Thunk driver can access it in Non-SMM mode.\r
@@ -676,6 +1130,16 @@ SmmBaseHelperMain (
   mSmmBaseHelperReady->FrameworkSmst = mFrameworkSmst;\r
   mSmmBaseHelperReady->ServiceEntry = SmmHandlerEntry;\r
 \r
+  //\r
+  // Register SMM Ready To Lock Protocol notification\r
+  //\r
+  Status = gSmst->SmmRegisterProtocolNotify (\r
+                    &gEfiSmmReadyToLockProtocolGuid,\r
+                    SmmReadyToLockEventNotify,\r
+                    &Registration\r
+                    );\r
+  ASSERT_EFI_ERROR (Status);\r
+\r
   ///\r
   /// Register SMM Base Helper services for SMM Base Thunk driver\r
   ///\r