]> git.proxmox.com Git - mirror_edk2.git/blobdiff - OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.c
OvmfPkg: Apply uncrustify changes
[mirror_edk2.git] / OvmfPkg / Library / BaseMemEncryptSevLib / X64 / VirtualMemory.c
index aed92127629f789542a1ab3afe6c9aa121154f55..2b759cafbfa7aa96946a018a992f5d22a1d7db2a 100644 (file)
 /** @file\r
 \r
-  Virtual Memory Management Services to set or clear the memory encryption bit\r
+  Virtual Memory Management Services to test an address range encryption state\r
 \r
-  Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>\r
-  Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>\r
+  Copyright (c) 2020, AMD Incorporated. All rights reserved.<BR>\r
 \r
-  This program and the accompanying materials are licensed and made available\r
-  under the terms and conditions of the BSD License which accompanies this\r
-  distribution.  The full text of the license may be found at\r
-  http://opensource.org/licenses/bsd-license.php\r
-\r
-  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS, WITHOUT\r
-  WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
-\r
-  Code is derived from MdeModulePkg/Core/DxeIplPeim/X64/VirtualMemory.c\r
+  SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
 #include <Library/CpuLib.h>\r
-#include <Register/Cpuid.h>\r
-#include <Register/Amd/Cpuid.h>\r
+#include <Library/MemEncryptSevLib.h>\r
 \r
 #include "VirtualMemory.h"\r
 \r
-STATIC BOOLEAN mAddressEncMaskChecked = FALSE;\r
-STATIC UINT64  mAddressEncMask;\r
-STATIC PAGE_TABLE_POOL   *mPageTablePool = NULL;\r
-\r
-typedef enum {\r
-   SetCBit,\r
-   ClearCBit\r
-} MAP_RANGE_MODE;\r
-\r
-/**\r
-  Get the memory encryption mask\r
-\r
-  @param[out]      EncryptionMask        contains the pte mask.\r
-\r
-**/\r
-STATIC\r
-UINT64\r
-GetMemEncryptionAddressMask (\r
-  VOID\r
-  )\r
-{\r
-  UINT64                            EncryptionMask;\r
-  CPUID_MEMORY_ENCRYPTION_INFO_EBX  Ebx;\r
-\r
-  if (mAddressEncMaskChecked) {\r
-    return mAddressEncMask;\r
-  }\r
-\r
-  //\r
-  // CPUID Fn8000_001F[EBX] Bit 0:5 (memory encryption bit position)\r
-  //\r
-  AsmCpuid (CPUID_MEMORY_ENCRYPTION_INFO, NULL, &Ebx.Uint32, NULL, NULL);\r
-  EncryptionMask = LShiftU64 (1, Ebx.Bits.PtePosBits);\r
-\r
-  mAddressEncMask = EncryptionMask & PAGING_1G_ADDRESS_MASK_64;\r
-  mAddressEncMaskChecked = TRUE;\r
-\r
-  return mAddressEncMask;\r
-}\r
-\r
-/**\r
-  Initialize a buffer pool for page table use only.\r
-\r
-  To reduce the potential split operation on page table, the pages reserved for\r
-  page table should be allocated in the times of PAGE_TABLE_POOL_UNIT_PAGES and\r
-  at the boundary of PAGE_TABLE_POOL_ALIGNMENT. So the page pool is always\r
-  initialized with number of pages greater than or equal to the given\r
-  PoolPages.\r
-\r
-  Once the pages in the pool are used up, this method should be called again to\r
-  reserve at least another PAGE_TABLE_POOL_UNIT_PAGES. Usually this won't\r
-  happen often in practice.\r
-\r
-  @param[in] PoolPages      The least page number of the pool to be created.\r
-\r
-  @retval TRUE    The pool is initialized successfully.\r
-  @retval FALSE   The memory is out of resource.\r
-**/\r
-STATIC\r
-BOOLEAN\r
-InitializePageTablePool (\r
-  IN  UINTN                           PoolPages\r
-  )\r
-{\r
-  VOID                      *Buffer;\r
-\r
-  //\r
-  // Always reserve at least PAGE_TABLE_POOL_UNIT_PAGES, including one page for\r
-  // header.\r
-  //\r
-  PoolPages += 1;   // Add one page for header.\r
-  PoolPages = ((PoolPages - 1) / PAGE_TABLE_POOL_UNIT_PAGES + 1) *\r
-              PAGE_TABLE_POOL_UNIT_PAGES;\r
-  Buffer = AllocateAlignedPages (PoolPages, PAGE_TABLE_POOL_ALIGNMENT);\r
-  if (Buffer == NULL) {\r
-    DEBUG ((DEBUG_ERROR, "ERROR: Out of aligned pages\r\n"));\r
-    return FALSE;\r
-  }\r
-\r
-  //\r
-  // Link all pools into a list for easier track later.\r
-  //\r
-  if (mPageTablePool == NULL) {\r
-    mPageTablePool = Buffer;\r
-    mPageTablePool->NextPool = mPageTablePool;\r
-  } else {\r
-    ((PAGE_TABLE_POOL *)Buffer)->NextPool = mPageTablePool->NextPool;\r
-    mPageTablePool->NextPool = Buffer;\r
-    mPageTablePool = Buffer;\r
-  }\r
-\r
-  //\r
-  // Reserve one page for pool header.\r
-  //\r
-  mPageTablePool->FreePages  = PoolPages - 1;\r
-  mPageTablePool->Offset = EFI_PAGES_TO_SIZE (1);\r
-\r
-  return TRUE;\r
-}\r
-\r
-/**\r
-  This API provides a way to allocate memory for page table.\r
-\r
-  This API can be called more than once to allocate memory for page tables.\r
-\r
-  Allocates the number of 4KB pages and returns a pointer to the allocated\r
-  buffer. The buffer returned is aligned on a 4KB boundary.\r
-\r
-  If Pages is 0, then NULL is returned.\r
-  If there is not enough memory remaining to satisfy the request, then NULL is\r
-  returned.\r
-\r
-  @param  Pages                 The number of 4 KB pages to allocate.\r
-\r
-  @return A pointer to the allocated buffer or NULL if allocation fails.\r
-\r
-**/\r
-STATIC\r
-VOID *\r
-EFIAPI\r
-AllocatePageTableMemory (\r
-  IN UINTN           Pages\r
-  )\r
-{\r
-  VOID                            *Buffer;\r
-\r
-  if (Pages == 0) {\r
-    return NULL;\r
-  }\r
-\r
-  //\r
-  // Renew the pool if necessary.\r
-  //\r
-  if (mPageTablePool == NULL ||\r
-      Pages > mPageTablePool->FreePages) {\r
-    if (!InitializePageTablePool (Pages)) {\r
-      return NULL;\r
-    }\r
-  }\r
-\r
-  Buffer = (UINT8 *)mPageTablePool + mPageTablePool->Offset;\r
-\r
-  mPageTablePool->Offset     += EFI_PAGES_TO_SIZE (Pages);\r
-  mPageTablePool->FreePages  -= Pages;\r
-\r
-  DEBUG ((\r
-    DEBUG_VERBOSE,\r
-    "%a:%a: Buffer=0x%Lx Pages=%ld\n",\r
-    gEfiCallerBaseName,\r
-    __FUNCTION__,\r
-    Buffer,\r
-    Pages\r
-    ));\r
-\r
-  return Buffer;\r
-}\r
-\r
-\r
-/**\r
-  Split 2M page to 4K.\r
-\r
-  @param[in]      PhysicalAddress       Start physical address the 2M page\r
-                                        covered.\r
-  @param[in, out] PageEntry2M           Pointer to 2M page entry.\r
-  @param[in]      StackBase             Stack base address.\r
-  @param[in]      StackSize             Stack size.\r
-\r
-**/\r
-STATIC\r
-VOID\r
-Split2MPageTo4K (\r
-  IN        PHYSICAL_ADDRESS               PhysicalAddress,\r
-  IN  OUT   UINT64                        *PageEntry2M,\r
-  IN        PHYSICAL_ADDRESS               StackBase,\r
-  IN        UINTN                          StackSize\r
-  )\r
-{\r
-  PHYSICAL_ADDRESS                  PhysicalAddress4K;\r
-  UINTN                             IndexOfPageTableEntries;\r
-  PAGE_TABLE_4K_ENTRY               *PageTableEntry, *PageTableEntry1;\r
-  UINT64                            AddressEncMask;\r
-\r
-  PageTableEntry = AllocatePageTableMemory(1);\r
-\r
-  PageTableEntry1 = PageTableEntry;\r
-\r
-  AddressEncMask = GetMemEncryptionAddressMask ();\r
-\r
-  ASSERT (PageTableEntry != NULL);\r
-  ASSERT (*PageEntry2M & AddressEncMask);\r
-\r
-  PhysicalAddress4K = PhysicalAddress;\r
-  for (IndexOfPageTableEntries = 0;\r
-       IndexOfPageTableEntries < 512;\r
-       (IndexOfPageTableEntries++,\r
-        PageTableEntry++,\r
-        PhysicalAddress4K += SIZE_4KB)) {\r
-    //\r
-    // Fill in the Page Table entries\r
-    //\r
-    PageTableEntry->Uint64 = (UINT64) PhysicalAddress4K | AddressEncMask;\r
-    PageTableEntry->Bits.ReadWrite = 1;\r
-    PageTableEntry->Bits.Present = 1;\r
-    if ((PhysicalAddress4K >= StackBase) &&\r
-        (PhysicalAddress4K < StackBase + StackSize)) {\r
-      //\r
-      // Set Nx bit for stack.\r
-      //\r
-      PageTableEntry->Bits.Nx = 1;\r
-    }\r
-  }\r
-\r
-  //\r
-  // Fill in 2M page entry.\r
-  //\r
-  *PageEntry2M = ((UINT64)(UINTN)PageTableEntry1 |\r
-                  IA32_PG_P | IA32_PG_RW | AddressEncMask);\r
-}\r
-\r
 /**\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
+  Returns the (updated) address range state based upon the page table\r
+  entry.\r
+\r
+  @param[in]  CurrentState            The current address range state\r
+  @param[in]  PageDirectoryEntry      The page table entry to check\r
+  @param[in]  AddressEncMask          The encryption mask\r
+\r
+  @retval MemEncryptSevAddressRangeUnencrypted  Address range is mapped\r
+                                                unencrypted\r
+  @retval MemEncryptSevAddressRangeEncrypted    Address range is mapped\r
+                                                encrypted\r
+  @retval MemEncryptSevAddressRangeMixed        Address range is mapped mixed\r
 **/\r
 STATIC\r
