From 30fdeb211887187e0b0c1a2c27ef8edc368cf444 Mon Sep 17 00:00:00 2001
From: donnaskiez <lachiehodges7785@gmail.com>
Date: Sun, 16 Jun 2024 21:42:13 +1000
Subject: [PATCH] fortnite

---
 driver/callbacks.c       |  2 ++
 driver/containers/tree.c | 63 +++++++++++++++++++++++++++++++++++++---
 driver/containers/tree.h |  8 +++++
 driver/io.c              |  3 +-
 4 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/driver/callbacks.c b/driver/callbacks.c
index 180847b..b55bc31 100644
--- a/driver/callbacks.c
+++ b/driver/callbacks.c
@@ -1220,6 +1220,8 @@ TimerObjectWorkItemRoutine(_In_ PDEVICE_OBJECT DeviceObject,
 
     DEBUG_VERBOSE("Integrity check timer callback invoked.");
 
+    RtlRbTreePrintCurrentStatistics(GetThreadTree());
+
     if (!ValidateOurDriversDispatchRoutines()) {
         DEBUG_VERBOSE("l");
     }
diff --git a/driver/containers/tree.c b/driver/containers/tree.c
index 2781fc2..3f28e33 100644
--- a/driver/containers/tree.c
+++ b/driver/containers/tree.c
@@ -46,10 +46,52 @@
  * https://www.kernel.org/doc/Documentation/rbtree.txt
  * https://github.com/torvalds/linux/blob/master/lib/rbtree.c
  * https://www.osronline.com/article.cfm%5Earticle=516.htm
- * https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/ntddk/ns-ntddk-_rtl_avl_table (for structure ideas)
- * 
+ * https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/ntddk/ns-ntddk-_rtl_avl_table
+ * (for structure ideas)
+ *
  */
 
+FORCEINLINE
+STATIC
+VOID
+RtlpRbTreeIncrementInsertionCount(_In_ PRB_TREE Tree)
+{
+    InterlockedIncrement(&Tree->insertion_count);
+}
+
+FORCEINLINE
+STATIC
+VOID
+RtlpRbTreeIncrementDeletionCount(_In_ PRB_TREE Tree)
+{
+    InterlockedIncrement(&Tree->deletion_count);
+}
+
+FORCEINLINE
+STATIC
+VOID
+RtlpRbTreeIncrementNodeCount(_In_ PRB_TREE Tree)
+{
+    InterlockedIncrement(&Tree->node_count);
+}
+
+FORCEINLINE
+STATIC
+VOID
+RtlpRbTreeDecrementNodeCount(_In_ PRB_TREE Tree)
+{
+    InterlockedDecrement(&Tree->node_count);
+}
+
+VOID
+RtlRbTreePrintCurrentStatistics(_In_ PRB_TREE Tree)
+{
+    DEBUG_VERBOSE("Tree: %llx", (UINT64)Tree);
+    DEBUG_VERBOSE("Node count: %lx", Tree->node_count);
+    DEBUG_VERBOSE("Insertion count: %lx", Tree->insertion_count);
+    DEBUG_VERBOSE("Deletion count: %lx", Tree->deletion_count);
+}
+
 /**
  * Initialises a caller allocated RB_TREE structure.
  *
@@ -95,7 +137,11 @@ RtlRbTreeCreate(_In_ RB_COMPARE Compare,
     if (!NT_SUCCESS(status))
         return status;
 
-    Tree->compare = Compare;
+    Tree->compare         = Compare;
+    Tree->deletion_count  = 0;
+    Tree->insertion_count = 0;
+    Tree->node_count      = 0;
+
     KeInitializeGuardedMutex(&Tree->lock);
 
     return STATUS_SUCCESS;
@@ -306,6 +352,9 @@ RtlRbTreeInsertNode(_In_ PRB_TREE Tree, _In_ PVOID Key)
         }
         else {
             ExFreeToLookasideListEx(&Tree->pool, node);
+
+            /* Since we allocate and free a node, no housekeeping regarding
+             * stats needs to be done. */
             return current->object;
         }
     }
@@ -320,6 +369,8 @@ RtlRbTreeInsertNode(_In_ PRB_TREE Tree, _In_ PVOID Key)
         parent->right = node;
 
     RtlpRbTreeFixupInsert(Tree, node);
+    RtlpRbTreeIncrementInsertionCount(Tree);
+    RtlpRbTreeIncrementNodeCount(Tree);
 
     return node->object;
 }
@@ -590,6 +641,9 @@ RtlRbTreeDeleteNode(_In_ PRB_TREE Tree, _In_ PVOID Key)
         RtlpRbTreeFixupDelete(Tree, child);
 
     ExFreeToLookasideListEx(&Tree->pool, target);
+
+    RtlpRbTreeIncrementDeletionCount(Tree);
+    RtlpRbTreeDecrementNodeCount(Tree);
 }
 
 /* Public API that is used to find the node object for an associated key. Should
@@ -661,13 +715,14 @@ RtlpPrintInOrder(PRB_TREE_NODE Node)
     RtlpPrintInOrder(Node->right);
 }
 
-/* assumes lock is held */
 VOID
 RtlRbTreeInOrderPrint(_In_ PRB_TREE Tree)
 {
     DEBUG_ERROR("*************************************************");
     DEBUG_ERROR("<><><><>STARTING IN ORDER PRINT <><><><><><");
+    RtlRbTreeAcquireLock(Tree);    
     RtlpPrintInOrder(Tree->root);
+    RtlRbTreeReleaselock(Tree);
     DEBUG_ERROR("<><><><>ENDING IN ORDER PRINT <><><><><><");
     DEBUG_ERROR("*************************************************");
 }
diff --git a/driver/containers/tree.h b/driver/containers/tree.h
index 49bf3c5..d5b4409 100644
--- a/driver/containers/tree.h
+++ b/driver/containers/tree.h
@@ -26,6 +26,11 @@ typedef struct _RB_TREE {
     LOOKASIDE_LIST_EX pool;
     UINT32            object_size;
     UINT32            active;
+
+    volatile UINT32 node_count;
+    volatile UINT32 insertion_count;
+    volatile UINT32 deletion_count;
+
 } RB_TREE, *PRB_TREE;
 
 typedef VOID (*RB_CALLBACK)(PRB_TREE_NODE Node);
@@ -72,4 +77,7 @@ RtlRbTreeReleaselock(_Inout_ PRB_TREE Tree)
     KeReleaseGuardedMutex(&Tree->lock);
 }
 
+VOID
+RtlRbTreePrintCurrentStatistics(_In_ PRB_TREE Tree);
+
 #endif
\ No newline at end of file
diff --git a/driver/io.c b/driver/io.c
index 1ff47f3..c3ed88c 100644
--- a/driver/io.c
+++ b/driver/io.c
@@ -1177,6 +1177,8 @@ DeviceClose(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp)
     UNREFERENCED_PARAMETER(DeviceObject);
     DEBUG_INFO("Handle to driver closed.");
 
+    /* This needs to be fixed lol, cos anyone can just open a handle whhich
+     * might not begin a session.*/
     SessionTerminate();
     UnregisterProcessObCallbacks();
     SharedMappingTerminate();
@@ -1191,7 +1193,6 @@ DeviceCreate(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp)
     PAGED_CODE();
     UNREFERENCED_PARAMETER(DeviceObject);
     DEBUG_INFO("Handle to driver opened.");
-
     IoCompleteRequest(Irp, IO_NO_INCREMENT);
     return Irp->IoStatus.Status;
 }
\ No newline at end of file