UefiCpuPkg/CpuDxe/CpuPageTable.c

   1 /** @file
   2   Page table management support.
   3
   4   Copyright (c) 2017 - 2019, Intel Corporation. All rights reserved.<BR>
   5   Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
   6
   7   SPDX-License-Identifier: BSD-2-Clause-Patent
   8
   9 **/
  10
  11 #include <Base.h>
  12 #include <Uefi.h>
  13 #include <Library/PeCoffGetEntryPointLib.h>
  14 #include <Library/SerialPortLib.h>
  15 #include <Library/SynchronizationLib.h>
  16 #include <Library/PrintLib.h>
  17 #include <Protocol/SmmBase2.h>
  18 #include <Register/Cpuid.h>
  19 #include <Register/Msr.h>
  20
  21 #include "CpuDxe.h"
  22 #include "CpuPageTable.h"
  23
  24 ///
  25 /// Page Table Entry
  26 ///
  27 #define IA32_PG_P                   BIT0
  28 #define IA32_PG_RW                  BIT1
  29 #define IA32_PG_U                   BIT2
  30 #define IA32_PG_WT                  BIT3
  31 #define IA32_PG_CD                  BIT4
  32 #define IA32_PG_A                   BIT5
  33 #define IA32_PG_D                   BIT6
  34 #define IA32_PG_PS                  BIT7
  35 #define IA32_PG_PAT_2M              BIT12
  36 #define IA32_PG_PAT_4K              IA32_PG_PS
  37 #define IA32_PG_PMNT                BIT62
  38 #define IA32_PG_NX                  BIT63
  39
  40 #define PAGE_ATTRIBUTE_BITS         (IA32_PG_D | IA32_PG_A | IA32_PG_U | IA32_PG_RW | IA32_PG_P)
  41 //
  42 // Bits 1, 2, 5, 6 are reserved in the IA32 PAE PDPTE
  43 // X64 PAE PDPTE does not have such restriction
  44 //
  45 #define IA32_PAE_PDPTE_ATTRIBUTE_BITS    (IA32_PG_P)
  46
  47 #define PAGE_PROGATE_BITS           (IA32_PG_NX | PAGE_ATTRIBUTE_BITS)
  48
  49 #define PAGING_4K_MASK  0xFFF
  50 #define PAGING_2M_MASK  0x1FFFFF
  51 #define PAGING_1G_MASK  0x3FFFFFFF
  52
  53 #define PAGING_PAE_INDEX_MASK  0x1FF
  54
  55 #define PAGING_4K_ADDRESS_MASK_64 0x000FFFFFFFFFF000ull
  56 #define PAGING_2M_ADDRESS_MASK_64 0x000FFFFFFFE00000ull
  57 #define PAGING_1G_ADDRESS_MASK_64 0x000FFFFFC0000000ull
  58
  59 #define MAX_PF_ENTRY_COUNT        10
  60 #define MAX_DEBUG_MESSAGE_LENGTH  0x100
  61 #define IA32_PF_EC_ID             BIT4
  62
  63 typedef enum {
  64   PageNone,
  65   Page4K,
  66   Page2M,
  67   Page1G,
  68 } PAGE_ATTRIBUTE;
  69
  70 typedef struct {
  71   PAGE_ATTRIBUTE   Attribute;
  72   UINT64           Length;
  73   UINT64           AddressMask;
  74 } PAGE_ATTRIBUTE_TABLE;
  75
  76 typedef enum {
  77   PageActionAssign,
  78   PageActionSet,
  79   PageActionClear,
  80 } PAGE_ACTION;
  81
  82 PAGE_ATTRIBUTE_TABLE mPageAttributeTable[] = {
  83   {Page4K,  SIZE_4KB, PAGING_4K_ADDRESS_MASK_64},
  84   {Page2M,  SIZE_2MB, PAGING_2M_ADDRESS_MASK_64},
  85   {Page1G,  SIZE_1GB, PAGING_1G_ADDRESS_MASK_64},
  86 };
  87
  88 PAGE_TABLE_POOL                   *mPageTablePool = NULL;
  89 BOOLEAN                           mPageTablePoolLock = FALSE;
  90 PAGE_TABLE_LIB_PAGING_CONTEXT     mPagingContext;
  91 EFI_SMM_BASE2_PROTOCOL            *mSmmBase2 = NULL;
  92
  93 //
  94 // Record the page fault exception count for one instruction execution.
  95 //
  96 UINTN                     *mPFEntryCount;
  97 UINT64                    *(*mLastPFEntryPointer)[MAX_PF_ENTRY_COUNT];
  98
  99 /**
 100  Check if current execution environment is in SMM mode or not, via
 101  EFI_SMM_BASE2_PROTOCOL.
 102
 103  This is necessary because of the fact that MdePkg\Library\SmmMemoryAllocationLib
 104  supports to free memory outside SMRAM. The library will call gBS->FreePool() or
 105  gBS->FreePages() and then SetMemorySpaceAttributes interface in turn to change
 106  memory paging attributes during free operation, if some memory related features
 107  are enabled (like Heap Guard).
 108
 109  This means that SetMemorySpaceAttributes() has chance to run in SMM mode. This
 110  will cause incorrect result because SMM mode always loads its own page tables,
 111  which are usually different from DXE. This function can be used to detect such
 112  situation and help to avoid further misoperations.
 113
 114   @retval TRUE    In SMM mode.
 115   @retval FALSE   Not in SMM mode.
 116 **/
 117 BOOLEAN
 118 IsInSmm (
 119   VOID
 120   )
 121 {
 122   BOOLEAN                 InSmm;
 123
 124   InSmm = FALSE;
 125   if (mSmmBase2 == NULL) {
 126     gBS->LocateProtocol (&gEfiSmmBase2ProtocolGuid, NULL, (VOID **)&mSmmBase2);
 127   }
 128
 129   if (mSmmBase2 != NULL) {
 130     mSmmBase2->InSmm (mSmmBase2, &InSmm);
 131   }
 132
 133   //
 134   // mSmmBase2->InSmm() can only detect if the caller is running in SMRAM
 135   // or from SMM driver. It cannot tell if the caller is running in SMM mode.
 136   // Check page table base address to guarantee that because SMM mode willl
 137   // load its own page table.
 138   //
 139   return (InSmm &&
 140           mPagingContext.ContextData.X64.PageTableBase != (UINT64)AsmReadCr3());
 141 }
 142
 143 /**
 144   Return current paging context.
 145
 146   @param[in,out]  PagingContext     The paging context.
 147 **/
 148 VOID
 149 GetCurrentPagingContext (
 150   IN OUT PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext
 151   )
 152 {
 153   UINT32                          RegEax;
 154   CPUID_EXTENDED_CPU_SIG_EDX      RegEdx;
 155   MSR_IA32_EFER_REGISTER          MsrEfer;
 156   IA32_CR4                        Cr4;
 157   IA32_CR0                        Cr0;
 158
 159   //
 160   // Don't retrieve current paging context from processor if in SMM mode.
 161   //
 162   if (!IsInSmm ()) {
 163     ZeroMem (&mPagingContext, sizeof(mPagingContext));
 164     if (sizeof(UINTN) == sizeof(UINT64)) {
 165       mPagingContext.MachineType = IMAGE_FILE_MACHINE_X64;
 166     } else {
 167       mPagingContext.MachineType = IMAGE_FILE_MACHINE_I386;
 168     }
 169
 170     Cr0.UintN = AsmReadCr0 ();
 171     Cr4.UintN = AsmReadCr4 ();
 172
 173     if (Cr0.Bits.PG != 0) {
 174       mPagingContext.ContextData.X64.PageTableBase = (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64);
 175     } else {
 176       mPagingContext.ContextData.X64.PageTableBase = 0;
 177     }
 178     if (Cr0.Bits.WP  != 0) {
 179       mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_WP_ENABLE;
 180     }
 181     if (Cr4.Bits.PSE != 0) {
 182       mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PSE;
 183     }
 184     if (Cr4.Bits.PAE != 0) {
 185       mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE;
 186     }
 187     if (Cr4.Bits.LA57 != 0) {
 188       mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_5_LEVEL;
 189     }
 190
 191     AsmCpuid (CPUID_EXTENDED_FUNCTION, &RegEax, NULL, NULL, NULL);
 192     if (RegEax >= CPUID_EXTENDED_CPU_SIG) {
 193       AsmCpuid (CPUID_EXTENDED_CPU_SIG, NULL, NULL, NULL, &RegEdx.Uint32);
 194
 195       if (RegEdx.Bits.NX != 0) {
 196         // XD supported
 197         MsrEfer.Uint64 = AsmReadMsr64(MSR_CORE_IA32_EFER);
 198         if (MsrEfer.Bits.NXE != 0) {
 199           // XD activated
 200           mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED;
 201         }
 202       }
 203
 204       if (RegEdx.Bits.Page1GB != 0) {
 205         mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAGE_1G_SUPPORT;
 206       }
 207     }
 208   }
 209
 210   //
 211   // This can avoid getting SMM paging context if in SMM mode. We cannot assume
 212   // SMM mode shares the same paging context as DXE.
 213   //
 214   CopyMem (PagingContext, &mPagingContext, sizeof (mPagingContext));
 215 }
 216
 217 /**
 218   Return length according to page attributes.
 219
 220   @param[in]  PageAttributes   The page attribute of the page entry.
 221
 222   @return The length of page entry.
 223 **/
 224 UINTN
 225 PageAttributeToLength (
 226   IN PAGE_ATTRIBUTE  PageAttribute
 227   )
 228 {
 229   UINTN  Index;
 230   for (Index = 0; Index < sizeof(mPageAttributeTable)/sizeof(mPageAttributeTable[0]); Index++) {
 231     if (PageAttribute == mPageAttributeTable[Index].Attribute) {
 232       return (UINTN)mPageAttributeTable[Index].Length;
 233     }
 234   }
 235   return 0;
 236 }
 237
 238 /**
 239   Return address mask according to page attributes.
 240
 241   @param[in]  PageAttributes   The page attribute of the page entry.
 242
 243   @return The address mask of page entry.
 244 **/
 245 UINTN
 246 PageAttributeToMask (
 247   IN PAGE_ATTRIBUTE  PageAttribute
 248   )
 249 {
 250   UINTN  Index;
 251   for (Index = 0; Index < sizeof(mPageAttributeTable)/sizeof(mPageAttributeTable[0]); Index++) {
 252     if (PageAttribute == mPageAttributeTable[Index].Attribute) {
 253       return (UINTN)mPageAttributeTable[Index].AddressMask;
 254     }
 255   }
 256   return 0;
 257 }
 258
 259 /**
 260   Return page table entry to match the address.
 261
 262   @param[in]  PagingContext     The paging context.
 263   @param[in]  Address           The address to be checked.
 264   @param[out] PageAttributes    The page attribute of the page entry.
 265
 266   @return The page entry.
 267 **/
 268 VOID *
 269 GetPageTableEntry (
 270   IN  PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext,
 271   IN  PHYSICAL_ADDRESS                  Address,
 272   OUT PAGE_ATTRIBUTE                    *PageAttribute
 273   )
 274 {
 275   UINTN                 Index1;
 276   UINTN                 Index2;
 277   UINTN                 Index3;
 278   UINTN                 Index4;
 279   UINTN                 Index5;
 280   UINT64                *L1PageTable;
 281   UINT64                *L2PageTable;
 282   UINT64                *L3PageTable;
 283   UINT64                *L4PageTable;
 284   UINT64                *L5PageTable;
 285   UINT64                AddressEncMask;
 286
 287   ASSERT (PagingContext != NULL);
 288
 289   Index5 = ((UINTN)RShiftU64 (Address, 48)) & PAGING_PAE_INDEX_MASK;
 290   Index4 = ((UINTN)RShiftU64 (Address, 39)) & PAGING_PAE_INDEX_MASK;
 291   Index3 = ((UINTN)Address >> 30) & PAGING_PAE_INDEX_MASK;
 292   Index2 = ((UINTN)Address >> 21) & PAGING_PAE_INDEX_MASK;
 293   Index1 = ((UINTN)Address >> 12) & PAGING_PAE_INDEX_MASK;
 294
 295   // Make sure AddressEncMask is contained to smallest supported address field.
 296   //
 297   AddressEncMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) & PAGING_1G_ADDRESS_MASK_64;
 298
 299   if (PagingContext->MachineType == IMAGE_FILE_MACHINE_X64) {
 300     if ((PagingContext->ContextData.X64.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_5_LEVEL) != 0) {
 301       L5PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.X64.PageTableBase;
 302       if (L5PageTable[Index5] == 0) {
 303         *PageAttribute = PageNone;
 304         return NULL;
 305       }
 306
 307       L4PageTable = (UINT64 *)(UINTN)(L5PageTable[Index5] & ~AddressEncMask & PAGING_4K_ADDRESS_MASK_64);
 308     } else {
 309       L4PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.X64.PageTableBase;
 310     }
 311     if (L4PageTable[Index4] == 0) {
 312       *PageAttribute = PageNone;
 313       return NULL;
 314     }
 315
 316     L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & ~AddressEncMask & PAGING_4K_ADDRESS_MASK_64);
 317   } else {
 318     ASSERT((PagingContext->ContextData.Ia32.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) != 0);
 319     L3PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.Ia32.PageTableBase;
 320   }
 321   if (L3PageTable[Index3] == 0) {
 322     *PageAttribute = PageNone;
 323     return NULL;
 324   }
 325   if ((L3PageTable[Index3] & IA32_PG_PS) != 0) {
 326     // 1G
 327     *PageAttribute = Page1G;
 328     return &L3PageTable[Index3];
 329   }
 330
 331   L2PageTable = (UINT64 *)(UINTN)(L3PageTable[Index3] & ~AddressEncMask & PAGING_4K_ADDRESS_MASK_64);
 332   if (L2PageTable[Index2] == 0) {
 333     *PageAttribute = PageNone;
 334     return NULL;
 335   }
 336   if ((L2PageTable[Index2] & IA32_PG_PS) != 0) {
 337     // 2M
 338     *PageAttribute = Page2M;
 339     return &L2PageTable[Index2];
 340   }
 341
 342   // 4k
 343   L1PageTable = (UINT64 *)(UINTN)(L2PageTable[Index2] & ~AddressEncMask & PAGING_4K_ADDRESS_MASK_64);
 344   if ((L1PageTable[Index1] == 0) && (Address != 0)) {
 345     *PageAttribute = PageNone;
 346     return NULL;
 347   }
 348   *PageAttribute = Page4K;
 349   return &L1PageTable[Index1];
 350 }
 351
 352 /**
 353   Return memory attributes of page entry.
 354
 355   @param[in]  PageEntry        The page entry.
 356
 357   @return Memory attributes of page entry.
 358 **/
 359 UINT64
 360 GetAttributesFromPageEntry (
 361   IN  UINT64                            *PageEntry
 362   )
 363 {
 364   UINT64  Attributes;
 365   Attributes = 0;
 366   if ((*PageEntry & IA32_PG_P) == 0) {
 367     Attributes |= EFI_MEMORY_RP;
 368   }
 369   if ((*PageEntry & IA32_PG_RW) == 0) {
 370     Attributes |= EFI_MEMORY_RO;
 371   }
 372   if ((*PageEntry & IA32_PG_NX) != 0) {
 373     Attributes |= EFI_MEMORY_XP;
 374   }
 375   return Attributes;
 376 }
 377
 378 /**
 379   Modify memory attributes of page entry.
 380
 381   @param[in]  PagingContext    The paging context.
 382   @param[in]  PageEntry        The page entry.
 383   @param[in]  Attributes       The bit mask of attributes to modify for the memory region.
 384   @param[in]  PageAction       The page action.
 385   @param[out] IsModified       TRUE means page table modified. FALSE means page table not modified.
 386 **/
 387 VOID
 388 ConvertPageEntryAttribute (
 389   IN  PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext,
 390   IN  UINT64                            *PageEntry,
 391   IN  UINT64                            Attributes,
 392   IN  PAGE_ACTION                       PageAction,
 393   OUT BOOLEAN                           *IsModified
 394   )
 395 {
 396   UINT64  CurrentPageEntry;
 397   UINT64  NewPageEntry;
 398
 399   CurrentPageEntry = *PageEntry;
 400   NewPageEntry = CurrentPageEntry;
 401   if ((Attributes & EFI_MEMORY_RP) != 0) {
 402     switch (PageAction) {
 403     case PageActionAssign:
 404     case PageActionSet:
 405       NewPageEntry &= ~(UINT64)IA32_PG_P;
 406       break;
 407     case PageActionClear:
 408       NewPageEntry |= IA32_PG_P;
 409       break;
 410     }
 411   } else {
 412     switch (PageAction) {
 413     case PageActionAssign:
 414       NewPageEntry |= IA32_PG_P;
 415       break;
 416     case PageActionSet:
 417     case PageActionClear:
 418       break;
 419     }
 420   }
 421   if ((Attributes & EFI_MEMORY_RO) != 0) {
 422     switch (PageAction) {
 423     case PageActionAssign:
 424     case PageActionSet:
 425       NewPageEntry &= ~(UINT64)IA32_PG_RW;
 426       break;
 427     case PageActionClear:
 428       NewPageEntry |= IA32_PG_RW;
 429       break;
 430     }
 431   } else {
 432     switch (PageAction) {
 433     case PageActionAssign:
 434       NewPageEntry |= IA32_PG_RW;
 435       break;
 436     case PageActionSet:
 437     case PageActionClear:
 438       break;
 439     }
 440   }
 441   if ((PagingContext->ContextData.Ia32.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED) != 0) {
 442     if ((Attributes & EFI_MEMORY_XP) != 0) {
 443       switch (PageAction) {
 444       case PageActionAssign:
 445       case PageActionSet:
 446         NewPageEntry |= IA32_PG_NX;
 447         break;
 448       case PageActionClear:
 449         NewPageEntry &= ~IA32_PG_NX;
 450         break;
 451       }
 452     } else {
 453       switch (PageAction) {
 454       case PageActionAssign:
 455         NewPageEntry &= ~IA32_PG_NX;
 456         break;
 457       case PageActionSet:
 458       case PageActionClear:
 459         break;
 460       }
 461     }
 462   }
 463   *PageEntry = NewPageEntry;
 464   if (CurrentPageEntry != NewPageEntry) {
 465     *IsModified = TRUE;
 466     DEBUG ((DEBUG_VERBOSE, "ConvertPageEntryAttribute 0x%lx", CurrentPageEntry));
 467     DEBUG ((DEBUG_VERBOSE, "->0x%lx\n", NewPageEntry));
 468   } else {
 469     *IsModified = FALSE;
 470   }
 471 }
 472
 473 /**
 474   This function returns if there is need to split page entry.
 475
 476   @param[in]  BaseAddress      The base address to be checked.
 477   @param[in]  Length           The length to be checked.
 478   @param[in]  PageEntry        The page entry to be checked.
 479   @param[in]  PageAttribute    The page attribute of the page entry.
 480
 481   @retval SplitAttributes on if there is need to split page entry.
 482 **/
 483 PAGE_ATTRIBUTE
 484 NeedSplitPage (
 485   IN  PHYSICAL_ADDRESS                  BaseAddress,
 486   IN  UINT64                            Length,
 487   IN  UINT64                            *PageEntry,
 488   IN  PAGE_ATTRIBUTE                    PageAttribute
 489   )
 490 {
 491   UINT64                PageEntryLength;
 492
 493   PageEntryLength = PageAttributeToLength (PageAttribute);
 494
 495   if (((BaseAddress & (PageEntryLength - 1)) == 0) && (Length >= PageEntryLength)) {
 496     return PageNone;
 497   }
 498
 499   if (((BaseAddress & PAGING_2M_MASK) != 0) || (Length < SIZE_2MB)) {
 500     return Page4K;
 501   }
 502
 503   return Page2M;
 504 }
 505
 506 /**
 507   This function splits one page entry to small page entries.
 508
 509   @param[in]  PageEntry         The page entry to be splitted.
 510   @param[in]  PageAttribute     The page attribute of the page entry.
 511   @param[in]  SplitAttribute    How to split the page entry.
 512   @param[in]  AllocatePagesFunc If page split is needed, this function is used to allocate more pages.
 513
 514   @retval RETURN_SUCCESS            The page entry is splitted.
 515   @retval RETURN_UNSUPPORTED        The page entry does not support to be splitted.
 516   @retval RETURN_OUT_OF_RESOURCES   No resource to split page entry.
 517 **/
 518 RETURN_STATUS
 519 SplitPage (
 520   IN  UINT64                            *PageEntry,
 521   IN  PAGE_ATTRIBUTE                    PageAttribute,
 522   IN  PAGE_ATTRIBUTE                    SplitAttribute,
 523   IN  PAGE_TABLE_LIB_ALLOCATE_PAGES     AllocatePagesFunc
 524   )
 525 {
 526   UINT64   BaseAddress;
 527   UINT64   *NewPageEntry;
 528   UINTN    Index;
 529   UINT64   AddressEncMask;
 530
 531   ASSERT (PageAttribute == Page2M || PageAttribute == Page1G);
 532
 533   ASSERT (AllocatePagesFunc != NULL);
 534
 535   // Make sure AddressEncMask is contained to smallest supported address field.
 536   //
 537   AddressEncMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) & PAGING_1G_ADDRESS_MASK_64;
 538
 539   if (PageAttribute == Page2M) {
 540     //
 541     // Split 2M to 4K
 542     //
 543     ASSERT (SplitAttribute == Page4K);
 544     if (SplitAttribute == Page4K) {
 545       NewPageEntry = AllocatePagesFunc (1);
 546       DEBUG ((DEBUG_VERBOSE, "Split - 0x%x\n", NewPageEntry));
 547       if (NewPageEntry == NULL) {
 548         return RETURN_OUT_OF_RESOURCES;
 549       }
 550       BaseAddress = *PageEntry & ~AddressEncMask & PAGING_2M_ADDRESS_MASK_64;
 551       for (Index = 0; Index < SIZE_4KB / sizeof(UINT64); Index++) {
 552         NewPageEntry[Index] = (BaseAddress + SIZE_4KB * Index) | AddressEncMask | ((*PageEntry) & PAGE_PROGATE_BITS);
 553       }
 554       (*PageEntry) = (UINT64)(UINTN)NewPageEntry | AddressEncMask | ((*PageEntry) & PAGE_ATTRIBUTE_BITS);
 555       return RETURN_SUCCESS;
 556     } else {
 557       return RETURN_UNSUPPORTED;
 558     }
 559   } else if (PageAttribute == Page1G) {
 560     //
 561     // Split 1G to 2M
 562     // No need support 1G->4K directly, we should use 1G->2M, then 2M->4K to get more compact page table.
 563     //
 564     ASSERT (SplitAttribute == Page2M || SplitAttribute == Page4K);
 565     if ((SplitAttribute == Page2M || SplitAttribute == Page4K)) {
 566       NewPageEntry = AllocatePagesFunc (1);
 567       DEBUG ((DEBUG_VERBOSE, "Split - 0x%x\n", NewPageEntry));
 568       if (NewPageEntry == NULL) {
 569         return RETURN_OUT_OF_RESOURCES;
 570       }
 571       BaseAddress = *PageEntry & ~AddressEncMask  & PAGING_1G_ADDRESS_MASK_64;
 572       for (Index = 0; Index < SIZE_4KB / sizeof(UINT64); Index++) {
 573         NewPageEntry[Index] = (BaseAddress + SIZE_2MB * Index) | AddressEncMask | IA32_PG_PS | ((*PageEntry) & PAGE_PROGATE_BITS);
 574       }
 575       (*PageEntry) = (UINT64)(UINTN)NewPageEntry | AddressEncMask | ((*PageEntry) & PAGE_ATTRIBUTE_BITS);
 576       return RETURN_SUCCESS;
 577     } else {
 578       return RETURN_UNSUPPORTED;
 579     }
 580   } else {
 581     return RETURN_UNSUPPORTED;
 582   }
 583 }
 584
 585 /**
 586  Check the WP status in CR0 register. This bit is used to lock or unlock write
 587  access to pages marked as read-only.
 588
 589   @retval TRUE    Write protection is enabled.
 590   @retval FALSE   Write protection is disabled.
 591 **/
 592 BOOLEAN
 593 IsReadOnlyPageWriteProtected (
 594   VOID
 595   )
 596 {
 597   IA32_CR0  Cr0;
 598   //
 599   // To avoid unforseen consequences, don't touch paging settings in SMM mode
 600   // in this driver.
 601   //
 602   if (!IsInSmm ()) {
 603     Cr0.UintN = AsmReadCr0 ();
 604     return (BOOLEAN) (Cr0.Bits.WP != 0);
 605   }
 606   return FALSE;
 607 }
 608
 609 /**
 610  Disable Write Protect on pages marked as read-only.
 611 **/
 612 VOID
 613 DisableReadOnlyPageWriteProtect (
 614   VOID
 615   )
 616 {
 617   IA32_CR0  Cr0;
 618   //
 619   // To avoid unforseen consequences, don't touch paging settings in SMM mode
 620   // in this driver.
 621   //
 622   if (!IsInSmm ()) {
 623     Cr0.UintN = AsmReadCr0 ();
 624     Cr0.Bits.WP = 0;
 625     AsmWriteCr0 (Cr0.UintN);
 626   }
 627 }
 628
 629 /**
 630  Enable Write Protect on pages marked as read-only.
 631 **/
 632 VOID
 633 EnableReadOnlyPageWriteProtect (
 634   VOID
 635   )
 636 {
 637   IA32_CR0  Cr0;
 638   //
 639   // To avoid unforseen consequences, don't touch paging settings in SMM mode
 640   // in this driver.
 641   //
 642   if (!IsInSmm ()) {
 643     Cr0.UintN = AsmReadCr0 ();
 644     Cr0.Bits.WP = 1;
 645     AsmWriteCr0 (Cr0.UintN);
 646   }
 647 }
 648
 649 /**
 650   This function modifies the page attributes for the memory region specified by BaseAddress and
 651   Length from their current attributes to the attributes specified by Attributes.
 652
 653   Caller should make sure BaseAddress and Length is at page boundary.
 654
 655   @param[in]  PagingContext     The paging context. NULL means get page table from current CPU context.
 656   @param[in]  BaseAddress       The physical address that is the start address of a memory region.
 657   @param[in]  Length            The size in bytes of the memory region.
 658   @param[in]  Attributes        The bit mask of attributes to modify for the memory region.
 659   @param[in]  PageAction        The page action.
 660   @param[in]  AllocatePagesFunc If page split is needed, this function is used to allocate more pages.
 661                                 NULL mean page split is unsupported.
 662   @param[out] IsSplitted        TRUE means page table splitted. FALSE means page table not splitted.
 663   @param[out] IsModified        TRUE means page table modified. FALSE means page table not modified.
 664
 665   @retval RETURN_SUCCESS           The attributes were modified for the memory region.
 666   @retval RETURN_ACCESS_DENIED     The attributes for the memory resource range specified by
 667                                    BaseAddress and Length cannot be modified.
 668   @retval RETURN_INVALID_PARAMETER Length is zero.
 669                                    Attributes specified an illegal combination of attributes that
 670                                    cannot be set together.
 671   @retval RETURN_OUT_OF_RESOURCES  There are not enough system resources to modify the attributes of
 672                                    the memory resource range.
 673   @retval RETURN_UNSUPPORTED       The processor does not support one or more bytes of the memory
 674                                    resource range specified by BaseAddress and Length.
 675                                    The bit mask of attributes is not support for the memory resource
 676                                    range specified by BaseAddress and Length.
 677 **/
 678 RETURN_STATUS
 679 ConvertMemoryPageAttributes (
 680   IN  PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext OPTIONAL,
 681   IN  PHYSICAL_ADDRESS                  BaseAddress,
 682   IN  UINT64                            Length,
 683   IN  UINT64                            Attributes,
 684   IN  PAGE_ACTION                       PageAction,
 685   IN  PAGE_TABLE_LIB_ALLOCATE_PAGES     AllocatePagesFunc OPTIONAL,
 686   OUT BOOLEAN                           *IsSplitted,  OPTIONAL
 687   OUT BOOLEAN                           *IsModified   OPTIONAL
 688   )
 689 {
 690   PAGE_TABLE_LIB_PAGING_CONTEXT     CurrentPagingContext;
 691   UINT64                            *PageEntry;
 692   PAGE_ATTRIBUTE                    PageAttribute;
 693   UINTN                             PageEntryLength;
 694   PAGE_ATTRIBUTE                    SplitAttribute;
 695   RETURN_STATUS                     Status;
 696   BOOLEAN                           IsEntryModified;
 697   BOOLEAN                           IsWpEnabled;
 698
 699   if ((BaseAddress & (SIZE_4KB - 1)) != 0) {
 700     DEBUG ((DEBUG_ERROR, "BaseAddress(0x%lx) is not aligned!\n", BaseAddress));
 701     return EFI_UNSUPPORTED;
 702   }
 703   if ((Length & (SIZE_4KB - 1)) != 0) {
 704     DEBUG ((DEBUG_ERROR, "Length(0x%lx) is not aligned!\n", Length));
 705     return EFI_UNSUPPORTED;
 706   }
 707   if (Length == 0) {
 708     DEBUG ((DEBUG_ERROR, "Length is 0!\n"));
 709     return RETURN_INVALID_PARAMETER;
 710   }
 711
 712   if ((Attributes & ~(EFI_MEMORY_RP | EFI_MEMORY_RO | EFI_MEMORY_XP)) != 0) {
 713     DEBUG ((DEBUG_ERROR, "Attributes(0x%lx) has unsupported bit\n", Attributes));
 714     return EFI_UNSUPPORTED;
 715   }
 716
 717   if (PagingContext == NULL) {
 718     GetCurrentPagingContext (&CurrentPagingContext);
 719   } else {
 720     CopyMem (&CurrentPagingContext, PagingContext, sizeof(CurrentPagingContext));
 721   }
 722   switch(CurrentPagingContext.MachineType) {
 723   case IMAGE_FILE_MACHINE_I386:
 724     if (CurrentPagingContext.ContextData.Ia32.PageTableBase == 0) {
 725       if (Attributes == 0) {
 726         return EFI_SUCCESS;
 727       } else {
 728         DEBUG ((DEBUG_ERROR, "PageTable is 0!\n"));
 729         return EFI_UNSUPPORTED;
 730       }
 731     }
 732     if ((CurrentPagingContext.ContextData.Ia32.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) == 0) {
 733       DEBUG ((DEBUG_ERROR, "Non-PAE Paging!\n"));
 734       return EFI_UNSUPPORTED;
 735     }
 736     if ((BaseAddress + Length) > BASE_4GB) {
 737       DEBUG ((DEBUG_ERROR, "Beyond 4GB memory in 32-bit mode!\n"));
 738       return EFI_UNSUPPORTED;
 739     }
 740     break;
 741   case IMAGE_FILE_MACHINE_X64:
 742     ASSERT (CurrentPagingContext.ContextData.X64.PageTableBase != 0);
 743     break;
 744   default:
 745     ASSERT(FALSE);
 746     return EFI_UNSUPPORTED;
 747     break;
 748   }
 749
 750 //  DEBUG ((DEBUG_ERROR, "ConvertMemoryPageAttributes(%x) - %016lx, %016lx, %02lx\n", IsSet, BaseAddress, Length, Attributes));
 751
 752   if (IsSplitted != NULL) {
 753     *IsSplitted = FALSE;
 754   }
 755   if (IsModified != NULL) {
 756     *IsModified = FALSE;
 757   }
 758   if (AllocatePagesFunc == NULL) {
 759     AllocatePagesFunc = AllocatePageTableMemory;
 760   }
 761
 762   //
 763   // Make sure that the page table is changeable.
 764   //
 765   IsWpEnabled = IsReadOnlyPageWriteProtected ();
 766   if (IsWpEnabled) {
 767     DisableReadOnlyPageWriteProtect ();
 768   }
 769
 770   //
 771   // Below logic is to check 2M/4K page to make sure we donot waist memory.
 772   //
 773   Status = EFI_SUCCESS;
 774   while (Length != 0) {
 775     PageEntry = GetPageTableEntry (&CurrentPagingContext, BaseAddress, &PageAttribute);
 776     if (PageEntry == NULL) {
 777       Status = RETURN_UNSUPPORTED;
 778       goto Done;
 779     }
 780     PageEntryLength = PageAttributeToLength (PageAttribute);
 781     SplitAttribute = NeedSplitPage (BaseAddress, Length, PageEntry, PageAttribute);
 782     if (SplitAttribute == PageNone) {
 783       ConvertPageEntryAttribute (&CurrentPagingContext, PageEntry, Attributes, PageAction, &IsEntryModified);
 784       if (IsEntryModified) {
 785         if (IsModified != NULL) {
 786           *IsModified = TRUE;
 787         }
 788       }
 789       //
 790       // Convert success, move to next
 791       //
 792       BaseAddress += PageEntryLength;
 793       Length -= PageEntryLength;
 794     } else {
 795       if (AllocatePagesFunc == NULL) {
 796         Status = RETURN_UNSUPPORTED;
 797         goto Done;
 798       }
 799       Status = SplitPage (PageEntry, PageAttribute, SplitAttribute, AllocatePagesFunc);
 800       if (RETURN_ERROR (Status)) {
 801         Status = RETURN_UNSUPPORTED;
 802         goto Done;
 803       }
 804       if (IsSplitted != NULL) {
 805         *IsSplitted = TRUE;
 806       }
 807       if (IsModified != NULL) {
 808         *IsModified = TRUE;
 809       }
 810       //
 811       // Just split current page
 812       // Convert success in next around
 813       //
 814     }
 815   }
 816
 817 Done:
 818   //
 819   // Restore page table write protection, if any.
 820   //
 821   if (IsWpEnabled) {
 822     EnableReadOnlyPageWriteProtect ();
 823   }
 824   return Status;
 825 }
 826
 827 /**
 828   This function assigns the page attributes for the memory region specified by BaseAddress and
 829   Length from their current attributes to the attributes specified by Attributes.
 830
 831   Caller should make sure BaseAddress and Length is at page boundary.
 832
 833   Caller need guarentee the TPL <= TPL_NOTIFY, if there is split page request.
 834
 835   @param[in]  PagingContext     The paging context. NULL means get page table from current CPU context.
 836   @param[in]  BaseAddress       The physical address that is the start address of a memory region.
 837   @param[in]  Length            The size in bytes of the memory region.
 838   @param[in]  Attributes        The bit mask of attributes to set for the memory region.
 839   @param[in]  AllocatePagesFunc If page split is needed, this function is used to allocate more pages.
 840                                 NULL mean page split is unsupported.
 841
 842   @retval RETURN_SUCCESS           The attributes were cleared for the memory region.
 843   @retval RETURN_ACCESS_DENIED     The attributes for the memory resource range specified by
 844                                    BaseAddress and Length cannot be modified.
 845   @retval RETURN_INVALID_PARAMETER Length is zero.
 846                                    Attributes specified an illegal combination of attributes that
 847                                    cannot be set together.
 848   @retval RETURN_OUT_OF_RESOURCES  There are not enough system resources to modify the attributes of
 849                                    the memory resource range.
 850   @retval RETURN_UNSUPPORTED       The processor does not support one or more bytes of the memory
 851                                    resource range specified by BaseAddress and Length.
 852                                    The bit mask of attributes is not support for the memory resource
 853                                    range specified by BaseAddress and Length.
 854 **/
 855 RETURN_STATUS
 856 EFIAPI
 857 AssignMemoryPageAttributes (
 858   IN  PAGE_TABLE_LIB_PAGING_CONTEXT     *PagingContext OPTIONAL,
 859   IN  PHYSICAL_ADDRESS                  BaseAddress,
 860   IN  UINT64                            Length,
 861   IN  UINT64                            Attributes,
 862   IN  PAGE_TABLE_LIB_ALLOCATE_PAGES     AllocatePagesFunc OPTIONAL
 863   )
 864 {
 865   RETURN_STATUS  Status;
 866   BOOLEAN        IsModified;
 867   BOOLEAN        IsSplitted;
 868
 869 //  DEBUG((DEBUG_INFO, "AssignMemoryPageAttributes: 0x%lx - 0x%lx (0x%lx)\n", BaseAddress, Length, Attributes));
 870   Status = ConvertMemoryPageAttributes (PagingContext, BaseAddress, Length, Attributes, PageActionAssign, AllocatePagesFunc, &IsSplitted, &IsModified);
 871   if (!EFI_ERROR(Status)) {
 872     if ((PagingContext == NULL) && IsModified) {
 873       //
 874       // Flush TLB as last step.
 875       //
 876       // Note: Since APs will always init CR3 register in HLT loop mode or do
 877       // TLB flush in MWAIT loop mode, there's no need to flush TLB for them
 878       // here.
 879       //
 880       CpuFlushTlb();
 881     }
 882   }
 883
 884   return Status;
 885 }
 886
 887 /**
 888  Check if Execute Disable feature is enabled or not.
 889 **/
 890 BOOLEAN
 891 IsExecuteDisableEnabled (
 892   VOID
 893   )
 894 {
 895   MSR_CORE_IA32_EFER_REGISTER    MsrEfer;
 896
 897   MsrEfer.Uint64 = AsmReadMsr64 (MSR_IA32_EFER);
 898   return (MsrEfer.Bits.NXE == 1);
 899 }
 900
 901 /**
 902   Update GCD memory space attributes according to current page table setup.
 903 **/
 904 VOID
 905 RefreshGcdMemoryAttributesFromPaging (
 906   VOID
 907   )
 908 {
 909   EFI_STATUS                          Status;
 910   UINTN                               NumberOfDescriptors;
 911   EFI_GCD_MEMORY_SPACE_DESCRIPTOR     *MemorySpaceMap;
 912   PAGE_TABLE_LIB_PAGING_CONTEXT       PagingContext;
 913   PAGE_ATTRIBUTE                      PageAttribute;
 914   UINT64                              *PageEntry;
 915   UINT64                              PageLength;
 916   UINT64                              MemorySpaceLength;
 917   UINT64                              Length;
 918   UINT64                              BaseAddress;
 919   UINT64                              PageStartAddress;
 920   UINT64                              Attributes;
 921   UINT64                              Capabilities;
 922   UINT64                              NewAttributes;
 923   UINTN                               Index;
 924
 925   //
 926   // Assuming that memory space map returned is sorted already; otherwise sort
 927   // them in the order of lowest address to highest address.
 928   //
 929   Status = gDS->GetMemorySpaceMap (&NumberOfDescriptors, &MemorySpaceMap);
 930   ASSERT_EFI_ERROR (Status);
 931
 932   GetCurrentPagingContext (&PagingContext);
 933
 934   Attributes      = 0;
 935   NewAttributes   = 0;
 936   BaseAddress     = 0;
 937   PageLength      = 0;
 938
 939   if (IsExecuteDisableEnabled ()) {
 940     Capabilities = EFI_MEMORY_RO | EFI_MEMORY_RP | EFI_MEMORY_XP;
 941   } else {
 942     Capabilities = EFI_MEMORY_RO | EFI_MEMORY_RP;
 943   }
 944
 945   for (Index = 0; Index < NumberOfDescriptors; Index++) {
 946     if (MemorySpaceMap[Index].GcdMemoryType == EfiGcdMemoryTypeNonExistent) {
 947       continue;
 948     }
 949
 950     //
 951     // Sync the actual paging related capabilities back to GCD service first.
 952     // As a side effect (good one), this can also help to avoid unnecessary
 953     // memory map entries due to the different capabilities of the same type
 954     // memory, such as multiple RT_CODE and RT_DATA entries in memory map,
 955     // which could cause boot failure of some old Linux distro (before v4.3).
 956     //
 957     Status = gDS->SetMemorySpaceCapabilities (
 958                     MemorySpaceMap[Index].BaseAddress,
 959                     MemorySpaceMap[Index].Length,
 960                     MemorySpaceMap[Index].Capabilities | Capabilities
 961                     );
 962     if (EFI_ERROR (Status)) {
 963       //
 964       // If we cannot udpate the capabilities, we cannot update its
 965       // attributes either. So just simply skip current block of memory.
 966       //
 967       DEBUG ((
 968         DEBUG_WARN,
 969         "Failed to update capability: [%lu] %016lx - %016lx (%016lx -> %016lx)\r\n",
 970         (UINT64)Index, MemorySpaceMap[Index].BaseAddress,
 971         MemorySpaceMap[Index].BaseAddress + MemorySpaceMap[Index].Length - 1,
 972         MemorySpaceMap[Index].Capabilities,
 973         MemorySpaceMap[Index].Capabilities | Capabilities
 974         ));
 975       continue;
 976     }
 977
 978     if (MemorySpaceMap[Index].BaseAddress >= (BaseAddress + PageLength)) {
 979       //
 980       // Current memory space starts at a new page. Resetting PageLength will
 981       // trigger a retrieval of page attributes at new address.
 982       //
 983       PageLength = 0;
 984     } else {
 985       //
 986       // In case current memory space is not adjacent to last one
 987       //
 988       PageLength -= (MemorySpaceMap[Index].BaseAddress - BaseAddress);
 989     }
 990
 991     //
 992     // Sync actual page attributes to GCD
 993     //
 994     BaseAddress       = MemorySpaceMap[Index].BaseAddress;
 995     MemorySpaceLength = MemorySpaceMap[Index].Length;
 996     while (MemorySpaceLength > 0) {
 997       if (PageLength == 0) {
 998         PageEntry = GetPageTableEntry (&PagingContext, BaseAddress, &PageAttribute);
 999         if (PageEntry == NULL) {
1000           break;
1001         }
1002
1003         //
1004         // Note current memory space might start in the middle of a page
1005         //
1006         PageStartAddress  = (*PageEntry) & (UINT64)PageAttributeToMask(PageAttribute);
1007         PageLength        = PageAttributeToLength (PageAttribute) - (BaseAddress - PageStartAddress);
1008         Attributes        = GetAttributesFromPageEntry (PageEntry);
1009       }
1010
1011       Length = MIN (PageLength, MemorySpaceLength);
1012       if (Attributes != (MemorySpaceMap[Index].Attributes &
1013                          EFI_MEMORY_PAGETYPE_MASK)) {
1014         NewAttributes = (MemorySpaceMap[Index].Attributes &
1015                          ~EFI_MEMORY_PAGETYPE_MASK) | Attributes;
1016         Status = gDS->SetMemorySpaceAttributes (
1017                         BaseAddress,
1018                         Length,
1019                         NewAttributes
1020                         );
1021         ASSERT_EFI_ERROR (Status);
1022         DEBUG ((
1023           DEBUG_VERBOSE,
1024           "Updated memory space attribute: [%lu] %016lx - %016lx (%016lx -> %016lx)\r\n",
1025           (UINT64)Index, BaseAddress, BaseAddress + Length - 1,
1026           MemorySpaceMap[Index].Attributes,
1027           NewAttributes
1028           ));
1029       }
1030
1031       PageLength        -= Length;
1032       MemorySpaceLength -= Length;
1033       BaseAddress       += Length;
1034     }
1035   }
1036
1037   FreePool (MemorySpaceMap);
1038 }
1039
1040 /**
1041   Initialize a buffer pool for page table use only.
1042
1043   To reduce the potential split operation on page table, the pages reserved for
1044   page table should be allocated in the times of PAGE_TABLE_POOL_UNIT_PAGES and
1045   at the boundary of PAGE_TABLE_POOL_ALIGNMENT. So the page pool is always
1046   initialized with number of pages greater than or equal to the given PoolPages.
1047
1048   Once the pages in the pool are used up, this method should be called again to
1049   reserve at least another PAGE_TABLE_POOL_UNIT_PAGES. Usually this won't happen
1050   often in practice.
1051
1052   @param[in] PoolPages      The least page number of the pool to be created.
1053
1054   @retval TRUE    The pool is initialized successfully.
1055   @retval FALSE   The memory is out of resource.
1056 **/
1057 BOOLEAN
1058 InitializePageTablePool (
1059   IN  UINTN                           PoolPages
1060   )
1061 {
1062   VOID                      *Buffer;
1063   BOOLEAN                   IsModified;
1064
1065   //
1066   // Do not allow re-entrance.
1067   //
1068   if (mPageTablePoolLock) {
1069     return FALSE;
1070   }
1071
1072   mPageTablePoolLock = TRUE;
1073   IsModified = FALSE;
1074
1075   //
1076   // Always reserve at least PAGE_TABLE_POOL_UNIT_PAGES, including one page for
1077   // header.
1078   //
1079   PoolPages += 1;   // Add one page for header.
1080   PoolPages = ((PoolPages - 1) / PAGE_TABLE_POOL_UNIT_PAGES + 1) *
1081               PAGE_TABLE_POOL_UNIT_PAGES;
1082   Buffer = AllocateAlignedPages (PoolPages, PAGE_TABLE_POOL_ALIGNMENT);
1083   if (Buffer == NULL) {
1084     DEBUG ((DEBUG_ERROR, "ERROR: Out of aligned pages\r\n"));
1085     goto Done;
1086   }
1087
1088   DEBUG ((
1089     DEBUG_INFO,
1090     "Paging: added %lu pages to page table pool\r\n",
1091     (UINT64)PoolPages
1092     ));
1093
1094   //
1095   // Link all pools into a list for easier track later.
1096   //
1097   if (mPageTablePool == NULL) {
1098     mPageTablePool = Buffer;
1099     mPageTablePool->NextPool = mPageTablePool;
1100   } else {
1101     ((PAGE_TABLE_POOL *)Buffer)->NextPool = mPageTablePool->NextPool;
1102     mPageTablePool->NextPool = Buffer;
1103     mPageTablePool = Buffer;
1104   }
1105
1106   //
1107   // Reserve one page for pool header.
1108   //
1109   mPageTablePool->FreePages  = PoolPages - 1;
1110   mPageTablePool->Offset = EFI_PAGES_TO_SIZE (1);
1111
1112   //
1113   // Mark the whole pool pages as read-only.
1114   //
1115   ConvertMemoryPageAttributes (
1116     NULL,
1117     (PHYSICAL_ADDRESS)(UINTN)Buffer,
1118     EFI_PAGES_TO_SIZE (PoolPages),
1119     EFI_MEMORY_RO,
1120     PageActionSet,
1121     AllocatePageTableMemory,
1122     NULL,
1123     &IsModified
1124     );
1125   ASSERT (IsModified == TRUE);
1126
1127 Done:
1128   mPageTablePoolLock = FALSE;
1129   return IsModified;
1130 }
1131
1132 /**
1133   This API provides a way to allocate memory for page table.
1134
1135   This API can be called more than once to allocate memory for page tables.
1136
1137   Allocates the number of 4KB pages and returns a pointer to the allocated
1138   buffer. The buffer returned is aligned on a 4KB boundary.
1139
1140   If Pages is 0, then NULL is returned.
1141   If there is not enough memory remaining to satisfy the request, then NULL is
1142   returned.
1143
1144   @param  Pages                 The number of 4 KB pages to allocate.
1145
1146   @return A pointer to the allocated buffer or NULL if allocation fails.
1147
1148 **/
1149 VOID *
1150 EFIAPI
1151 AllocatePageTableMemory (
1152   IN UINTN           Pages
1153   )
1154 {
1155   VOID                            *Buffer;
1156
1157   if (Pages == 0) {
1158     return NULL;
1159   }
1160
1161   //
1162   // Renew the pool if necessary.
1163   //
1164   if (mPageTablePool == NULL ||
1165       Pages > mPageTablePool->FreePages) {
1166     if (!InitializePageTablePool (Pages)) {
1167       return NULL;
1168     }
1169   }
1170
1171   Buffer = (UINT8 *)mPageTablePool + mPageTablePool->Offset;
1172
1173   mPageTablePool->Offset     += EFI_PAGES_TO_SIZE (Pages);
1174   mPageTablePool->FreePages  -= Pages;
1175
1176   return Buffer;
1177 }
1178
1179 /**
1180   Special handler for #DB exception, which will restore the page attributes
1181   (not-present). It should work with #PF handler which will set pages to
1182   'present'.
1183
1184   @param ExceptionType  Exception type.
1185   @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
1186
1187 **/
1188 VOID
1189 EFIAPI
1190 DebugExceptionHandler (
1191   IN EFI_EXCEPTION_TYPE   ExceptionType,
1192   IN EFI_SYSTEM_CONTEXT   SystemContext
1193   )
1194 {
1195   UINTN     CpuIndex;
1196   UINTN     PFEntry;
1197   BOOLEAN   IsWpEnabled;
1198
1199   MpInitLibWhoAmI (&CpuIndex);
1200
1201   //
1202   // Clear last PF entries
1203   //
1204   IsWpEnabled = IsReadOnlyPageWriteProtected ();
1205   if (IsWpEnabled) {
1206     DisableReadOnlyPageWriteProtect ();
1207   }
1208
1209   for (PFEntry = 0; PFEntry < mPFEntryCount[CpuIndex]; PFEntry++) {
1210     if (mLastPFEntryPointer[CpuIndex][PFEntry] != NULL) {
1211       *mLastPFEntryPointer[CpuIndex][PFEntry] &= ~(UINT64)IA32_PG_P;
1212     }
1213   }
1214
1215   if (IsWpEnabled) {
1216     EnableReadOnlyPageWriteProtect ();
1217   }
1218
1219   //
1220   // Reset page fault exception count for next page fault.
1221   //
1222   mPFEntryCount[CpuIndex] = 0;
1223
1224   //
1225   // Flush TLB
1226   //
1227   CpuFlushTlb ();
1228
1229   //
1230   // Clear TF in EFLAGS
1231   //
1232   if (mPagingContext.MachineType == IMAGE_FILE_MACHINE_I386) {
1233     SystemContext.SystemContextIa32->Eflags &= (UINT32)~BIT8;
1234   } else {
1235     SystemContext.SystemContextX64->Rflags &= (UINT64)~BIT8;
1236   }
1237 }
1238
1239 /**
1240   Special handler for #PF exception, which will set the pages which caused
1241   #PF to be 'present'. The attribute of those pages should be restored in
1242   the subsequent #DB handler.
1243
1244   @param ExceptionType  Exception type.
1245   @param SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
1246
1247 **/
1248 VOID
1249 EFIAPI
1250 PageFaultExceptionHandler (
1251   IN EFI_EXCEPTION_TYPE   ExceptionType,
1252   IN EFI_SYSTEM_CONTEXT   SystemContext
1253   )
1254 {
1255   EFI_STATUS                      Status;
1256   UINT64                          PFAddress;
1257   PAGE_TABLE_LIB_PAGING_CONTEXT   PagingContext;
1258   PAGE_ATTRIBUTE                  PageAttribute;
1259   UINT64                          Attributes;
1260   UINT64                          *PageEntry;
1261   UINTN                           Index;
1262   UINTN                           CpuIndex;
1263   UINTN                           PageNumber;
1264   BOOLEAN                         NonStopMode;
1265
1266   PFAddress = AsmReadCr2 () & ~EFI_PAGE_MASK;
1267   if (PFAddress < BASE_4KB) {
1268     NonStopMode = NULL_DETECTION_NONSTOP_MODE ? TRUE : FALSE;
1269   } else {
1270     NonStopMode = HEAP_GUARD_NONSTOP_MODE ? TRUE : FALSE;
1271   }
1272
1273   if (NonStopMode) {
1274     MpInitLibWhoAmI (&CpuIndex);
1275     GetCurrentPagingContext (&PagingContext);
1276     //
1277     // Memory operation cross page boundary, like "rep mov" instruction, will
1278     // cause infinite loop between this and Debug Trap handler. We have to make
1279     // sure that current page and the page followed are both in PRESENT state.
1280     //
1281     PageNumber = 2;
1282     while (PageNumber > 0) {
1283       PageEntry = GetPageTableEntry (&PagingContext, PFAddress, &PageAttribute);
1284       ASSERT(PageEntry != NULL);
1285
1286       if (PageEntry != NULL) {
1287         Attributes = GetAttributesFromPageEntry (PageEntry);
1288         if ((Attributes & EFI_MEMORY_RP) != 0) {
1289           Attributes &= ~EFI_MEMORY_RP;
1290           Status = AssignMemoryPageAttributes (&PagingContext, PFAddress,
1291                                                EFI_PAGE_SIZE, Attributes, NULL);
1292           if (!EFI_ERROR(Status)) {
1293             Index = mPFEntryCount[CpuIndex];
1294             //
1295             // Re-retrieve page entry because above calling might update page
1296             // table due to table split.
1297             //
1298             PageEntry = GetPageTableEntry (&PagingContext, PFAddress, &PageAttribute);
1299             mLastPFEntryPointer[CpuIndex][Index++] = PageEntry;
1300             mPFEntryCount[CpuIndex] = Index;
1301           }
1302         }
1303       }
1304
1305       PFAddress += EFI_PAGE_SIZE;
1306       --PageNumber;
1307     }
1308   }
1309
1310   //
1311   // Initialize the serial port before dumping.
1312   //
1313   SerialPortInitialize ();
1314   //
1315   // Display ExceptionType, CPU information and Image information
1316   //
1317   DumpCpuContext (ExceptionType, SystemContext);
1318   if (NonStopMode) {
1319     //
1320     // Set TF in EFLAGS
1321     //
1322     if (mPagingContext.MachineType == IMAGE_FILE_MACHINE_I386) {
1323       SystemContext.SystemContextIa32->Eflags |= (UINT32)BIT8;
1324     } else {
1325       SystemContext.SystemContextX64->Rflags |= (UINT64)BIT8;
1326     }
1327   } else {
1328     CpuDeadLoop ();
1329   }
1330 }
1331
1332 /**
1333   Initialize the Page Table lib.
1334 **/
1335 VOID
1336 InitializePageTableLib (
1337   VOID
1338   )
1339 {
1340   PAGE_TABLE_LIB_PAGING_CONTEXT     CurrentPagingContext;
1341
1342   GetCurrentPagingContext (&CurrentPagingContext);
1343
1344   //
1345   // Reserve memory of page tables for future uses, if paging is enabled.
1346   //
1347   if (CurrentPagingContext.ContextData.X64.PageTableBase != 0 &&
1348       (CurrentPagingContext.ContextData.Ia32.Attributes &
1349        PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE) != 0) {
1350     DisableReadOnlyPageWriteProtect ();
1351     InitializePageTablePool (1);
1352     EnableReadOnlyPageWriteProtect ();
1353   }
1354
1355   if (HEAP_GUARD_NONSTOP_MODE || NULL_DETECTION_NONSTOP_MODE) {
1356     mPFEntryCount = (UINTN *)AllocateZeroPool (sizeof (UINTN) * mNumberOfProcessors);
1357     ASSERT (mPFEntryCount != NULL);
1358
1359     mLastPFEntryPointer = (UINT64 *(*)[MAX_PF_ENTRY_COUNT])
1360                           AllocateZeroPool (sizeof (mLastPFEntryPointer[0]) * mNumberOfProcessors);
1361     ASSERT (mLastPFEntryPointer != NULL);
1362   }
1363
1364   DEBUG ((DEBUG_INFO, "CurrentPagingContext:\n", CurrentPagingContext.MachineType));
1365   DEBUG ((DEBUG_INFO, "  MachineType   - 0x%x\n", CurrentPagingContext.MachineType));
1366   DEBUG ((DEBUG_INFO, "  PageTableBase - 0x%x\n", CurrentPagingContext.ContextData.X64.PageTableBase));
1367   DEBUG ((DEBUG_INFO, "  Attributes    - 0x%x\n", CurrentPagingContext.ContextData.X64.Attributes));
1368
1369   return ;
1370 }
1371