talloc: allow replacement allocator

This allows us to both allocators which handle failure themselves, and
allocators which insert failures.

3 files changed, 88 insertions, 5 deletions

diff --git a/ccan/talloc/talloc.c b/ccan/talloc/talloc.c
index 63423a50..0c6bcda9 100644
--- a/ccan/talloc/talloc.c
+++ b/ccan/talloc/talloc.c
@@ -82,6 +82,10 @@
 static void *null_context;
 static pid_t *autofree_context;
 
+static void *(*tc_malloc)(size_t size) = malloc;
+static void (*tc_free)(void *ptr) = free;
+static void *(*tc_realloc)(void *ptr, size_t size) = realloc;
+
 static void *(*tc_external_realloc)(const void *parent, void *ptr, size_t size);
 static void (*tc_lock)(const void *ctx);
 static void (*tc_unlock)(void);
@@ -283,7 +287,7 @@ static inline void *__talloc(const void *context, size_t size)
 		}
 	}
 
-	tc = (struct talloc_chunk *)malloc(TC_HDR_SIZE+size);
+	tc = (struct talloc_chunk *)tc_malloc(TC_HDR_SIZE+size);
 alloc_done:
 	return init_talloc(parent, tc, size, external);
 }
@@ -571,7 +575,7 @@ static inline int _talloc_free(const void *ptr)
 	if (unlikely(tc->flags & TALLOC_FLAG_EXT_ALLOC))
 		tc_external_realloc(oldparent, tc, 0);
 	else
-		free(tc);
+		tc_free(tc);
 
 	return 0;
 }
@@ -923,13 +927,13 @@ void *_talloc_realloc(const void *context, void *ptr, size_t size, const char *n
 		tc->flags |= TALLOC_FLAG_FREE;
 
 #if ALWAYS_REALLOC
-		new_ptr = malloc(size + TC_HDR_SIZE);
+		new_ptr = tc_malloc(size + TC_HDR_SIZE);
 		if (new_ptr) {
 			memcpy(new_ptr, tc, tc->size + TC_HDR_SIZE);
-			free(tc);
+			tc_free(tc);
 		}
 #else
-		new_ptr = realloc(tc, size + TC_HDR_SIZE);
+		new_ptr = tc_realloc(tc, size + TC_HDR_SIZE);
 #endif
 	}
 
@@ -1613,6 +1617,15 @@ int talloc_is_parent(const void *context, const void *ptr)
 	return ret;
 }
 
+void talloc_set_allocator(void *(*malloc)(size_t size),
+			  void (*free)(void *ptr),
+			  void *(*realloc)(void *ptr, size_t size))
+{
+	tc_malloc = malloc;
+	tc_free = free;
+	tc_realloc = realloc;
+}
+
 void *talloc_add_external(const void *ctx,
 			  void *(*realloc)(const void *, void *, size_t),
 			  void (*lock)(const void *p),
diff --git a/ccan/talloc/talloc.h b/ccan/talloc/talloc.h
index d27c6892..5263b9fe 100644
--- a/ccan/talloc/talloc.h
+++ b/ccan/talloc/talloc.h
@@ -958,6 +958,19 @@ size_t talloc_get_size(const void *ctx);
 void *talloc_find_parent_byname(const void *ctx, const char *name);
 
 /**
+ * talloc_set_allocator - set the allocations function(s) for talloc.
+ * @malloc: the malloc function
+ * @free: the free function
+ * @realloc: the realloc function
+ *
+ * Instead of using the standard malloc, free and realloc, talloc will use
+ * these replacements.  @realloc will never be called with size 0 or ptr NULL.
+ */
+void talloc_set_allocator(void *(*malloc)(size_t size),
+			  void (*free)(void *ptr),
+			  void *(*realloc)(void *ptr, size_t size));
+
+/**
  * talloc_add_external - create an externally allocated node
  * @ctx: the parent
  * @realloc: the realloc() equivalent
diff --git a/ccan/talloc/test/run-set_allocator.c b/ccan/talloc/test/run-set_allocator.c
new file mode 100644
index 00000000..e485c623
--- /dev/null
+++ b/ccan/talloc/test/run-set_allocator.c
@@ -0,0 +1,57 @@
+#include <ccan/failtest/failtest_override.h>
+#include <ccan/talloc/talloc.c>
+#include <stdbool.h>
+#include <ccan/tap/tap.h>
+#include <ccan/failtest/failtest.h>
+
+static unsigned my_malloc_count, my_free_count, my_realloc_count;
+
+static void *my_malloc(size_t size)
+{
+	my_malloc_count++;
+	return malloc(size);
+}
+
+static void my_free(void *ptr)
+{
+	my_free_count++;
+	free(ptr);
+}
+
+static void *my_realloc(void *ptr, size_t size)
+{
+	my_realloc_count++;
+	ok1(ptr);
+	ok1(size);
+	return realloc(ptr, size);
+}
+
+int main(int argc, char *argv[])
+{
+	int *p1, *p2;
+
+	plan_tests(14);
+	failtest_init(argc, argv);
+	talloc_set_allocator(my_malloc, my_free, my_realloc);
+	p1 = talloc_array(NULL, int, 10);
+	ok1(my_malloc_count == 1);
+	ok1(my_free_count == 0);
+	ok1(my_realloc_count == 0);
+
+	p1 = talloc_realloc(NULL, p1, int, 10000);
+	ok1(my_malloc_count == 1);
+	ok1(my_free_count == 0);
+	ok1(my_realloc_count == 1);
+
+	p2 = talloc(p1, int);
+	ok1(my_malloc_count == 2);
+	ok1(my_free_count == 0);
+	ok1(my_realloc_count == 1);
+
+	talloc_free(p1);
+	ok1(my_malloc_count == 2);
+	ok1(my_free_count == 2);
+	ok1(my_realloc_count == 1);
+
+	failtest_exit(exit_status());
+}