[PATCH] MdeModulePkg/PciBus: Disable BME of all devices when entering RT

[Ruiyu Ni <ruiyu.ni@intel.com>]

The patch assumes IOMMU protections are disabled after PciBus
disables the BMT bit in Command register.
It ensures all DMA transactions are protected by IOMMU.

Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Michael Turner <michael.turner@microsoft.com>
Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Cc: Michael D Kinney <michael.d.kinney@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
---
 MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h           |  2 +
 MdeModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf      |  3 +
 MdeModulePkg/Bus/Pci/PciBusDxe/PciDeviceSupport.c | 86 +++++++++++++++++++++++
 3 files changed, 91 insertions(+)

diff --git a/MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h b/MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h
index 55eb3a5a80..79b5b71082 100644
--- a/MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h
+++ b/MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h
@@ -18,6 +18,8 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 #include <PiDxe.h>
 
+#include <Guid/EventGroup.h>
+
 #include <Protocol/LoadedImage.h>
 #include <Protocol/PciHostBridgeResourceAllocation.h>
 #include <Protocol/PciIo.h>
diff --git a/MdeModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf b/MdeModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf
index 97608bfcf2..d5b8fab3ca 100644
--- a/MdeModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf
+++ b/MdeModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf
@@ -80,6 +80,9 @@ [LibraryClasses]
   DebugLib
   PeCoffLib
 
+[Guids]
+  gEfiEventExitBootServicesGuid                   ## SOMETIMES_CONSUMES ## Event
+
 [Protocols]
   gEfiPciHotPlugRequestProtocolGuid               ## SOMETIMES_PRODUCES
   gEfiPciIoProtocolGuid                           ## BY_START
diff --git a/MdeModulePkg/Bus/Pci/PciBusDxe/PciDeviceSupport.c b/MdeModulePkg/Bus/Pci/PciBusDxe/PciDeviceSupport.c
index 97bb971a59..b5530a13d1 100644
--- a/MdeModulePkg/Bus/Pci/PciBusDxe/PciDeviceSupport.c
+++ b/MdeModulePkg/Bus/Pci/PciBusDxe/PciDeviceSupport.c
@@ -21,6 +21,72 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 LIST_ENTRY  mPciDevicePool;
 
 /**
+ Disable Bus Master Enable bit in all devices in the list.
+
+ @param Devices  A device list.
+**/
+VOID
+DisableBmeOnTree (
+  IN LIST_ENTRY      *Devices
+  )
+{
+  LIST_ENTRY      *Link;
+  PCI_IO_DEVICE   *PciIoDevice;
+  UINT16           Command;
+
+  for ( Link = GetFirstNode (Devices)
+      ; !IsNull (Devices, Link)
+      ; Link = GetNextNode (Devices, Link)
+      ) {
+    PciIoDevice = PCI_IO_DEVICE_FROM_LINK (Link);
+    //
+    // Turn off all children's Bus Master, if any
+    //
+    DisableBmeOnTree (&PciIoDevice->ChildList);
+
+    //
+    // If this is a device that supports BME, disable BME on this device.
+    //
+    if ((PciIoDevice->Supports & EFI_PCI_IO_ATTRIBUTE_BUS_MASTER) != 0) {
+      PCI_READ_COMMAND_REGISTER(PciIoDevice, &Command);
+      if ((Command & EFI_PCI_COMMAND_BUS_MASTER) != 0) {
+        Command &= ~EFI_PCI_COMMAND_BUS_MASTER;
+        PCI_SET_COMMAND_REGISTER (PciIoDevice, Command);
+        DEBUG ((
+          EFI_D_INFO,"  %02x   %02x      %02x         %04x\n",
+          PciIoDevice->BusNumber, PciIoDevice->DeviceNumber, PciIoDevice->FunctionNumber,
+          Command
+          ));
+      }
+    }
+  }
+}
+
+/**
+  Exit Boot Services Event notification handler.
+
+  Disable Bus Master on any that were enabled during BDS.
+
+  @param[in]  Event     Event whose notification function is being invoked.
+  @param[in]  Context   Pointer to the notification function's context.
+
+**/
+VOID
+EFIAPI
+OnExitBootServices (
+  IN      EFI_EVENT                 Event,
+  IN      VOID                      *Context
+  )
+{
+  DEBUG ((
+    DEBUG_INFO,
+    "PciBus: Disable Bus Master of all devices...\n"
+    "  Bus# Device# Function#  NewCommand\n"
+    ));
+  DisableBmeOnTree(&mPciDevicePool);
+}
+
+/**
   Initialize the PCI devices pool.
 
 **/
@@ -29,7 +95,27 @@ InitializePciDevicePool (
   VOID
   )
 {
+  EFI_EVENT   ExitBootServicesEvent;
+  EFI_STATUS  Status;
+
   InitializeListHead (&mPciDevicePool);
+
+  //
+  // DisableBME on ExitBootServices should be synchonized with any IOMMU ExitBootServices routine.
+  // DisableBME should be run before the IOMMU protections are disabled.
+  // One way to do this is to ensure that the IOMMU ExitBootServices callback runs at TPL_CALLBACK.
+  //
+  Status = gBS->CreateEventEx (
+                  EVT_NOTIFY_SIGNAL,
+                  TPL_NOTIFY,
+                  OnExitBootServices,
+                  NULL,
+                  &gEfiEventExitBootServicesGuid,
+                  &ExitBootServicesEvent
+                  );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "PciBus: Unable to hook ExitBootServices event - %r\n", Status));
+  }
 }
 
 /**
-- 
2.12.2.windows.2