[PATCH v3 12/15] OvmfPkg/MemEncryptSevLib: Address range encryption state interface

["Lendacky, Thomas" <thomas.lendacky@amd.com>]

From: Tom Lendacky <thomas.lendacky@amd.com>

BZ: https://bugzilla.tianocore.org/show_bug.cgi?id=3108

Update the MemEncryptSevLib library to include an interface that can
report the encryption state on a range of memory. The values will
represent the range as being unencrypted, encrypted, a mix of unencrypted
and encrypted, and error (e.g. ranges that aren't mapped).

Cc: Jordan Justen <jordan.l.justen@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Ard Biesheuvel <ard.biesheuvel@arm.com>
Cc: Brijesh Singh <brijesh.singh@amd.com>
Acked-by: Laszlo Ersek <lersek@redhat.com>
Signed-off-by: Tom Lendacky <thomas.lendacky@amd.com>
---
 OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf   |   1 +
 OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf   |   1 +
 OvmfPkg/Library/BaseMemEncryptSevLib/SecMemEncryptSevLib.inf   |   1 +
 OvmfPkg/Include/Library/MemEncryptSevLib.h                     |  33 ++++
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.h       |  35 +++-
 OvmfPkg/Library/BaseMemEncryptSevLib/Ia32/MemEncryptSevLib.c   |  31 ++-
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/MemEncryptSevLib.c    |  32 ++-
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiDxeVirtualMemory.c |  19 +-
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/SecVirtualMemory.c    |  20 ++
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.c       | 207 ++++++++++++++++++++
 10 files changed, 368 insertions(+), 12 deletions(-)

diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf b/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
index 4480e4cc7c89..8e3b8ddd5a95 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
@@ -36,6 +36,7 @@ [Sources]
 [Sources.X64]
   X64/MemEncryptSevLib.c
   X64/PeiDxeVirtualMemory.c
+  X64/VirtualMemory.c
   X64/VirtualMemory.h
 
 [Sources.IA32]
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
index 0697f1dab502..03a78c32df28 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
@@ -36,6 +36,7 @@ [Sources]
 [Sources.X64]
   X64/MemEncryptSevLib.c
   X64/PeiDxeVirtualMemory.c
+  X64/VirtualMemory.c
   X64/VirtualMemory.h
 
 [Sources.IA32]
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/SecMemEncryptSevLib.inf b/OvmfPkg/Library/BaseMemEncryptSevLib/SecMemEncryptSevLib.inf
index 7cd0111fe47b..279c38bfbc2c 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/SecMemEncryptSevLib.inf
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/SecMemEncryptSevLib.inf
@@ -35,6 +35,7 @@ [Sources]
 [Sources.X64]
   X64/MemEncryptSevLib.c
   X64/SecVirtualMemory.c
+  X64/VirtualMemory.c
   X64/VirtualMemory.h
 
 [Sources.IA32]
diff --git a/OvmfPkg/Include/Library/MemEncryptSevLib.h b/OvmfPkg/Include/Library/MemEncryptSevLib.h
index 872abe6725dc..ec470b8d0363 100644
--- a/OvmfPkg/Include/Library/MemEncryptSevLib.h
+++ b/OvmfPkg/Include/Library/MemEncryptSevLib.h
@@ -33,6 +33,16 @@ typedef struct _SEC_SEV_ES_WORK_AREA {
   UINT64   EncryptionMask;
 } SEC_SEV_ES_WORK_AREA;
 
+//
+// Memory encryption address range states.
+//
+typedef enum {
+  MemEncryptSevAddressRangeUnencrypted,
+  MemEncryptSevAddressRangeEncrypted,
+  MemEncryptSevAddressRangeMixed,
+  MemEncryptSevAddressRangeError,
+} MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE;
+
 /**
   Returns a boolean to indicate whether SEV-ES is enabled.
 
@@ -147,4 +157,27 @@ MemEncryptSevGetEncryptionMask (
   VOID
   );
 
+/**
+  Returns the encryption state of the specified virtual address range.
+
+  @param[in]  Cr3BaseAddress          Cr3 Base Address (if zero then use
+                                      current CR3)
+  @param[in]  BaseAddress             Base address to check
+  @param[in]  Length                  Length of virtual address range
+
+  @retval MemEncryptSevAddressRangeUnencrypted  Address range is mapped
+                                                unencrypted
+  @retval MemEncryptSevAddressRangeEncrypted    Address range is mapped
+                                                encrypted
+  @retval MemEncryptSevAddressRangeMixed        Address range is mapped mixed
+  @retval MemEncryptSevAddressRangeError        Address range is not mapped
+**/
+MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE
+EFIAPI
+MemEncryptSevGetAddressRangeState (
+  IN PHYSICAL_ADDRESS         Cr3BaseAddress,
+  IN PHYSICAL_ADDRESS         BaseAddress,
+  IN UINTN                    Length
+  );
+
 #endif // _MEM_ENCRYPT_SEV_LIB_H_
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.h b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.h
index 26d26cd922a4..996f94f07ebb 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.h
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.h
@@ -3,7 +3,7 @@
   Virtual Memory Management Services to set or clear the memory encryption bit
 
   Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>
-  Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
+  Copyright (c) 2017 - 2020, AMD Incorporated. All rights reserved.<BR>
 
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -178,7 +178,17 @@ typedef struct {
   UINTN           FreePages;
 } PAGE_TABLE_POOL;
 
+/**
+  Return the pagetable memory encryption mask.
 
+  @return  The pagetable memory encryption mask.
+
+**/
+UINT64
+EFIAPI
+InternalGetMemEncryptionAddressMask (
+  VOID
+  );
 
 /**
   This function clears memory encryption bit for the memory region specified by
@@ -234,4 +244,27 @@ InternalMemEncryptSevSetMemoryEncrypted (
   IN  BOOLEAN                 Flush
   );
 
+/**
+  Returns the encryption state of the specified virtual address range.
+
+  @param[in]  Cr3BaseAddress          Cr3 Base Address (if zero then use
+                                      current CR3)
+  @param[in]  BaseAddress             Base address to check
+  @param[in]  Length                  Length of virtual address range
+
+  @retval MemEncryptSevAddressRangeUnencrypted  Address range is mapped
+                                                unencrypted
+  @retval MemEncryptSevAddressRangeEncrypted    Address range is mapped
+                                                encrypted
+  @retval MemEncryptSevAddressRangeMixed        Address range is mapped mixed
+  @retval MemEncryptSevAddressRangeError        Address range is not mapped
+**/
+MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE
+EFIAPI
+InternalMemEncryptSevGetAddressRangeState (
+  IN PHYSICAL_ADDRESS         Cr3BaseAddress,
+  IN PHYSICAL_ADDRESS         BaseAddress,
+  IN UINTN                    Length
+  );
+
 #endif
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/Ia32/MemEncryptSevLib.c b/OvmfPkg/Library/BaseMemEncryptSevLib/Ia32/MemEncryptSevLib.c
index b4f6e5738e6e..12a5bf495bd7 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/Ia32/MemEncryptSevLib.c
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/Ia32/MemEncryptSevLib.c
@@ -2,7 +2,7 @@
 
   Secure Encrypted Virtualization (SEV) library helper function
 
-  Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
+  Copyright (c) 2017 - 2020, AMD Incorporated. All rights reserved.<BR>
 
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -82,3 +82,32 @@ MemEncryptSevSetPageEncMask (
   //
   return RETURN_UNSUPPORTED;
 }
+
+/**
+  Returns the encryption state of the specified virtual address range.
+
+  @param[in]  Cr3BaseAddress          Cr3 Base Address (if zero then use
+                                      current CR3)
+  @param[in]  BaseAddress             Base address to check
+  @param[in]  Length                  Length of virtual address range
+
+  @retval MemEncryptSevAddressRangeUnencrypted  Address range is mapped
+                                                unencrypted
+  @retval MemEncryptSevAddressRangeEncrypted    Address range is mapped
+                                                encrypted
+  @retval MemEncryptSevAddressRangeMixed        Address range is mapped mixed
+  @retval MemEncryptSevAddressRangeError        Address range is not mapped
+**/
+MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE
+EFIAPI
+MemEncryptSevGetAddressRangeState (
+  IN PHYSICAL_ADDRESS         Cr3BaseAddress,
+  IN PHYSICAL_ADDRESS         BaseAddress,
+  IN UINTN                    Length
+  )
+{
+  //
+  // Memory is always encrypted in 32-bit mode
+  //
+  return MemEncryptSevAddressRangeEncrypted;
+}
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/MemEncryptSevLib.c b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/MemEncryptSevLib.c
index cf0921e21464..4fea6a6be0ac 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/MemEncryptSevLib.c
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/MemEncryptSevLib.c
@@ -2,7 +2,7 @@
 
   Secure Encrypted Virtualization (SEV) library helper function
 
-  Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
+  Copyright (c) 2017 - 2020, AMD Incorporated. All rights reserved.<BR>
 
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -88,3 +88,33 @@ MemEncryptSevSetPageEncMask (
            Flush
            );
 }
