opt: add allocator setting.

Good for tal usage.

Signed-off-by: Rusty Russell <rusty@rustcorp.com.au>

7 files changed, 402 insertions, 13 deletions

diff --git a/ccan/opt/helpers.c b/ccan/opt/helpers.c
index e7ab273c..25929bd5 100644
--- a/ccan/opt/helpers.c
+++ b/ccan/opt/helpers.c
@@ -134,7 +134,7 @@ char *opt_usage_and_exit(const char *extra)
 	char *usage = opt_usage(opt_argv0, extra);
 	printf("%s", usage);
 	/* Don't have valgrind complain! */
-	free(usage);
+	opt_alloc.free(usage);
 	opt_free_table();
 	exit(0);
 }
diff --git a/ccan/opt/opt.c b/ccan/opt/opt.c
index 94eb0d5a..85de935a 100644
--- a/ccan/opt/opt.c
+++ b/ccan/opt/opt.c
@@ -12,6 +12,9 @@
 struct opt_table *opt_table;
 unsigned int opt_count, opt_num_short, opt_num_short_arg, opt_num_long;
 const char *opt_argv0;
+struct opt_alloc opt_alloc = {
+	malloc, realloc, free
+};
 
 /* Returns string after first '-'. */
 static const char *first_name(const char *names, unsigned *len)
@@ -150,7 +153,8 @@ static void check_opt(const struct opt_table *entry)
 
 static void add_opt(const struct opt_table *entry)
 {
-	opt_table = realloc(opt_table, sizeof(opt_table[0]) * (opt_count+1));
+	opt_table = opt_alloc.realloc(opt_table,
+				      sizeof(opt_table[0]) * (opt_count+1));
 	opt_table[opt_count++] = *entry;
 }
 
