]> 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 347e83d01f61b5d47aac6bf5e562e814677e0fcb..08ac649b3b1857fb8630ede6d3b0a9a1659d6093 100644 (file)
 #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
@@ -39,6 +45,8 @@ typedef struct {
   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
@@ -57,9 +65,19 @@ typedef struct {
 \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
@@ -88,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
@@ -124,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
@@ -134,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
@@ -160,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_SMM_CPU_SAVE_STATE),\r
-                    (VOID **)&FrameworkSmst->CpuSaveState\r
-                    );\r
-  ASSERT_EFI_ERROR (Status);\r
-  ZeroMem (FrameworkSmst->CpuSaveState, gSmst->NumberOfCpus * sizeof (EFI_SMM_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
@@ -181,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
@@ -196,149 +617,72 @@ 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
+  } while (Status == EFI_BUFFER_TOO_SMALL);\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
-\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
@@ -351,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
@@ -413,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
@@ -436,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
-                          (UINTN)sizeof (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
@@ -490,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
-                          (UINTN)sizeof (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
@@ -534,12 +854,16 @@ CallbackThunk (
 **/\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
@@ -547,36 +871,35 @@ 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
@@ -590,11 +913,15 @@ HelperAllocatePool (
   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
@@ -607,9 +934,50 @@ HelperFreePool (
   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
@@ -618,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
@@ -649,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
@@ -671,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
@@ -688,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 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
@@ -714,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