+
+/**
+  Returns the encryption state of the specified virtual address range.
+
+  @param[in]  Cr3BaseAddress          Cr3 Base Address (if zero then use
+                                      current CR3)
+  @param[in]  BaseAddress             Base address to check
+  @param[in]  Length                  Length of virtual address range
+
+  @retval MemEncryptSevAddressRangeUnencrypted  Address range is mapped
+                                                unencrypted
+  @retval MemEncryptSevAddressRangeEncrypted    Address range is mapped
+                                                encrypted
+  @retval MemEncryptSevAddressRangeMixed        Address range is mapped mixed
+  @retval MemEncryptSevAddressRangeError        Address range is not mapped
+**/
+MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE
+EFIAPI
+MemEncryptSevGetAddressRangeState (
+  IN PHYSICAL_ADDRESS         Cr3BaseAddress,
+  IN PHYSICAL_ADDRESS         BaseAddress,
+  IN UINTN                    Length
+  )
+{
+  return InternalMemEncryptSevGetAddressRangeState (
+           Cr3BaseAddress,
+           BaseAddress,
+           Length
+           );
+}
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiDxeVirtualMemory.c b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiDxeVirtualMemory.c
index 3a5bab657bd7..d3455e812bd1 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiDxeVirtualMemory.c
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/PeiDxeVirtualMemory.c
@@ -28,14 +28,14 @@ typedef enum {
 } MAP_RANGE_MODE;
 
 /**
-  Get the memory encryption mask
+  Return the pagetable memory encryption mask.
 
-  @param[out]      EncryptionMask        contains the pte mask.
+  @return  The pagetable memory encryption mask.
 
 **/
