X-Git-Url: https://git.ozlabs.org/?p=ccan;a=blobdiff_plain;f=ccan%2Ftal%2Ftal.c;h=5ab6cb682cd5349b060b8731e99d9f18c45bb15e;hp=88bcb5ab6135257529eda0e999df024507c0386d;hb=932d65dd6537250e617516749f03a00fea3b34f6;hpb=84b29c907c7608a024f65c4ef6803a1f5bb80ccd

diff --git a/ccan/tal/tal.c b/ccan/tal/tal.c
index 88bcb5ab..5ab6cb68 100644
--- a/ccan/tal/tal.c
+++ b/ccan/tal/tal.c
@@ -3,6 +3,7 @@
 #include <ccan/compiler/compiler.h>
 #include <ccan/hash/hash.h>
 #include <ccan/list/list.h>
+#include <ccan/take/take.h>
 #include <assert.h>
 #include <stdio.h>
 #include <stdarg.h>
@@ -103,12 +104,14 @@ static struct group *next_group(struct group *group)
 	return list_entry(group->list.n.next, struct group, list.n);
 }
 
-static bool atexit_set = false;
+static bool initialized = false;
+
 /* This means valgrind can see leaks. */
-static void unlink_null(void)
+static void tal_cleanup(void)
 {
 	struct group *i, *next;
 
+	/* Unlink null_parent. */
 	for (i = next_group(&null_parent.c.group);
 	     i != &null_parent.c.group;
 	     i = next) {
@@ -116,6 +119,15 @@ static void unlink_null(void)
 		freefn(i);
 	}
 	null_parent.c.group.first_child = NULL;
+
+	/* Cleanup any taken pointers. */
+	take_cleanup();
+}
+
+/* For allocation failures inside ccan/take */
+static void take_alloc_failed(const void *p)
+{
+	tal_free(p);
 }
 
 #ifndef NDEBUG
@@ -332,9 +344,10 @@ static bool add_child(struct tal_hdr *parent, struct tal_hdr *child)
 	if (unlikely(!group->first_child)) {
 		assert(group == &children->group);
 		/* This hits on first child appended to null parent. */
-		if (unlikely(!atexit_set)) {
-			atexit(unlink_null);
-			atexit_set = true;
+		if (unlikely(!initialized)) {
+			atexit(tal_cleanup);
+			take_allocfail(take_alloc_failed);
+			initialized = true;
 		}
 		/* Link group into this child, make it the first one. */
 		group->hdr.next = child->prop;
@@ -480,18 +493,17 @@ static struct tal_hdr *remove_node(struct tal_hdr *t)
 	return NULL;
 }
 
-void tal_free(const tal_t *ctx)
+void *tal_free(const tal_t *ctx)
 {
-        struct tal_hdr *t;
-	int saved_errno = errno;
-
-        if (!ctx)
-                return;
-
-        t = debug_tal(to_tal_hdr(ctx));
-        remove_node(t);
-        del_tree(t);
-	errno = saved_errno;
+        if (ctx) {
+		struct tal_hdr *t;
+		int saved_errno = errno;
+		t = debug_tal(to_tal_hdr(ctx));
+		remove_node(t);
+		del_tree(t);
+		errno = saved_errno;
+	}
+	return NULL;
 }
 
 void *tal_steal_(const tal_t *new_parent, const tal_t *ctx)
@@ -705,18 +717,26 @@ bool tal_resize_(tal_t **ctxp, size_t size)
 
 char *tal_strdup(const tal_t *ctx, const char *p)
 {
-	return tal_dup(ctx, char, p, strlen(p)+1, 0);
+	/* We have to let through NULL for take(). */
+	return tal_dup(ctx, char, p, p ? strlen(p) + 1: 1, 0);
 }
 
 char *tal_strndup(const tal_t *ctx, const char *p, size_t n)
 {
+	size_t len;
 	char *ret;
 
-	if (strlen(p) < n)
-		n = strlen(p);
-	ret = tal_dup(ctx, char, p, n, 1);
+	/* We have to let through NULL for take(). */
+	if (likely(p)) {
+		len = strlen(p);
+		if (len > n)
+			len = n;
+	} else
+		len = n;
+
+	ret = tal_dup(ctx, char, p, len, 1);
 	if (ret)
-		ret[n] = '\0';
+		ret[len] = '\0';
 	return ret;
 }
 
@@ -728,18 +748,18 @@ void *tal_dup_(const tal_t *ctx, const void *p, size_t n, size_t extra,
 	/* Beware overflow! */
 	if (n + extra < n || n + extra + sizeof(struct tal_hdr) < n) {
 		call_error("dup size overflow");
-		if (ctx == TAL_TAKE)
+		if (taken(p))
 			tal_free(p);
 		return NULL;
 	}
 
-	if (ctx == TAL_TAKE) {
+	if (taken(p)) {
 		if (unlikely(!p))
 			return NULL;
-		if (!tal_resize_((void **)&p, n + extra)) {
-			tal_free(p);
-			return NULL;
-		}
+		if (unlikely(!tal_resize_((void **)&p, n + extra)))
+			return tal_free(p);
+		if (unlikely(!tal_steal(ctx, p)))
+			return tal_free(p);
 		return (void *)p;
 	}
 	ret = tal_alloc_(ctx, n + extra, false, label);
@@ -762,15 +782,16 @@ char *tal_asprintf(const tal_t *ctx, const char *fmt, ...)
 
 char *tal_vasprintf(const tal_t *ctx, const char *fmt, va_list ap)
 {
-	size_t max = strlen(fmt) * 2;
+	size_t max;
 	char *buf;
 	int ret;
 
-	if (ctx == TAL_TAKE)
-		buf = tal_arr(tal_parent(fmt), char, max);
-	else
-		buf = tal_arr(ctx, char, max);
+	if (!fmt && taken(fmt))
+		return NULL;
 
+	/* A decent guess to start. */
+	max = strlen(fmt) * 2;
+	buf = tal_arr(ctx, char, max);
 	while (buf) {
 		va_list ap2;
 
@@ -780,12 +801,10 @@ char *tal_vasprintf(const tal_t *ctx, const char *fmt, va_list ap)
 
 		if (ret < max)
 			break;
-		if (!tal_resize(&buf, max *= 2)) {
-			tal_free(buf);
-			buf = NULL;
-		}
+		if (!tal_resize(&buf, max *= 2))
+			buf = tal_free(buf);
 	}
-	if (ctx == TAL_TAKE)
+	if (taken(fmt))
 		tal_free(fmt);
 	return buf;
 }