[PATCH V2 3/4] OvmfPkg/IoMmuDxe: Add SEV support for reserved shared memory

["Min Xu" <min.m.xu@intel.com>]

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

Add support to use the reserved shared memory within the IoMmu library.
This improves boot times for all SEV guests, with SEV-SNP benefiting the
most as it avoids the page state change call to the hypervisor.

Cc: Erdem Aktas <erdemaktas@google.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Min Xu <min.m.xu@intel.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
Acked-by: Jiewen Yao <Jiewen.yao@intel.com>
Signed-off-by: Tom Lendacky <thomas.lendacky@amd.com>
---
 OvmfPkg/IoMmuDxe/CcIoMmu.c     | 81 +++++++++++++++++-----------------
 OvmfPkg/IoMmuDxe/IoMmuBuffer.c | 54 ++++++++++++++++++-----
 2 files changed, 83 insertions(+), 52 deletions(-)

diff --git a/OvmfPkg/IoMmuDxe/CcIoMmu.c b/OvmfPkg/IoMmuDxe/CcIoMmu.c
index 1479af469881..e5cbf037c50d 100644
--- a/OvmfPkg/IoMmuDxe/CcIoMmu.c
+++ b/OvmfPkg/IoMmuDxe/CcIoMmu.c
@@ -223,30 +223,33 @@ IoMmuMap (
       goto FreeMapInfo;
   }
 
