对驱动型键盘记录程序Klog的修改

 对驱动型键盘记录程序Klog的修改    vocanicy 2007/09/27

　　Windows上的各种软件都做过，就是没接触做驱动，真是惭愧。最近下载了Clandestiny写的驱动型键盘记录程序的源码，在学习的过程中发现了几点问题。这些问题经过小的改动就可以解决，我相信是作者有意留在源码中的。
　　本人初次接触驱动程序开发，下面这些小的改动，也让我欣赏了很多次Windows的蓝屏。在这我说说修改的思路。

1、内存泄漏
　　Klog的OnReadCompletion IRP完成处理函数中获取按键信息KEY_DATA，然后为其分配4K内存，并加入键盘记录的链表。Klog另外有一个专门负责将链表中的记录写入文件的线程。但是这个线程将键盘记录取出后，并没有将其内存释放，导致内存泄漏。

由于OnReadCompletion的IRQL为DISPATCH_LEVEL，原程序中使用ExAllocatePool分配NonPagedPool内存。这里我将其替换成ExAllocatePoolWithTag后，就可以在PoolTag工具中看到Tag内存的分配情况。每次按键之后，内存都重新分配，但是原内存并未释放。
KEY_DATA* kData = (KEY_DATA*)ExAllocatePoolWithTag(NonPagedPool, sizeof(KEY_DATA), (ULONG)'GOLK');

解决方法：在线程将按键记录从链表中取出写入文件后，调用ExFreePool(kData)将其释放。

2、内存碎片
　　第1点中提到的OnReadCompletion的IRQL为DISPATCH_LEVEL，因此原程序中使用ExAllocatePool分配NonPagedPool内存。DDK的文档中有说明使用ExAllocatePool分配的内存是以PAGE_SIZE对齐的。Windows中的PAGE_SIZE大小一般为4K。也就是说每按一个键，KLog都将分配4K的NonPagedPool内存，而实际上每个按键记录KEY_DATA结构只需要2个字节。更重要的是DDK文档中提到NonPagedPool内存为稀有资源（scarce resource）。这可以说是非常巨大的浪费。（如果没有修改第一点提到的内存泄漏，这个程序不知道能运行多长时间。4K×击键次数）
注：原程序中的这个缺点，作者是有提到的。

解决方法：在驱动的DriverEntry中预先分配好内存，并自己建立一个内存队列。这样可以避免频繁分配释放内存。下面是我写的键盘循环队列的结构。这里用到一个小技巧，使结构大小对齐到PAGE_SIZE。

#define HDR_SIZE        (sizeof(KSEMAPHORE) + sizeof(KSPIN_LOCK) + sizeof(int) * 2)
#define KEY_QUEUE_COUNT        ((PAGE_SIZE - HDR_SIZE) / sizeof(KEY_DATA))
struct KEY_QUEUE
{
        KSEMAPHORE KeySem;
        KSPIN_LOCK KeyLock;
       
        int Start;      // 队列头
        int End;        // 队列尾
       
        KEY_DATA Data[KEY_QUEUE_COUNT];
};

typedef KEY_QUEUE* PKEY_QUEUE;

3、驱动卸载
　　Klog中为了截获键盘的按键信息，为IRP_MJ_READ操作的IRP设置了OnReadCompletion完成处理。OnReadCompletion需要等到某个按键按下之后才被调用。原KLog程序的Unload函数中设置一个等待循环，等到所有的IRP处理完成之后才退出。这就造成了一个问题。如果驱动卸载时，没有按键操作，系统将被挂起直到某个按键被按下。
注：这个缺点作者在注释中也有说明，因为驱动一般情况下不会被卸载，这个问题就显得不重要了。

解决方法：我的解决方法是在IoSetCompletionRoutine的同时，为IRP再调用IoSetCancelRoutine指定一个取消处理函数。在驱动卸载时，对Pending的IRP调用IoCancelIrp，迫使IRP立刻返回。