[PATCH] MdeModulePkg/NonDiscoverablePciDeviceDxe: connect handles on registration

["Ard Biesheuvel" <ard.biesheuvel@arm.com>]

The way EDK2 invokes the UEFI driver model assumes that PCI I/O
protocol instances exist for all PCI I/O controllers in the system.

For instance, UefiBootManagerLib connects the short-form USB device
path of the console input by looking for PCI I/O controllers that
have the 'USB host controller' class code, and passing each one to
ConnectController(), using the device path as the 'RemainingDevicePath'
argument.

For true PCI I/O protocol instances produced by the PCI root bridge
driver, this works fine, since it always enumerates the PCIe hierarchy
exhaustively. However, for platform devices that are wired to PCI class
drivers using the non-discoverable PCIe driver, this breaks down, due
to the fact that the PCI I/O protocol instance does not exist unless the
non-discoverable device protocol handle is connected first.

So let's connect these handles non-recursively as soon as they appear.

Cc: Hao A Wu <hao.a.wu@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Signed-off-by: Ard Biesheuvel <ard.biesheuvel@arm.com>
---
 MdeModulePkg/Bus/Pci/NonDiscoverablePciDeviceDxe/NonDiscoverablePciDeviceDxe.c | 43 ++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/MdeModulePkg/Bus/Pci/NonDiscoverablePciDeviceDxe/NonDiscoverablePciDeviceDxe.c b/MdeModulePkg/Bus/Pci/NonDiscoverablePciDeviceDxe/NonDiscoverablePciDeviceDxe.c
index 5c93e2a7663c..a14c06e7f4e1 100644
--- a/MdeModulePkg/Bus/Pci/NonDiscoverablePciDeviceDxe/NonDiscoverablePciDeviceDxe.c
+++ b/MdeModulePkg/Bus/Pci/NonDiscoverablePciDeviceDxe/NonDiscoverablePciDeviceDxe.c
@@ -15,6 +15,8 @@
 STATIC UINTN               mUniqueIdCounter = 0;
 EFI_CPU_ARCH_PROTOCOL      *mCpu;
 
+STATIC VOID                *mProtocolNotifyRegistration;
+
 //
 // We only support the following device types
 //
@@ -250,6 +252,43 @@ STATIC EFI_DRIVER_BINDING_PROTOCOL gDriverBinding = {
   NULL
 };
 
+STATIC
+VOID
+EFIAPI
+NonDiscoverablePciDeviceProtocolNotify (
+  IN EFI_EVENT            Event,
+  IN VOID                 *Context
+  )
+{
+  EFI_STATUS        Status;
+  EFI_HANDLE        *Handles;
+  UINTN             HandleCount;
+  UINTN             Index;
+
+  Status = gBS->LocateHandleBuffer (ByRegisterNotify, NULL,
+                  mProtocolNotifyRegistration, &HandleCount, &Handles);
+  if (EFI_ERROR (Status)) {
+    if (Status != EFI_NOT_FOUND) {
+      DEBUG ((DEBUG_WARN, "%a: LocateHandleBuffer() failed - %r\n",
+        __FUNCTION__, Status));
+      }
+    return;
+  }
+
+  for (Index = 0; Index < HandleCount; Index++) {
+    //
+    // Connect each newly registered gEdkiiNonDiscoverableDeviceProtocolGuid
+    // instance non-recursively to this driver specifically. This ensures that
+    // PCI I/O instances exist for each, regardless of whether ConnectAll() is
+    // used at any point.
+    //
+    Status = gBS->ConnectController (Handles[Index], gImageHandle, NULL, FALSE);
+    DEBUG ((DEBUG_VERBOSE, "%a: ConnectController () returned %r\n",
+      __FUNCTION__, Status));
+  }
+  gBS->FreePool (Handles);
+}
+
 /**
   Entry point of this driver.
 
@@ -272,6 +311,10 @@ NonDiscoverablePciDeviceDxeEntryPoint (
   Status = gBS->LocateProtocol (&gEfiCpuArchProtocolGuid, NULL, (VOID **)&mCpu);
   ASSERT_EFI_ERROR(Status);
 
+  EfiCreateProtocolNotifyEvent (&gEdkiiNonDiscoverableDeviceProtocolGuid,
+    TPL_CALLBACK, NonDiscoverablePciDeviceProtocolNotify, NULL,
+    &mProtocolNotifyRegistration);
+
   return EfiLibInstallDriverBindingComponentName2 (
            ImageHandle,
            SystemTable,
-- 
2.17.1