From ffa9875fa6a2fe05369df3af523752305fb055c3 Mon Sep 17 00:00:00 2001
From: Hendiadyoin1 <leon.a@serenityos.org>
Date: Sat, 27 Apr 2024 13:43:13 +0200
Subject: [PATCH] LibDeviceTree: Fix stale pointers in phandle and parent
 handling

We need to set them after we've created the full tree, as otherwise the
HashMap containing the nodes may reallocate and invalidate the pointers.
---
 .../Libraries/LibDeviceTree/DeviceTree.cpp    | 28 ++++++++----
 Userland/Libraries/LibDeviceTree/DeviceTree.h | 45 +++++++++++++++++++
 2 files changed, 65 insertions(+), 8 deletions(-)

diff --git a/Userland/Libraries/LibDeviceTree/DeviceTree.cpp b/Userland/Libraries/LibDeviceTree/DeviceTree.cpp
index 51cfef64572..f29951ee2a6 100644
--- a/Userland/Libraries/LibDeviceTree/DeviceTree.cpp
+++ b/Userland/Libraries/LibDeviceTree/DeviceTree.cpp
@@ -39,14 +39,7 @@ ErrorOr<NonnullOwnPtr<DeviceTree>> DeviceTree::parse(ReadonlyBytes flattened_dev
                 current_node = current_node->parent();
                 return IterationDecision::Continue;
             },
-            .on_property = [&device_tree, &current_node](StringView name, ReadonlyBytes value) -> ErrorOr<IterationDecision> {
-                DeviceTreeProperty property { value };
-
-                if (name == "phandle"sv) {
-                    auto phandle = property.as<u32>();
-                    TRY(device_tree->set_phandle(phandle, current_node));
-                }
-
+            .on_property = [&current_node](StringView name, ReadonlyBytes value) -> ErrorOr<IterationDecision> {
                 TRY(current_node->properties().try_set(name, DeviceTreeProperty { value }));
                 return IterationDecision::Continue;
             },
@@ -58,6 +51,25 @@ ErrorOr<NonnullOwnPtr<DeviceTree>> DeviceTree::parse(ReadonlyBytes flattened_dev
             },
         }));
 
+    // FIXME: While growing the a nodes children map, we might have reallocated it's storage
+    //        breaking the parent pointers of the children, so we need to fix them here
+    auto fix_parent = [](auto self, DeviceTreeNodeView& node) -> void {
+        for (auto& [name, child] : node.children()) {
+            child.m_parent = &node;
+            self(self, child);
+        }
+    };
+
+    fix_parent(fix_parent, *device_tree);
+    // Note: For the same reason as above, we need to postpone setting the phandles until the tree is fully built
+    TRY(device_tree->for_each_node([&device_tree]([[maybe_unused]] StringView name, DeviceTreeNodeView& node) -> ErrorOr<RecursionDecision> {
+        if (auto phandle = node.get_property("phandle"sv); phandle.has_value()) {
+            auto phandle_value = phandle.value().as<u32>();
+            TRY(device_tree->set_phandle(phandle_value, &node));
+        }
+        return RecursionDecision::Recurse;
+    }));
+
     return device_tree;
 }
 
diff --git a/Userland/Libraries/LibDeviceTree/DeviceTree.h b/Userland/Libraries/LibDeviceTree/DeviceTree.h
index b50cba50029..b4ccf8670c4 100644
--- a/Userland/Libraries/LibDeviceTree/DeviceTree.h
+++ b/Userland/Libraries/LibDeviceTree/DeviceTree.h
@@ -13,6 +13,7 @@
 #include <AK/IterationDecision.h>
 #include <AK/MemoryStream.h>
 #include <AK/Optional.h>
+#include <AK/RecursionDecision.h>
 #include <AK/Span.h>
 
 namespace DeviceTree {
@@ -168,6 +169,28 @@ public:
     //        bonus points if it could automatically recurse in the tree under some conditions,
     //        like "simple-bus" or "pci-bridge" nodes
 
+    auto for_each_node(CallableAs<ErrorOr<RecursionDecision>, StringView, DeviceTreeNodeView const&> auto callback) const
+    {
+        auto iterate = [&](auto self, StringView name, DeviceTreeNodeView const& node) -> ErrorOr<RecursionDecision> {
+            auto result = TRY(callback(name, node));
+
+            if (result == RecursionDecision::Recurse) {
+                for (auto const& [name, child] : node.children()) {
+                    auto child_result = TRY(self(self, name, child));
+
+                    if (child_result == RecursionDecision::Break)
+                        return RecursionDecision::Break;
+                }
+
+                return RecursionDecision::Continue;
+            }
+
+            return result;
+        };
+
+        return iterate(iterate, "/"sv, *this);
+    }
+
     DeviceTreeNodeView const* phandle(u32 phandle) const
     {
         if (phandle >= m_phandles.size())
@@ -184,6 +207,28 @@ private:
     {
     }
 
+    auto for_each_node(CallableAs<ErrorOr<RecursionDecision>, StringView, DeviceTreeNodeView&> auto callback)
+    {
+        auto iterate = [&](auto self, StringView name, DeviceTreeNodeView& node) -> ErrorOr<RecursionDecision> {
+            auto result = TRY(callback(name, node));
+
+            if (result == RecursionDecision::Recurse) {
+                for (auto& [name, child] : node.children()) {
+                    auto child_result = TRY(self(self, name, child));
+
+                    if (child_result == RecursionDecision::Break)
+                        return RecursionDecision::Break;
+                }
+
+                return RecursionDecision::Continue;
+            }
+
+            return result;
+        };
+
+        return iterate(iterate, "/"sv, *this);
+    }
+
     ErrorOr<void> set_phandle(u32 phandle, DeviceTreeNodeView* node)
     {
         if (m_phandles.size() > phandle && m_phandles[phandle] != nullptr)