[PATCH V2 27/28] OvmfPkg: Update IoMmuDxe to support TDX

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

RFC： https://bugzilla.tianocore.org/show_bug.cgi?id=3429

The IOMMU protocol driver provides capabilities to set a DMA access
attribute and methods to allocate, free, map and unmap the DMA memory
for the PCI Bus devices.

The current IoMmuDxe driver supports DMA operations inside SEV guest.
To support DMA operation in TDX guest, mIoMmuType is added to determine
if it is Legac guest, SEV guest or TDX guest.

Due to security reasons all DMA operations inside the SEV/TDX guest must
be performed on shared pages. The IOMMU protocol driver for the SEV/TDX
guest uses a bounce buffer to map guest DMA buffer to shared pages in
order to provide the support for DMA operations inside SEV/TDX guest.

The call of SEV or TDX specific function to set/clear EncMask/SharedBit
is determined by mIoMmuType.

Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Jordan Justen <jordan.l.justen@intel.com>
Cc: Brijesh Singh <brijesh.singh@amd.com>
Cc: Erdem Aktas <erdemaktas@google.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Signed-off-by: Min Xu <min.m.xu@intel.com>
---
 OvmfPkg/IoMmuDxe/AmdSevIoMmu.c | 134 ++++++++++++++++++++++++---------
 OvmfPkg/IoMmuDxe/AmdSevIoMmu.h |  12 +++
 OvmfPkg/IoMmuDxe/IoMmuDxe.c    |   4 +-
 OvmfPkg/IoMmuDxe/IoMmuDxe.inf  |   1 +
 OvmfPkg/OvmfPkgX64.dsc         |   2 +
 5 files changed, 117 insertions(+), 36 deletions(-)

