This commit is contained in:
donnaskiez 2024-06-16 21:42:13 +10:00
parent 424f1e9a63
commit 30fdeb2118
4 changed files with 71 additions and 5 deletions

View file

@ -1220,6 +1220,8 @@ TimerObjectWorkItemRoutine(_In_ PDEVICE_OBJECT DeviceObject,
DEBUG_VERBOSE("Integrity check timer callback invoked.");
RtlRbTreePrintCurrentStatistics(GetThreadTree());
if (!ValidateOurDriversDispatchRoutines()) {
DEBUG_VERBOSE("l");
}

View file

@ -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("*************************************************");
}

View file

@ -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

View file

@ -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;
}