-  if (CC_GUEST_IS_SEV (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
-    //
-    // Clear the memory encryption mask on the plaintext buffer.
-    //
-    Status = MemEncryptSevClearPageEncMask (
-               0,
-               MapInfo->PlainTextAddress,
-               MapInfo->NumberOfPages
-               );
-  } else if (CC_GUEST_IS_TDX (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+  if (MapInfo->ReservedMemBitmap == 0) {
     //
     // Set the memory shared bit.
     // If MapInfo->ReservedMemBitmap is 0, it means the bounce buffer is not allocated
     // from the pre-allocated shared memory, so it must be converted to shared memory here.
     //
-    if (MapInfo->ReservedMemBitmap == 0) {
+    if (CC_GUEST_IS_SEV (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+      //
+      // Clear the memory encryption mask on the plaintext buffer.
+      //
+      Status = MemEncryptSevClearPageEncMask (
+                 0,
+                 MapInfo->PlainTextAddress,
+                 MapInfo->NumberOfPages
+                 );
+    } else if (CC_GUEST_IS_TDX (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+      //
+      // Set the memory shared bit.
+      //
       Status = MemEncryptTdxSetPageSharedBit (
                  0,
                  MapInfo->PlainTextAddress,
                  MapInfo->NumberOfPages
                  );
+    } else {
+      ASSERT (FALSE);
     }
-  } else {
-    ASSERT (FALSE);
   }
 
   ASSERT_EFI_ERROR (Status);
@@ -396,30 +399,30 @@ IoMmuUnmapWorker (
       break;
   }
 
-  if (CC_GUEST_IS_SEV (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
-    //
-    // Restore the memory encryption mask on the area we used to hold the
-    // plaintext.
-    //
-    Status = MemEncryptSevSetPageEncMask (
-               0,
-               MapInfo->PlainTextAddress,
-               MapInfo->NumberOfPages
-               );
-  } else if (CC_GUEST_IS_TDX (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
-    //
-    // Restore the memory shared bit mask on the area we used to hold the
-    // plaintext.
-    //
-    if (MapInfo->ReservedMemBitmap == 0) {
+  if (MapInfo->ReservedMemBitmap == 0) {
+    if (CC_GUEST_IS_SEV (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+      //
+      // Restore the memory encryption mask on the area we used to hold the
+      // plaintext.
+      //
+      Status = MemEncryptSevSetPageEncMask (
+                 0,
+                 MapInfo->PlainTextAddress,
+                 MapInfo->NumberOfPages
+                 );
+    } else if (CC_GUEST_IS_TDX (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+      //
+      // Restore the memory shared bit mask on the area we used to hold the
+      // plaintext.
+      //
       Status = MemEncryptTdxClearPageSharedBit (
                  0,
                  MapInfo->PlainTextAddress,
                  MapInfo->NumberOfPages
                  );
+    } else {
+      ASSERT (FALSE);
     }
-  } else {
-    ASSERT (FALSE);
   }
 
   ASSERT_EFI_ERROR (Status);
@@ -924,16 +927,14 @@ InstallIoMmuProtocol (
   }
 
   //
-  // Currently only Tdx guest support Reserved shared memory for DMA operation.
+  // For CC guests, use reserved shared memory for DMA operation.
   //
-  if (CC_GUEST_IS_TDX (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
-    mReservedSharedMemSupported = TRUE;
-    Status                      = IoMmuInitReservedSharedMem ();
-    if (EFI_ERROR (Status)) {
-      mReservedSharedMemSupported = FALSE;
-    } else {
-      DEBUG ((DEBUG_INFO, "%a: Feature of reserved memory for DMA is supported.\n", __FUNCTION__));
-    }
+  mReservedSharedMemSupported = TRUE;
+  Status                      = IoMmuInitReservedSharedMem ();
+  if (EFI_ERROR (Status)) {
+    mReservedSharedMemSupported = FALSE;
+  } else {
+    DEBUG ((DEBUG_INFO, "%a: Feature of reserved memory for DMA is supported.\n", __FUNCTION__));
   }
 
   return EFI_SUCCESS;
diff --git a/OvmfPkg/IoMmuDxe/IoMmuBuffer.c b/OvmfPkg/IoMmuDxe/IoMmuBuffer.c
index 1e77d8a57402..3139d10f4c2d 100644
--- a/OvmfPkg/IoMmuDxe/IoMmuBuffer.c
+++ b/OvmfPkg/IoMmuDxe/IoMmuBuffer.c
@@ -9,7 +9,9 @@
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
 #include <Library/MemoryAllocationLib.h>
+#include <Library/MemEncryptSevLib.h>
 #include <Library/MemEncryptTdxLib.h>
+#include <Library/PcdLib.h>
 #include <Library/UefiBootServicesTableLib.h>
 #include "IoMmuInternal.h"
 
@@ -139,6 +141,7 @@ IoMmuInitReservedSharedMem (
   UINTN                     TotalPages;
   IOMMU_RESERVED_MEM_RANGE  *MemRange;
   EFI_PHYSICAL_ADDRESS      PhysicalAddress;
+  UINT64                    SharedAddress;
 
   if (!mReservedSharedMemSupported) {
     return EFI_UNSUPPORTED;
@@ -163,12 +166,25 @@ IoMmuInitReservedSharedMem (
     MemRange->StartAddressOfMemRange = PhysicalAddress;
 
     for (Index2 = 0; Index2 < MemRange->Slots; Index2++) {
-      Status = MemEncryptTdxSetPageSharedBit (
-                 0,
-                 (UINT64)(UINTN)(MemRange->StartAddressOfMemRange + Index2 * SIZE_OF_MEM_RANGE (MemRange) + MemRange->HeaderSize),
-                 EFI_SIZE_TO_PAGES (MemRange->DataSize)
-                 );
-      ASSERT (!EFI_ERROR (Status));
+      SharedAddress = (UINT64)(UINTN)(MemRange->StartAddressOfMemRange + Index2 * SIZE_OF_MEM_RANGE (MemRange) + MemRange->HeaderSize);
+
+      if (CC_GUEST_IS_SEV (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+        Status = MemEncryptSevClearPageEncMask (
+                   0,
+                   SharedAddress,
+                   EFI_SIZE_TO_PAGES (MemRange->DataSize)
+                   );
+        ASSERT (!EFI_ERROR (Status));
+      } else if (CC_GUEST_IS_TDX (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+        Status = MemEncryptTdxSetPageSharedBit (
+                   0,
+                   SharedAddress,
+                   EFI_SIZE_TO_PAGES (MemRange->DataSize)
+                   );
+        ASSERT (!EFI_ERROR (Status));
+      } else {
+        ASSERT (FALSE);
+      }
     }
 
     PhysicalAddress += (MemRange->Slots * SIZE_OF_MEM_RANGE (MemRange));
@@ -190,16 +206,30 @@ IoMmuReleaseReservedSharedMem (
   EFI_STATUS                Status;
   UINT32                    Index1, Index2;
   IOMMU_RESERVED_MEM_RANGE  *MemRange;
+  UINT64                    SharedAddress;
 
   for (Index1 = 0; Index1 < ARRAY_SIZE (mReservedMemRanges); Index1++) {
     MemRange = &mReservedMemRanges[Index1];
     for (Index2 = 0; Index2 < MemRange->Slots; Index2++) {
-      Status = MemEncryptTdxClearPageSharedBit (
-                 0,
-                 (UINT64)(UINTN)(MemRange->StartAddressOfMemRange + Index2 * SIZE_OF_MEM_RANGE (MemRange) + MemRange->HeaderSize),
-                 EFI_SIZE_TO_PAGES (MemRange->DataSize)
-                 );
-      ASSERT (!EFI_ERROR (Status));
+      SharedAddress = (UINT64)(UINTN)(MemRange->StartAddressOfMemRange + Index2 * SIZE_OF_MEM_RANGE (MemRange) + MemRange->HeaderSize);
+
+      if (CC_GUEST_IS_SEV (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+        Status = MemEncryptSevSetPageEncMask (
+                   0,
+                   SharedAddress,
+                   EFI_SIZE_TO_PAGES (MemRange->DataSize)
+                   );
+        ASSERT (!EFI_ERROR (Status));
+      } else if (CC_GUEST_IS_TDX (PcdGet64 (PcdConfidentialComputingGuestAttr))) {
+        Status = MemEncryptTdxClearPageSharedBit (
+                   0,
+                   SharedAddress,
+                   EFI_SIZE_TO_PAGES (MemRange->DataSize)
+                   );
+        ASSERT (!EFI_ERROR (Status));
+      } else {
+        ASSERT (FALSE);
+      }
     }
   }
 
-- 
2.29.2.windows.2