? ? ? ? 這篇我們看一個”容錯“”節(jié)省“的實例。一下是一個Win32API的聲明（轉(zhuǎn)載請指明出處）

LONG WINAPI RegEnumKeyEx(__in         HKEY hKey,__in         DWORD dwIndex,__out        LPTSTR lpName,__inout      LPDWORD lpcName,__reserved   LPDWORD lpReserved,__inout      LPTSTR lpClass,__inout_opt  LPDWORD lpcClass,__out_opt    PFILETIME lpftLastWriteTime
);

節(jié)省 ? ? ? ?

? ? ? ?這個函數(shù)底層是使用ZwEnumerateKey，使用過該函數(shù)的同學(xué)應(yīng)該知道，該函數(shù)根據(jù)傳入的KEY_INFORMATION_CLASS不同而查詢該項不同結(jié)構(gòu)體的數(shù)據(jù)。我們來看兩個不同結(jié)構(gòu)體

typedef struct _KEY_BASIC_INFORMATION {LARGE_INTEGER LastWriteTime;ULONG         TitleIndex;ULONG         NameLength;WCHAR         Name[1];
} KEY_BASIC_INFORMATION, *PKEY_BASIC_INFORMATION;

typedef struct _KEY_NODE_INFORMATION {LARGE_INTEGER LastWriteTime;ULONG         TitleIndex;ULONG         ClassOffset;ULONG         ClassLength;ULONG         NameLength;WCHAR         Name[1];
} KEY_NODE_INFORMATION, *PKEY_NODE_INFORMATION;

? ? ? ?可見到NODE比BASIC多出ClassOffset和ClassLength和藏在Name中的Class信息。RegEnumKeyEx要獲取的信息中是可以通過是否為NULL來定的，如果你不想獲取Class的信息，可以將lpClass和lpcClass指定為NULL。那么Reactos中如何實現(xiàn)的呢？我們先思考下，如果我們在NtEnumerateKey中一股腦兒使用KEY_NODE_INFORMATION去查詢，是很簡單，但是是不是有點浪費呢（也許用戶不用查詢Class信息）？如果使用KEY_BASIC_INFORMATION，則如果用戶要查詢Class信息，則我們將查詢不到。或許你想到了——特殊情況特殊處理，是的Reactos就是如此做的。下面看一段我修改后的代碼

