[NTOS:IO][NTOS:PNP] Add PiActionAddBootDevices device action

During the boot process, it makes possible to initalize the driver's
devices right after the driver is loaded. Moreover, this way one can be
sure that all critical devices are initialized before the
IopMarkBootPartition call (because we explicitly call the driver's
AddDevice routine now, after each driver is loaded)

CORE-7826

ntoskrnl/include/internal/io.h

@@ -533,7 +533,8 @@ typedef enum _DEVICE_ACTION
 {
     PiActionEnumDeviceTree,
     PiActionEnumRootDevices,
-    PiActionResetDevice
+    PiActionResetDevice,
+    PiActionAddBootDevices
 } DEVICE_ACTION;
 
 //

ntoskrnl/io/iomgr/driver.c

@@ -738,14 +738,16 @@ LdrProcessDriverModule(PLDR_DATA_TABLE_ENTRY LdrEntry,
     return STATUS_SUCCESS;
 }
 
+PDEVICE_OBJECT
+IopGetDeviceObjectFromDeviceInstance(PUNICODE_STRING DeviceInstance);
+
 /*
  * IopInitializeBuiltinDriver
  *
  * Initialize a driver that is already loaded in memory.
  */
 CODE_SEG("INIT")
-NTSTATUS
-NTAPI
+BOOLEAN
 IopInitializeBuiltinDriver(IN PLDR_DATA_TABLE_ENTRY BootLdrEntry)
 {
     PDRIVER_OBJECT DriverObject;
@@ -769,7 +771,7 @@ IopInitializeBuiltinDriver(IN PLDR_DATA_TABLE_ENTRY BootLdrEntry)
                                    TAG_IO);
     if (Buffer == NULL)
     {
-        return STATUS_INSUFFICIENT_RESOURCES;
+        return FALSE;
     }
 
     RtlCopyMemory(Buffer, ModuleName->Buffer, ModuleName->Length);
@@ -795,7 +797,7 @@ IopInitializeBuiltinDriver(IN PLDR_DATA_TABLE_ENTRY BootLdrEntry)
     ExFreePoolWithTag(Buffer, TAG_IO);
     if (!Success)
     {
-        return STATUS_INSUFFICIENT_RESOURCES;
+        return FALSE;
     }
 
     FileExtension = wcsrchr(ServiceName.Buffer, L'.');
@@ -813,7 +815,7 @@ IopInitializeBuiltinDriver(IN PLDR_DATA_TABLE_ENTRY BootLdrEntry)
     RegistryPath.Buffer = ExAllocatePoolWithTag(PagedPool, RegistryPath.MaximumLength, TAG_IO);
     if (RegistryPath.Buffer == NULL)
     {
-        return STATUS_INSUFFICIENT_RESOURCES;
+        return FALSE;
     }
     RtlAppendUnicodeToString(&RegistryPath, ServicesKeyName);
     RtlAppendUnicodeStringToString(&RegistryPath, &ServiceName);
@@ -824,7 +826,7 @@ IopInitializeBuiltinDriver(IN PLDR_DATA_TABLE_ENTRY BootLdrEntry)
     RtlFreeUnicodeString(&RegistryPath);
     if (!NT_SUCCESS(Status))
     {
-        return Status;
+        return FALSE;
     }
 
     /* Lookup the new Ldr entry in PsLoadedModuleList */
@@ -851,18 +853,94 @@ IopInitializeBuiltinDriver(IN PLDR_DATA_TABLE_ENTRY BootLdrEntry)
                                        serviceHandle,
                                        &DriverObject,
                                        &driverEntryStatus);
-    ZwClose(serviceHandle);
 
     if (!NT_SUCCESS(Status))
     {
         DPRINT1("Driver '%wZ' load failed, status (%x)\n", ModuleName, Status);
-        return Status;
+        return FALSE;
     }
 