-STATIC
 UINT64
-GetMemEncryptionAddressMask (
+EFIAPI
+InternalGetMemEncryptionAddressMask (
   VOID
   )
 {
@@ -200,7 +200,7 @@ Split2MPageTo4K (
 
   PageTableEntry1 = PageTableEntry;
 
-  AddressEncMask = GetMemEncryptionAddressMask ();
+  AddressEncMask = InternalGetMemEncryptionAddressMask ();
 
   ASSERT (PageTableEntry != NULL);
   ASSERT (*PageEntry2M & AddressEncMask);
@@ -286,7 +286,7 @@ SetPageTablePoolReadOnly (
   LevelSize[3] = SIZE_1GB;
   LevelSize[4] = SIZE_512GB;
 
-  AddressEncMask  = GetMemEncryptionAddressMask();
+  AddressEncMask  = InternalGetMemEncryptionAddressMask();
   PageTable       = (UINT64 *)(UINTN)PageTableBase;
   PoolUnitSize    = PAGE_TABLE_POOL_UNIT_SIZE;
 
@@ -431,7 +431,7 @@ Split1GPageTo2M (
 
   PageDirectoryEntry = AllocatePageTableMemory(1);
 
-  AddressEncMask = GetMemEncryptionAddressMask ();
+  AddressEncMask = InternalGetMemEncryptionAddressMask ();
   ASSERT (PageDirectoryEntry != NULL);
   ASSERT (*PageEntry1G & AddressEncMask);
   //
@@ -485,7 +485,7 @@ SetOrClearCBit(
 {
   UINT64      AddressEncMask;
 
-  AddressEncMask = GetMemEncryptionAddressMask ();
+  AddressEncMask = InternalGetMemEncryptionAddressMask ();
 
   if (Mode == SetCBit) {
     *PageTablePointer |= AddressEncMask;
@@ -527,6 +527,7 @@ DisableReadOnlyPageWriteProtect (
 /**
  Enable Write Protect on pages marked as read-only.
 **/
+STATIC
 VOID
 EnableReadOnlyPageWriteProtect (
   VOID
@@ -605,7 +606,7 @@ SetMemoryEncDec (
   //
   // Check if we have a valid memory encryption mask
   //
-  AddressEncMask = GetMemEncryptionAddressMask ();
+  AddressEncMask = InternalGetMemEncryptionAddressMask ();
   if (!AddressEncMask) {
     return RETURN_ACCESS_DENIED;
   }
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SecVirtualMemory.c b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SecVirtualMemory.c
index 5c337ea0b820..bca5e3febb1b 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SecVirtualMemory.c
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/SecVirtualMemory.c
@@ -13,6 +13,26 @@
 
 #include "VirtualMemory.h"
 
+/**
+  Return the pagetable memory encryption mask.
+
+  @return  The pagetable memory encryption mask.
+
+**/
+UINT64
+EFIAPI
+InternalGetMemEncryptionAddressMask (
+  VOID
+  )
+{
+  UINT64                            EncryptionMask;
+
+  EncryptionMask = MemEncryptSevGetEncryptionMask ();
+  EncryptionMask &= PAGING_1G_ADDRESS_MASK_64;
+
+  return EncryptionMask;
+}
+
 /**
   This function clears memory encryption bit for the memory region specified by
   PhysicalAddress and Length from the current page table context.
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.c b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.c
new file mode 100644
index 000000000000..36aabcf556a7
--- /dev/null
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.c
@@ -0,0 +1,207 @@
+/** @file
+
+  Virtual Memory Management Services to test an address range encryption state
+
+  Copyright (c) 2020, AMD Incorporated. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <Library/CpuLib.h>
+#include <Library/MemEncryptSevLib.h>
+
+#include "VirtualMemory.h"
+
+/**
+  Returns the (updated) address range state based upon the page table
+  entry.
+
+  @param[in]  CurrentState            The current address range state
+  @param[in]  PageDirectoryEntry      The page table entry to check
+  @param[in]  AddressEncMask          The encryption mask
+
+  @retval MemEncryptSevAddressRangeUnencrypted  Address range is mapped
+                                                unencrypted
+  @retval MemEncryptSevAddressRangeEncrypted    Address range is mapped
+                                                encrypted
+  @retval MemEncryptSevAddressRangeMixed        Address range is mapped mixed
+**/
+STATIC
+MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE
+UpdateAddressState (
+  IN MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE  CurrentState,
+  IN UINT64                               PageDirectoryEntry,
+  IN UINT64                               AddressEncMask
+  )
+{
+  if (CurrentState == MemEncryptSevAddressRangeEncrypted) {
+    if ((PageDirectoryEntry & AddressEncMask) == 0) {
+      CurrentState = MemEncryptSevAddressRangeMixed;
+    }
+  } else if (CurrentState == MemEncryptSevAddressRangeUnencrypted) {
+    if ((PageDirectoryEntry & AddressEncMask) != 0) {
+      CurrentState = MemEncryptSevAddressRangeMixed;
+    }
+  } else if (CurrentState == MemEncryptSevAddressRangeError) {
+    //
+    // First address check, set initial state
+    //
+    if ((PageDirectoryEntry & AddressEncMask) == 0) {
+      CurrentState = MemEncryptSevAddressRangeUnencrypted;
+    } else {
+      CurrentState = MemEncryptSevAddressRangeEncrypted;
+    }
+  }
+
+  return CurrentState;
+}
+
+/**
+  Returns the encryption state of the specified virtual address range.
+
+  @param[in]  Cr3BaseAddress          Cr3 Base Address (if zero then use
+                                      current CR3)
+  @param[in]  BaseAddress             Base address to check
+  @param[in]  Length                  Length of virtual address range
+
+  @retval MemEncryptSevAddressRangeUnencrypted  Address range is mapped
+                                                unencrypted
+  @retval MemEncryptSevAddressRangeEncrypted    Address range is mapped
+                                                encrypted
+  @retval MemEncryptSevAddressRangeMixed        Address range is mapped mixed
+  @retval MemEncryptSevAddressRangeError        Address range is not mapped
+**/
+MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE
+EFIAPI
+InternalMemEncryptSevGetAddressRangeState (
+  IN PHYSICAL_ADDRESS  Cr3BaseAddress,
+  IN PHYSICAL_ADDRESS  BaseAddress,
+  IN UINTN             Length
+  )
+{
+  PAGE_MAP_AND_DIRECTORY_POINTER       *PageMapLevel4Entry;
+  PAGE_MAP_AND_DIRECTORY_POINTER       *PageUpperDirectoryPointerEntry;
+  PAGE_MAP_AND_DIRECTORY_POINTER       *PageDirectoryPointerEntry;
+  PAGE_TABLE_1G_ENTRY                  *PageDirectory1GEntry;
+  PAGE_TABLE_ENTRY                     *PageDirectory2MEntry;
+  PAGE_TABLE_4K_ENTRY                  *PageTableEntry;
+  UINT64                               AddressEncMask;
+  UINT64                               PgTableMask;
+  PHYSICAL_ADDRESS                     Address;
+  PHYSICAL_ADDRESS                     AddressEnd;
+  MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE  State;
+
+  //
+  // If Cr3BaseAddress is not specified then read the current CR3
+  //
+  if (Cr3BaseAddress == 0) {
+    Cr3BaseAddress = AsmReadCr3();
+  }
+
+  AddressEncMask = MemEncryptSevGetEncryptionMask ();
+  AddressEncMask &= PAGING_1G_ADDRESS_MASK_64;
+
+  PgTableMask = AddressEncMask | EFI_PAGE_MASK;
+
+  State = MemEncryptSevAddressRangeError;
+
+  //
+  // Encryption is on a page basis, so start at the beginning of the
+  // virtual address page boundary and walk page-by-page.
+  //
+  Address = (PHYSICAL_ADDRESS) (UINTN) BaseAddress & ~EFI_PAGE_MASK;
+  AddressEnd = (PHYSICAL_ADDRESS)
+                 (UINTN) (BaseAddress + Length);
+
+  while (Address < AddressEnd) {
+    PageMapLevel4Entry = (VOID*) (Cr3BaseAddress & ~PgTableMask);
+    PageMapLevel4Entry += PML4_OFFSET (Address);
+    if (!PageMapLevel4Entry->Bits.Present) {
+      return MemEncryptSevAddressRangeError;
+    }
+
+    PageDirectory1GEntry = (VOID *) (
+                             (PageMapLevel4Entry->Bits.PageTableBaseAddress <<
+                              12) & ~PgTableMask
+                             );
+    PageDirectory1GEntry += PDP_OFFSET (Address);
+    if (!PageDirectory1GEntry->Bits.Present) {
+      return MemEncryptSevAddressRangeError;
+    }
+
+    //
+    // If the MustBe1 bit is not 1, it's not actually a 1GB entry
+    //
+    if (PageDirectory1GEntry->Bits.MustBe1) {
+      //
+      // Valid 1GB page
+      //
+      State = UpdateAddressState (
+                State,
+                PageDirectory1GEntry->Uint64,
+                AddressEncMask
+                );
+
+      Address += BIT30;
+      continue;
+    }
+
+    //
+    // Actually a PDP
+    //
+    PageUpperDirectoryPointerEntry =
+      (PAGE_MAP_AND_DIRECTORY_POINTER *) PageDirectory1GEntry;
+    PageDirectory2MEntry =
+      (VOID *) (
+        (PageUpperDirectoryPointerEntry->Bits.PageTableBaseAddress <<
+         12) & ~PgTableMask
+        );
+    PageDirectory2MEntry += PDE_OFFSET (Address);
+    if (!PageDirectory2MEntry->Bits.Present) {
+      return MemEncryptSevAddressRangeError;
+    }
+
+    //
+    // If the MustBe1 bit is not a 1, it's not a 2MB entry
+    //
+    if (PageDirectory2MEntry->Bits.MustBe1) {
+      //
+      // Valid 2MB page
+      //
+      State = UpdateAddressState (
+                State,
+                PageDirectory2MEntry->Uint64,
+                AddressEncMask
+                );
+
+      Address += BIT21;
+      continue;
+    }
+
+    //
+    // Actually a PMD
+    //
+    PageDirectoryPointerEntry =
+      (PAGE_MAP_AND_DIRECTORY_POINTER *)PageDirectory2MEntry;
+    PageTableEntry =
+      (VOID *)(
+        (PageDirectoryPointerEntry->Bits.PageTableBaseAddress <<
+         12) & ~PgTableMask
+        );
+    PageTableEntry += PTE_OFFSET (Address);
+    if (!PageTableEntry->Bits.Present) {
+      return MemEncryptSevAddressRangeError;
+    }
+
+    State = UpdateAddressState (
+              State,
+              PageTableEntry->Uint64,
+              AddressEncMask
+              );
+
+    Address += EFI_PAGE_SIZE;
+  }
+
+  return State;
+}
-- 
2.30.0