@@ -214,7 +218,7 @@ bool opt_early_parse(int argc, char *argv[],
 {
 	int ret;
 	unsigned off = 0;
-	char **tmpargv = malloc(sizeof(argv[0]) * (argc + 1));
+	char **tmpargv = opt_alloc.alloc(sizeof(argv[0]) * (argc + 1));
 
 	/* We could avoid a copy and skip instead, but this is simple. */
 	memcpy(tmpargv, argv, sizeof(argv[0]) * (argc + 1));
@@ -224,7 +228,7 @@ bool opt_early_parse(int argc, char *argv[],
 
 	while ((ret = parse_one(&argc, tmpargv, OPT_EARLY, &off, errlog)) == 1);
 
-	free(tmpargv);
+	opt_alloc.free(tmpargv);
 
 	/* parse_one returns 0 on finish, -1 on error */
 	return (ret == 0);
@@ -232,7 +236,7 @@ bool opt_early_parse(int argc, char *argv[],
 
 void opt_free_table(void)
 {
-	free(opt_table);
+	opt_alloc.free(opt_table);
 	opt_table = NULL;
 	opt_count = opt_num_short = opt_num_short_arg = opt_num_long = 0;
 }
@@ -260,7 +264,16 @@ void opt_log_stderr_exit(const char *fmt, ...)
 
 char *opt_invalid_argument(const char *arg)
 {
-	char *str = malloc(sizeof("Invalid argument '%s'") + strlen(arg));
+	char *str = opt_alloc.alloc(sizeof("Invalid argument '%s'") + strlen(arg));
 	sprintf(str, "Invalid argument '%s'", arg);
 	return str;
 }
+
+void opt_set_alloc(void *(*allocfn)(size_t size),
+		   void *(*reallocfn)(void *ptr, size_t size),
+		   void (*freefn)(void *ptr))
+{
+	opt_alloc.alloc = allocfn;
+	opt_alloc.realloc = reallocfn;
+	opt_alloc.free = freefn;
+}
diff --git a/ccan/opt/opt.h b/ccan/opt/opt.h
index 8bbc8f51..6c3cbee1 100644
--- a/ccan/opt/opt.h
+++ b/ccan/opt/opt.h
@@ -21,7 +21,7 @@ struct opt_table;
  *
  * If the @cb returns non-NULL, opt_parse() will stop parsing, use the
  * returned string to form an error message for errlog(), free() the
- * string and return false.
+ * string (or see opt_set_alloc) and return false.
  *
  * Any number of equivalent short or long options can be listed in @names,
  * separated by '|'.  Short options are a single hyphen followed by a single
@@ -60,7 +60,7 @@ struct opt_table;
  *
  * If the @cb returns non-NULL, opt_parse() will stop parsing, use the
  * returned string to form an error message for errlog(), free() the
- * string and return false.
+ * string (or see opt_set_alloc) and return false.
  *
  * See Also:
  *	OPT_WITHOUT_ARG()
@@ -159,7 +159,7 @@ void opt_register_table(const struct opt_table *table, const char *desc);
  *
  * If the @cb returns non-NULL, opt_parse() will stop parsing, use the
  * returned string to form an error message for errlog(), free() the
- * string and return false.
+ * string (or see opt_set_alloc) and return false.
  */
 #define opt_register_noarg(names, cb, arg, desc)			\
 	_opt_register((names), OPT_CB_NOARG((cb), 0, (arg)), (arg), (desc))
@@ -182,7 +182,7 @@ void opt_register_table(const struct opt_table *table, const char *desc);
  *
  * If the @cb returns non-NULL, opt_parse() will stop parsing, use the
  * returned string to form an error message for errlog(), free() the
- * string and return false.
+ * string (or see opt_set_alloc) and return false.
  *
  * Example:
  * static char *explode(const char *optarg, void *unused)
@@ -297,6 +297,19 @@ bool opt_early_parse(int argc, char *argv[],
 void opt_free_table(void);
 
 /**
+ * opt_set_alloc - set alloc/realloc/free function for opt to use.
+ * @allocfn: allocator function
+ * @reallocfn: reallocator function, ptr may be NULL, size never 0.
+ * @freefn: free function
+ *
+ * By default opt uses malloc/realloc/free, and simply crashes if they fail.
+ * You can set your own variants here.
+ */
+void opt_set_alloc(void *(*allocfn)(size_t size),
+		   void *(*reallocfn)(void *ptr, size_t size),
+		   void (*freefn)(void *ptr));
+
+/**
  * opt_log_stderr - print message to stderr.
  * @fmt: printf-style format.
  *
diff --git a/ccan/opt/parse.c b/ccan/opt/parse.c
index 5b2cec8b..14c7afbb 100644
--- a/ccan/opt/parse.c
+++ b/ccan/opt/parse.c
@@ -116,7 +116,7 @@ int parse_one(int *argc, char *argv[], enum opt_type is_early, unsigned *offset,
 
 	if (problem) {
 		parse_err(errlog, argv[0], o, len, problem);
-		free(problem);
+		opt_alloc.free(problem);
 		return -1;
 	}
 
diff --git a/ccan/opt/private.h b/ccan/opt/private.h
index 12c56931..8d7815c8 100644
--- a/ccan/opt/private.h
+++ b/ccan/opt/private.h
@@ -14,6 +14,13 @@ const char *next_sopt(const char *names, unsigned *i);
 const char *first_lopt(unsigned *i, unsigned *len);
 const char *next_lopt(const char *p, unsigned *i, unsigned *len);
 
+struct opt_alloc {
+	void *(*alloc)(size_t size);
+	void *(*realloc)(void *ptr, size_t size);
+	void (*free)(void *ptr);
+};
+extern struct opt_alloc opt_alloc;
+
 int parse_one(int *argc, char *argv[], enum opt_type is_early, unsigned *offset,
 	      void (*errlog)(const char *fmt, ...));
 
diff --git a/ccan/opt/test/run-set_alloc.c b/ccan/opt/test/run-set_alloc.c
new file mode 100644
index 00000000..b30a77d8
--- /dev/null
+++ b/ccan/opt/test/run-set_alloc.c
@@ -0,0 +1,356 @@
+#include <ccan/tap/tap.h>
+#include <stdlib.h>
+
+/* Make sure we override these! */
+static void *no_malloc(size_t size)
+{
+	abort();
+}
+static void *no_realloc(void *p, size_t size)
+{
+	abort();
+}
+static void no_free(void *p)
+{
+	abort();
+}
+#define malloc no_malloc
+#define realloc no_realloc
+#define free no_free
+
+#include <ccan/opt/opt.c>
+#include <ccan/opt/usage.c>
+#include <ccan/opt/helpers.c>
+#include <ccan/opt/parse.c>
+#include "utils.h"
+
+#undef malloc
+#undef realloc
+#undef free
+
+static unsigned int alloc_count, realloc_count, free_count;
+static void *ptrs[100];
+
+static void **find_ptr(void *p)
+{
+	unsigned int i;
+
+	for (i = 0; i < 100; i++)
+		if (ptrs[i] == p)
+			return ptrs + i;
+	return NULL;
+}
+
+static void *allocfn(size_t size)
+{
+	alloc_count++;
+	return *find_ptr(NULL) = malloc(size);
+}
+
+static void *reallocfn(void *ptr, size_t size)
+{
+	realloc_count++;
+	if (!ptr)
+		alloc_count++;
+
+	return *find_ptr(ptr) = realloc(ptr, size);
+}
+
+static void freefn(void *ptr)
+{
+	free_count++;
+	free(ptr);
+	*find_ptr(ptr) = NULL;
+}
+
+int main(int argc, char *argv[])
+{
+	const char *myname = argv[0];
+
+	plan_tests(220);
+
+	opt_set_alloc(allocfn, reallocfn, freefn);
+
+	/* Simple short arg.*/
+	opt_register_noarg("-a", test_noarg, NULL, "All");
+	ok1(parse_args(&argc, &argv, "-a", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 1);
+
+	/* Simple long arg. */
+	opt_register_noarg("--aaa", test_noarg, NULL, "AAAAll");
+	ok1(parse_args(&argc, &argv, "--aaa", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 2);
+
+	/* Both long and short args. */
+	opt_register_noarg("--aaa|-a", test_noarg, NULL, "AAAAAAll");
+	ok1(parse_args(&argc, &argv, "--aaa", "-a", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 4);
+
+	/* Extra arguments preserved. */
+	ok1(parse_args(&argc, &argv, "--aaa", "-a", "extra", "args", NULL));
+	ok1(argc == 3);
+	ok1(argv[0] == myname);
+	ok1(strcmp(argv[1], "extra") == 0);
+	ok1(strcmp(argv[2], "args") == 0);
+	ok1(test_cb_called == 6);
+
+	/* Malformed versions. */
+	ok1(!parse_args(&argc, &argv, "--aaa=arg", NULL));
+	ok1(strstr(err_output, ": --aaa: doesn't allow an argument"));
+	ok1(!parse_args(&argc, &argv, "--aa", NULL));
+	ok1(strstr(err_output, ": --aa: unrecognized option"));
+	ok1(!parse_args(&argc, &argv, "--aaargh", NULL));
+	ok1(strstr(err_output, ": --aaargh: unrecognized option"));
+
+	/* Argument variants. */
+	reset_options();
+	test_cb_called = 0;
+	opt_register_arg("-a|--aaa", test_arg, NULL, "aaa", "AAAAAAll");
+	ok1(parse_args(&argc, &argv, "--aaa", "aaa", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(test_cb_called == 1);
+
+	ok1(parse_args(&argc, &argv, "--aaa=aaa", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(test_cb_called == 2);
+
+	ok1(parse_args(&argc, &argv, "-a", "aaa", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(test_cb_called == 3);
+
+	/* Malformed versions. */
+	ok1(!parse_args(&argc, &argv, "-a", NULL));
+	ok1(strstr(err_output, ": -a: requires an argument"));
+	ok1(!parse_args(&argc, &argv, "--aaa", NULL));
+	ok1(strstr(err_output, ": --aaa: requires an argument"));
+	ok1(!parse_args(&argc, &argv, "--aa", NULL));
+	ok1(strstr(err_output, ": --aa: unrecognized option"));
+	ok1(!parse_args(&argc, &argv, "--aaargh", NULL));
+	ok1(strstr(err_output, ": --aaargh: unrecognized option"));
+
+	/* Now, tables. */
+	/* Short table: */
+	reset_options();
+	test_cb_called = 0;
+	opt_register_table(short_table, NULL);
+	ok1(parse_args(&argc, &argv, "-a", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 1);
+	/* This one needs an arg. */
+	ok1(parse_args(&argc, &argv, "-b", NULL) == false);
+	ok1(test_cb_called == 1);
+	ok1(parse_args(&argc, &argv, "-b", "b", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 2);
+
+	/* Long table: */
+	reset_options();
+	test_cb_called = 0;
+	opt_register_table(long_table, NULL);
+	ok1(parse_args(&argc, &argv, "--ddd", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 1);
+	/* This one needs an arg. */
+	ok1(parse_args(&argc, &argv, "--eee", NULL) == false);
+	ok1(test_cb_called == 1);
+	ok1(parse_args(&argc, &argv, "--eee", "eee", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 2);
+
+	/* Short and long, both. */
+	reset_options();
+	test_cb_called = 0;
+	opt_register_table(long_and_short_table, NULL);
+	ok1(parse_args(&argc, &argv, "-g", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 1);
+	ok1(parse_args(&argc, &argv, "--ggg", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 2);
+	/* This one needs an arg. */
+	ok1(parse_args(&argc, &argv, "-h", NULL) == false);
+	ok1(test_cb_called == 2);
+	ok1(parse_args(&argc, &argv, "-h", "hhh", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 3);
+	ok1(parse_args(&argc, &argv, "--hhh", NULL) == false);
+	ok1(test_cb_called == 3);
+	ok1(parse_args(&argc, &argv, "--hhh", "hhh", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 4);
+
+	/* Those will all work as tables. */
+	test_cb_called = 0;
+	reset_options();
+	opt_register_table(subtables, NULL);
+	ok1(parse_args(&argc, &argv, "-a", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 1);
+	/* This one needs an arg. */
+	ok1(parse_args(&argc, &argv, "-b", NULL) == false);
+	ok1(test_cb_called == 1);
+	ok1(parse_args(&argc, &argv, "-b", "b", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 2);
+
+	ok1(parse_args(&argc, &argv, "--ddd", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 3);
+	/* This one needs an arg. */
+	ok1(parse_args(&argc, &argv, "--eee", NULL) == false);
+	ok1(test_cb_called == 3);
+	ok1(parse_args(&argc, &argv, "--eee", "eee", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 4);
+
+	/* Short and long, both. */
+	ok1(parse_args(&argc, &argv, "-g", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 5);
+	ok1(parse_args(&argc, &argv, "--ggg", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 6);
+	/* This one needs an arg. */
+	ok1(parse_args(&argc, &argv, "-h", NULL) == false);
+	ok1(test_cb_called == 6);
+	ok1(parse_args(&argc, &argv, "-h", "hhh", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 7);
+	ok1(parse_args(&argc, &argv, "--hhh", NULL) == false);
+	ok1(test_cb_called == 7);
+	ok1(parse_args(&argc, &argv, "--hhh", "hhh", NULL));
+	ok1(argc == 1);
+	ok1(argv[0] == myname);
+	ok1(argv[1] == NULL);
+	ok1(test_cb_called == 8);
+
+	/* Now the tricky one: -? must not be confused with an unknown option */
+	test_cb_called = 0;
+	reset_options();
+
+	/* glibc's getopt does not handle ? with arguments. */
+	opt_register_noarg("-?", test_noarg, NULL, "Help");
+	ok1(parse_args(&argc, &argv, "-?", NULL));
+	ok1(test_cb_called == 1);
+	ok1(parse_args(&argc, &argv, "-a", NULL) == false);
+	ok1(test_cb_called == 1);
+	ok1(strstr(err_output, ": -a: unrecognized option"));
+	ok1(parse_args(&argc, &argv, "--aaaa", NULL) == false);
+	ok1(test_cb_called == 1);
+	ok1(strstr(err_output, ": --aaaa: unrecognized option"));
+
+	test_cb_called = 0;
+	reset_options();
+
+	/* Corner cases involving short arg parsing weirdness. */
+	opt_register_noarg("-a|--aaa", test_noarg, NULL, "a");
+	opt_register_arg("-b|--bbb", test_arg, NULL, "bbb", "b");
+	opt_register_arg("-c|--ccc", test_arg, NULL, "aaa", "c");
+	/* -aa == -a -a */
+	ok1(parse_args(&argc, &argv, "-aa", NULL));
+	ok1(test_cb_called == 2);
+	ok1(parse_args(&argc, &argv, "-aab", NULL) == false);
+	ok1(test_cb_called == 4);
+	ok1(strstr(err_output, ": -b: requires an argument"));
+	ok1(parse_args(&argc, &argv, "-bbbb", NULL));
+	ok1(test_cb_called == 5);
+	ok1(parse_args(&argc, &argv, "-aabbbb", NULL));
+	ok1(test_cb_called == 8);
+	ok1(parse_args(&argc, &argv, "-aabbbb", "-b", "bbb", NULL));
+	ok1(test_cb_called == 12);
+	ok1(parse_args(&argc, &argv, "-aabbbb", "--bbb", "bbb", NULL));
+	ok1(test_cb_called == 16);
+	ok1(parse_args(&argc, &argv, "-aabbbb", "--bbb=bbb", NULL));
+	ok1(test_cb_called == 20);
+	ok1(parse_args(&argc, &argv, "-aacaaa", NULL));
+	ok1(test_cb_called == 23);
+	ok1(parse_args(&argc, &argv, "-aacaaa", "-a", NULL));
+	ok1(test_cb_called == 27);
+	ok1(parse_args(&argc, &argv, "-aacaaa", "--bbb", "bbb", "-aacaaa",
+		       NULL));
+	ok1(test_cb_called == 34);
+
+	test_cb_called = 0;
+	reset_options();
+
+	/* -- and POSIXLY_CORRECT */
+	opt_register_noarg("-a|--aaa", test_noarg, NULL, "a");
+	ok1(parse_args(&argc, &argv, "-a", "--", "-a", NULL));
+	ok1(test_cb_called == 1);
+	ok1(argc == 2);
+	ok1(strcmp(argv[1], "-a") == 0);
+	ok1(!argv[2]);
+
+	unsetenv("POSIXLY_CORRECT");
+	ok1(parse_args(&argc, &argv, "-a", "somearg", "-a", "--", "-a", NULL));
+	ok1(test_cb_called == 3);
+	ok1(argc == 3);
+	ok1(strcmp(argv[1], "somearg") == 0);
+	ok1(strcmp(argv[2], "-a") == 0);
+	ok1(!argv[3]);
+
+	setenv("POSIXLY_CORRECT", "1", 1);
+	ok1(parse_args(&argc, &argv, "-a", "somearg", "-a", "--", "-a", NULL));
+	ok1(test_cb_called == 4);
+	ok1(argc == 5);
+	ok1(strcmp(argv[1], "somearg") == 0);
+	ok1(strcmp(argv[2], "-a") == 0);
+	ok1(strcmp(argv[3], "--") == 0);
+	ok1(strcmp(argv[4], "-a") == 0);
+	ok1(!argv[5]);
+
+	/* We should have tested each one at least once! */
+	ok1(realloc_count);
+	ok1(alloc_count);
+	ok1(free_count);
+
+	ok1(free_count < alloc_count);
+	reset_options();
+	ok1(free_count == alloc_count);
+
+	/* parse_args allocates argv */
+	free(argv);
+	return exit_status();
+}
diff --git a/ccan/opt/usage.c b/ccan/opt/usage.c
index 7c1971c2..1142fb85 100644
--- a/ccan/opt/usage.c
+++ b/ccan/opt/usage.c
@@ -60,7 +60,7 @@ static char *add_str_len(char *base, size_t *len, size_t *max,
 			 const char *str, size_t slen)
 {
 	if (slen >= *max - *len)
-		base = realloc(base, *max = (*max * 2 + slen + 1));
+		base = opt_alloc.realloc(base, *max = (*max * 2 + slen + 1));
 	memcpy(base + *len, str, slen);
 	*len += slen;
 	return base;
@@ -74,7 +74,7 @@ static char *add_str(char *base, size_t *len, size_t *max, const char *str)
 static char *add_indent(char *base, size_t *len, size_t *max, size_t indent)
 {
 	if (indent >= *max - *len)
-		base = realloc(base, *max = (*max * 2 + indent + 1));
+		base = opt_alloc.realloc(base, *max = (*max * 2 + indent + 1));
 	memset(base + *len, ' ', indent);
 	*len += indent;
 	return base;