]> git.proxmox.com Git - mirror_edk2.git/blobdiff - OvmfPkg/IoMmuDxe/AmdSevIoMmu.c
OvmfPkg/PlatformDxe: list "Platform.h" in the INF file
[mirror_edk2.git] / OvmfPkg / IoMmuDxe / AmdSevIoMmu.c
index d899b0ab9e4105a6f0d4168dcc4f5aad4958d568..6b7df8d8aa3ddd5a2b866a34c765f4dec3a3ebef 100644 (file)
 \r
 #include "AmdSevIoMmu.h"\r
 \r
+#define MAP_INFO_SIG SIGNATURE_64 ('M', 'A', 'P', '_', 'I', 'N', 'F', 'O')\r
+\r
 typedef struct {\r
+  UINT64                                    Signature;\r
+  LIST_ENTRY                                Link;\r
   EDKII_IOMMU_OPERATION                     Operation;\r
   UINTN                                     NumberOfBytes;\r
   UINTN                                     NumberOfPages;\r
@@ -28,7 +32,51 @@ typedef struct {
   EFI_PHYSICAL_ADDRESS                      PlainTextAddress;\r
 } MAP_INFO;\r
 \r
-#define NO_MAPPING             (VOID *) (UINTN) -1\r
+//\r
+// List of the MAP_INFO structures that have been set up by IoMmuMap() and not\r
+// yet torn down by IoMmuUnmap(). The list represents the full set of mappings\r
+// currently in effect.\r
+//\r
+STATIC LIST_ENTRY mMapInfos = INITIALIZE_LIST_HEAD_VARIABLE (mMapInfos);\r
+\r
+#define COMMON_BUFFER_SIG SIGNATURE_64 ('C', 'M', 'N', 'B', 'U', 'F', 'F', 'R')\r
+\r
+//\r
+// ASCII names for EDKII_IOMMU_OPERATION constants, for debug logging.\r
+//\r
+STATIC CONST CHAR8 * CONST\r
+mBusMasterOperationName[EdkiiIoMmuOperationMaximum] = {\r
+  "Read",\r
+  "Write",\r
+  "CommonBuffer",\r
+  "Read64",\r
+  "Write64",\r
+  "CommonBuffer64"\r
+};\r
+\r
+//\r
+// The following structure enables Map() and Unmap() to perform in-place\r
+// decryption and encryption, respectively, for BusMasterCommonBuffer[64]\r
+// operations, without dynamic memory allocation or release.\r
+//\r
+// Both COMMON_BUFFER_HEADER and COMMON_BUFFER_HEADER.StashBuffer are allocated\r
+// by AllocateBuffer() and released by FreeBuffer().\r
+//\r
+#pragma pack (1)\r
+typedef struct {\r
+  UINT64 Signature;\r
+\r
+  //\r
+  // Always allocated from EfiBootServicesData type memory, and always\r
+  // encrypted.\r
+  //\r
+  VOID *StashBuffer;\r
+\r
+  //\r
+  // Followed by the actual common buffer, starting at the next page.\r
+  //\r
+} COMMON_BUFFER_HEADER;\r
+#pragma pack ()\r
 \r
 /**\r
   Provides the controller-specific addresses required to access system memory\r
@@ -72,98 +120,140 @@ IoMmuMap (
   )\r
 {\r
   EFI_STATUS                                        Status;\r
-  EFI_PHYSICAL_ADDRESS                              PhysicalAddress;\r
   MAP_INFO                                          *MapInfo;\r
-  EFI_PHYSICAL_ADDRESS                              DmaMemoryTop;\r
   EFI_ALLOCATE_TYPE                                 AllocateType;\r
+  COMMON_BUFFER_HEADER                              *CommonBufferHeader;\r
+  VOID                                              *DecryptionSource;\r
+\r
+  DEBUG ((\r
+    DEBUG_VERBOSE,\r
+    "%a: Operation=%a Host=0x%p Bytes=0x%Lx\n",\r
+    __FUNCTION__,\r
+    ((Operation >= 0 &&\r
+      Operation < ARRAY_SIZE (mBusMasterOperationName)) ?\r
+     mBusMasterOperationName[Operation] :\r
+     "Invalid"),\r
+    HostAddress,\r
+    (UINT64)((NumberOfBytes == NULL) ? 0 : *NumberOfBytes)\r
+    ));\r
 \r
   if (HostAddress == NULL || NumberOfBytes == NULL || DeviceAddress == NULL ||\r
       Mapping == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
-  //\r
-  // Make sure that Operation is valid\r
-  //\r
-  if ((UINT32) Operation >= EdkiiIoMmuOperationMaximum) {\r
-    return EFI_INVALID_PARAMETER;\r
-  }\r
-  PhysicalAddress = (EFI_PHYSICAL_ADDRESS) (UINTN) HostAddress;\r
-\r
-  DmaMemoryTop = (UINTN)-1;\r
-  AllocateType = AllocateAnyPages;\r
-\r
-  if (((Operation != EdkiiIoMmuOperationBusMasterRead64 &&\r
-        Operation != EdkiiIoMmuOperationBusMasterWrite64 &&\r
-        Operation != EdkiiIoMmuOperationBusMasterCommonBuffer64)) &&\r
-      ((PhysicalAddress + *NumberOfBytes) > SIZE_4GB)) {\r
-    //\r
-    // If the root bridge or the device cannot handle performing DMA above\r
-    // 4GB but any part of the DMA transfer being mapped is above 4GB, then\r
-    // map the DMA transfer to a buffer below 4GB.\r
-    //\r
-    DmaMemoryTop = SIZE_4GB - 1;\r
-    AllocateType = AllocateMaxAddress;\r
-\r
-    if (Operation == EdkiiIoMmuOperationBusMasterCommonBuffer ||\r
-        Operation == EdkiiIoMmuOperationBusMasterCommonBuffer64) {\r
-        //\r
-        // Common Buffer operations can not be remapped.  If the common buffer\r
-        // if above 4GB, then it is not possible to generate a mapping, so\r
-        // return an error.\r
-        //\r
-        return EFI_UNSUPPORTED;\r
-    }\r
-  }\r
-\r
-  //\r
-  // CommandBuffer was allocated by us (AllocateBuffer) and is already in\r
-  // unencryted buffer so no need to create bounce buffer\r
-  //\r
-  if (Operation == EdkiiIoMmuOperationBusMasterCommonBuffer ||\r
-      Operation == EdkiiIoMmuOperationBusMasterCommonBuffer64) {\r
-    *Mapping = NO_MAPPING;\r
-    *DeviceAddress = PhysicalAddress;\r
-\r
-    return EFI_SUCCESS;\r
-  }\r
-\r
   //\r
   // Allocate a MAP_INFO structure to remember the mapping when Unmap() is\r
   // called later.\r
   //\r
   MapInfo = AllocatePool (sizeof (MAP_INFO));\r
   if (MapInfo == NULL) {\r
-    *NumberOfBytes = 0;\r
-    return EFI_OUT_OF_RESOURCES;\r
+    Status = EFI_OUT_OF_RESOURCES;\r
+    goto Failed;\r
   }\r
 \r
   //\r
-  // Initialize the MAP_INFO structure\r
+  // Initialize the MAP_INFO structure, except the PlainTextAddress field\r
   //\r
+  ZeroMem (&MapInfo->Link, sizeof MapInfo->Link);\r
+  MapInfo->Signature         = MAP_INFO_SIG;\r
   MapInfo->Operation         = Operation;\r
   MapInfo->NumberOfBytes     = *NumberOfBytes;\r
   MapInfo->NumberOfPages     = EFI_SIZE_TO_PAGES (MapInfo->NumberOfBytes);\r
-  MapInfo->CryptedAddress    = PhysicalAddress;\r
-  MapInfo->PlainTextAddress  = DmaMemoryTop;\r
+  MapInfo->CryptedAddress    = (UINTN)HostAddress;\r
 \r
   //\r
-  // Allocate a buffer to map the transfer to.\r
+  // In the switch statement below, we point "MapInfo->PlainTextAddress" to the\r
+  // plaintext buffer, according to Operation. We also set "DecryptionSource".\r
   //\r
-  Status = gBS->AllocatePages (\r
-                  AllocateType,\r
-                  EfiBootServicesData,\r
-                  MapInfo->NumberOfPages,\r
-                  &MapInfo->PlainTextAddress\r
-                  );\r
-  if (EFI_ERROR (Status)) {\r
-    FreePool (MapInfo);\r
-    *NumberOfBytes = 0;\r
-    return Status;\r
+  MapInfo->PlainTextAddress = MAX_ADDRESS;\r
+  AllocateType = AllocateAnyPages;\r
+  DecryptionSource = (VOID *)(UINTN)MapInfo->CryptedAddress;\r
+  switch (Operation) {\r
+  //\r
+  // For BusMasterRead[64] and BusMasterWrite[64] operations, a bounce buffer\r
+  // is necessary regardless of whether the original (crypted) buffer crosses\r
+  // the 4GB limit or not -- we have to allocate a separate plaintext buffer.\r
+  // The only variable is whether the plaintext buffer should be under 4GB.\r
+  //\r
+  case EdkiiIoMmuOperationBusMasterRead:\r
+  case EdkiiIoMmuOperationBusMasterWrite:\r
+    MapInfo->PlainTextAddress = BASE_4GB - 1;\r
+    AllocateType = AllocateMaxAddress;\r
+    //\r
+    // fall through\r
+    //\r
+  case EdkiiIoMmuOperationBusMasterRead64:\r
+  case EdkiiIoMmuOperationBusMasterWrite64:\r
+    //\r
+    // Allocate the implicit plaintext bounce buffer.\r
+    //\r
+    Status = gBS->AllocatePages (\r
+                    AllocateType,\r
+                    EfiBootServicesData,\r
+                    MapInfo->NumberOfPages,\r
+                    &MapInfo->PlainTextAddress\r
+                    );\r
+    if (EFI_ERROR (Status)) {\r
+      goto FreeMapInfo;\r
+    }\r
+    break;\r
+\r
+  //\r
+  // For BusMasterCommonBuffer[64] operations, a to-be-plaintext buffer and a\r
+  // stash buffer (for in-place decryption) have been allocated already, with\r
+  // AllocateBuffer(). We only check whether the address of the to-be-plaintext\r
+  // buffer is low enough for the requested operation.\r
+  //\r
+  case EdkiiIoMmuOperationBusMasterCommonBuffer:\r
+    if ((MapInfo->CryptedAddress > BASE_4GB) ||\r
+        (EFI_PAGES_TO_SIZE (MapInfo->NumberOfPages) >\r
+         BASE_4GB - MapInfo->CryptedAddress)) {\r
+      //\r
+      // CommonBuffer operations cannot be remapped. If the common buffer is\r
+      // above 4GB, then it is not possible to generate a mapping, so return an\r
+      // error.\r
+      //\r
+      Status = EFI_UNSUPPORTED;\r
+      goto FreeMapInfo;\r
+    }\r
+    //\r
+    // fall through\r
+    //\r
+  case EdkiiIoMmuOperationBusMasterCommonBuffer64:\r
+    //\r
+    // The buffer at MapInfo->CryptedAddress comes from AllocateBuffer().\r
+    //\r
+    MapInfo->PlainTextAddress = MapInfo->CryptedAddress;\r
+    //\r
+    // Stash the crypted data.\r
+    //\r
+    CommonBufferHeader = (COMMON_BUFFER_HEADER *)(\r
+                           (UINTN)MapInfo->CryptedAddress - EFI_PAGE_SIZE\r
+                           );\r
+    ASSERT (CommonBufferHeader->Signature == COMMON_BUFFER_SIG);\r
+    CopyMem (\r
+      CommonBufferHeader->StashBuffer,\r
+      (VOID *)(UINTN)MapInfo->CryptedAddress,\r
+      MapInfo->NumberOfBytes\r
+      );\r
+    //\r
+    // Point "DecryptionSource" to the stash buffer so that we decrypt\r
+    // it to the original location, after the switch statement.\r
+    //\r
+    DecryptionSource = CommonBufferHeader->StashBuffer;\r
+    break;\r
+\r
+  default:\r
+    //\r
+    // Operation is invalid\r
+    //\r
+    Status = EFI_INVALID_PARAMETER;\r
+    goto FreeMapInfo;\r
   }\r
 \r
   //\r
-  // Clear the memory encryption mask from the device buffer\r
+  // Clear the memory encryption mask on the plaintext buffer.\r
   //\r
   Status = MemEncryptSevClearPageEncMask (\r
              0,\r
@@ -171,50 +261,71 @@ IoMmuMap (
              MapInfo->NumberOfPages,\r
              TRUE\r
              );\r
-  ASSERT_EFI_ERROR(Status);\r
+  ASSERT_EFI_ERROR (Status);\r
+  if (EFI_ERROR (Status)) {\r
+    CpuDeadLoop ();\r
+  }\r
 \r
   //\r
   // If this is a read operation from the Bus Master's point of view,\r
   // then copy the contents of the real buffer into the mapped buffer\r
   // so the Bus Master can read the contents of the real buffer.\r
   //\r
+  // For BusMasterCommonBuffer[64] operations, the CopyMem() below will decrypt\r
+  // the original data (from the stash buffer) back to the original location.\r
+  //\r
   if (Operation == EdkiiIoMmuOperationBusMasterRead ||\r
-      Operation == EdkiiIoMmuOperationBusMasterRead64) {\r
+      Operation == EdkiiIoMmuOperationBusMasterRead64 ||\r
+      Operation == EdkiiIoMmuOperationBusMasterCommonBuffer ||\r
+      Operation == EdkiiIoMmuOperationBusMasterCommonBuffer64) {\r
     CopyMem (\r
       (VOID *) (UINTN) MapInfo->PlainTextAddress,\r
-      (VOID *) (UINTN) MapInfo->CryptedAddress,\r
+      DecryptionSource,\r
       MapInfo->NumberOfBytes\r
       );\r
   }\r
 \r
   //\r
-  // The DeviceAddress is the address of the maped buffer below 4GB\r
+  // Track all MAP_INFO structures.\r
   //\r
-  *DeviceAddress = MapInfo->PlainTextAddress;\r
-\r
+  InsertHeadList (&mMapInfos, &MapInfo->Link);\r
   //\r
-  // Return a pointer to the MAP_INFO structure in Mapping\r
+  // Populate output parameters.\r
   //\r
+  *DeviceAddress = MapInfo->PlainTextAddress;\r
   *Mapping       = MapInfo;\r
 \r
   DEBUG ((\r
     DEBUG_VERBOSE,\r
-    "%a PlainText 0x%Lx Crypted 0x%Lx Pages 0x%Lx Bytes 0x%Lx\n",\r
+    "%a: Mapping=0x%p Device(PlainText)=0x%Lx Crypted=0x%Lx Pages=0x%Lx\n",\r
     __FUNCTION__,\r
+    MapInfo,\r
     MapInfo->PlainTextAddress,\r
     MapInfo->CryptedAddress,\r
-    (UINT64)MapInfo->NumberOfPages,\r
-    (UINT64)MapInfo->NumberOfBytes\r
+    (UINT64)MapInfo->NumberOfPages\r
     ));\r
 \r
   return EFI_SUCCESS;\r
+\r
+FreeMapInfo:\r
+  FreePool (MapInfo);\r
+\r
+Failed:\r
+  *NumberOfBytes = 0;\r
+  return Status;\r
 }\r
 \r
 /**\r
   Completes the Map() operation and releases any corresponding resources.\r
 \r
+  This is an internal worker function that only extends the Map() API with\r
+  the MemoryMapLocked parameter.\r
+\r
   @param  This                  The protocol instance pointer.\r
   @param  Mapping               The mapping value returned from Map().\r
+  @param  MemoryMapLocked       The function is executing on the stack of\r
+                                gBS->ExitBootServices(); changes to the UEFI\r
+                                memory map are forbidden.\r
 \r
   @retval EFI_SUCCESS           The range was unmapped.\r
   @retval EFI_INVALID_PARAMETER Mapping is not a value that was returned by\r
@@ -222,56 +333,82 @@ IoMmuMap (
   @retval EFI_DEVICE_ERROR      The data was not committed to the target system\r
                                 memory.\r
 **/\r
+STATIC\r
 EFI_STATUS\r
 EFIAPI\r
-IoMmuUnmap (\r
+IoMmuUnmapWorker (\r
   IN  EDKII_IOMMU_PROTOCOL                     *This,\r
-  IN  VOID                                     *Mapping\r
+  IN  VOID                                     *Mapping,\r
+  IN  BOOLEAN                                  MemoryMapLocked\r
   )\r
 {\r
   MAP_INFO                 *MapInfo;\r
   EFI_STATUS               Status;\r
+  COMMON_BUFFER_HEADER     *CommonBufferHeader;\r
+  VOID                     *EncryptionTarget;\r
+\r
+  DEBUG ((\r
+    DEBUG_VERBOSE,\r
+    "%a: Mapping=0x%p MemoryMapLocked=%d\n",\r
+    __FUNCTION__,\r
+    Mapping,\r
+    MemoryMapLocked\r
+    ));\r
 \r
   if (Mapping == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
+  MapInfo = (MAP_INFO *)Mapping;\r
+\r
   //\r
-  // See if the Map() operation associated with this Unmap() required a mapping\r
-  // buffer. If a mapping buffer was not required, then this function simply\r
-  // buffer. If a mapping buffer was not required, then this function simply\r
+  // set CommonBufferHeader to suppress incorrect compiler/analyzer warnings\r
   //\r
-  if (Mapping == NO_MAPPING) {\r
-    return EFI_SUCCESS;\r
-  }\r
-\r
-  MapInfo = (MAP_INFO *)Mapping;\r
+  CommonBufferHeader = NULL;\r
 \r
   //\r
-  // If this is a write operation from the Bus Master's point of view,\r
-  // then copy the contents of the mapped buffer into the real buffer\r
-  // so the processor can read the contents of the real buffer.\r
+  // For BusMasterWrite[64] operations and BusMasterCommonBuffer[64] operations\r
+  // we have to encrypt the results, ultimately to the original place (i.e.,\r
+  // "MapInfo->CryptedAddress").\r
+  //\r
+  // For BusMasterCommonBuffer[64] operations however, this encryption has to\r
+  // land in-place, so divert the encryption to the stash buffer first.\r
   //\r
-  if (MapInfo->Operation == EdkiiIoMmuOperationBusMasterWrite ||\r
-      MapInfo->Operation == EdkiiIoMmuOperationBusMasterWrite64) {\r
+  EncryptionTarget = (VOID *)(UINTN)MapInfo->CryptedAddress;\r
+\r
+  switch (MapInfo->Operation) {\r
+  case EdkiiIoMmuOperationBusMasterCommonBuffer:\r
+  case EdkiiIoMmuOperationBusMasterCommonBuffer64:\r
+    ASSERT (MapInfo->PlainTextAddress == MapInfo->CryptedAddress);\r
+\r
+    CommonBufferHeader = (COMMON_BUFFER_HEADER *)(\r
+                           (UINTN)MapInfo->PlainTextAddress - EFI_PAGE_SIZE\r
+                           );\r
+    ASSERT (CommonBufferHeader->Signature == COMMON_BUFFER_SIG);\r
+    EncryptionTarget = CommonBufferHeader->StashBuffer;\r
+    //\r
+    // fall through\r
+    //\r
+\r
+  case EdkiiIoMmuOperationBusMasterWrite:\r
+  case EdkiiIoMmuOperationBusMasterWrite64:\r
     CopyMem (\r
-      (VOID *) (UINTN) MapInfo->CryptedAddress,\r
+      EncryptionTarget,\r
       (VOID *) (UINTN) MapInfo->PlainTextAddress,\r
       MapInfo->NumberOfBytes\r
       );\r
+    break;\r
+\r
+  default:\r
+    //\r
+    // nothing to encrypt after BusMasterRead[64] operations\r
+    //\r
+    break;\r
   }\r
 \r
-  DEBUG ((\r
-    DEBUG_VERBOSE,\r
-    "%a PlainText 0x%Lx Crypted 0x%Lx Pages 0x%Lx Bytes 0x%Lx\n",\r
-    __FUNCTION__,\r
-    MapInfo->PlainTextAddress,\r
-    MapInfo->CryptedAddress,\r
-    (UINT64)MapInfo->NumberOfPages,\r
-    (UINT64)MapInfo->NumberOfBytes\r
-    ));\r
   //\r
-  // Restore the memory encryption mask\r
+  // Restore the memory encryption mask on the area we used to hold the\r
+  // plaintext.\r
   //\r
   Status = MemEncryptSevSetPageEncMask (\r
              0,\r
@@ -279,20 +416,74 @@ IoMmuUnmap (
              MapInfo->NumberOfPages,\r
              TRUE\r
              );\r
-  ASSERT_EFI_ERROR(Status);\r
-  ZeroMem (\r
-    (VOID*)(UINTN)MapInfo->PlainTextAddress,\r
-    EFI_PAGES_TO_SIZE (MapInfo->NumberOfPages)\r
-    );\r
+  ASSERT_EFI_ERROR (Status);\r
+  if (EFI_ERROR (Status)) {\r
+    CpuDeadLoop ();\r
+  }\r
 \r
   //\r
-  // Free the mapped buffer and the MAP_INFO structure.\r
+  // For BusMasterCommonBuffer[64] operations, copy the stashed data to the\r
+  // original (now encrypted) location.\r
   //\r
-  gBS->FreePages (MapInfo->PlainTextAddress, MapInfo->NumberOfPages);\r
-  FreePool (Mapping);\r
+  // For all other operations, fill the late bounce buffer (which existed as\r
+  // plaintext at some point) with zeros, and then release it (unless the UEFI\r
+  // memory map is locked).\r
+  //\r
+  if (MapInfo->Operation == EdkiiIoMmuOperationBusMasterCommonBuffer ||\r
+      MapInfo->Operation == EdkiiIoMmuOperationBusMasterCommonBuffer64) {\r
+    CopyMem (\r
+      (VOID *)(UINTN)MapInfo->CryptedAddress,\r
+      CommonBufferHeader->StashBuffer,\r
+      MapInfo->NumberOfBytes\r
+      );\r
+  } else {\r
+    ZeroMem (\r
+      (VOID *)(UINTN)MapInfo->PlainTextAddress,\r
+      EFI_PAGES_TO_SIZE (MapInfo->NumberOfPages)\r
+      );\r
+    if (!MemoryMapLocked) {\r
+      gBS->FreePages (MapInfo->PlainTextAddress, MapInfo->NumberOfPages);\r
+    }\r
+  }\r
+\r
+  //\r
+  // Forget the MAP_INFO structure, then free it (unless the UEFI memory map is\r
+  // locked).\r
+  //\r
+  RemoveEntryList (&MapInfo->Link);\r
+  if (!MemoryMapLocked) {\r
+    FreePool (MapInfo);\r
+  }\r
+\r
   return EFI_SUCCESS;\r
 }\r
 \r
+/**\r
+  Completes the Map() operation and releases any corresponding resources.\r
+\r
+  @param  This                  The protocol instance pointer.\r
+  @param  Mapping               The mapping value returned from Map().\r
+\r
+  @retval EFI_SUCCESS           The range was unmapped.\r
+  @retval EFI_INVALID_PARAMETER Mapping is not a value that was returned by\r
+                                Map().\r
+  @retval EFI_DEVICE_ERROR      The data was not committed to the target system\r
+                                memory.\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+IoMmuUnmap (\r
+  IN  EDKII_IOMMU_PROTOCOL                     *This,\r
+  IN  VOID                                     *Mapping\r
+  )\r
+{\r
+  return IoMmuUnmapWorker (\r
+           This,\r
+           Mapping,\r
+           FALSE    // MemoryMapLocked\r
+           );\r
+}\r
+\r
 /**\r
   Allocates pages that are suitable for an OperationBusMasterCommonBuffer or\r
   OperationBusMasterCommonBuffer64 mapping.\r
@@ -328,6 +519,18 @@ IoMmuAllocateBuffer (
 {\r
   EFI_STATUS                Status;\r
   EFI_PHYSICAL_ADDRESS      PhysicalAddress;\r
+  VOID                      *StashBuffer;\r
+  UINTN                     CommonBufferPages;\r
+  COMMON_BUFFER_HEADER      *CommonBufferHeader;\r
+\r
+  DEBUG ((\r
+    DEBUG_VERBOSE,\r
+    "%a: MemoryType=%u Pages=0x%Lx Attributes=0x%Lx\n",\r
+    __FUNCTION__,\r
+    (UINT32)MemoryType,\r
+    (UINT64)Pages,\r
+    Attributes\r
+    ));\r
 \r
   //\r
   // Validate Attributes\r
@@ -352,6 +555,32 @@ IoMmuAllocateBuffer (
     return EFI_INVALID_PARAMETER;\r
   }\r
 \r
+  //\r
+  // We'll need a header page for the COMMON_BUFFER_HEADER structure.\r
+  //\r
+  if (Pages > MAX_UINTN - 1) {\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+  CommonBufferPages = Pages + 1;\r
+\r
+  //\r
+  // Allocate the stash in EfiBootServicesData type memory.\r
+  //\r
+  // Map() will temporarily save encrypted data in the stash for\r
+  // BusMasterCommonBuffer[64] operations, so the data can be decrypted to the\r
+  // original location.\r
+  //\r
+  // Unmap() will temporarily save plaintext data in the stash for\r
+  // BusMasterCommonBuffer[64] operations, so the data can be encrypted to the\r
+  // original location.\r
+  //\r
+  // StashBuffer always resides in encrypted memory.\r
+  //\r
+  StashBuffer = AllocatePages (Pages);\r
+  if (StashBuffer == NULL) {\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+\r
   PhysicalAddress = (UINTN)-1;\r
   if ((Attributes & EDKII_IOMMU_ATTRIBUTE_DUAL_ADDRESS_CYCLE) == 0) {\r
     //\r
@@ -362,26 +591,32 @@ IoMmuAllocateBuffer (
   Status = gBS->AllocatePages (\r
                   AllocateMaxAddress,\r
                   MemoryType,\r
-                  Pages,\r
+                  CommonBufferPages,\r
                   &PhysicalAddress\r
                   );\r
-  if (!EFI_ERROR (Status)) {\r
-    *HostAddress = (VOID *) (UINTN) PhysicalAddress;\r
-\r
-    //\r
-    // Clear memory encryption mask\r
-    //\r
-    Status = MemEncryptSevClearPageEncMask (0, PhysicalAddress, Pages, TRUE);\r
-    ASSERT_EFI_ERROR(Status);\r
+  if (EFI_ERROR (Status)) {\r
+    goto FreeStashBuffer;\r
   }\r
 \r
+  CommonBufferHeader = (VOID *)(UINTN)PhysicalAddress;\r
+  PhysicalAddress += EFI_PAGE_SIZE;\r
+\r
+  CommonBufferHeader->Signature = COMMON_BUFFER_SIG;\r
+  CommonBufferHeader->StashBuffer = StashBuffer;\r
+\r
+  *HostAddress = (VOID *)(UINTN)PhysicalAddress;\r
+\r
   DEBUG ((\r
     DEBUG_VERBOSE,\r
-    "%a Address 0x%Lx Pages 0x%Lx\n",\r
+    "%a: Host=0x%Lx Stash=0x%p\n",\r
     __FUNCTION__,\r
     PhysicalAddress,\r
-    (UINT64)Pages\r
+    StashBuffer\r
     ));\r
+  return EFI_SUCCESS;\r
+\r
+FreeStashBuffer:\r
+  FreePages (StashBuffer, Pages);\r
   return Status;\r
 }\r
 \r
@@ -406,28 +641,41 @@ IoMmuFreeBuffer (
   IN  VOID                                     *HostAddress\r
   )\r
 {\r
-  EFI_STATUS  Status;\r
-\r
-  //\r
-  // Set memory encryption mask\r
-  //\r
-  Status = MemEncryptSevSetPageEncMask (\r
-             0,\r
-             (EFI_PHYSICAL_ADDRESS)(UINTN)HostAddress,\r
-             Pages,\r
-             TRUE\r
-             );\r
-  ASSERT_EFI_ERROR(Status);\r
-  ZeroMem (HostAddress, EFI_PAGES_TO_SIZE (Pages));\r
+  UINTN                CommonBufferPages;\r
+  COMMON_BUFFER_HEADER *CommonBufferHeader;\r
 \r
   DEBUG ((\r
     DEBUG_VERBOSE,\r
-    "%a Address 0x%Lx Pages 0x%Lx\n",\r
+    "%a: Host=0x%p Pages=0x%Lx\n",\r
     __FUNCTION__,\r
-    (UINT64)(UINTN)HostAddress,\r
+    HostAddress,\r
     (UINT64)Pages\r
     ));\r
-  return gBS->FreePages ((EFI_PHYSICAL_ADDRESS) (UINTN) HostAddress, Pages);\r
+\r
+  CommonBufferPages = Pages + 1;\r
+  CommonBufferHeader = (COMMON_BUFFER_HEADER *)(\r
+                         (UINTN)HostAddress - EFI_PAGE_SIZE\r
+                         );\r
+\r
+  //\r
+  // Check the signature.\r
+  //\r
+  ASSERT (CommonBufferHeader->Signature == COMMON_BUFFER_SIG);\r
+  if (CommonBufferHeader->Signature != COMMON_BUFFER_SIG) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  //\r
+  // Free the stash buffer. This buffer was always encrypted, so no need to\r
+  // zero it.\r
+  //\r
+  FreePages (CommonBufferHeader->StashBuffer, Pages);\r
+\r
+  //\r
+  // Release the common buffer itself. Unmap() has re-encrypted it in-place, so\r
+  // no need to zero it.\r
+  //\r
+  return gBS->FreePages ((UINTN)CommonBufferHeader, CommonBufferPages);\r
 }\r
 \r
 \r
@@ -500,6 +748,107 @@ EDKII_IOMMU_PROTOCOL  mAmdSev = {
   IoMmuFreeBuffer,\r
 };\r
 \r
+/**\r
+  Notification function that is queued when gBS->ExitBootServices() signals the\r
+  EFI_EVENT_GROUP_EXIT_BOOT_SERVICES event group. This function signals another\r
+  event, received as Context, and returns.\r
+\r
+  Signaling an event in this context is safe. The UEFI spec allows\r
+  gBS->SignalEvent() to return EFI_SUCCESS only; EFI_OUT_OF_RESOURCES is not\r
+  listed, hence memory is not allocated. The edk2 implementation also does not\r
+  release memory (and we only have to care about the edk2 implementation\r
+  because EDKII_IOMMU_PROTOCOL is edk2-specific anyway).\r
+\r
+  @param[in] Event          Event whose notification function is being invoked.\r
+                            Event is permitted to request the queueing of this\r
+                            function at TPL_CALLBACK or TPL_NOTIFY task\r
+                            priority level.\r
+\r
+  @param[in] EventToSignal  Identifies the EFI_EVENT to signal. EventToSignal\r
+                            is permitted to request the queueing of its\r
+                            notification function only at TPL_CALLBACK level.\r
+**/\r
+STATIC\r
+VOID\r
+EFIAPI\r
+AmdSevExitBoot (\r
+  IN EFI_EVENT Event,\r
+  IN VOID      *EventToSignal\r
+  )\r
+{\r
+  //\r
+  // (1) The NotifyFunctions of all the events in\r
+  //     EFI_EVENT_GROUP_EXIT_BOOT_SERVICES will have been queued before\r
+  //     AmdSevExitBoot() is entered.\r
+  //\r
+  // (2) AmdSevExitBoot() is executing minimally at TPL_CALLBACK.\r
+  //\r
+  // (3) AmdSevExitBoot() has been queued in unspecified order relative to the\r
+  //     NotifyFunctions of all the other events in\r
+  //     EFI_EVENT_GROUP_EXIT_BOOT_SERVICES whose NotifyTpl is the same as\r
+  //     Event's.\r
+  //\r
+  // Consequences:\r
+  //\r
+  // - If Event's NotifyTpl is TPL_CALLBACK, then some other NotifyFunctions\r
+  //   queued at TPL_CALLBACK may be invoked after AmdSevExitBoot() returns.\r
+  //\r
+  // - If Event's NotifyTpl is TPL_NOTIFY, then some other NotifyFunctions\r
+  //   queued at TPL_NOTIFY may be invoked after AmdSevExitBoot() returns; plus\r
+  //   *all* NotifyFunctions queued at TPL_CALLBACK will be invoked strictly\r
+  //   after all NotifyFunctions queued at TPL_NOTIFY, including\r
+  //   AmdSevExitBoot(), have been invoked.\r
+  //\r
+  // - By signaling EventToSignal here, whose NotifyTpl is TPL_CALLBACK, we\r
+  //   queue EventToSignal's NotifyFunction after the NotifyFunctions of *all*\r
+  //   events in EFI_EVENT_GROUP_EXIT_BOOT_SERVICES.\r
+  //\r
+  DEBUG ((DEBUG_VERBOSE, "%a\n", __FUNCTION__));\r
+  gBS->SignalEvent (EventToSignal);\r
+}\r
+\r
+/**\r
+  Notification function that is queued after the notification functions of all\r
+  events in the EFI_EVENT_GROUP_EXIT_BOOT_SERVICES event group. The same memory\r
+  map restrictions apply.\r
+\r
+  This function unmaps all currently existing IOMMU mappings.\r
+\r
+  @param[in] Event    Event whose notification function is being invoked. Event\r
+                      is permitted to request the queueing of this function\r
+                      only at TPL_CALLBACK task priority level.\r
+\r
+  @param[in] Context  Ignored.\r
+**/\r
+STATIC\r
+VOID\r
+EFIAPI\r
+AmdSevUnmapAllMappings (\r
+  IN EFI_EVENT Event,\r
+  IN VOID      *Context\r
+  )\r
+{\r
+  LIST_ENTRY *Node;\r
+  LIST_ENTRY *NextNode;\r
+  MAP_INFO   *MapInfo;\r
+\r
+  DEBUG ((DEBUG_VERBOSE, "%a\n", __FUNCTION__));\r
+\r
+  //\r
+  // All drivers that had set up IOMMU mappings have halted their respective\r
+  // controllers by now; tear down the mappings.\r
+  //\r
+  for (Node = GetFirstNode (&mMapInfos); Node != &mMapInfos; Node = NextNode) {\r
+    NextNode = GetNextNode (&mMapInfos, Node);\r
+    MapInfo = CR (Node, MAP_INFO, Link, MAP_INFO_SIG);\r
+    IoMmuUnmapWorker (\r
+      &mAmdSev, // This\r
+      MapInfo,  // Mapping\r
+      TRUE      // MemoryMapLocked\r
+      );\r
+  }\r
+}\r
+\r
 /**\r
   Initialize Iommu Protocol.\r
 \r
@@ -511,13 +860,57 @@ AmdSevInstallIoMmuProtocol (
   )\r
 {\r
   EFI_STATUS  Status;\r
+  EFI_EVENT   UnmapAllMappingsEvent;\r
+  EFI_EVENT   ExitBootEvent;\r
   EFI_HANDLE  Handle;\r
 \r
+  //\r
+  // Create the "late" event whose notification function will tear down all\r
+  // left-over IOMMU mappings.\r
+  //\r
+  Status = gBS->CreateEvent (\r
+                  EVT_NOTIFY_SIGNAL,      // Type\r
+                  TPL_CALLBACK,           // NotifyTpl\r
+                  AmdSevUnmapAllMappings, // NotifyFunction\r
+                  NULL,                   // NotifyContext\r
+                  &UnmapAllMappingsEvent  // Event\r
+                  );\r
+  if (EFI_ERROR (Status)) {\r
+    return Status;\r
+  }\r
+\r
+  //\r
+  // Create the event whose notification function will be queued by\r
+  // gBS->ExitBootServices() and will signal the event created above.\r
+  //\r
+  Status = gBS->CreateEvent (\r
+                  EVT_SIGNAL_EXIT_BOOT_SERVICES, // Type\r
+                  TPL_CALLBACK,                  // NotifyTpl\r
+                  AmdSevExitBoot,                // NotifyFunction\r
+                  UnmapAllMappingsEvent,         // NotifyContext\r
+                  &ExitBootEvent                 // Event\r
+                  );\r
+  if (EFI_ERROR (Status)) {\r
+    goto CloseUnmapAllMappingsEvent;\r
+  }\r
+\r
   Handle = NULL;\r
   Status = gBS->InstallMultipleProtocolInterfaces (\r
                   &Handle,\r
                   &gEdkiiIoMmuProtocolGuid, &mAmdSev,\r
                   NULL\r
                   );\r
+  if (EFI_ERROR (Status)) {\r
+    goto CloseExitBootEvent;\r
+  }\r
+\r
+  return EFI_SUCCESS;\r
+\r
+CloseExitBootEvent:\r
+  gBS->CloseEvent (ExitBootEvent);\r
+\r
+CloseUnmapAllMappingsEvent:\r
+  gBS->CloseEvent (UnmapAllMappingsEvent);\r
+\r
   return Status;\r
 }\r