#include <ntifs.h>
#include <windef.h>
/*
内核中的二叉树之splay link tree。
发现内核中有操作以下树的函数。
splay link tree。
Adelson-Velsky/Landis (AVL) trees。
以后可以用AVL函数遍历进程的地址空间了，可能还需要知道数据结构（树的结构，进程对象，寄存器）。
splay link tree的配合函数。
RtlRandom
ExInitializeFastMutex
此文仅仅示例splay link tree的用法。
更多信息：
http://msdn.microsoft.com/en-us/library/windows/hardware/ff553327(v=vs.85).aspx
http://msdn.microsoft.com/en-us/library/windows/hardware/ff553345(v=vs.85).aspx
made by correy
made at 2014.08.22
email:kouleguan at hotmail dot com
homepage:http://correy.webs.com
*/
#define TAG 'tset' //test
typedef struct _BinTree {
RTL_SPLAY_LINKS rsl;//注意这个不是指针，这个和下面的成员是并排的，且是第一个成员保持不变。
int i;
} BinTree;
typedef BinTree *PBinTree;
BinTree g_bt;
DRIVER_UNLOAD Unload;
VOID Unload(__in PDRIVER_OBJECT DriverObject)
{
}
VOID traversal(IN BinTree * t)
{
if ( t != 0)
{
KdPrint(("%d\n",t->i));
traversal((BinTree *) t->rsl.LeftChild);
traversal((BinTree *) t->rsl.RightChild);
}
}
BOOL InsertChild(IN BinTree * g , IN BinTree * t )
{
BOOL B = FALSE ;
BinTree * /*PRTL_SPLAY_LINKS*/ p ;
if (RtlLeftChild(&g->rsl) == 0)
{
RtlInsertAsLeftChild(&g->rsl, &t->rsl);
return TRUE ;
}
if (RtlRightChild(&g->rsl) == 0)
{
RtlInsertAsRightChild(&g->rsl, &t->rsl);
return TRUE ;
}
p = (BinTree *)RtlRightChild(&g->rsl); //这个在前，下面的在后，这样遍历出的结果是顺序的。
if (InsertChild(p, t ) == TRUE )
{
return TRUE ;
}
p = (BinTree *)RtlLeftChild(&g->rsl);
if (InsertChild(p, t ) == TRUE )
{
return TRUE ;
}
return B ;
}
BOOL DeleteOne()
{
BOOL B = FALSE ;
return B ;
}
BOOL DeleteALL()
{
BOOL B = FALSE ;
return B ;
}
DRIVER_INITIALIZE DriverEntry;
NTSTATUS DriverEntry( __in struct _DRIVER_OBJECT * DriverObject, __in PUNICODE_STRING RegistryPath)
{
NTSTATUS status = STATUS_UNSUCCESSFUL;
BOOLEAN B = FALSE;
PRTL_SPLAY_LINKS t;
int i;
KdBreakPoint();
DriverObject->DriverUnload = Unload;
//RtlInitializeSplayLinks(&g_bt.rsl);//设置上级为自己，其实源码也是这样的。
//RtlInsertAsRightChild(&g_bt.rsl, &g_bt1.rsl);//设置上级的右孩子为自己，同时设置自己的上级为第一个参数，其实源码也是这样的。
//RtlInsertAsLeftChild(&g_bt.rsl, &g_bt2.rsl);//同上，不再解释。
//B = RtlIsRoot(&g_bt.rsl);//有源码。
//t = RtlParent(&g_bt1.rsl);//有源码。
//t = RtlRightChild(&g_bt.rsl);//有源码。
//t = RtlLeftChild(&g_bt.rsl);//有源码。
//B = RtlIsLeftChild//有源码。
//B = RtlIsRightChild //有源码。
//RtlDeleteNoSplay(&g_bt.rsl, &t);//没有源码。
//t = RtlSubtreeSuccessor(&g_bt.rsl);//没有源码。
//t = RtlSubtreePredecessor(&g_bt.rsl);//没有源码。
//t = RtlRealSuccessor(&g_bt.rsl);//没有源码。
//t = RtlRealPredecessor(&g_bt.rsl);//没有源码。
//t = RtlSplay(&g_bt.rsl);//没有源码。
//t = RtlDelete(&g_bt.rsl);//没有源码。保持平衡
//创建
g_bt.i = 10;
RtlInitializeSplayLinks(&g_bt.rsl);
//插入
for (i = 0 ; i < 10; i++)
{
BinTree * btt = (BinTree *)ExAllocatePoolWithTag(NonPagedPool, sizeof(BinTree), TAG); //这里不能用全局变量，否则栈溢出。
if (btt == NULL) {
break;
}
RtlZeroMemory(btt, sizeof(BinTree));
btt->i = i;
InsertChild(&g_bt, btt);
}
//遍历
traversal(&g_bt);
//删除
return status;//STATUS_SUCCESS
}
--------------------------------------------------------------------------------------------------------------------------------
/*
By default, the operating system uses splay trees to implement generic tables.
Under some circumstances, operations on a splay tree will make the tree deep and narrow and might even turn it into a straight line.
Very deep trees degrade the performance of searches.
You can ensure a more balanced, shallower tree implementation of generic tables by using Adelson-Velsky/Landis (AVL) trees.
If you want to configure the generic table routines to use AVL trees instead of splay trees in your driver, insert the following define statement in a common header file before including Ntddk.h:
*/
//#define RTL_USE_AVL_TABLES 0
/*
If RTL_USE_AVL_TABLES is not defined, you must use the AVL form of the generic table routines.
For example, use the RtlInitializeGenericTableAvl routine instead of RtlInitializeGenericTable.
RtlInitializeGenericTableAvl returns an initialized RTL_AVL_TABLE table structure in the buffer to which the Table parameter points.
In the call to RtlInitializeGenericTableAvl, the caller must pass a PRTL_AVL_COMPARE_ROUTINE-typed comparison callback routine,
a PRTL_AVL_ALLOCATE_ROUTINE-typed allocation callback routine, and a PRTL_AVL_FREE_ROUTINE-typed deallocation callback routine rather than the similar PRTL_GENERIC_Xxx-typed routines.
*/
#include <ntifs.h>
#include <windef.h>
#define TAG 'tset' //test
/*
功能：内核中的splay trees和Adelson-Velsky/Landis (AVL) trees的用法示例。
内核中的splay trees和Adelson-Velsky/Landis (AVL) trees知道久矣！一直没有机会去做实验。
微软好像建议用AVL trees，不建议用splay trees。个人感觉AVL trees比splay trees快。
不过会用了，那就想用哪个就用哪个。
我们平常应该使用最多的应该是：RtlGetElementGenericTable，这个可能很快，不用再自己遍历链表和匹配了，
其次是RtlGetElementGenericTable，
RtlEnumerateGenericTable和RtlEnumerateGenericTableWithoutSplaying应该是用的最少的。
但是RtlDeleteElementGenericTable和RtlInsertElementGenericTable和RtlInitializeGenericTable是必须的。
示例工程：
1.Windows 8 Driver Samples\ClassPnP Storage Class Driver Library，注意：7600.16385.1\src\storage\class\classpnp没有。
2.Windows 8 Driver Samples\AvScan File System Minifilter Driver
3.https://github.com/desowin/usbpcap/blob/master/USBPcapDriver/USBPcapTables.c
Callers of RtlInitializeGenericTable must be running at IRQL <= DISPATCH_LEVEL.
Note that if Rtl...GenericTable routines are to be used at IRQL DISPATCH_LEVEL, the CompareRoutine, AllocateRoutine, and FreeRoutine must all be nonpageable code,
and the AllocateRoutine should allocate memory from nonpaged pool.
RtlDeleteElementGenericTableAvl
RtlEnumerateGenericTableAvl
RtlEnumerateGenericTableLikeADirectory
RtlEnumerateGenericTableWithoutSplayingAvl
RtlGetElementGenericTableAvl
RtlInsertElementGenericTableAvl
RtlInsertElementGenericTableFullAvl
RtlIsGenericTableEmptyAvl
RtlLookupElementGenericTableAvl
RtlLookupElementGenericTableFullAvl
RtlLookupFirstMatchingElementGenericTableAvl
RtlNumberGenericTableElementsAvl
made by correy
made at 2015.11.12
homepage:http://correy.webs.com
*/
#ifdef RTL_USE_AVL_TABLES
RTL_AVL_TABLE g_Table = {0}; //编译时RtlInitializeGenericTable第一个参数类型不对。
#else
RTL_GENERIC_TABLE g_Table = {0};//WDK 7600.16385.1的文档上没有公开这个结构（搞得好乖神秘的），MSDN上公开了。
#endif
/*
有的说这个结构不可过小，可是经测试也没事。
有的在结构的头部添加：RTL_SPLAY_LINKS和LIST_ENTRY两成员。
*/
typedef struct _content{
int i;
//...
}content, *Pcontent;
#define BUFFERSIZE (sizeof(RTL_SPLAY_LINKS) + sizeof(LIST_ENTRY) + sizeof(content))
FAST_MUTEX g_FastMutex;
RTL_GENERIC_COMPARE_RESULTS CompareRoutine (__in struct _RTL_GENERIC_TABLE *Table, __in PVOID FirstStruct, __in PVOID SecondStruct)
/*
The caller-supplied CompareRoutine is called before the AllocateRoutine to locate an appropriate location at which a new element should be inserted.
The CompareRoutine also is called before the FreeRoutine to locate an element to be deleted.
Arguments:
Table A pointer to the generic table.
FirstStruct A pointer to the first item to be compared.
SecondStruct A pointer to the second item to be compared.
The CompareRoutine must strictly track the ordering of all elements in the generic table so that it can identify any particular element.
The caller-defined structure for element data usually includes a member whose value is unique and can be used as a sorting key.
All Rtl...GenericTable routines that call the CompareRoutine take a buffer pointer as a parameter, which is passed in turn to the CompareRoutine.
The buffer contains a caller-supplied key value to be matched by the CompareRoutine to the key of the element that is being searched for.
Given two such key values, the CompareRoutine returns GenericLessThan, GenericGreaterThan, or GenericEqual.
RtlInsertElementGenericTable的第一次调用不会走这里，可能发现表是空的。
*/
{
Pcontent p1 = (Pcontent)FirstStruct;
Pcontent p2 = (Pcontent)SecondStruct;
if (p1->i < p2->i)
{
return GenericLessThan;
}
else if (p1->i > p2->i)
{
return GenericGreaterThan;
}
else
{
return GenericEqual;
}
}
PVOID AllocateRoutine (__in struct _RTL_GENERIC_TABLE *Table, __in CLONG ByteSize)
/*
调用RtlInsertElementGenericTable会走这里。
Arguments:
Table A pointer to the generic table.
ByteSize The number of bytes to allocate.
For each new element, the AllocateRoutine is called to allocate memory for caller-supplied data plus some additional memory for use by the Rtl...GenericTable routines.
Note that because of this "additional memory," caller-supplied routines must not access the first (sizeof(RTL_SPLAY_LINKS) + sizeof(LIST_ENTRY)) bytes of any element in the generic table.
*/
{
return ExAllocatePoolWithTag(NonPagedPool, ByteSize, TAG);
}
VOID FreeRoutine (__in struct _RTL_GENERIC_TABLE *Table, __in PVOID Buffer)
/*
调用RtlDeleteElementGenericTable会走这里。
Arguments:
Table A pointer to the generic table.
Buffer A pointer to the element that is being deleted.
Rtl...GenericTable routines call the FreeRoutine to deallocate memory for elements to be deleted from the generic table.
The FreeRoutine is the opposite of the AllocateRoutine.
*/
{
ExFreePoolWithTag(Buffer, TAG);
}
DRIVER_UNLOAD Unload;
VOID Unload(__in PDRIVER_OBJECT DriverObject)
{
NTSTATUS status = STATUS_UNSUCCESSFUL;
}
void InitializeGenericTable()
{
ExInitializeFastMutex(&g_FastMutex);
RtlInitializeGenericTable(&g_Table, (PRTL_GENERIC_COMPARE_ROUTINE)CompareRoutine, (PRTL_GENERIC_ALLOCATE_ROUTINE)AllocateRoutine, (PRTL_GENERIC_FREE_ROUTINE)FreeRoutine, NULL);
}
void InsertElementGenericTable()
{
PVOID p = NULL;
PVOID Buffer = NULL;
BOOLEAN NewElement = FALSE;
int i = 0;
Buffer = ExAllocatePoolWithTag(NonPagedPool, sizeof(content), TAG);
ASSERT(Buffer != NULL);
for ( ; i < 9; i++)
{
Pcontent pc = (Pcontent)Buffer;
RtlZeroMemory(Buffer, sizeof(content));
pc->i = i;//设置内容，因为不准重复的元素。
ExAcquireFastMutex(&g_FastMutex);
p = RtlInsertElementGenericTable(&g_Table, Buffer, sizeof(content), &NewElement);
ASSERT(p != NULL);
//ASSERT(NewElement == TRUE );
ExReleaseFastMutex(&g_FastMutex);
}
ExFreePoolWithTag(Buffer, TAG);
}
void GetElementGenericTable()
{
PVOID p = NULL;
ULONG I = 0;
Pcontent pc = NULL;
ExAcquireFastMutex(&g_FastMutex);
I = RtlNumberGenericTableElements(&g_Table) - 1;
p = RtlGetElementGenericTable(&g_Table, I);
ASSERT(p != NULL);
pc = (Pcontent)p;
ExReleaseFastMutex(&g_FastMutex);
}
void LookupElementGenericTable()
{
PVOID p = NULL;
content test = {0};
Pcontent pc = NULL;
ExAcquireFastMutex(&g_FastMutex);
test.i = 1;
p = RtlLookupElementGenericTable(&g_Table, &test);
ASSERT(p != NULL);
pc = (Pcontent)p;
ASSERT(pc->i == 1);
ExReleaseFastMutex(&g_FastMutex);
}
void EnumerateGenericTable()
{
PVOID p = NULL;
ExAcquireFastMutex(&g_FastMutex);
for (p = RtlEnumerateGenericTable ( &g_Table, TRUE ); p != NULL; p = RtlEnumerateGenericTable ( &g_Table, FALSE ))
{
// Process the element pointed to by p
Pcontent pc = (Pcontent)p;
KdPrint(("%d.\r\n", pc->i));
}
ExReleaseFastMutex(&g_FastMutex);
}
void EnumerateGenericTableWithoutSplaying()
{
PVOID ptr = NULL;
PVOID RestartKey = NULL;
ExAcquireFastMutex(&g_FastMutex);
for (ptr = RtlEnumerateGenericTableWithoutSplaying(&g_Table, &RestartKey); ptr != NULL; ptr = RtlEnumerateGenericTableWithoutSplaying(&g_Table, &RestartKey))
{
// Process the element pointed to by ptr
Pcontent pc = (Pcontent)ptr;
KdPrint(("%d.\r\n", pc->i));
}
ExReleaseFastMutex(&g_FastMutex);
}
void EnumerateGenericTableEx()
{
ULONG i = 0;
ExAcquireFastMutex(&g_FastMutex);
for ( ; i < RtlNumberGenericTableElements(&g_Table); i++)
{
// Process the element pointed to by p
Pcontent pc = (Pcontent)RtlGetElementGenericTable(&g_Table, i);
KdPrint(("%d.\r\n", pc->i));
}
ExReleaseFastMutex(&g_FastMutex);
}
void DeleteElementGenericTable1()
/*
功能：删除整个表。
参考：https://msdn.microsoft.com/en-us/library/windows/hardware/ff552243(v=vs.85).aspx
因为：RtlEnumerateGenericTable returns a pointer to the next element，所以这没有走完就退出了，结果是没有删除完毕。
要做还得改进。
*/
{
PVOID p = NULL;
ExAcquireFastMutex(&g_FastMutex);
for (p = RtlEnumerateGenericTable ( &g_Table, TRUE );
p != NULL;
p = RtlEnumerateGenericTable ( &g_Table, FALSE ))
{
// Process the element pointed to by p
BOOLEAN B = RtlDeleteElementGenericTable(&g_Table, p);
ASSERT(B != TRUE );
}
ExReleaseFastMutex(&g_FastMutex);
}
void DeleteElementGenericTable2()
/*
功能：删除整个表。
参考：https://msdn.microsoft.com/en-us/library/windows/hardware/ff552247(v=vs.85).aspx
原因：同上，要做还有很多改进的余地。
*/
{
PVOID ptr = NULL;
PVOID RestartKey = NULL;
ExAcquireFastMutex(&g_FastMutex);
for (ptr = RtlEnumerateGenericTableWithoutSplaying(&g_Table, &RestartKey);
ptr != NULL;
ptr = RtlEnumerateGenericTableWithoutSplaying(&g_Table, &RestartKey))
{
// Process the element pointed to by ptr
BOOLEAN B = RtlDeleteElementGenericTable(&g_Table, ptr);
ASSERT(B != TRUE );
}
ExReleaseFastMutex(&g_FastMutex);
}
void DeleteElementGenericTable3()
/*
功能：删除整个表。
摘自：Windows 8 Driver Samples\AvScan File System Minifilter Driver\C++\filter\avscan.c
*/
{
ExAcquireFastMutex(&g_FastMutex);
while (!RtlIsGenericTableEmpty( &g_Table ) )
{
PVOID entry = RtlGetElementGenericTable(&g_Table, 0);
RtlDeleteElementGenericTable(&g_Table, entry);
}
ExReleaseFastMutex(&g_FastMutex);
}
DRIVER_INITIALIZE DriverEntry;
NTSTATUS DriverEntry( __in struct _DRIVER_OBJECT * DriverObject, __in PUNICODE_STRING RegistryPath)
{
NTSTATUS status = STATUS_UNSUCCESSFUL;
KdBreakPoint();
DriverObject->DriverUnload = Unload;
InitializeGenericTable();
InsertElementGenericTable();
GetElementGenericTable();
LookupElementGenericTable();
EnumerateGenericTable();
EnumerateGenericTableWithoutSplaying();
EnumerateGenericTableEx();
//DeleteElementGenericTable1();
//DeleteElementGenericTable2();
DeleteElementGenericTable3();
return status;
}
