From 4513806a7ec7f8c4da746a7ffde952d9f170a07f Mon Sep 17 00:00:00 2001
From: Gustavo Moreira <gmoreira@gmail.com>
Date: Thu, 15 Aug 2024 07:55:45 -0700
Subject: [PATCH] Move RBTree to the rb_root extension object for better
 integration

---
 .../framework/symbols/linux/__init__.py       | 34 +------------------
 .../symbols/linux/extensions/__init__.py      | 30 +++++++++++++++-
 2 files changed, 30 insertions(+), 34 deletions(-)

diff --git a/volatility3/framework/symbols/linux/__init__.py b/volatility3/framework/symbols/linux/__init__.py
index 525492141..c2eb3d791 100644
--- a/volatility3/framework/symbols/linux/__init__.py
+++ b/volatility3/framework/symbols/linux/__init__.py
@@ -46,6 +46,7 @@ class LinuxKernelIntermedSymbols(intermed.IntermediateSymbolTable):
         # Might not exist in older kernels or the current symbols
         self.optional_set_type_class("mount", extensions.mount)
         self.optional_set_type_class("mnt_namespace", extensions.mnt_namespace)
+        self.optional_set_type_class("rb_root", extensions.rb_root)
 
         # Network
         self.set_type_class("net", extensions.net)
@@ -425,36 +426,3 @@ class LinuxUtilities(interfaces.configuration.VersionableInterface):
         kernel = context.modules[kernel_module_name]
 
         return kernel
-
-
-class RBTree(object):
-    """Simple Red-Black tree abstraction"""
-
-    def __init__(self, root):
-        self.root = root
-
-    def _walk_nodes(self, root_node) -> Iterator[int]:
-        """Traverses the Red-Black tree from the root node and yields a pointer to each
-        node in this tree.
-
-        Args:
-            root_node: A Red-Black tree node from which to start descending
-
-        Yields:
-            A pointer to every node descending from the specified root node
-        """
-        if not root_node:
-            return
-
-        yield root_node
-        yield from self._walk_nodes(root_node.rb_left)
-        yield from self._walk_nodes(root_node.rb_right)
-
-    def get_nodes(self) -> Iterator[int]:
-        """Yields a pointer to each node in the Red-Black tree
-
-        Yields:
-            A pointer to every node in the Red-Black tree
-        """
-
-        yield from self._walk_nodes(root_node=self.root.rb_node)
diff --git a/volatility3/framework/symbols/linux/extensions/__init__.py b/volatility3/framework/symbols/linux/extensions/__init__.py
index 645f493cb..c18cd093d 100644
--- a/volatility3/framework/symbols/linux/extensions/__init__.py
+++ b/volatility3/framework/symbols/linux/extensions/__init__.py
@@ -1326,7 +1326,7 @@ class mnt_namespace(objects.StructType):
             vmlinux = linux.LinuxUtilities.get_module_from_volobj_type(
                 self._context, self
             )
-            for node in linux.RBTree(self.mounts).get_nodes():
+            for node in self.mounts.get_nodes():
                 mnt = linux.LinuxUtilities.container_of(
                     node, "mount", "mnt_list", vmlinux
                 )
@@ -1925,3 +1925,31 @@ class inode(objects.StructType):
             The inode's file mode string
         """
         return stat.filemode(self.i_mode)
+
+
+class rb_root(objects.StructType):
+    def _walk_nodes(self, root_node) -> Iterator[int]:
+        """Traverses the Red-Black tree from the root node and yields a pointer to each
+        node in this tree.
+
+        Args:
+            root_node: A Red-Black tree node from which to start descending
+
+        Yields:
+            A pointer to every node descending from the specified root node
+        """
+        if not root_node:
+            return
+
+        yield root_node
+        yield from self._walk_nodes(root_node.rb_left)
+        yield from self._walk_nodes(root_node.rb_right)
+
+    def get_nodes(self) -> Iterator[int]:
+        """Yields a pointer to each node in the Red-Black tree
+
+        Yields:
+            A pointer to every node in the Red-Black tree
+        """
+
+        yield from self._walk_nodes(root_node=self.rb_node)