[gschneider] 44566: Alexander Yastrebov <menone7 AT gmail DOT com>: - Simplify and fix error checking in RegQuery/SetEx functions - Add handling for missing termination characters - Fixes one advapi32 winetest See issue #4993 for more details.

Author: gschneider
Date: Sun Dec 13 17:25:50 2009
New Revision: 44566

URL: http://svn.reactos.org/svn/reactos?rev=44566&view=rev
Log:
Alexander Yastrebov <menone7 AT gmail DOT com>: 
- Simplify and fix error checking in RegQuery/SetEx functions
- Add handling for missing termination characters
- Fixes one advapi32 winetest
See issue #4993 for more details.

Modified:
    trunk/reactos/dll/win32/advapi32/reg/reg.c

Modified: trunk/reactos/dll/win32/advapi32/reg/reg.c
URL: http://svn.reactos.org/svn/reactos/trunk/reactos/dll/win32/advapi32/reg/reg....
==============================================================================
--- trunk/reactos/dll/win32/advapi32/reg/reg.c [iso-8859-1] (original)
+++ trunk/reactos/dll/win32/advapi32/reg/reg.c [iso-8859-1] Sun Dec 13 17:25:50 2009
@@ -3989,10 +3989,11 @@
 {
     UNICODE_STRING ValueName;
     UNICODE_STRING ValueData;
-    ANSI_STRING AnsiString;
     LONG ErrorCode;
     DWORD Length;
     DWORD Type;
+    NTSTATUS Status;
+    ULONG Index;

TRACE("hKey 0x%X  lpValueName %s  lpData 0x%X  lpcbData %d\n",
           hKey, lpValueName, lpData, lpcbData ? *lpcbData : 0);
@@ -4041,16 +4042,23 @@
         ErrorCode == ERROR_MORE_DATA)
     {

-        if ((Type == REG_SZ) || (Type == REG_MULTI_SZ) || (Type == REG_EXPAND_SZ))
+        if (is_string(Type))
         {
             if (ErrorCode == ERROR_SUCCESS && ValueData.Buffer != NULL)
             {
-                RtlInitAnsiString(&AnsiString, NULL);
-                AnsiString.Buffer = (LPSTR)lpData;
-                AnsiString.MaximumLength = *lpcbData;
-                ValueData.Length = Length;
-                ValueData.MaximumLength = ValueData.Length + sizeof(WCHAR);
-                RtlUnicodeStringToAnsiString(&AnsiString, &ValueData, FALSE);
+                Status = RtlUnicodeToMultiByteN((PCHAR)lpData, *lpcbData, &Index, (PWCHAR)ValueData.Buffer, Length);
+                if (NT_SUCCESS(Status))
+                {
+                    PCHAR szData = (PCHAR)lpData;
+                    if(&szData[Index] < (PCHAR)(lpData + *lpcbData))
+                    {
+                        szData[Index] = '\0';
+                    }
+                }
+                else
+                {
+                    ErrorCode = RtlNtStatusToDosError(Status);
+                }
             }

Length = Length / sizeof(WCHAR);
@@ -4135,7 +4143,7 @@

status = NtQueryValueKey( hkey, &name_str, KeyValuePartialInformation,
                               buffer, total_size, &total_size );
-    if (status && status != STATUS_BUFFER_OVERFLOW) goto done;
+    if (!NT_SUCCESS(status) && status != STATUS_BUFFER_OVERFLOW) goto done;

if (data)
     {
@@ -4153,12 +4161,12 @@
                                       buf_ptr, total_size, &total_size );
         }

-        if (!status)
+        if (NT_SUCCESS(status))
         {
             memcpy( data, 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 (total_size - info_size <= *count-sizeof(WCHAR) && is_string(info->Type))
+            if (is_string(info->Type) && total_size - info_size <= *count-sizeof(WCHAR))
             {
                 WCHAR *ptr = (WCHAR *)(data + total_size - info_size);
                 if (ptr > (WCHAR *)data && ptr[-1]) *ptr = 0;
@@ -4743,12 +4751,16 @@
     LONG ErrorCode;
     LPBYTE pData;
     DWORD DataSize;
-
-    if (lpValueName != NULL &&
-        strlen(lpValueName) != 0)
-    {
-        RtlCreateUnicodeStringFromAsciiz(&ValueName,
-                                         (PSTR)lpValueName);
+    NTSTATUS Status;
+
+    /* Convert SubKey name to Unicode */
+    if (lpValueName != NULL && lpValueName[0] != '\0')
+    {
+        BOOL bConverted;
+        bConverted = RtlCreateUnicodeStringFromAsciiz(&ValueName,
+                                                  (PSTR)lpValueName);
+        if(!bConverted)
+            return ERROR_NOT_ENOUGH_MEMORY;
     }
     else
     {
@@ -4757,32 +4769,32 @@

pValueName = (LPWSTR)ValueName.Buffer;

-    if (((dwType == REG_SZ) ||
-         (dwType == REG_MULTI_SZ) ||
-         (dwType == REG_EXPAND_SZ)) &&
-        (cbData != 0))
-    {
-        /* NT adds one if the caller forgot the NULL-termination character */
-        if (lpData[cbData - 1] != '\0')
-        {
-            cbData++;
-        }
-
-        RtlInitAnsiString(&AnsiString,
-                          NULL);
+
+    if (is_string(dwType) && (cbData != 0))
+    {
+        /* Convert ANSI string Data to Unicode */
+        /* If last character NOT zero then increment length */
+        LONG bNoNulledStr = ((lpData[cbData-1] != '\0') ? 1 : 0);
         AnsiString.Buffer = (PSTR)lpData;
-        AnsiString.Length = cbData - 1;
-        AnsiString.MaximumLength = cbData;
-        RtlAnsiStringToUnicodeString(&Data,
+        AnsiString.Length = cbData + bNoNulledStr;
+        AnsiString.MaximumLength = cbData + bNoNulledStr;
+        Status = RtlAnsiStringToUnicodeString(&Data,
                                      &AnsiString,
                                      TRUE);
+
+        if (!NT_SUCCESS(Status))
+        {
+            if (pValueName != NULL)
+                RtlFreeUnicodeString(&ValueName);
+
+            return RtlNtStatusToDosError(Status);
+        }
         pData = (LPBYTE)Data.Buffer;
         DataSize = cbData * sizeof(WCHAR);
     }
     else
     {
-        RtlInitUnicodeString(&Data,
-                             NULL);
+        Data.Buffer = NULL;
         pData = (LPBYTE)lpData;
         DataSize = cbData;
     }
@@ -4793,19 +4805,12 @@
                                dwType,
                                pData,
                                DataSize);
+
     if (pValueName != NULL)
-    {
-        RtlFreeHeap(ProcessHeap,
-                    0,
-                    ValueName.Buffer);
-    }
+        RtlFreeUnicodeString(&ValueName);

if (Data.Buffer != NULL)
-    {
-        RtlFreeHeap(ProcessHeap,
-                    0,
-                    Data.Buffer);
-    }
+        RtlFreeUnicodeString(&Data);

return ErrorCode;
 }
@@ -4847,13 +4852,16 @@
     }
     pValueName = &ValueName;

-    if (((dwType == REG_SZ) ||
-         (dwType == REG_MULTI_SZ) ||
-         (dwType == REG_EXPAND_SZ)) &&
-        (cbData != 0) && (*(((PWCHAR)lpData) + (cbData / sizeof(WCHAR)) - 1) != L'\0'))
-    {
-        /* NT adds one if the caller forgot the NULL-termination character */
-        cbData += sizeof(WCHAR);
+    if (is_string(dwType) && (cbData != 0))
+    {
+        PWSTR pwsData = (PWSTR)lpData;
+
+        if((pwsData[cbData / sizeof(WCHAR) - 1] != L'\0') &&
+            (pwsData[cbData / sizeof(WCHAR)] == L'\0'))
+        {
+            /* Increment length if last character is not zero and next is zero */
+            cbData += sizeof(WCHAR);
+        }
     }

Status = NtSetValueKey(KeyHandle,