]> 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 e7976b9c0c795285aae357d22dd53c642e912dc1..08ac649b3b1857fb8630ede6d3b0a9a1659d6093 100644 (file)
@@ -25,6 +25,8 @@
 #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
@@ -69,6 +71,13 @@ EFI_SMM_BASE_HELPER_READY_PROTOCOL *mSmmBaseHelperReady;
 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
@@ -97,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
@@ -133,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
@@ -164,6 +584,7 @@ ConstructFrameworkSmst (
   FrameworkSmst->Hdr.Revision = EFI_SMM_SYSTEM_TABLE_REVISION;\r
   CopyGuid (&FrameworkSmst->EfiSmmCpuIoGuid, &mEfiSmmCpuIoGuid);\r
 \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
@@ -360,44 +781,22 @@ 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 < mNumberOfProcessors; 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
@@ -422,16 +821,11 @@ CallbackThunk (
   ///\r
   /// Save CPU Save States in case any of them was modified\r
   ///\r
-  for (CpuIndex = 0; CpuIndex < mNumberOfProcessors; 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