1 files changed, 40 insertions, 19 deletions
diff --git a/alloc.c b/alloc.c
index 55c3020..2b19c48 100644
--- a/alloc.c
+++ b/alloc.c
@@ -31,35 +31,56 @@
 #include "fmacros.h"
 #include "alloc.h"
 #include <string.h>
+#include <stdlib.h>
 
-void *hi_malloc(size_t size) {
-    void *ptr = malloc(size);
-    if (ptr == NULL)
-        HIREDIS_OOM_HANDLER;
+hiredisAllocFuncs hiredisAllocFns = {
+    .malloc = malloc,
+    .calloc = calloc,
+    .realloc = realloc,
+    .strdup = strdup,
+    .free = free,
+};
+
+/* Override hiredis' allocators with ones supplied by the user */
+hiredisAllocFuncs hiredisSetAllocators(hiredisAllocFuncs *override) {
+    hiredisAllocFuncs orig = hiredisAllocFns;
 
-    return ptr;
+    hiredisAllocFns = *override;
+
+    return orig;
 }
 
-void *hi_calloc(size_t nmemb, size_t size) {
-    void *ptr = calloc(nmemb, size);
-    if (ptr == NULL)
-        HIREDIS_OOM_HANDLER;
+/* Reset allocators to use libc defaults */
+void hiredisResetAllocators(void) {
+    hiredisAllocFns = (hiredisAllocFuncs) {
+        .malloc = malloc,
+        .calloc = calloc,
+        .realloc = realloc,
+        .strdup = strdup,
+        .free = free,
+    };
+}
+
+#ifdef _WIN32
 
-    return ptr;
+void *hi_malloc(size_t size) {
+    return hiredisAllocFns.malloc(size);
 }
 
-void *hi_realloc(void *ptr, size_t size) {
-    void *newptr = realloc(ptr, size);
-    if (newptr == NULL)
-        HIREDIS_OOM_HANDLER;
+void *hi_calloc(size_t nmemb, size_t size) {
+    return hiredisAllocFns.calloc(nmemb, size);
+}
 
-    return newptr;
+void *hi_realloc(void *ptr, size_t size) {
+    return hiredisAllocFns.realloc(ptr, size);
 }
 
 char *hi_strdup(const char *str) {
-    char *newstr = strdup(str);
-    if (newstr == NULL)
-        HIREDIS_OOM_HANDLER;
+    return hiredisAllocFns.strdup(str);
+}
 
-    return newstr;
+void hi_free(void *ptr) {
+    hiredisAllocFns.free(ptr);
 }
+
+#endif