lpKeyInfo = HeapAlloc( GetProcessHeap(), 0, BufferSize );if ( NULL == lpKeyInfo ) {// 分配失敗lRes = STATUS_MEMORY_NOT_ALLOCATED;break;}ULONG ResultSize = 0;// 查詢不同類型的結(jié)構(gòu)體取決于lpClass是否需要查詢lRes = GETORIFUNC(EnumerateKey)(hKeyHandle,dwIndex,lpClass ? KeyNodeInformation : KeyBasicInformation,lpKeyInfo,BufferSize,&ResultSize );CHECKRESULT(lRes);// 定義一個聯(lián)合體typedef union {KEY_NODE_INFORMATION Node;KEY_BASIC_INFORMATION Basic;} un;un* KeyInfo = (un*)lpKeyInfo;if ( NT_FAILED(lRes) ) {break;}else {if ( NULL == lpClass ) {// 如果lpClass不需要查詢，則使用KEY_BASIC_INFORMATION結(jié)構(gòu)體中數(shù)據(jù)if ( KeyInfo->Basic.NameLength > NameLength ) {// 承載的空間不夠lRes = STATUS_BUFFER_OVERFLOW;}else {RtlCopyMemory( lpName, KeyInfo->Basic.Name, KeyInfo->Basic.NameLength );// 返回的長度不包含結(jié)尾符*lpcName = (DWORD)( KeyInfo->Basic.NameLength / sizeof(WCHAR) );// 設(shè)置結(jié)尾符lpName[*lpcName] = (WCHAR)0;}}else {// 如果lpClass需要查詢，則使用KEY_NODE_INFORMATION結(jié)構(gòu)體中數(shù)據(jù)if ( KeyInfo->Node.NameLength > NameLength || KeyInfo->Node.ClassLength > ClassLength ) {// 承載的空間不夠lRes = STATUS_BUFFER_OVERFLOW;}else {// 拷貝數(shù)據(jù)到內(nèi)存中RtlCopyMemory( lpName, KeyInfo->Node.Name, KeyInfo->Node.NameLength );// 設(shè)置返回的大小，不包含結(jié)尾符*lpcName = KeyInfo->Node.NameLength / sizeof(WCHAR);// 設(shè)置結(jié)尾符lpName[*lpcName] = (WCHAR)0;// 拷貝數(shù)據(jù)到內(nèi)存中RtlCopyMemory( lpClass,(PVOID)((ULONG_PTR)KeyInfo->Node.Name + KeyInfo->Node.ClassOffset),KeyInfo->Node.ClassLength );// 設(shè)置返回的大小，不包含結(jié)尾符*lpcClass = (DWORD)(KeyInfo->Node.ClassLength / sizeof(WCHAR));// 設(shè)置結(jié)尾符lpClass[*lpcClass] = (WCHAR)0;}}if ( lRes == STATUS_SUCCESS && NULL != lpftLastWriteTime ) {if ( lpClass == NULL ) {// 如果lpClass需要查詢，則使用KEY_NODE_INFORMATION結(jié)構(gòu)體中數(shù)據(jù)lpftLastWriteTime->dwLowDateTime = KeyInfo->Basic.LastWriteTime.u.LowPart;lpftLastWriteTime->dwHighDateTime = KeyInfo->Basic.LastWriteTime.u.HighPart;}else {// 如果lpClass需要查詢，則使用KEY_NODE_INFORMATION結(jié)構(gòu)體中數(shù)據(jù)lpftLastWriteTime->dwLowDateTime = KeyInfo->Node.LastWriteTime.u.LowPart;lpftLastWriteTime->dwHighDateTime = KeyInfo->Node.LastWriteTime.u.HighPart;}}}

? ? ? ? Reactos使用了聯(lián)合體解決了這個問題。

容錯。

? ? ? ? 我們寫的API，往往會接受調(diào)用方傳入的一些數(shù)據(jù)。如果這個數(shù)據(jù)是個很大的且沒有固定結(jié)構(gòu)的數(shù)據(jù)時，那么就要非常注意這個空間的大小了。如RegEnumKeyEx函數(shù)就接受了兩個用戶傳入的空間及其大小。

? ? ? ? 在我們重寫的RegEnumKey中對用戶傳入的數(shù)據(jù)進行填充前，我們要先準(zhǔn)確無誤地獲取數(shù)據(jù)，而用戶傳入的空間和大小我們不能用，因為我們不知道他對不對，于是我們要先分配一個適合大小的空間，調(diào)用NtEnumerateKey得到數(shù)據(jù)后再對用戶傳入空間大小進行判斷，對空間進行填充。但是這個空間大小如何定義呢？有一種辦法就是不斷試錯，通過ResultLength參數(shù)得到適合的空間大小。

? ? ? ? 但是是不是很費呢？是的，Reactos對RegEnumKey的實現(xiàn)則是利用用戶傳入的空間大小，而沒有用其傳入的空間，這樣一旦空間過小，會快速發(fā)現(xiàn)，而不用等數(shù)據(jù)都查完了才發(fā)現(xiàn)用戶傳入的空間太小。但是現(xiàn)在存在一個問題，如果用戶傳入的空間大小特別大，實際用不到這么大的數(shù)據(jù)，那怎么辦？難道我們也要聽從用戶分配一個巨大的內(nèi)存空間么？不是，我們看看Reactos的做法（我修改后的代碼）

LPVOID lpKeyInfo = NULL;do { ULONG NameLength = 0;if ( *lpcName > 0 ) {NameLength = min( *lpcName - 1, MAX_PATH ) * sizeof(WCHAR);}else {NameLength = 0;}ULONG BufferSize = 0;ULONG ClassLength = 0;if ( lpClass ) {if ( *lpcClass > 0 ) {ClassLength = min( *lpcClass - 1, MAX_PATH ) * sizeof(WCHAR);}else {ClassLength = 0;}// +3 再& ~3是為了讓大小按4字節(jié)對齊且取其上限// 如果存在lpClass，則要將ClassLength的長度算進去// 如果存在lpClass，則要使用KEY_NODE_INFORMATION結(jié)構(gòu)體BufferSize = ( ( sizeof(KEY_NODE_INFORMATION) + NameLength + 3 ) & ~3 ) + ClassLength;}else {// 如果不存在lpClass，則使用KEY_BASIC_INFORMATION結(jié)構(gòu)體BufferSize = sizeof(KEY_BASIC_INFORMATION) + NameLength;}// 在堆上分配一段適合大小的空間lpKeyInfo = HeapAlloc( GetProcessHeap(), 0, BufferSize );

? ? ? 它在別人傳入的空間-1和260之間選了一個最小值。如果調(diào)用方傳了一個巨大的空間大小，我們也就分配260個WCHAR的大小。可能有人問：那么如果Class和KeyNamed的長度就是長于260呢？好問題！Reactos系統(tǒng)中Class和KeyName的最大長度就是260，何來長于260的名字呢？我在我電腦上剛做了實驗，將某鍵名改成250個1，Regedit就會報錯，說名字太長。

? ? ? 這種容錯還用在RegQueryValueEx的實現(xiàn)中，以下列出我修改后的部分代碼

NTSTATUS WINAPI OriRegQueryValueEx(HANDLE KeyHandle,LPCWSTR lpValueName,LPDWORD lpReserved,LPDWORD lpType,LPBYTE lpData,LPDWORD lpcbData )
{NTSTATUS lRes = STATUS_INVALID_PARAMETER;char buffer[256] = {0};char *buf_ptr = buffer;do {KEY_VALUE_PARTIAL_INFORMATION *info = (KEY_VALUE_PARTIAL_INFORMATION *)buffer;// KEY_VALUE_PARTIAL_INFORMATION結(jié)構(gòu)的必要結(jié)構(gòu)體長度static const int info_size = offsetof( KEY_VALUE_PARTIAL_INFORMATION, Data );// 參數(shù)判斷if ( ( NULL != lpData && NULL == lpcbData ) || lpReserved ) {return STATUS_INVALID_PARAMETER;}UNICODE_STRING name_str;RtlInitUnicodeString_( &name_str, lpValueName );// 取一段比較合理的空間大小，這樣避免用戶傳入過大的空間大小DWORD total_size = 0;if ( NULL != lpData ) {total_size = min( sizeof(buffer), *lpcbData + info_size );}else {total_size = info_size;if ( NULL != lpcbData ) {*lpcbData = 0;}}/* this matches Win9x behaviour - NT sets *type to a random value */if ( lpType ) {*lpType = REG_NONE;}
……

? ? ? ??因為Value的長度理論上是可以超過260的，于是有以下處理

lRes = GETORIFUNC(QueryValueKey)(KeyHandle,&name_str, KeyValuePartialInformation,buffer, total_size, &total_size );if ( NT_FAILED(lRes) && lRes != STATUS_BUFFER_OVERFLOW ) {break;}if ( NULL == lpData ) {lRes = STATUS_SUCCESS;}else {while ( lRes == STATUS_BUFFER_OVERFLOW && total_size - info_size <= *lpcbData ) {// 空間不足，且需要的空間大小比用戶傳入的小// 釋放之前分配的內(nèi)存if ( buf_ptr != buffer ) {HeapFree( GetProcessHeap(), 0, buf_ptr );buf_ptr = NULL;}// 重新分配內(nèi)存buf_ptr = (char*)HeapAlloc( GetProcessHeap(), 0, total_size );if ( NULL == buf_ptr ) {return STATUS_NO_MEMORY;}info = (KEY_VALUE_PARTIAL_INFORMATION *)buf_ptr;lRes = GETORIFUNC(QueryValueKey)( KeyHandle, &name_str, KeyValuePartialInformation,buf_ptr, total_size, &total_size );}if ( NT_SUCCESS(lRes) ) {memcpy( lpData, buf_ptr + info_size, total_size - info_size );/* if the type is REG_SZ and data is not 0-terminated* and there is enough space in the buffer NT appends a \0 */if ( is_string(info->Type) && total_size - info_size <= *lpcbData - sizeof(WCHAR) ) {WCHAR *ptr = (WCHAR *)( lpData + total_size - info_size );if (ptr > (WCHAR *)lpData && ptr[-1]) {*ptr = 0;}}}else if ( lRes != STATUS_BUFFER_OVERFLOW ) {break;}}if ( NULL != lpType) {*lpType = info->Type;}if ( NULL != lpcbData) {*lpcbData = total_size - info_size;}} while (0);

? ? ? ??在空間不夠的情況下，會在堆上分配更多的空間。

總結(jié)

以上是生活随笔為你收集整理的一种注册表沙箱的思路、实现——研究Reactos中注册表函数的实现3的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。