[ADVAPI32]

[reactos.git] / reactos / dll / win32 / advapi32 / reg / reg.c

diff --git a/reactos/dll/win32/advapi32/reg/reg.c b/reactos/dll/win32/advapi32/reg/reg.c
index 26036dc..cfc8ba1 100644 (file)
--- a/reactos/dll/win32/advapi32/reg/reg.c
+++ b/reactos/dll/win32/advapi32/reg/reg.c
@@ -2793,14 +2793,14 @@ RegEnumValueA(HKEY hKey,
 
     status = NtEnumerateValueKey( KeyHandle, index, KeyValueFullInformation,
                                   buffer, total_size, &total_size );
-    if (status && status != STATUS_BUFFER_OVERFLOW) goto done;
+    if (status && (status != STATUS_BUFFER_OVERFLOW) && (status != STATUS_BUFFER_TOO_SMALL)) goto done;
 
     /* we need to fetch the contents for a string type even if not requested,
      * because we need to compute the length of the ASCII string. */
     if (value || data || is_string(info->Type))
     {
         /* retry with a dynamically allocated buffer */
-        while (status == STATUS_BUFFER_OVERFLOW)
+        while ((status == STATUS_BUFFER_OVERFLOW) || (status == STATUS_BUFFER_TOO_SMALL))
         {
             if (buf_ptr != buffer) HeapFree( GetProcessHeap(), 0, buf_ptr );
             if (!(buf_ptr = HeapAlloc( GetProcessHeap(), 0, total_size )))
@@ -2819,14 +2819,14 @@ RegEnumValueA(HKEY hKey,
         {
             ULONG len;
             RtlUnicodeToMultiByteSize( &len, (WCHAR *)(buf_ptr + info->DataOffset),
-                                       total_size - info->DataOffset );
+                                       info->DataLength );
             if (data && len)
             {
                 if (len > *count) status = STATUS_BUFFER_OVERFLOW;
                 else
                 {
                     RtlUnicodeToMultiByteN( (PCHAR)data, len, NULL, (WCHAR *)(buf_ptr + info->DataOffset),
-                                            total_size - info->DataOffset );
+                                            info->DataLength );
                     /* 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 (len < *count && data[len-1]) data[len] = 0;
@@ -2836,8 +2836,8 @@ RegEnumValueA(HKEY hKey,
         }
         else if (data)
         {
-            if (total_size - info->DataOffset > *count) status = STATUS_BUFFER_OVERFLOW;
-            else memcpy( data, buf_ptr + info->DataOffset, total_size - info->DataOffset );
+            if (info->DataLength > *count) status = STATUS_BUFFER_OVERFLOW;
+            else memcpy( data, buf_ptr + info->DataOffset, info->DataLength );
         }
 
         if (value && !status)
@@ -2928,12 +2928,12 @@ RegEnumValueW(HKEY hKey,
 
     status = NtEnumerateValueKey( KeyHandle, index, KeyValueFullInformation,
                                   buffer, total_size, &total_size );
-    if (status && status != STATUS_BUFFER_OVERFLOW) goto done;
+    if (status && (status != STATUS_BUFFER_OVERFLOW) && (status != STATUS_BUFFER_TOO_SMALL)) goto done;
 
     if (value || data)
     {
         /* retry with a dynamically allocated buffer */
-        while (status == STATUS_BUFFER_OVERFLOW)
+        while ((status == STATUS_BUFFER_OVERFLOW) || (status == STATUS_BUFFER_TOO_SMALL))
         {
             if (buf_ptr != buffer) HeapFree( GetProcessHeap(), 0, buf_ptr );
             if (!(buf_ptr = HeapAlloc( GetProcessHeap(), 0, total_size )))
@@ -2962,17 +2962,17 @@ RegEnumValueW(HKEY hKey,
 
         if (data)
         {
-            if (total_size - info->DataOffset > *count)
+            if (info->DataLength > *count)
             {
                 status = STATUS_BUFFER_OVERFLOW;
                 goto overflow;
             }
-            memcpy( data, buf_ptr + info->DataOffset, total_size - info->DataOffset );
-            if (total_size - info->DataOffset <= *count-sizeof(WCHAR) && is_string(info->Type))
+            memcpy( data, buf_ptr + info->DataOffset, info->DataLength );
+            if (is_string(info->Type) && info->DataLength <= *count - sizeof(WCHAR))
             {
                 /* if the type is REG_SZ and data is not 0-terminated
                  * and there is enough space in the buffer NT appends a \0 */
-                WCHAR *ptr = (WCHAR *)(data + total_size - info->DataOffset);
+                WCHAR *ptr = (WCHAR *)(data + info->DataLength);
                 if (ptr > (WCHAR *)data && ptr[-1]) *ptr = 0;
             }
         }
@@ -3988,7 +3988,7 @@ RegQueryValueExA(HKEY hKey,
                  LPDWORD lpcbData)
 {
     UNICODE_STRING ValueName;
-    UNICODE_STRING ValueData;
+    LPWSTR lpValueBuffer;
     LONG ErrorCode;
     DWORD Length;
     DWORD Type;
@@ -4003,39 +4003,42 @@ RegQueryValueExA(HKEY hKey,
         return ERROR_INVALID_PARAMETER;
     }
 
+    Length = (lpcbData == NULL || lpData == NULL) ? 0 : *lpcbData * sizeof(WCHAR);
+
     if (lpData)
     {
-        ValueData.Length = 0;
-        ValueData.MaximumLength = (*lpcbData + 1) * sizeof(WCHAR);
-        ValueData.Buffer = RtlAllocateHeap(ProcessHeap,
-                                           0,
-                                           ValueData.MaximumLength);
-        if (!ValueData.Buffer)
+        lpValueBuffer = RtlAllocateHeap(ProcessHeap,
+                                        0,
+                                        Length + sizeof(WCHAR));
+        if (!lpValueBuffer)
         {
             return ERROR_OUTOFMEMORY;
         }
     }
     else
     {
-        ValueData.Buffer = NULL;
-        ValueData.Length = 0;
-        ValueData.MaximumLength = 0;
+        lpValueBuffer = NULL;
 
         if (lpcbData)
             *lpcbData = 0;
     }
 
-    RtlCreateUnicodeStringFromAsciiz(&ValueName,
-                                     (LPSTR)lpValueName);
+    if(!RtlCreateUnicodeStringFromAsciiz(&ValueName,
+                                         (LPSTR)lpValueName))
+    {
+        ERR("RtlCreateUnicodeStringFromAsciiz failed!\n");
+        ErrorCode = ERROR_OUTOFMEMORY;
+        goto cleanup;
+    }
 
-    Length = (lpcbData == NULL) ? 0 : *lpcbData * sizeof(WCHAR);
     ErrorCode = RegQueryValueExW(hKey,
                                  ValueName.Buffer,
                                  lpReserved,
                                  &Type,
-                                 (lpData == NULL) ? NULL : (LPBYTE)ValueData.Buffer,
+                                 (LPBYTE)lpValueBuffer,
                                  &Length);
     TRACE("ErrorCode %lu\n", ErrorCode);
+
     RtlFreeUnicodeString(&ValueName);
 
     if (ErrorCode == ERROR_SUCCESS ||
@@ -4044,9 +4047,9 @@ RegQueryValueExA(HKEY hKey,
 
         if (is_string(Type))
         {
-            if (ErrorCode == ERROR_SUCCESS && ValueData.Buffer != NULL)
+            if (ErrorCode == ERROR_SUCCESS && lpValueBuffer != NULL)
             {
-                Status = RtlUnicodeToMultiByteN((PCHAR)lpData, *lpcbData, &Index, (PWCHAR)ValueData.Buffer, Length);
+                Status = RtlUnicodeToMultiByteN((PCHAR)lpData, *lpcbData, &Index, (PWCHAR)lpValueBuffer, Length);
                 if (NT_SUCCESS(Status))
                 {
                     PCHAR szData = (PCHAR)lpData;
@@ -4063,7 +4066,7 @@ RegQueryValueExA(HKEY hKey,
 
             Length = Length / sizeof(WCHAR);
         }
-        else if (ErrorCode == ERROR_SUCCESS && ValueData.Buffer != NULL)
+        else if (ErrorCode == ERROR_SUCCESS && lpValueBuffer != NULL)
         {
             if (*lpcbData < Length)
             {
@@ -4071,7 +4074,7 @@ RegQueryValueExA(HKEY hKey,
             }
             else
             {
-                RtlMoveMemory(lpData, ValueData.Buffer, Length);
+                RtlMoveMemory(lpData, lpValueBuffer, Length);
             }
         }
 
@@ -4086,9 +4089,10 @@ RegQueryValueExA(HKEY hKey,
         *lpType = Type;
     }
 
-    if (ValueData.Buffer != NULL)
+cleanup:
+    if (lpValueBuffer != NULL)
     {
-        RtlFreeHeap(ProcessHeap, 0, ValueData.Buffer);
+        RtlFreeHeap(ProcessHeap, 0, lpValueBuffer);
     }
 
     return ErrorCode;