Add IoMmu support to NonDiscoverablePciDeviceDxe

Uses IoMmuLib to do IoMmu protocol mappings.

Author: eeshanl

MdeModulePkg/Bus/Pci/NonDiscoverablePciDeviceDxe/NonDiscoverablePciDeviceDxe.inf

@@ -33,6 +33,7 @@
   UefiBootServicesTableLib
   UefiDriverEntryPoint
   UefiLib
+  IoMmuLib                                      ## MU_CHANGE
 
 [Protocols]
   gEfiPciIoProtocolGuid                         ## BY_START

MdeModulePkg/Bus/Pci/NonDiscoverablePciDeviceDxe/NonDiscoverablePciDeviceIo.c

@@ -15,11 +15,14 @@
 
 #include <Protocol/PciRootBridgeIo.h>
 
+#include <Library/IoMmuLib.h> // MU_CHANGE
+
 typedef struct {
   EFI_PHYSICAL_ADDRESS             AllocAddress;
   VOID                             *HostAddress;
   EFI_PCI_IO_PROTOCOL_OPERATION    Operation;
   UINTN                            NumberOfBytes;
+  VOID                             *IoMmuContext; // MU_CHANGE
 } NON_DISCOVERABLE_PCI_DEVICE_MAP_INFO;
 
 /**
@@ -1278,6 +1281,13 @@ NonCoherentPciIoMap (
   EFI_GCD_MEMORY_SPACE_DESCRIPTOR       GcdDescriptor;
   BOOLEAN                               Bounce;
 
+  // MU_CHANGE [BEGIN]
+  VOID                   *IoMmuHostAddress;
+  EDKII_IOMMU_OPERATION  IoMmuOperation;
+  UINT64                 IoMmuAttribute;
+
+  // MU_CHANGE [END]
+
   AllocAddress = NULL; // MS_CHANGE for vs2017
 
   if ((HostAddress   == NULL) ||
@@ -1303,6 +1313,7 @@ NonCoherentPciIoMap (
   MapInfo->HostAddress   = HostAddress;
   MapInfo->Operation     = Operation;
   MapInfo->NumberOfBytes = *NumberOfBytes;
+  MapInfo->IoMmuContext  = NULL; // MU_CHANGE
 
   Dev = NON_DISCOVERABLE_PCI_DEVICE_FROM_PCI_IO (This);
 
@@ -1372,10 +1383,12 @@ NonCoherentPciIoMap (
       gBS->CopyMem (AllocAddress, HostAddress, *NumberOfBytes);
     }
 
-    *DeviceAddress = MapInfo->AllocAddress;
+    *DeviceAddress   = MapInfo->AllocAddress;
+    IoMmuHostAddress = (VOID *)(UINTN)MapInfo->AllocAddress; // MU_CHANGE
   } else {
     MapInfo->AllocAddress = 0;
     *DeviceAddress        = (EFI_PHYSICAL_ADDRESS)(UINTN)HostAddress;
+    IoMmuHostAddress      = HostAddress; // MU_CHANGE
 
     //
     // We are not using a bounce buffer: the mapping is sufficiently
@@ -1396,6 +1409,60 @@ NonCoherentPciIoMap (
   }
 
   *Mapping = MapInfo;
+
+  // MU_CHANGE [BEGIN]
+  switch (Operation) {
+    case EfiPciIoOperationBusMasterRead:
+      IoMmuOperation = EdkiiIoMmuOperationBusMasterRead;
+      IoMmuAttribute = EDKII_IOMMU_ACCESS_READ;
+      break;
+
+    case EfiPciIoOperationBusMasterWrite:
+      IoMmuOperation = EdkiiIoMmuOperationBusMasterWrite;
+      IoMmuAttribute = EDKII_IOMMU_ACCESS_WRITE;
+      break;
+
+    case EfiPciIoOperationBusMasterCommonBuffer:
+      IoMmuOperation = EdkiiIoMmuOperationBusMasterCommonBuffer;
+      IoMmuAttribute = EDKII_IOMMU_ACCESS_READ | EDKII_IOMMU_ACCESS_WRITE;
+      break;
+
+    default:
+      DEBUG ((DEBUG_ERROR, "%a - Invalid operation %d\n", __func__, Operation));
+      ASSERT (FALSE);
+      return EFI_INVALID_PARAMETER;
+  }
+
+  Status = IoMmuMap (
+             IoMmuOperation,
+             IoMmuHostAddress,
+             NumberOfBytes,
+             DeviceAddress,
+             &MapInfo->IoMmuContext
+             );
+  if (Status != EFI_UNSUPPORTED) {
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "%a - IoMmuMap failed.\n", __func__));
+      ASSERT (FALSE);
+      return Status;
+    }
+  }
+
+  Status = IoMmuSetAttribute (
+             NULL,
+             MapInfo->IoMmuContext,
+             IoMmuAttribute
+             );
+  if (Status != EFI_UNSUPPORTED) {
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "%a - IoMmuSetAttribute failed.\n", __func__));
+      ASSERT (FALSE);
+      return Status;
+    }
+  }
+
+  // MU_CHANGE [END]
+
   return EFI_SUCCESS;
 
 FreeMapInfo:
@@ -1422,12 +1489,26 @@ NonCoherentPciIoUnmap (
   )
 {
   NON_DISCOVERABLE_PCI_DEVICE_MAP_INFO  *MapInfo;
+  EFI_STATUS                            Status; // MU_CHANGE
 
   if (Mapping == NULL) {
     return EFI_DEVICE_ERROR;
   }
 
   MapInfo = Mapping;
+
+  // MU_CHANGE [BEGIN]
+  Status = IoMmuUnmap (MapInfo->IoMmuContext);
+  if (Status != EFI_UNSUPPORTED) {
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "%a - IoMmuUnmap failed.\n", __func__));
+      ASSERT (FALSE);
+      return Status;
+    }
+  }
+
+  // MU_CHANGE [END]
+
   if (MapInfo->AllocAddress != 0) {
     //
     // We are using a bounce buffer: copy back the data if necessary,