[PATCH 2/2] reg: Properly handle HeapReAlloc() failure in the query operation

[Hugh McMaster]

Signed-off-by: Hugh McMaster <hugh.mcmaster at outlook.com>
---
 programs/reg/reg.c | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/programs/reg/reg.c b/programs/reg/reg.c
index 59c8eda732..30e7fb1c27 100644
--- a/programs/reg/reg.c
+++ b/programs/reg/reg.c
@@ -654,8 +654,15 @@ static int query_value(HKEY key, WCHAR *value_name, WCHAR *path, BOOL recurse)
         rc = RegQueryValueExW(key, value_name, NULL, &type, data, &data_size);
         if (rc == ERROR_MORE_DATA)
         {
+            BYTE *tmp = HeapReAlloc(GetProcessHeap(), 0, data, data_size);
+            if (!tmp)
+            {
+                ERR("Failed to increase data memory allocation\n");
+                HeapFree(GetProcessHeap(), 0, data);
+                return 1;
+            }
+            data = tmp;
             max_data_bytes = data_size;
-            data = HeapReAlloc(GetProcessHeap(), 0, data, max_data_bytes);
         }
         else break;
     }
@@ -763,8 +770,16 @@ static int query_all(HKEY key, WCHAR *path, BOOL recurse)
         {
             if (data_size > max_data_bytes)
             {
+                BYTE *tmp = HeapReAlloc(GetProcessHeap(), 0, data, data_size);
+                if (!tmp)
+                {
+                    ERR("Failed to increase data memory allocation\n");
+                    HeapFree(GetProcessHeap(), 0, data);
+                    HeapFree(GetProcessHeap(), 0, value_name);
+                    return 1;
+                }
+                data = tmp;
                 max_data_bytes = data_size;
-                data = HeapReAlloc(GetProcessHeap(), 0, data, max_data_bytes);
             }
             else
             {
-- 
2.13.2