+/**\r
+ Set one page of page table pool memory to be read-only.\r
+\r
+ @param[in] PageTableBase Base address of page table (CR3).\r
+ @param[in] Address Start address of a page to be set as read-only.\r
+ @param[in] Level4Paging Level 4 paging flag.\r
+\r
+**/\r
+VOID\r
+SetPageTablePoolReadOnly (\r
+ IN UINTN PageTableBase,\r
+ IN EFI_PHYSICAL_ADDRESS Address,\r
+ IN BOOLEAN Level4Paging\r
+ )\r
+{\r
+ UINTN Index;\r
+ UINTN EntryIndex;\r
+ UINT64 AddressEncMask;\r
+ EFI_PHYSICAL_ADDRESS PhysicalAddress;\r
+ UINT64 *PageTable;\r
+ UINT64 *NewPageTable;\r
+ UINT64 PageAttr;\r
+ UINT64 LevelSize[5];\r
+ UINT64 LevelMask[5];\r
+ UINTN LevelShift[5];\r
+ UINTN Level;\r
+ UINT64 PoolUnitSize;\r
+\r
+ ASSERT (PageTableBase != 0);\r
+\r
+ //\r
+ // Since the page table is always from page table pool, which is always\r
+ // located at the boundary of PcdPageTablePoolAlignment, we just need to\r
+ // set the whole pool unit to be read-only.\r
+ //\r
+ Address = Address & PAGE_TABLE_POOL_ALIGN_MASK;\r
+\r
+ LevelShift[1] = PAGING_L1_ADDRESS_SHIFT;\r
+ LevelShift[2] = PAGING_L2_ADDRESS_SHIFT;\r
+ LevelShift[3] = PAGING_L3_ADDRESS_SHIFT;\r
+ LevelShift[4] = PAGING_L4_ADDRESS_SHIFT;\r
+\r
+ LevelMask[1] = PAGING_4K_ADDRESS_MASK_64;\r
+ LevelMask[2] = PAGING_2M_ADDRESS_MASK_64;\r
+ LevelMask[3] = PAGING_1G_ADDRESS_MASK_64;\r
+ LevelMask[4] = PAGING_1G_ADDRESS_MASK_64;\r
+\r
+ LevelSize[1] = SIZE_4KB;\r
+ LevelSize[2] = SIZE_2MB;\r
+ LevelSize[3] = SIZE_1GB;\r
+ LevelSize[4] = SIZE_512GB;\r
+\r
+ AddressEncMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) &\r
+ PAGING_1G_ADDRESS_MASK_64;\r
+ PageTable = (UINT64 *)(UINTN)PageTableBase;\r
+ PoolUnitSize = PAGE_TABLE_POOL_UNIT_SIZE;\r
+\r
+ for (Level = (Level4Paging) ? 4 : 3; Level > 0; --Level) {\r
+ Index = ((UINTN)RShiftU64 (Address, LevelShift[Level]));\r
+ Index &= PAGING_PAE_INDEX_MASK;\r
+\r
+ PageAttr = PageTable[Index];\r
+ if ((PageAttr & IA32_PG_PS) == 0) {\r
+ //\r
+ // Go to next level of table.\r
+ //\r
+ PageTable = (UINT64 *)(UINTN)(PageAttr & ~AddressEncMask &\r
+ PAGING_4K_ADDRESS_MASK_64);\r
+ continue;\r
+ }\r
+\r
+ if (PoolUnitSize >= LevelSize[Level]) {\r
+ //\r
+ // Clear R/W bit if current page granularity is not larger than pool unit\r
+ // size.\r
+ //\r
+ if ((PageAttr & IA32_PG_RW) != 0) {\r
+ while (PoolUnitSize > 0) {\r
+ //\r
+ // PAGE_TABLE_POOL_UNIT_SIZE and PAGE_TABLE_POOL_ALIGNMENT are fit in\r
+ // one page (2MB). Then we don't need to update attributes for pages\r
+ // crossing page directory. ASSERT below is for that purpose.\r
+ //\r
+ ASSERT (Index < EFI_PAGE_SIZE/sizeof (UINT64));\r
+\r
+ PageTable[Index] &= ~(UINT64)IA32_PG_RW;\r
+ PoolUnitSize -= LevelSize[Level];\r
+\r
+ ++Index;\r
+ }\r
+ }\r
+\r
+ break;\r
+ } else {\r
+ //\r
+ // The smaller granularity of page must be needed.\r
+ //\r
+ ASSERT (Level > 1);\r
+\r
+ NewPageTable = AllocatePageTableMemory (1);\r
+ ASSERT (NewPageTable != NULL);\r
+\r
+ PhysicalAddress = PageAttr & LevelMask[Level];\r
+ for (EntryIndex = 0;\r
+ EntryIndex < EFI_PAGE_SIZE/sizeof (UINT64);\r
+ ++EntryIndex)\r
+ {\r
+ NewPageTable[EntryIndex] = PhysicalAddress | AddressEncMask |\r
+ IA32_PG_P | IA32_PG_RW;\r
+ if (Level > 2) {\r
+ NewPageTable[EntryIndex] |= IA32_PG_PS;\r
+ }\r
+\r
+ PhysicalAddress += LevelSize[Level - 1];\r
+ }\r
+\r
+ PageTable[Index] = (UINT64)(UINTN)NewPageTable | AddressEncMask |\r
+ IA32_PG_P | IA32_PG_RW;\r
+ PageTable = NewPageTable;\r
+ }\r
+ }\r
+}\r
+\r
+/**\r
+ Prevent the memory pages used for page table from been overwritten.\r
+\r
+ @param[in] PageTableBase Base address of page table (CR3).\r
+ @param[in] Level4Paging Level 4 paging flag.\r
+\r
+**/\r
+VOID\r
+EnablePageTableProtection (\r
+ IN UINTN PageTableBase,\r
+ IN BOOLEAN Level4Paging\r
+ )\r
+{\r
+ PAGE_TABLE_POOL *HeadPool;\r
+ PAGE_TABLE_POOL *Pool;\r
+ UINT64 PoolSize;\r
+ EFI_PHYSICAL_ADDRESS Address;\r
+\r
+ if (mPageTablePool == NULL) {\r
+ return;\r
+ }\r
+\r
+ //\r
+ // No need to clear CR0.WP since PageTableBase has't been written to CR3 yet.\r
+ // SetPageTablePoolReadOnly might update mPageTablePool. It's safer to\r
+ // remember original one in advance.\r
+ //\r
+ HeadPool = mPageTablePool;\r
+ Pool = HeadPool;\r
+ do {\r
+ Address = (EFI_PHYSICAL_ADDRESS)(UINTN)Pool;\r
+ PoolSize = Pool->Offset + EFI_PAGES_TO_SIZE (Pool->FreePages);\r
+\r
+ //\r
+ // The size of one pool must be multiple of PAGE_TABLE_POOL_UNIT_SIZE, which\r
+ // is one of page size of the processor (2MB by default). Let's apply the\r
+ // protection to them one by one.\r
+ //\r
+ while (PoolSize > 0) {\r
+ SetPageTablePoolReadOnly (PageTableBase, Address, Level4Paging);\r
+ Address += PAGE_TABLE_POOL_UNIT_SIZE;\r
+ PoolSize -= PAGE_TABLE_POOL_UNIT_SIZE;\r
+ }\r
+\r
+ Pool = Pool->NextPool;\r
+ } while (Pool != HeadPool);\r
+\r
+ //\r
+ // Enable write protection, after page table attribute updated.\r
+ //\r
+ AsmWriteCr0 (AsmReadCr0 () | CR0_WP);\r
+}\r
+\r