+VOID\r
+PageFaultHandlerHook (\r
+ VOID\r
+ );\r
+\r
+VOID\r
+ReadCpuSaveState (\r
+ UINTN CpuIndex,\r
+ 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
+VOID\r
+WriteCpuSaveState (\r
+ UINTN CpuIndex,\r
+ 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
+VOID\r
+ReadWriteCpuStatePage (\r
+ UINT64 PageAddress,\r
+ 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
+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
+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
+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
+UINT64 *\r
+InitCpuStatePageTable (\r
+ 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
+VOID\r
+HookCpuStateMemory (\r
+ 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