-VOID\r
-SetPageTablePoolReadOnly (\r
-  IN  UINTN                             PageTableBase,\r
-  IN  EFI_PHYSICAL_ADDRESS              Address,\r
-  IN  BOOLEAN                           Level4Paging\r
+MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE\r
+UpdateAddressState (\r
+  IN MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE  CurrentState,\r
+  IN UINT64                               PageDirectoryEntry,\r
+  IN UINT64                               AddressEncMask\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  = GetMemEncryptionAddressMask() &\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
+  if (CurrentState == MemEncryptSevAddressRangeEncrypted) {\r
+    if ((PageDirectoryEntry & AddressEncMask) == 0) {\r
+      CurrentState = MemEncryptSevAddressRangeMixed;\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
-\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
-        NewPageTable[EntryIndex] = PhysicalAddress  | AddressEncMask |\r
-                                   IA32_PG_P | IA32_PG_RW;\r
-        if (Level > 2) {\r
-          NewPageTable[EntryIndex] |= IA32_PG_PS;\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
+  } else if (CurrentState == MemEncryptSevAddressRangeUnencrypted) {\r
+    if ((PageDirectoryEntry & AddressEncMask) != 0) {\r
+      CurrentState = MemEncryptSevAddressRangeMixed;\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
-STATIC\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
-  // 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
+  } else if (CurrentState == MemEncryptSevAddressRangeError) {\r
     //\r
-    // The size of one pool must be multiple of PAGE_TABLE_POOL_UNIT_SIZE,\r
-    // which is one of page size of the processor (2MB by default). Let's apply\r
-    // the protection to them one by one.\r
+    // First address check, set initial state\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
-\r
-\r
-/**\r
-  Split 1G page to 2M.\r
-\r
-  @param[in]      PhysicalAddress       Start physical address the 1G page\r
-                                        covered.\r
-  @param[in, out] PageEntry1G           Pointer to 1G page entry.\r
-  @param[in]      StackBase             Stack base address.\r
-  @param[in]      StackSize             Stack size.\r
-\r
-**/\r
-STATIC\r
-VOID\r
-Split1GPageTo2M (\r
-  IN          PHYSICAL_ADDRESS               PhysicalAddress,\r
-  IN  OUT     UINT64                         *PageEntry1G,\r
-  IN          PHYSICAL_ADDRESS               StackBase,\r
-  IN          UINTN                          StackSize\r
-  )\r
-{\r
-  PHYSICAL_ADDRESS                  PhysicalAddress2M;\r
-  UINTN                             IndexOfPageDirectoryEntries;\r
-  PAGE_TABLE_ENTRY                  *PageDirectoryEntry;\r
-  UINT64                            AddressEncMask;\r
-\r
-  PageDirectoryEntry = AllocatePageTableMemory(1);\r
-\r
-  AddressEncMask = GetMemEncryptionAddressMask ();\r
-  ASSERT (PageDirectoryEntry != NULL);\r
-  ASSERT (*PageEntry1G & GetMemEncryptionAddressMask ());\r
-  //\r
-  // Fill in 1G page entry.\r
-  //\r
-  *PageEntry1G = ((UINT64)(UINTN)PageDirectoryEntry |\r
-                  IA32_PG_P | IA32_PG_RW | AddressEncMask);\r
-\r
-  PhysicalAddress2M = PhysicalAddress;\r
-  for (IndexOfPageDirectoryEntries = 0;\r
-       IndexOfPageDirectoryEntries < 512;\r
-       (IndexOfPageDirectoryEntries++,\r
-        PageDirectoryEntry++,\r
-        PhysicalAddress2M += SIZE_2MB)) {\r
-    if ((PhysicalAddress2M < StackBase + StackSize) &&\r
-        ((PhysicalAddress2M + SIZE_2MB) > StackBase)) {\r
-      //\r
-      // Need to split this 2M page that covers stack range.\r
-      //\r
-      Split2MPageTo4K (\r
-        PhysicalAddress2M,\r
-        (UINT64 *)PageDirectoryEntry,\r
-        StackBase,\r
-        StackSize\r
-        );\r
+    if ((PageDirectoryEntry & AddressEncMask) == 0) {\r
+      CurrentState = MemEncryptSevAddressRangeUnencrypted;\r
     } else {\r
-      //\r
-      // Fill in the Page Directory entries\r
-      //\r
-      PageDirectoryEntry->Uint64 = (UINT64) PhysicalAddress2M | AddressEncMask;\r
-      PageDirectoryEntry->Bits.ReadWrite = 1;\r
-      PageDirectoryEntry->Bits.Present = 1;\r
-      PageDirectoryEntry->Bits.MustBe1 = 1;\r
+      CurrentState = MemEncryptSevAddressRangeEncrypted;\r
     }\r
   }\r
-}\r
-\r
-\r
-/**\r
-  Set or Clear the memory encryption bit\r
-\r
-  @param[in]      PagetablePoint        Page table entry pointer (PTE).\r
-  @param[in]      Mode                  Set or Clear encryption bit\r
-\r
-**/\r
-STATIC VOID\r
-SetOrClearCBit(\r
-  IN   OUT     UINT64*            PageTablePointer,\r
-  IN           MAP_RANGE_MODE     Mode\r
-  )\r
-{\r
-  UINT64      AddressEncMask;\r
-\r
-  AddressEncMask = GetMemEncryptionAddressMask ();\r
-\r
-  if (Mode == SetCBit) {\r
-    *PageTablePointer |= AddressEncMask;\r
-  } else {\r
-    *PageTablePointer &= ~AddressEncMask;\r
-  }\r
-\r
-}\r
-\r
-/**\r
- Check the WP status in CR0 register. This bit is used to lock or unlock write\r
- access to pages marked as read-only.\r
-\r
-  @retval TRUE    Write protection is enabled.\r
-  @retval FALSE   Write protection is disabled.\r
-**/\r
-STATIC\r
-BOOLEAN\r
-IsReadOnlyPageWriteProtected (\r
-  VOID\r
-  )\r
-{\r
-  return ((AsmReadCr0 () & BIT16) != 0);\r
-}\r
 \r
-\r
-/**\r
- Disable Write Protect on pages marked as read-only.\r
-**/\r
-STATIC\r
-VOID\r
-DisableReadOnlyPageWriteProtect (\r
-  VOID\r
-  )\r
-{\r
-  AsmWriteCr0 (AsmReadCr0() & ~BIT16);\r
+  return CurrentState;\r
 }\r
 \r
 /**\r
- Enable Write Protect on pages marked as read-only.\r
-**/\r
-VOID\r
-EnableReadOnlyPageWriteProtect (\r
-  VOID\r
-  )\r
-{\r
-  AsmWriteCr0 (AsmReadCr0() | BIT16);\r
-}\r
-\r
-\r
-/**\r
-  This function either sets or clears memory encryption bit for the memory\r
-  region specified by PhysicalAddress and Length from the current page table\r
-  context.\r
-\r
-  The function iterates through the PhysicalAddress one page at a time, and set\r
-  or clears the memory encryption mask in the page table. If it encounters\r
-  that a given physical address range is part of large page then it attempts to\r
-  change the attribute at one go (based on size), otherwise it splits the\r
-  large pages into smaller (e.g 2M page into 4K pages) and then try to set or\r
-  clear the encryption bit on the smallest page size.\r
+  Returns the encryption state of the specified virtual address range.\r
 \r
   @param[in]  Cr3BaseAddress          Cr3 Base Address (if zero then use\r
                                       current CR3)\r
-  @param[in]  PhysicalAddress         The physical address that is the start\r
-                                      address of a memory region.\r
-  @param[in]  Length                  The length of memory region\r
-  @param[in]  Mode                    Set or Clear mode\r
-  @param[in]  CacheFlush              Flush the caches before applying the\r
-                                      encryption mask\r
-\r
-  @retval RETURN_SUCCESS              The attributes were cleared for the\r
-                                      memory region.\r
-  @retval RETURN_INVALID_PARAMETER    Number of pages is zero.\r
-  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute\r
-                                      is not supported\r
+  @param[in]  BaseAddress             Base address to check\r
+  @param[in]  Length                  Length of virtual address range\r
+\r
+  @retval MemEncryptSevAddressRangeUnencrypted  Address range is mapped\r
+                                                unencrypted\r
+  @retval MemEncryptSevAddressRangeEncrypted    Address range is mapped\r
+                                                encrypted\r
+  @retval MemEncryptSevAddressRangeMixed        Address range is mapped mixed\r
+  @retval MemEncryptSevAddressRangeError        Address range is not mapped\r
 **/\r
-\r
-STATIC\r
-RETURN_STATUS\r
+MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE\r
 EFIAPI\r
-SetMemoryEncDec (\r
-  IN    PHYSICAL_ADDRESS         Cr3BaseAddress,\r
-  IN    PHYSICAL_ADDRESS         PhysicalAddress,\r
-  IN    UINTN                    Length,\r
-  IN    MAP_RANGE_MODE           Mode,\r
-  IN    BOOLEAN                  CacheFlush\r
+InternalMemEncryptSevGetAddressRangeState (\r
+  IN PHYSICAL_ADDRESS  Cr3BaseAddress,\r
+  IN PHYSICAL_ADDRESS  BaseAddress,\r
+  IN UINTN             Length\r
   )\r
 {\r
-  PAGE_MAP_AND_DIRECTORY_POINTER *PageMapLevel4Entry;\r
-  PAGE_MAP_AND_DIRECTORY_POINTER *PageUpperDirectoryPointerEntry;\r
-  PAGE_MAP_AND_DIRECTORY_POINTER *PageDirectoryPointerEntry;\r
-  PAGE_TABLE_1G_ENTRY            *PageDirectory1GEntry;\r
-  PAGE_TABLE_ENTRY               *PageDirectory2MEntry;\r
-  PAGE_TABLE_4K_ENTRY            *PageTableEntry;\r
-  UINT64                         PgTableMask;\r
-  UINT64                         AddressEncMask;\r
-  BOOLEAN                        IsWpEnabled;\r
-  RETURN_STATUS                  Status;\r
-\r
-  DEBUG ((\r
-    DEBUG_VERBOSE,\r
-    "%a:%a: Cr3Base=0x%Lx Physical=0x%Lx Length=0x%Lx Mode=%a CacheFlush=%u\n",\r
-    gEfiCallerBaseName,\r
-    __FUNCTION__,\r
-    Cr3BaseAddress,\r
-    PhysicalAddress,\r
-    (UINT64)Length,\r
-    (Mode == SetCBit) ? "Encrypt" : "Decrypt",\r
-    (UINT32)CacheFlush\r
-    ));\r
+  PAGE_MAP_AND_DIRECTORY_POINTER       *PageMapLevel4Entry;\r
+  PAGE_MAP_AND_DIRECTORY_POINTER       *PageUpperDirectoryPointerEntry;\r
+  PAGE_MAP_AND_DIRECTORY_POINTER       *PageDirectoryPointerEntry;\r
+  PAGE_TABLE_1G_ENTRY                  *PageDirectory1GEntry;\r
+  PAGE_TABLE_ENTRY                     *PageDirectory2MEntry;\r
+  PAGE_TABLE_4K_ENTRY                  *PageTableEntry;\r
+  UINT64                               AddressEncMask;\r
+  UINT64                               PgTableMask;\r
+  PHYSICAL_ADDRESS                     Address;\r
+  PHYSICAL_ADDRESS                     AddressEnd;\r
+  MEM_ENCRYPT_SEV_ADDRESS_RANGE_STATE  State;\r
 \r
   //\r
-  // Check if we have a valid memory encryption mask\r
+  // If Cr3BaseAddress is not specified then read the current CR3\r
   //\r
-  AddressEncMask = GetMemEncryptionAddressMask ();\r
-  if (!AddressEncMask) {\r
-    return RETURN_ACCESS_DENIED;\r
+  if (Cr3BaseAddress == 0) {\r
+    Cr3BaseAddress = AsmReadCr3 ();\r
   }\r
 \r
-  PgTableMask = AddressEncMask | EFI_PAGE_MASK;\r
+  AddressEncMask  = MemEncryptSevGetEncryptionMask ();\r
+  AddressEncMask &= PAGING_1G_ADDRESS_MASK_64;\r
 \r
-  if (Length == 0) {\r
-    return RETURN_INVALID_PARAMETER;\r
-  }\r
+  PgTableMask = AddressEncMask | EFI_PAGE_MASK;\r
 \r
-  //\r
-  // We are going to change the memory encryption attribute from C=0 -> C=1 or\r
-  // vice versa Flush the caches to ensure that data is written into memory\r
-  // with correct C-bit\r
-  //\r
-  if (CacheFlush) {\r
-    WriteBackInvalidateDataCacheRange((VOID*) (UINTN)PhysicalAddress, Length);\r
-  }\r
+  State = MemEncryptSevAddressRangeError;\r
 \r
   //\r
-  // Make sure that the page table is changeable.\r
+  // Encryption is on a page basis, so start at the beginning of the\r
+  // virtual address page boundary and walk page-by-page.\r
   //\r
-  IsWpEnabled = IsReadOnlyPageWriteProtected ();\r
-  if (IsWpEnabled) {\r
-    DisableReadOnlyPageWriteProtect ();\r
-  }\r
-\r
-  Status = EFI_SUCCESS;\r
-\r
-  while (Length)\r
-  {\r
-    //\r
-    // If Cr3BaseAddress is not specified then read the current CR3\r
-    //\r
-    if (Cr3BaseAddress == 0) {\r
-      Cr3BaseAddress = AsmReadCr3();\r
-    }\r
+  Address    = (PHYSICAL_ADDRESS)(UINTN)BaseAddress & ~EFI_PAGE_MASK;\r
+  AddressEnd = (PHYSICAL_ADDRESS)\r
+               (UINTN)(BaseAddress + Length);\r
 \r
-    PageMapLevel4Entry = (VOID*) (Cr3BaseAddress & ~PgTableMask);\r
-    PageMapLevel4Entry += PML4_OFFSET(PhysicalAddress);\r
+  while (Address < AddressEnd) {\r
+    PageMapLevel4Entry  = (VOID *)(Cr3BaseAddress & ~PgTableMask);\r
+    PageMapLevel4Entry += PML4_OFFSET (Address);\r
     if (!PageMapLevel4Entry->Bits.Present) {\r
-      DEBUG ((\r
-        DEBUG_ERROR,\r
-        "%a:%a: bad PML4 for Physical=0x%Lx\n",\r
-        gEfiCallerBaseName,\r
-        __FUNCTION__,\r
-        PhysicalAddress\r
-        ));\r
-      Status = RETURN_NO_MAPPING;\r
-      goto Done;\r
+      return MemEncryptSevAddressRangeError;\r
     }\r
 \r
     PageDirectory1GEntry = (VOID *)(\r
-                             (PageMapLevel4Entry->Bits.PageTableBaseAddress <<\r
-                              12) & ~PgTableMask\r
-                             );\r
-    PageDirectory1GEntry += PDP_OFFSET(PhysicalAddress);\r
+                                    (PageMapLevel4Entry->Bits.PageTableBaseAddress <<\r
+                                     12) & ~PgTableMask\r
+                                    );\r
+    PageDirectory1GEntry += PDP_OFFSET (Address);\r
     if (!PageDirectory1GEntry->Bits.Present) {\r
-      DEBUG ((\r
-        DEBUG_ERROR,\r
-        "%a:%a: bad PDPE for Physical=0x%Lx\n",\r
-        gEfiCallerBaseName,\r
-        __FUNCTION__,\r
-        PhysicalAddress\r
-        ));\r
-      Status = RETURN_NO_MAPPING;\r
-      goto Done;\r
+      return MemEncryptSevAddressRangeError;\r
     }\r
 \r
     //\r
@@ -687,208 +136,72 @@ SetMemoryEncDec (
     if (PageDirectory1GEntry->Bits.MustBe1) {\r
       //\r
       // Valid 1GB page\r
-      // If we have at least 1GB to go, we can just update this entry\r
-      //\r
-      if (!(PhysicalAddress & (BIT30 - 1)) && Length >= BIT30) {\r
-        SetOrClearCBit(&PageDirectory1GEntry->Uint64, Mode);\r
-        DEBUG ((\r
-          DEBUG_VERBOSE,\r
-          "%a:%a: updated 1GB entry for Physical=0x%Lx\n",\r
-          gEfiCallerBaseName,\r
-          __FUNCTION__,\r
-          PhysicalAddress\r
-          ));\r
-        PhysicalAddress += BIT30;\r
-        Length -= BIT30;\r
-      } else {\r
-        //\r
-        // We must split the page\r
-        //\r
-        DEBUG ((\r
-          DEBUG_VERBOSE,\r
-          "%a:%a: splitting 1GB page for Physical=0x%Lx\n",\r
-          gEfiCallerBaseName,\r
-          __FUNCTION__,\r
-          PhysicalAddress\r
-          ));\r
-        Split1GPageTo2M (\r
-          (UINT64)PageDirectory1GEntry->Bits.PageTableBaseAddress << 30,\r
-          (UINT64 *)PageDirectory1GEntry,\r
-          0,\r
-          0\r
-          );\r
-        continue;\r
-      }\r
-    } else {\r
       //\r
-      // Actually a PDP\r
-      //\r
-      PageUpperDirectoryPointerEntry =\r
-        (PAGE_MAP_AND_DIRECTORY_POINTER *)PageDirectory1GEntry;\r
-      PageDirectory2MEntry =\r
-        (VOID *)(\r
-          (PageUpperDirectoryPointerEntry->Bits.PageTableBaseAddress <<\r
-           12) & ~PgTableMask\r
-          );\r
-      PageDirectory2MEntry += PDE_OFFSET(PhysicalAddress);\r
-      if (!PageDirectory2MEntry->Bits.Present) {\r
-        DEBUG ((\r
-          DEBUG_ERROR,\r
-          "%a:%a: bad PDE for Physical=0x%Lx\n",\r
-          gEfiCallerBaseName,\r
-          __FUNCTION__,\r
-          PhysicalAddress\r
-          ));\r
-        Status = RETURN_NO_MAPPING;\r
-        goto Done;\r
-      }\r
-      //\r
-      // If the MustBe1 bit is not a 1, it's not a 2MB entry\r
-      //\r
-      if (PageDirectory2MEntry->Bits.MustBe1) {\r
-        //\r
-        // Valid 2MB page\r
-        // If we have at least 2MB left to go, we can just update this entry\r
-        //\r
-        if (!(PhysicalAddress & (BIT21-1)) && Length >= BIT21) {\r
-          SetOrClearCBit (&PageDirectory2MEntry->Uint64, Mode);\r
-          PhysicalAddress += BIT21;\r
-          Length -= BIT21;\r
-        } else {\r
-          //\r
-          // We must split up this page into 4K pages\r
-          //\r
-          DEBUG ((\r
-            DEBUG_VERBOSE,\r
-            "%a:%a: splitting 2MB page for Physical=0x%Lx\n",\r
-            gEfiCallerBaseName,\r
-            __FUNCTION__,\r
-            PhysicalAddress\r
-            ));\r
-          Split2MPageTo4K (\r
-            (UINT64)PageDirectory2MEntry->Bits.PageTableBaseAddress << 21,\r
-            (UINT64 *)PageDirectory2MEntry,\r
-            0,\r
-            0\r
-            );\r
-          continue;\r
-        }\r
-      } else {\r
-        PageDirectoryPointerEntry =\r
-          (PAGE_MAP_AND_DIRECTORY_POINTER *)PageDirectory2MEntry;\r
-        PageTableEntry =\r
-          (VOID *)(\r
-            (PageDirectoryPointerEntry->Bits.PageTableBaseAddress <<\r
-             12) & ~PgTableMask\r
-            );\r
-        PageTableEntry += PTE_OFFSET(PhysicalAddress);\r
-        if (!PageTableEntry->Bits.Present) {\r
-          DEBUG ((\r
-            DEBUG_ERROR,\r
-            "%a:%a: bad PTE for Physical=0x%Lx\n",\r
-            gEfiCallerBaseName,\r
-            __FUNCTION__,\r
-            PhysicalAddress\r
-            ));\r
-          Status = RETURN_NO_MAPPING;\r
-          goto Done;\r
-        }\r
-        SetOrClearCBit (&PageTableEntry->Uint64, Mode);\r
-        PhysicalAddress += EFI_PAGE_SIZE;\r
-        Length -= EFI_PAGE_SIZE;\r
-      }\r
-    }\r
-  }\r
-\r
-  //\r
-  // Protect the page table by marking the memory used for page table to be\r
-  // read-only.\r
-  //\r
-  if (IsWpEnabled) {\r
-    EnablePageTableProtection ((UINTN)PageMapLevel4Entry, TRUE);\r
-  }\r
-\r
-  //\r
-  // Flush TLB\r
-  //\r
-  CpuFlushTlb();\r
+      State = UpdateAddressState (\r
+                State,\r
+                PageDirectory1GEntry->Uint64,\r
+                AddressEncMask\r
+                );\r
 \r
-Done:\r
-  //\r
-  // Restore page table write protection, if any.\r
-  //\r
-  if (IsWpEnabled) {\r
-    EnableReadOnlyPageWriteProtect ();\r
-  }\r
-\r
-  return Status;\r
-}\r
+      Address += BIT30;\r
+      continue;\r
+    }\r
 \r
-/**\r
-  This function clears memory encryption bit for the memory region specified by\r
-  PhysicalAddress and length from the current page table context.\r
+    //\r
+    // Actually a PDP\r
+    //\r
+    PageUpperDirectoryPointerEntry =\r
+      (PAGE_MAP_AND_DIRECTORY_POINTER *)PageDirectory1GEntry;\r
+    PageDirectory2MEntry =\r
+      (VOID *)(\r
+               (PageUpperDirectoryPointerEntry->Bits.PageTableBaseAddress <<\r
+                12) & ~PgTableMask\r
+               );\r
+    PageDirectory2MEntry += PDE_OFFSET (Address);\r
+    if (!PageDirectory2MEntry->Bits.Present) {\r
+      return MemEncryptSevAddressRangeError;\r
+    }\r
 \r
-  @param[in]  PhysicalAddress         The physical address that is the start\r
-                                      address of a memory region.\r
-  @param[in]  Length                  The length of memory region\r
-  @param[in]  Flush                   Flush the caches before applying the\r
-                                      encryption mask\r
+    //\r
+    // If the MustBe1 bit is not a 1, it's not a 2MB entry\r
+    //\r
+    if (PageDirectory2MEntry->Bits.MustBe1) {\r
+      //\r
+      // Valid 2MB page\r
+      //\r
+      State = UpdateAddressState (\r
+                State,\r
+                PageDirectory2MEntry->Uint64,\r
+                AddressEncMask\r
+                );\r
 \r
-  @retval RETURN_SUCCESS              The attributes were cleared for the\r
-                                      memory region.\r
-  @retval RETURN_INVALID_PARAMETER    Number of pages is zero.\r
-  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute\r
-                                      is not supported\r
-**/\r
-RETURN_STATUS\r
-EFIAPI\r
-InternalMemEncryptSevSetMemoryDecrypted (\r
-  IN  PHYSICAL_ADDRESS        Cr3BaseAddress,\r
-  IN  PHYSICAL_ADDRESS        PhysicalAddress,\r
-  IN  UINTN                   Length,\r
-  IN  BOOLEAN                 Flush\r
-  )\r
-{\r
+      Address += BIT21;\r
+      continue;\r
+    }\r
 \r
-  return SetMemoryEncDec (\r
-           Cr3BaseAddress,\r
-           PhysicalAddress,\r
-           Length,\r
-           ClearCBit,\r
-           Flush\r
-           );\r
-}\r
+    //\r
+    // Actually a PMD\r
+    //\r
+    PageDirectoryPointerEntry =\r
+      (PAGE_MAP_AND_DIRECTORY_POINTER *)PageDirectory2MEntry;\r
+    PageTableEntry =\r
+      (VOID *)(\r
+               (PageDirectoryPointerEntry->Bits.PageTableBaseAddress <<\r
+                12) & ~PgTableMask\r
+               );\r
+    PageTableEntry += PTE_OFFSET (Address);\r
+    if (!PageTableEntry->Bits.Present) {\r
+      return MemEncryptSevAddressRangeError;\r
+    }\r
 \r
-/**\r
-  This function sets memory encryption bit for the memory region specified by\r
-  PhysicalAddress and length from the current page table context.\r
+    State = UpdateAddressState (\r
+              State,\r
+              PageTableEntry->Uint64,\r
+              AddressEncMask\r
+              );\r
 \r
-  @param[in]  PhysicalAddress         The physical address that is the start\r
-                                      address of a memory region.\r
-  @param[in]  Length                  The length of memory region\r
-  @param[in]  Flush                   Flush the caches before applying the\r
-                                      encryption mask\r
+    Address += EFI_PAGE_SIZE;\r
+  }\r
 \r
-  @retval RETURN_SUCCESS              The attributes were cleared for the\r
-                                      memory region.\r
-  @retval RETURN_INVALID_PARAMETER    Number of pages is zero.\r
-  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute\r
-                                      is not supported\r
-**/\r
-RETURN_STATUS\r
-EFIAPI\r
-InternalMemEncryptSevSetMemoryEncrypted (\r
-  IN  PHYSICAL_ADDRESS        Cr3BaseAddress,\r
-  IN  PHYSICAL_ADDRESS        PhysicalAddress,\r
-  IN  UINTN                   Length,\r
-  IN  BOOLEAN                 Flush\r
-  )\r
-{\r
-  return SetMemoryEncDec (\r
-           Cr3BaseAddress,\r
-           PhysicalAddress,\r
-           Length,\r
-           SetCBit,\r
-           Flush\r
-           );\r
+  return State;\r
 }\r