對驅動型鍵盤記錄程序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立刻返回。