diff --git a/OvmfPkg/IoMmuDxe/AmdSevIoMmu.c b/OvmfPkg/IoMmuDxe/AmdSevIoMmu.c
index b30628078f73..6c9570677cdc 100644
--- a/OvmfPkg/IoMmuDxe/AmdSevIoMmu.c
+++ b/OvmfPkg/IoMmuDxe/AmdSevIoMmu.c
@@ -1,9 +1,9 @@
 /** @file
 
   The protocol provides support to allocate, free, map and umap a DMA buffer
-  for bus master (e.g PciHostBridge). When SEV is enabled, the DMA operations
-  must be performed on unencrypted buffer hence we use a bounce buffer to map
-  the guest buffer into an unencrypted DMA buffer.
+  for bus master (e.g PciHostBridge). When SEV or TDX is enabled, the DMA
+  operations must be performed on unencrypted buffer hence we use a bounce
+  buffer to map the guest buffer into an unencrypted DMA buffer.
 
   Copyright (c) 2017, AMD Inc. All rights reserved.<BR>
   Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
@@ -14,6 +14,12 @@
 
 #include "AmdSevIoMmu.h"
 
+#define IO_MMU_LEGACY  0x0
+#define IO_MMU_SEV     0x01
+#define IO_MMU_TDX     0x02
+
+UINTN mIoMmuType = IO_MMU_LEGACY;
+
 #define MAP_INFO_SIG SIGNATURE_64 ('M', 'A', 'P', '_', 'I', 'N', 'F', 'O')
 
 typedef struct {
@@ -74,7 +80,7 @@ typedef struct {
 
 /**
   Provides the controller-specific addresses required to access system memory
-  from a DMA bus master. On SEV guest, the DMA operations must be performed on
+  from a DMA bus master. On SEV/TDX guest, the DMA operations must be performed on
   shared buffer hence we allocate a bounce buffer to map the HostAddress to a
   DeviceAddress. The Encryption attribute is removed from the DeviceAddress
   buffer.
@@ -246,14 +252,29 @@ IoMmuMap (
     goto FreeMapInfo;
   }
 
-  //
-  // Clear the memory encryption mask on the plaintext buffer.
-  //
-  Status = MemEncryptSevClearPageEncMask (
-             0,
-             MapInfo->PlainTextAddress,
-             MapInfo->NumberOfPages
-             );
+  if (mIoMmuType == IO_MMU_SEV) {
+    //
+    // Clear the memory encryption mask on the plaintext buffer.
+    //
+    Status = MemEncryptSevClearPageEncMask (
+               0,
+               MapInfo->PlainTextAddress,
+               MapInfo->NumberOfPages
+               );
+  } else if (mIoMmuType == IO_MMU_TDX) {
+    //
+    // Set the memory shared bit.
+    //
+    Status = MemEncryptTdxSetPageSharedBit (
+               0,
+               MapInfo->PlainTextAddress,
+               MapInfo->NumberOfPages
+               );
+
+  } else {
+    ASSERT (FALSE);
+  }
+
   ASSERT_EFI_ERROR (Status);
   if (EFI_ERROR (Status)) {
     CpuDeadLoop ();
@@ -399,15 +420,30 @@ IoMmuUnmapWorker (
     break;
   }
 
-  //
-  // Restore the memory encryption mask on the area we used to hold the
-  // plaintext.
-  //
-  Status = MemEncryptSevSetPageEncMask (
-             0,
-             MapInfo->PlainTextAddress,
-             MapInfo->NumberOfPages
-             );
+  if (mIoMmuType == IO_MMU_SEV) {
+    //
+    // Restore the memory encryption mask on the area we used to hold the
+    // plaintext.
+    //
+    Status = MemEncryptSevSetPageEncMask (
+               0,
+               MapInfo->PlainTextAddress,
+               MapInfo->NumberOfPages
+               );
+  } else if (mIoMmuType == IO_MMU_TDX) {
+    //
+    // 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);
+  }
+
   ASSERT_EFI_ERROR (Status);
   if (EFI_ERROR (Status)) {
     CpuDeadLoop ();
@@ -731,7 +767,7 @@ IoMmuSetAttribute (
   return EFI_UNSUPPORTED;
 }
 
-EDKII_IOMMU_PROTOCOL  mAmdSev = {
+EDKII_IOMMU_PROTOCOL  mIoMmu = {
   EDKII_IOMMU_PROTOCOL_REVISION,
   IoMmuSetAttribute,
   IoMmuMap,
@@ -763,7 +799,7 @@ EDKII_IOMMU_PROTOCOL  mAmdSev = {
 STATIC
 VOID
 EFIAPI
-AmdSevExitBoot (
+IoMmuExitBoot (
   IN EFI_EVENT Event,
   IN VOID      *EventToSignal
   )
@@ -771,11 +807,11 @@ AmdSevExitBoot (
   //
   // (1) The NotifyFunctions of all the events in
   //     EFI_EVENT_GROUP_EXIT_BOOT_SERVICES will have been queued before
-  //     AmdSevExitBoot() is entered.
+  //     IoMmuExitBoot() is entered.
   //
-  // (2) AmdSevExitBoot() is executing minimally at TPL_CALLBACK.
+  // (2) IoMmuExitBoot() is executing minimally at TPL_CALLBACK.
   //
-  // (3) AmdSevExitBoot() has been queued in unspecified order relative to the
+  // (3) IoMmuExitBoot() has been queued in unspecified order relative to the
   //     NotifyFunctions of all the other events in
   //     EFI_EVENT_GROUP_EXIT_BOOT_SERVICES whose NotifyTpl is the same as
   //     Event's.
@@ -783,13 +819,13 @@ AmdSevExitBoot (
   // Consequences:
   //
   // - If Event's NotifyTpl is TPL_CALLBACK, then some other NotifyFunctions
-  //   queued at TPL_CALLBACK may be invoked after AmdSevExitBoot() returns.
+  //   queued at TPL_CALLBACK may be invoked after IoMmuExitBoot() returns.
   //
   // - If Event's NotifyTpl is TPL_NOTIFY, then some other NotifyFunctions
-  //   queued at TPL_NOTIFY may be invoked after AmdSevExitBoot() returns; plus
+  //   queued at TPL_NOTIFY may be invoked after IoMmuExitBoot() returns; plus
   //   *all* NotifyFunctions queued at TPL_CALLBACK will be invoked strictly
   //   after all NotifyFunctions queued at TPL_NOTIFY, including
-  //   AmdSevExitBoot(), have been invoked.
+  //   IoMmuExitBoot(), have been invoked.
   //
   // - By signaling EventToSignal here, whose NotifyTpl is TPL_CALLBACK, we
   //   queue EventToSignal's NotifyFunction after the NotifyFunctions of *all*
@@ -815,7 +851,7 @@ AmdSevExitBoot (
 STATIC
 VOID
 EFIAPI
-AmdSevUnmapAllMappings (
+IoMmuUnmapAllMappings (
   IN EFI_EVENT Event,
   IN VOID      *Context
   )
@@ -834,7 +870,7 @@ AmdSevUnmapAllMappings (
     NextNode = GetNextNode (&mMapInfos, Node);
     MapInfo = CR (Node, MAP_INFO, Link, MAP_INFO_SIG);
     IoMmuUnmapWorker (
-      &mAmdSev, // This
+      &mIoMmu,  // This
       MapInfo,  // Mapping
       TRUE      // MemoryMapLocked
       );
@@ -847,7 +883,7 @@ AmdSevUnmapAllMappings (
 **/
 EFI_STATUS
 EFIAPI
