一种注册表沙箱的思路、实现——研究Reactos中注册表函数的实现3

这篇我们看一个”容错“”节省“的实例。一下是一个Win32API的声明（转载请指明出处）

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
);

节省

这个函数底层是使用ZwEnumerateKey，使用过该函数的同学应该知道，该函数根据传入的KEY_INFORMATION_CLASS不同而查询该项不同结构体的数据。我们来看两个不同结构体

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中如何实现的呢？我们先思考下，如果我们在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;// 查询不同类型的结构体取决于lpClass是否需要查询lRes = GETORIFUNC(EnumerateKey)(hKeyHandle,dwIndex,lpClass ? KeyNodeInformation : KeyBasicInformation,lpKeyInfo,BufferSize,&ResultSize );CHECKRESULT(lRes);// 定义一个联合体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结构体中数据if ( KeyInfo->Basic.NameLength > NameLength ) {// 承载的空间不够lRes = STATUS_BUFFER_OVERFLOW;}else {RtlCopyMemory( lpName, KeyInfo->Basic.Name, KeyInfo->Basic.NameLength );// 返回的长度不包含结尾符*lpcName = (DWORD)( KeyInfo->Basic.NameLength / sizeof(WCHAR) );// 设置结尾符lpName[*lpcName] = (WCHAR)0;}}else {// 如果lpClass需要查询，则使用KEY_NODE_INFORMATION结构体中数据if ( KeyInfo->Node.NameLength > NameLength || KeyInfo->Node.ClassLength > ClassLength ) {// 承载的空间不够lRes = STATUS_BUFFER_OVERFLOW;}else {// 拷贝数据到内存中RtlCopyMemory( lpName, KeyInfo->Node.Name, KeyInfo->Node.NameLength );// 设置返回的大小，不包含结尾符*lpcName = KeyInfo->Node.NameLength / sizeof(WCHAR);// 设置结尾符lpName[*lpcName] = (WCHAR)0;// 拷贝数据到内存中RtlCopyMemory( lpClass,(PVOID)((ULONG_PTR)KeyInfo->Node.Name + KeyInfo->Node.ClassOffset),KeyInfo->Node.ClassLength );// 设置返回的大小，不包含结尾符*lpcClass = (DWORD)(KeyInfo->Node.ClassLength / sizeof(WCHAR));// 设置结尾符lpClass[*lpcClass] = (WCHAR)0;}}if ( lRes == STATUS_SUCCESS && NULL != lpftLastWriteTime ) {if ( lpClass == NULL ) {// 如果lpClass需要查询，则使用KEY_NODE_INFORMATION结构体中数据lpftLastWriteTime->dwLowDateTime = KeyInfo->Basic.LastWriteTime.u.LowPart;lpftLastWriteTime->dwHighDateTime = KeyInfo->Basic.LastWriteTime.u.HighPart;}else {// 如果lpClass需要查询，则使用KEY_NODE_INFORMATION结构体中数据lpftLastWriteTime->dwLowDateTime = KeyInfo->Node.LastWriteTime.u.LowPart;lpftLastWriteTime->dwHighDateTime = KeyInfo->Node.LastWriteTime.u.HighPart;}}}

Reactos使用了联合体解决了这个问题。

容错。

        我们写的API，往往会接受调用方传入的一些数据。如果这个数据是个很大的且没有固定结构的数据时，那么就要非常注意这个空间的大小了。如RegEnumKeyEx函数就接受了两个用户传入的空间及其大小。

        在我们重写的RegEnumKey中对用户传入的数据进行填充前，我们要先准确无误地获取数据，而用户传入的空间和大小我们不能用，因为我们不知道他对不对，于是我们要先分配一个适合大小的空间，调用NtEnumerateKey得到数据后再对用户传入空间大小进行判断，对空间进行填充。但是这个空间大小如何定义呢？有一种办法就是不断试错，通过ResultLength参数得到适合的空间大小。

        但是是不是很费呢？是的，Reactos对RegEnumKey的实现则是利用用户传入的空间大小，而没有用其传入的空间，这样一旦空间过小，会快速发现，而不用等数据都查完了才发现用户传入的空间太小。但是现在存在一个问题，如果用户传入的空间大小特别大，实际用不到这么大的数据，那怎么办？难道我们也要听从用户分配一个巨大的内存空间么？不是，我们看看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字节对齐且取其上限// 如果存在lpClass，则要将ClassLength的长度算进去// 如果存在lpClass，则要使用KEY_NODE_INFORMATION结构体BufferSize = ( ( sizeof(KEY_NODE_INFORMATION) + NameLength + 3 ) & ~3 ) + ClassLength;}else {// 如果不存在lpClass，则使用KEY_BASIC_INFORMATION结构体BufferSize = sizeof(KEY_BASIC_INFORMATION) + NameLength;}// 在堆上分配一段适合大小的空间lpKeyInfo = HeapAlloc( GetProcessHeap(), 0, BufferSize );

它在别人传入的空间-1和260之间选了一个最小值。如果调用方传了一个巨大的空间大小，我们也就分配260个WCHAR的大小。可能有人问：那么如果Class和KeyNamed的长度就是长于260呢？好问题！Reactos系统中Class和KeyName的最大长度就是260，何来长于260的名字呢？我在我电脑上刚做了实验，将某键名改成250个1，Regedit就会报错，说名字太长。

      这种容错还用在RegQueryValueEx的实现中，以下列出我修改后的部分代码

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结构的必要结构体长度static const int info_size = offsetof( KEY_VALUE_PARTIAL_INFORMATION, Data );// 参数判断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 ) {// 空间不足，且需要的空间大小比用户传入的小// 释放之前分配的内存if ( buf_ptr != buffer ) {HeapFree( GetProcessHeap(), 0, buf_ptr );buf_ptr = NULL;}// 重新分配内存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);

在空间不够的情况下，会在堆上分配更多的空间。