[NTOS:PNP] Avoid recursion in IopTraverseDeviceTree(Node). CORE-17215

diff --git a/ntoskrnl/io/pnpmgr/devnode.c b/ntoskrnl/io/pnpmgr/devnode.c
index 6f5e029..a068cd7 100644 (file)
--- a/ntoskrnl/io/pnpmgr/devnode.c
+++ b/ntoskrnl/io/pnpmgr/devnode.c
@@ -321,56 +321,36 @@ IopFreeDeviceNode(
 
 static
 NTSTATUS
-IopTraverseDeviceTreeNode(
+IopFindNextDeviceNodeForTraversal(
     _In_ PDEVICETREE_TRAVERSE_CONTEXT Context)
 {
-    PDEVICE_NODE ParentDeviceNode;
-    PDEVICE_NODE ChildDeviceNode;
-    PDEVICE_NODE NextDeviceNode;
-    NTSTATUS Status;
-
-    /* Copy context data so we don't overwrite it in subsequent calls to this function */
-    ParentDeviceNode = Context->DeviceNode;
-
-    /* HACK: Keep a reference to the PDO so we can keep traversing the tree
-     * if the device is deleted. In a perfect world, children would have to be
-     * deleted before their parents, and we'd restart the traversal after
-     * deleting a device node. */
-    ObReferenceObject(ParentDeviceNode->PhysicalDeviceObject);
-
-    /* Call the action routine */
-    Status = (Context->Action)(ParentDeviceNode, Context->Context);
-    if (!NT_SUCCESS(Status))
+    /* If we have a child, simply go down the tree */
+    if (Context->DeviceNode->Child != NULL)
     {
-        ObDereferenceObject(ParentDeviceNode->PhysicalDeviceObject);
-        return Status;
+        ASSERT(Context->DeviceNode->Child->Parent == Context->DeviceNode);
+        Context->DeviceNode = Context->DeviceNode->Child;
+        return STATUS_SUCCESS;
     }
 
-    /* Traversal of all children nodes */
-    for (ChildDeviceNode = ParentDeviceNode->Child;
-         ChildDeviceNode != NULL;
-         ChildDeviceNode = NextDeviceNode)
+    while (Context->DeviceNode != Context->FirstDeviceNode)
     {
-        /* HACK: We need this reference to ensure we can get Sibling below. */
-        ObReferenceObject(ChildDeviceNode->PhysicalDeviceObject);
-
-        /* Pass the current device node to the action routine */
-        Context->DeviceNode = ChildDeviceNode;
-
-        Status = IopTraverseDeviceTreeNode(Context);
-        if (!NT_SUCCESS(Status))
+        /* All children processed -- go sideways */
+        if (Context->DeviceNode->Sibling != NULL)
         {
-            ObDereferenceObject(ChildDeviceNode->PhysicalDeviceObject);
-            ObDereferenceObject(ParentDeviceNode->PhysicalDeviceObject);
-            return Status;
+            ASSERT(Context->DeviceNode->Sibling->Parent == Context->DeviceNode->Parent);
+            Context->DeviceNode = Context->DeviceNode->Sibling;
+            return STATUS_SUCCESS;
         }
 
-        NextDeviceNode = ChildDeviceNode->Sibling;
-        ObDereferenceObject(ChildDeviceNode->PhysicalDeviceObject);
+        /* We're the last sibling -- go back up */
+        ASSERT(Context->DeviceNode->Parent->LastChild == Context->DeviceNode);
+        Context->DeviceNode = Context->DeviceNode->Parent;
+
+        /* We already visited the parent and all its children, so keep looking */
     }
 
-    ObDereferenceObject(ParentDeviceNode->PhysicalDeviceObject);
-    return Status;
+    /* Done with all children of the start node -- stop enumeration */
+    return STATUS_UNSUCCESSFUL;
 }
 
 NTSTATUS
@@ -378,6 +358,7 @@ IopTraverseDeviceTree(
     _In_ PDEVICETREE_TRAVERSE_CONTEXT Context)
 {
     NTSTATUS Status;
+    PDEVICE_NODE DeviceNode;
 
     DPRINT("Context 0x%p\n", Context);
 
@@ -387,8 +368,32 @@ IopTraverseDeviceTree(
     /* Start from the specified device node */
     Context->DeviceNode = Context->FirstDeviceNode;
 
-    /* Recursively traverse the device tree */
-    Status = IopTraverseDeviceTreeNode(Context);
+    /* Traverse the device tree */
+    do
+    {
+        DeviceNode = Context->DeviceNode;
+
+        /* HACK: Keep a reference to the PDO so we can keep traversing the tree
+         * if the device is deleted. In a perfect world, children would have to be
+         * deleted before their parents, and we'd restart the traversal after
+         * deleting a device node. */
+        ObReferenceObject(DeviceNode->PhysicalDeviceObject);
+
+        /* Call the action routine */
+        Status = (Context->Action)(DeviceNode, Context->Context);
+        if (NT_SUCCESS(Status))
+        {
+            /* Find next device node */
+            ASSERT(Context->DeviceNode == DeviceNode);
+            Status = IopFindNextDeviceNodeForTraversal(Context);
+        }
+
+        /* We need to either abort or make progress */
+        ASSERT(!NT_SUCCESS(Status) || Context->DeviceNode != DeviceNode);
+
+        ObDereferenceObject(DeviceNode->PhysicalDeviceObject);
+    } while (NT_SUCCESS(Status));
+
     if (Status == STATUS_UNSUCCESSFUL)
     {
         /* The action routine just wanted to terminate the traversal with status