From e58aa474329ef63a64150b7625c56b48a00e31da Mon Sep 17 00:00:00 2001
From: "Ni, Ray" <ray.ni@intel.com>
Date: Thu, 1 Aug 2019 17:58:26 +0800
Subject: [PATCH] UefiCpuPkg/CpuDxe: Support parsing 5-level page table

REF: https://bugzilla.tianocore.org/show_bug.cgi?id=2008

Signed-off-by: Ray Ni <ray.ni@intel.com>
Reviewed-by: Eric Dong <eric.dong@intel.com>
Regression-tested-by: Laszlo Ersek <lersek@redhat.com>
Signed-off-by: Eric Dong <eric.dong@intel.com>
---
 UefiCpuPkg/CpuDxe/CpuPageTable.c | 18 +++++++++++++++++-
 UefiCpuPkg/CpuDxe/CpuPageTable.h |  3 ++-
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/UefiCpuPkg/CpuDxe/CpuPageTable.c b/UefiCpuPkg/CpuDxe/CpuPageTable.c
index 16a2528b55..36ce90d66c 100644
--- a/UefiCpuPkg/CpuDxe/CpuPageTable.c
+++ b/UefiCpuPkg/CpuDxe/CpuPageTable.c
@@ -184,6 +184,9 @@ GetCurrentPagingContext (
     if (Cr4.Bits.PAE != 0) {
       mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE;
     }
+    if (Cr4.Bits.LA57 != 0) {
+      mPagingContext.ContextData.Ia32.Attributes |= PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_5_LEVEL;
+    }
 
     AsmCpuid (CPUID_EXTENDED_FUNCTION, &RegEax, NULL, NULL, NULL);
     if (RegEax >= CPUID_EXTENDED_CPU_SIG) {
@@ -273,14 +276,17 @@ GetPageTableEntry (
   UINTN                 Index2;
   UINTN                 Index3;
   UINTN                 Index4;
+  UINTN                 Index5;
   UINT64                *L1PageTable;
   UINT64                *L2PageTable;
   UINT64                *L3PageTable;
   UINT64                *L4PageTable;
+  UINT64                *L5PageTable;
   UINT64                AddressEncMask;
 
   ASSERT (PagingContext != NULL);
 
+  Index5 = ((UINTN)RShiftU64 (Address, 48)) & PAGING_PAE_INDEX_MASK;
   Index4 = ((UINTN)RShiftU64 (Address, 39)) & PAGING_PAE_INDEX_MASK;
   Index3 = ((UINTN)Address >> 30) & PAGING_PAE_INDEX_MASK;
   Index2 = ((UINTN)Address >> 21) & PAGING_PAE_INDEX_MASK;
@@ -291,7 +297,17 @@ GetPageTableEntry (
   AddressEncMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) & PAGING_1G_ADDRESS_MASK_64;
 
   if (PagingContext->MachineType == IMAGE_FILE_MACHINE_X64) {
-    L4PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.X64.PageTableBase;
+    if ((PagingContext->ContextData.X64.Attributes & PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_5_LEVEL) != 0) {
+      L5PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.X64.PageTableBase;
+      if (L5PageTable[Index5] == 0) {
+        *PageAttribute = PageNone;
+        return NULL;
+      }
+
+      L4PageTable = (UINT64 *)(UINTN)(L5PageTable[Index5] & ~AddressEncMask & PAGING_4K_ADDRESS_MASK_64);
+    } else {
+      L4PageTable = (UINT64 *)(UINTN)PagingContext->ContextData.X64.PageTableBase;
+    }
     if (L4PageTable[Index4] == 0) {
       *PageAttribute = PageNone;
       return NULL;
diff --git a/UefiCpuPkg/CpuDxe/CpuPageTable.h b/UefiCpuPkg/CpuDxe/CpuPageTable.h
index 02d62f2b14..f845956f73 100644
--- a/UefiCpuPkg/CpuDxe/CpuPageTable.h
+++ b/UefiCpuPkg/CpuDxe/CpuPageTable.h
@@ -1,7 +1,7 @@
 /** @file
   Page table management header file.
 
-  Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2017 - 2019, Intel Corporation. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -14,6 +14,7 @@
 #define PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PSE              BIT0
 #define PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE              BIT1
 #define PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAGE_1G_SUPPORT  BIT2
+#define PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_5_LEVEL          BIT3
 #define PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_WP_ENABLE        BIT30
 #define PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_XD_ACTIVATED     BIT31
 // Other bits are reserved for future use
-- 
2.39.2