ntdll: Fix the return status of NtQueryValueKey for a too small buffer.

dlls/kernel32/except.c

@@ -220,7 +220,7 @@ static BOOL	start_debugger(PEXCEPTION_POINTERS epointers, HANDLE hEvent)
 
         RtlInitUnicodeString( &nameW, DebuggerW );
         if (NtQueryValueKey( hDbgConf, &nameW, KeyValuePartialInformation,
-                             NULL, 0, &format_size ) == STATUS_BUFFER_OVERFLOW)
+                             NULL, 0, &format_size ) == STATUS_BUFFER_TOO_SMALL)
         {
             char *data = HeapAlloc(GetProcessHeap(), 0, format_size);
             NtQueryValueKey( hDbgConf, &nameW, KeyValuePartialInformation,

dlls/kernel32/path.c

@@ -182,7 +182,7 @@ static BOOL add_boot_rename_entry( LPCWSTR source, LPCWSTR dest, DWORD flags )
 
     /* First we check if the key exists and if so how many bytes it already contains. */
     if (NtQueryValueKey( Reboot, &nameW, KeyValuePartialInformation,
-                         NULL, 0, &DataSize ) == STATUS_BUFFER_OVERFLOW)
+                         NULL, 0, &DataSize ) == STATUS_BUFFER_TOO_SMALL)
     {
         if (!(Buffer = HeapAlloc( GetProcessHeap(), 0, DataSize + len1 + len2 + sizeof(WCHAR) )))
             goto Quit;

dlls/ntdll/reg.c

@@ -477,7 +477,7 @@ NTSTATUS WINAPI NtQueryValueKey( HANDLE handle, const UNICODE_STRING *name,
 {
     NTSTATUS ret;
     UCHAR *data_ptr;
-    unsigned int fixed_size = 0;
+    unsigned int fixed_size = 0, min_size = 0;
 
     TRACE( "(%p,%s,%d,%p,%d)\n", handle, debugstr_us(name), info_class, info, length );
 
@@ -489,30 +489,26 @@ NTSTATUS WINAPI NtQueryValueKey( HANDLE handle, const UNICODE_STRING *name,
     case KeyValueBasicInformation:
     {
         KEY_VALUE_BASIC_INFORMATION *basic_info = info;
-        if (FIELD_OFFSET(KEY_VALUE_BASIC_INFORMATION, Name) < length)
-        {
-            memcpy(basic_info->Name, name->Buffer,
-                   min(length - FIELD_OFFSET(KEY_VALUE_BASIC_INFORMATION, Name), name->Length));
-        }
-        fixed_size = FIELD_OFFSET(KEY_VALUE_BASIC_INFORMATION, Name) + name->Length;
+        min_size = FIELD_OFFSET(KEY_VALUE_BASIC_INFORMATION, Name);
+        fixed_size = min_size + name->Length;
+        if (min_size < length)
+            memcpy(basic_info->Name, name->Buffer, min(length - min_size, name->Length));
         data_ptr = NULL;
         break;
     }
     case KeyValueFullInformation:
     {
         KEY_VALUE_FULL_INFORMATION *full_info = info;
-        if (FIELD_OFFSET(KEY_VALUE_FULL_INFORMATION, Name) < length)
-        {
-            memcpy(full_info->Name, name->Buffer,
-                   min(length - FIELD_OFFSET(KEY_VALUE_FULL_INFORMATION, Name), name->Length));
-        }
+        min_size = FIELD_OFFSET(KEY_VALUE_FULL_INFORMATION, Name);
+        fixed_size = min_size + name->Length;
+        if (min_size < length)
+            memcpy(full_info->Name, name->Buffer, min(length - min_size, name->Length));
         data_ptr = (UCHAR *)full_info->Name + name->Length;
-        fixed_size = (char *)data_ptr - (char *)info;
         break;
     }
     case KeyValuePartialInformation:
+        min_size = fixed_size = FIELD_OFFSET(KEY_VALUE_PARTIAL_INFORMATION, Data);
         data_ptr = ((KEY_VALUE_PARTIAL_INFORMATION *)info)->Data;
-        fixed_size = (char *)data_ptr - (char *)info;
         break;
     default:
         FIXME( "Information class %d not implemented\n", info_class );
@@ -529,7 +525,8 @@ NTSTATUS WINAPI NtQueryValueKey( HANDLE handle, const UNICODE_STRING *name,
             copy_key_value_info( info_class, info, length, reply->type,
                                  name->Length, reply->total );
             *result_len = fixed_size + (info_class == KeyValueBasicInformation ? 0 : reply->total);
-            if (length < *result_len) ret = STATUS_BUFFER_OVERFLOW;
+            if (length < min_size) ret = STATUS_BUFFER_TOO_SMALL;
+            else if (length < *result_len) ret = STATUS_BUFFER_OVERFLOW;
         }
     }
     SERVER_END_REQ;

dlls/ntdll/tests/reg.c

@@ -480,7 +480,7 @@ static void test_NtQueryValueKey(void)
     KEY_VALUE_BASIC_INFORMATION *basic_info;
     KEY_VALUE_PARTIAL_INFORMATION *partial_info;
     KEY_VALUE_FULL_INFORMATION *full_info;
-    DWORD len;
+    DWORD len, expected;
 
     pRtlCreateUnicodeStringFromAsciiz(&ValName, "deletetest");
 
@@ -554,7 +554,7 @@ static void test_NtQueryValueKey(void)
     pRtlCreateUnicodeStringFromAsciiz(&ValName, "stringtest");
 
     status = pNtQueryValueKey(key, &ValName, KeyValuePartialInformation, NULL, 0, &len);
-    todo_wine ok(status == STATUS_BUFFER_TOO_SMALL, "NtQueryValueKey should have returned STATUS_BUFFER_TOO_SMALL instead of 0x%08x\n", status);
+    ok(status == STATUS_BUFFER_TOO_SMALL, "NtQueryValueKey should have returned STATUS_BUFFER_TOO_SMALL instead of 0x%08x\n", status);
     partial_info = HeapAlloc(GetProcessHeap(), 0, len+1);
     memset(partial_info, 0xbd, len+1);
     status = pNtQueryValueKey(key, &ValName, KeyValuePartialInformation, partial_info, len, &len);
@@ -564,6 +564,21 @@ static void test_NtQueryValueKey(void)
     ok(partial_info->DataLength == STR_TRUNC_SIZE, "NtQueryValueKey returned wrong DataLength %d\n", partial_info->DataLength);
     ok(!memcmp(partial_info->Data, stringW, STR_TRUNC_SIZE), "incorrect Data returned\n");
     ok(*(partial_info->Data+STR_TRUNC_SIZE) == 0xbd, "string overflowed %02x\n", *(partial_info->Data+STR_TRUNC_SIZE));
+
+    expected = len;
+    status = pNtQueryValueKey(key, &ValName, KeyValuePartialInformation, partial_info, 0, &len);
+    ok(status == STATUS_BUFFER_TOO_SMALL, "NtQueryValueKey wrong status 0x%08x\n", status);
+    ok(len == expected, "NtQueryValueKey wrong len %u\n", len);
+    status = pNtQueryValueKey(key, &ValName, KeyValuePartialInformation, partial_info, 1, &len);
+    ok(status == STATUS_BUFFER_TOO_SMALL, "NtQueryValueKey wrong status 0x%08x\n", status);
+    ok(len == expected, "NtQueryValueKey wrong len %u\n", len);
+    status = pNtQueryValueKey(key, &ValName, KeyValuePartialInformation, partial_info, FIELD_OFFSET(KEY_VALUE_PARTIAL_INFORMATION, Data) - 1, &len);
+    ok(status == STATUS_BUFFER_TOO_SMALL, "NtQueryValueKey wrong status 0x%08x\n", status);
+    ok(len == expected, "NtQueryValueKey wrong len %u\n", len);
+    status = pNtQueryValueKey(key, &ValName, KeyValuePartialInformation, partial_info, FIELD_OFFSET(KEY_VALUE_PARTIAL_INFORMATION, Data), &len);
+    ok(status == STATUS_BUFFER_OVERFLOW, "NtQueryValueKey wrong status 0x%08x\n", status);
+    ok(len == expected, "NtQueryValueKey wrong len %u\n", len);
+
     HeapFree(GetProcessHeap(), 0, partial_info);
 
     pRtlFreeUnicodeString(&ValName);