-AmdSevInstallIoMmuProtocol (
+InstallIoMmuProtocol (
   VOID
   )
 {
@@ -863,7 +899,7 @@ AmdSevInstallIoMmuProtocol (
   Status = gBS->CreateEvent (
                   EVT_NOTIFY_SIGNAL,      // Type
                   TPL_CALLBACK,           // NotifyTpl
-                  AmdSevUnmapAllMappings, // NotifyFunction
+                  IoMmuUnmapAllMappings,  // NotifyFunction
                   NULL,                   // NotifyContext
                   &UnmapAllMappingsEvent  // Event
                   );
@@ -878,7 +914,7 @@ AmdSevInstallIoMmuProtocol (
   Status = gBS->CreateEvent (
                   EVT_SIGNAL_EXIT_BOOT_SERVICES, // Type
                   TPL_CALLBACK,                  // NotifyTpl
-                  AmdSevExitBoot,                // NotifyFunction
+                  IoMmuExitBoot,                 // NotifyFunction
                   UnmapAllMappingsEvent,         // NotifyContext
                   &ExitBootEvent                 // Event
                   );
@@ -889,7 +925,7 @@ AmdSevInstallIoMmuProtocol (
   Handle = NULL;
   Status = gBS->InstallMultipleProtocolInterfaces (
                   &Handle,
-                  &gEdkiiIoMmuProtocolGuid, &mAmdSev,
+                  &gEdkiiIoMmuProtocolGuid, &mIoMmu,
                   NULL
                   );
   if (EFI_ERROR (Status)) {
@@ -906,3 +942,31 @@ CloseUnmapAllMappingsEvent:
 
   return Status;
 }
+
+/**
+  Initialize Iommu Protocol for Intel TDX.
+
+**/
+EFI_STATUS
+EFIAPI
+IntelTdxInstallIoMmuProtocol (
+  VOID
+  )
+{
+  mIoMmuType = IO_MMU_TDX;
+  return InstallIoMmuProtocol ();
+}
+
+/**
+  Initialize Iommu Protocol for Intel TDX.
+
+**/
+EFI_STATUS
+EFIAPI
+AmdSevInstallIoMmuProtocol (
+  VOID
+  )
+{
+  mIoMmuType = IO_MMU_SEV;
+  return InstallIoMmuProtocol ();
+}
diff --git a/OvmfPkg/IoMmuDxe/AmdSevIoMmu.h b/OvmfPkg/IoMmuDxe/AmdSevIoMmu.h
index 8244f28b57fd..768d18028198 100644
--- a/OvmfPkg/IoMmuDxe/AmdSevIoMmu.h
+++ b/OvmfPkg/IoMmuDxe/AmdSevIoMmu.h
@@ -21,6 +21,7 @@
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
 #include <Library/MemEncryptSevLib.h>
+#include <Library/MemEncryptTdxLib.h>
 #include <Library/MemoryAllocationLib.h>
 #include <Library/UefiBootServicesTableLib.h>
 
@@ -35,4 +36,15 @@ AmdSevInstallIoMmuProtocol (
   VOID
   );
 
+/**
+  Install IOMMU protocol to provide the DMA support for PciHostBridge and
+  MemEncryptSevLib.
+
+**/
+EFI_STATUS
+EFIAPI
+IntelTdxInstallIoMmuProtocol (
+  VOID
+  );
+
 #endif
diff --git a/OvmfPkg/IoMmuDxe/IoMmuDxe.c b/OvmfPkg/IoMmuDxe/IoMmuDxe.c
index 13df8ba874c5..698229b16bfa 100644
--- a/OvmfPkg/IoMmuDxe/IoMmuDxe.c
+++ b/OvmfPkg/IoMmuDxe/IoMmuDxe.c
@@ -22,11 +22,13 @@ IoMmuDxeEntryPoint (
   EFI_HANDLE    Handle;
 
   //
-  // When SEV is enabled, install IoMmu protocol otherwise install the
+  // When SEV or TDX is enabled, install IoMmu protocol otherwise install the
   // placeholder protocol so that other dependent module can run.
   //
   if (MemEncryptSevIsEnabled ()) {
     Status = AmdSevInstallIoMmuProtocol ();
+  } else if (MemEncryptTdxIsEnabled ()) {
+    Status = IntelTdxInstallIoMmuProtocol ();
   } else {
     Handle = NULL;
 
diff --git a/OvmfPkg/IoMmuDxe/IoMmuDxe.inf b/OvmfPkg/IoMmuDxe/IoMmuDxe.inf
index 2ebd74e5558c..52ad6f2efdb7 100644
--- a/OvmfPkg/IoMmuDxe/IoMmuDxe.inf
+++ b/OvmfPkg/IoMmuDxe/IoMmuDxe.inf
@@ -32,6 +32,7 @@
   BaseMemoryLib
   DebugLib
   MemEncryptSevLib
+  MemEncryptTdxLib
   MemoryAllocationLib
   UefiBootServicesTableLib
   UefiDriverEntryPoint
diff --git a/OvmfPkg/OvmfPkgX64.dsc b/OvmfPkg/OvmfPkgX64.dsc
index 2c4a6613b1ea..455e901c2eb8 100644
--- a/OvmfPkg/OvmfPkgX64.dsc
+++ b/OvmfPkg/OvmfPkgX64.dsc
@@ -180,6 +180,8 @@
   VirtioLib|OvmfPkg/Library/VirtioLib/VirtioLib.inf
   LoadLinuxLib|OvmfPkg/Library/LoadLinuxLib/LoadLinuxLib.inf
   MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
+  MemEncryptTdxLib|OvmfPkg/Library/BaseMemEncryptTdxLib/BaseMemEncryptTdxLib.inf
+
 !if $(SMM_REQUIRE) == FALSE
   LockBoxLib|OvmfPkg/Library/LockBoxLib/LockBoxBaseLib.inf
 !endif
-- 
2.29.2.windows.2