+    // The driver has been loaded, now check if where are any PDOs
+    // for that driver, and queue AddDevice call for them.
+    // The check is possible because HKLM/SYSTEM/CCS/Services/<ServiceName>/Enum directory
+    // is populated upon a new device arrival based on a (critical) device database
+    
+    // Legacy drivers may add devices inside DriverEntry.
+    // We're lazy and always assume that they are doing so
+    BOOLEAN deviceAdded = (_Bool)(DriverObject->Flags & DRVO_LEGACY_DRIVER);
+
+    HANDLE enumServiceHandle;
+    UNICODE_STRING enumName = RTL_CONSTANT_STRING(L"Enum");
+
+    Status = IopOpenRegistryKeyEx(&enumServiceHandle, serviceHandle, &enumName, KEY_READ);
+    ZwClose(serviceHandle);
+
+    if (NT_SUCCESS(Status))
+    {
+        UINT32 instanceCount = 0;
+        PKEY_VALUE_FULL_INFORMATION kvInfo;
+        Status = IopGetRegistryValue(enumServiceHandle, L"Count", &kvInfo);
+        if (!NT_SUCCESS(Status))
+        {
+            goto Cleanup;
+        }
+        if (kvInfo->Type != REG_DWORD)
+        {
+            ExFreePool(kvInfo);
+            goto Cleanup;
+        }
+
+        RtlMoveMemory(&instanceCount,
+                      (PVOID)((ULONG_PTR)kvInfo + kvInfo->DataOffset),
+                      sizeof(UINT32));
+        ExFreePool(kvInfo);
+
+        DPRINT("Processing %u instances for %wZ module\n", instanceCount, ModuleName);
+
+        for (UINT32 i = 0; i < instanceCount; i++)
+        {
+            WCHAR num[11];
+            UNICODE_STRING instancePath;
+            RtlStringCchPrintfW(num, sizeof(num), L"%u", i);
+
+            Status = IopGetRegistryValue(enumServiceHandle, num, &kvInfo);
+            if (!NT_SUCCESS(Status))
+            {
+                continue;
+            }
+            if (kvInfo->Type != REG_SZ || kvInfo->DataLength == 0)
+            {
+                ExFreePool(kvInfo);
+                continue;
+            }
+
+            instancePath.Length = kvInfo->DataLength - sizeof(WCHAR),
+            instancePath.MaximumLength = kvInfo->DataLength,
+            instancePath.Buffer = ExAllocatePoolWithTag(NonPagedPool,
+                                                        instancePath.MaximumLength,
+                                                        TAG_IO);
+            if (instancePath.Buffer)
+            {
+                RtlMoveMemory(instancePath.Buffer,
+                          (PVOID)((ULONG_PTR)kvInfo + kvInfo->DataOffset),
+                          instancePath.Length);
+
+                PDEVICE_OBJECT pdo = IopGetDeviceObjectFromDeviceInstance(&instancePath);
+                PiQueueDeviceAction(pdo, PiActionAddBootDevices, NULL, NULL);
+                ObDereferenceObject(pdo);
+                deviceAdded = TRUE;
+            }
+
+            ExFreePool(kvInfo);            
+        }
+
+        ZwClose(enumServiceHandle);
+    }
+Cleanup:
     /* Remove extra reference from IopInitializeDriverModule */
     ObDereferenceObject(DriverObject);
 
-    return Status;
+    return deviceAdded;
 }
 
 /*
@@ -1028,13 +1106,14 @@ IopInitializeBootDrivers(VOID)
             LdrEntry = DriverInfo->DataTableEntry->LdrEntry;
 
             /* Initialize it */
-            IopInitializeBuiltinDriver(LdrEntry);
-
-            /* Start the devices found by a driver (if any) */
-            PiQueueDeviceAction(IopRootDeviceNode->PhysicalDeviceObject,
-                                PiActionEnumRootDevices,
-                                NULL,
-                                NULL);
+            if (IopInitializeBuiltinDriver(LdrEntry))
+            {
+                // it does not make sense to enumerate the tree if there are no new devices added
+                PiQueueDeviceAction(IopRootDeviceNode->PhysicalDeviceObject,
+                                    PiActionEnumRootDevices,
+                                    NULL,
+                                    NULL);
+            }           
 
             /* Next entry */
             NextEntry = NextEntry->Flink;

ntoskrnl/io/pnpmgr/devaction.c

@@ -2480,6 +2480,8 @@ ActionToStr(
             return "PiActionEnumRootDevices";
         case PiActionResetDevice:
             return "PiActionResetDevice";
+        case PiActionAddBootDevices:
+            return "PiActionAddBootDevices";
         default:
             return "(request unknown)";
     }
@@ -2517,6 +2519,15 @@ PipDeviceActionWorker(
 
         switch (Request->Action)
         {
+            case PiActionAddBootDevices:
+            {
+                if (deviceNode->State == DeviceNodeInitialized && 
+                    !(deviceNode->Flags & DNF_HAS_PROBLEM))
+                {
+                    status = PiCallDriverAddDevice(deviceNode, PnPBootDriversInitialized);
+                }
+                break;
+            }
             case PiActionEnumRootDevices:
             case PiActionEnumDeviceTree:
                 deviceNode->Flags |= DNF_REENUMERATE;
@@ -2589,7 +2600,7 @@ PiQueueDeviceAction(
     KeAcquireSpinLock(&IopDeviceActionLock, &OldIrql);
     InsertTailList(&IopDeviceActionRequestList, &Request->RequestListEntry);
 
-    if (Action == PiActionEnumRootDevices)
+    if (Action == PiActionEnumRootDevices || Action == PiActionAddBootDevices)
     {
         ASSERT(!IopDeviceActionInProgress);