[PATCH 4/5] reg: Use a helper function to allocate and re-allocate memory

[Hugh McMaster]

Signed-off-by: Hugh McMaster <hugh.mcmaster at outlook.com>
---
 programs/reg/reg.c  | 76 ++++++++++++++++++++++++-----------------------------
 programs/reg/reg.h  |  1 +
 programs/reg/reg.rc |  1 +
 3 files changed, 36 insertions(+), 42 deletions(-)

diff --git a/programs/reg/reg.c b/programs/reg/reg.c
index 573c75040a..6d48dece05 100644
--- a/programs/reg/reg.c
+++ b/programs/reg/reg.c
@@ -17,9 +17,12 @@
  */
 
 #include <windows.h>
+#include <errno.h>
+#include <stdlib.h>
+
 #include <wine/unicode.h>
 #include <wine/debug.h>
-#include <errno.h>
+
 #include "reg.h"
 
 #define ARRAY_SIZE(A) (sizeof(A)/sizeof(*A))
@@ -174,6 +177,26 @@ static BOOL ask_confirm(unsigned int msgid, WCHAR *reg_info)
     }
 }
 
+#define CHECK_ENOUGH_MEMORY(p) \
+if (!(p)) \
+{ \
+    output_message(STRING_OUT_OF_MEMORY, __FILE__, __LINE__); \
+    exit(1); \
+}
+
+static void *resize_buffer(void *buf, size_t count)
+{
+    void *new_buf;
+
+    if (buf)
+        new_buf = HeapReAlloc(GetProcessHeap(), 0, buf, count);
+    else
+        new_buf = HeapAlloc(GetProcessHeap(), 0, count);
+
+    CHECK_ENOUGH_MEMORY(new_buf);
+    return new_buf;
+}
+
 static inline BOOL path_rootname_cmp(const WCHAR *input_path, const WCHAR *rootkey_name)
 {
     DWORD length = strlenW(rootkey_name);
@@ -615,13 +638,9 @@ static WCHAR *build_subkey_path(WCHAR *path, DWORD path_len, WCHAR *subkey_name,
     WCHAR *subkey_path;
     WCHAR fmt[] = {'%','s','\\','%','s',0};
 
-    subkey_path = HeapAlloc(GetProcessHeap(), 0, (path_len + subkey_len + 2) * sizeof(WCHAR));
-    if (!subkey_path)
-    {
-        ERR("Failed to allocate memory for subkey_path\n");
-        return NULL;
-    }
+    subkey_path = resize_buffer(NULL, (path_len + subkey_len + 2) * sizeof(WCHAR));
     sprintfW(subkey_path, fmt, path, subkey_name);
+
     return subkey_path;
 }
 
@@ -641,12 +660,7 @@ static int query_value(HKEY key, WCHAR *value_name, WCHAR *path, BOOL recurse)
     WCHAR *subkey_name, *subkey_path;
     HKEY subkey;
 
-    data = HeapAlloc(GetProcessHeap(), 0, max_data_bytes);
-    if (!data)
-    {
-        ERR("Failed to allocate memory for data\n");
-        return 1;
-    }
+    data = resize_buffer(NULL, max_data_bytes);
 
     for (;;)
     {
@@ -655,7 +669,7 @@ static int query_value(HKEY key, WCHAR *value_name, WCHAR *path, BOOL recurse)
         if (rc == ERROR_MORE_DATA)
         {
             max_data_bytes = data_size;
-            data = HeapReAlloc(GetProcessHeap(), 0, data, max_data_bytes);
+            data = resize_buffer(data, max_data_bytes);
         }
         else break;
     }
@@ -685,12 +699,7 @@ static int query_value(HKEY key, WCHAR *value_name, WCHAR *path, BOOL recurse)
         return 0;
     }
 
-    subkey_name = HeapAlloc(GetProcessHeap(), 0, MAX_SUBKEY_LEN * sizeof(WCHAR));
-    if (!subkey_name)
-    {
-        ERR("Failed to allocate memory for subkey_name\n");
-        return 1;
-    }
+    subkey_name = resize_buffer(NULL, MAX_SUBKEY_LEN * sizeof(WCHAR));
 
     path_len = strlenW(path);
 
@@ -733,20 +742,8 @@ static int query_all(HKEY key, WCHAR *path, BOOL recurse)
 
     output_string(fmt, path);
 
-    value_name = HeapAlloc(GetProcessHeap(), 0, max_value_len * sizeof(WCHAR));
-    if (!value_name)
-    {
-        ERR("Failed to allocate memory for value_name\n");
-        return 1;
-    }
-
-    data = HeapAlloc(GetProcessHeap(), 0, max_data_bytes);
-    if (!data)
-    {
-        HeapFree(GetProcessHeap(), 0, value_name);
-        ERR("Failed to allocate memory for data\n");
-        return 1;
-    }
+    value_name = resize_buffer(NULL, max_value_len * sizeof(WCHAR));
+    data = resize_buffer(NULL, max_data_bytes);
 
     i = 0;
     for (;;)
@@ -764,12 +761,12 @@ static int query_all(HKEY key, WCHAR *path, BOOL recurse)
             if (data_size > max_data_bytes)
             {
                 max_data_bytes = data_size;
-                data = HeapReAlloc(GetProcessHeap(), 0, data, max_data_bytes);
+                data = resize_buffer(data, max_data_bytes);
             }
             else
             {
                 max_value_len *= 2;
-                value_name = HeapReAlloc(GetProcessHeap(), 0, value_name, max_value_len * sizeof(WCHAR));
+                value_name = resize_buffer(value_name, max_value_len * sizeof(WCHAR));
             }
         }
         else break;
@@ -781,12 +778,7 @@ static int query_all(HKEY key, WCHAR *path, BOOL recurse)
     if (i || recurse)
         output_string(newlineW);
 
-    subkey_name = HeapAlloc(GetProcessHeap(), 0, MAX_SUBKEY_LEN * sizeof(WCHAR));
-    if (!subkey_name)
-    {
-        ERR("Failed to allocate memory for subkey_name\n");
-        return 1;
-    }
+    subkey_name = resize_buffer(NULL, MAX_SUBKEY_LEN * sizeof(WCHAR));
 
     path_len = strlenW(path);
 
diff --git a/programs/reg/reg.h b/programs/reg/reg.h
index 1142278a19..212e8d2f2e 100644
--- a/programs/reg/reg.h
+++ b/programs/reg/reg.h
@@ -52,3 +52,4 @@
 #define STRING_REG_HELP         129
 #define STRING_FUNC_HELP        130
 #define STRING_VALUE_NOT_SET    131
+#define STRING_OUT_OF_MEMORY    132
diff --git a/programs/reg/reg.rc b/programs/reg/reg.rc
index ae69cb7adf..e76c08e84b 100644
--- a/programs/reg/reg.rc
+++ b/programs/reg/reg.rc
@@ -62,4 +62,5 @@ STRINGTABLE
     STRING_REG_HELP, "Type \"REG /?\" for help.\n"
     STRING_FUNC_HELP, "Type \"REG %1 /?\" for help.\n"
     STRING_VALUE_NOT_SET, "(value not set)"
+    STRING_OUT_OF_MEMORY, "reg: Out of memory! (%1!S!, line %2!u!)\n"
 }
-- 
2.13.2