utils: use a cleaner pattern to capture wally allocations.

We force use of tal_wally_start/tal_wally_end around every wally
allocation, and with "end" make the caller choose where to reparent
everything.

This is particularly powerful where we allocate a tx or a psbt: we
want that tx or psbt to be the parent of the other allocations, so
this way we can reparent the tx or psbt, then reparent everything
else onto it.

Implementing psbt_finalize (which uses a behavior flag antipattern)
was tricky, so I ended up splitting that into 'psbt_finalize' and
'psbt_final_tx', which I think also makes the callers clearer.

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

bitcoin/base58.c

@@ -24,12 +24,14 @@ static char *to_base58(const tal_t *ctx, u8 version,
 	buf[0] = version;
 	memcpy(buf + 1, rmd, sizeof(*rmd));
 
+	tal_wally_start();
 	if (wally_base58_from_bytes((const unsigned char *) buf,
 				    total_length, BASE58_FLAG_CHECKSUM, &out)
 	    != WALLY_OK)
 		out = NULL;
+	tal_wally_end(tal_steal(ctx, out));
 
-	return tal_steal(ctx, out);
+	return out;
 }
 
 char *bitcoin_to_base58(const tal_t *ctx, const struct chainparams *chainparams,

bitcoin/psbt.c

@@ -27,16 +27,15 @@ static struct wally_psbt *init_psbt(const tal_t *ctx, size_t num_inputs, size_t
 	int wally_err;
 	struct wally_psbt *psbt;
 
+	tal_wally_start();
 	if (is_elements(chainparams))
 		wally_err = wally_psbt_elements_init_alloc(0, num_inputs, num_outputs, 0, &psbt);
 	else
 		wally_err = wally_psbt_init_alloc(0, num_inputs, num_outputs, 0, &psbt);
 	assert(wally_err == WALLY_OK);
-	tal_steal(ctx, psbt);
-	/* If both of these are zero, no sub-allocations were done */
-	if (num_inputs || num_outputs)
-		tal_gather_wally(psbt);
 	tal_add_destructor(psbt, psbt_destroy);
+	tal_wally_end(tal_steal(ctx, psbt));
+
 	return psbt;
 }
 
@@ -46,16 +45,20 @@ struct wally_psbt *create_psbt(const tal_t *ctx, size_t num_inputs, size_t num_o
 	struct wally_tx *wtx;
 	struct wally_psbt *psbt;
 
+	tal_wally_start();
 	if (wally_tx_init_alloc(WALLY_TX_VERSION_2, locktime, num_inputs, num_outputs, &wtx) != WALLY_OK)
 		abort();
-	tal_gather_wally(tal_steal(NULL, wtx));
+	/* wtx is freed below */
+	tal_wally_end(NULL);
 
 	psbt = init_psbt(ctx, num_inputs, num_outputs);
 
+	tal_wally_start();
 	wally_err = wally_psbt_set_global_tx(psbt, wtx);
 	assert(wally_err == WALLY_OK);
+	tal_wally_end(psbt);
+
 	wally_tx_free(wtx);
-	tal_gather_wally(psbt);
 	return psbt;
 }
 
@@ -66,6 +69,7 @@ struct wally_psbt *new_psbt(const tal_t *ctx, const struct wally_tx *wtx)
 
 	psbt = init_psbt(ctx, wtx->num_inputs, wtx->num_outputs);
 
+	tal_wally_start();
 	/* Set directly: avoids psbt checks for non-NULL scripts/witnesses */
 	wally_err = wally_tx_clone_alloc(wtx, 0, &psbt->tx);
 	assert(wally_err == WALLY_OK);
@@ -90,7 +94,7 @@ struct wally_psbt *new_psbt(const tal_t *ctx, const struct wally_tx *wtx)
 		}
 	}
 
-	tal_gather_wally(tal_steal(ctx, psbt));
+	tal_wally_end(psbt);
 	return psbt;
 }
 
@@ -109,9 +113,10 @@ struct wally_psbt_input *psbt_add_input(struct wally_psbt *psbt,
 	const u32 flags = WALLY_PSBT_FLAG_NON_FINAL; /* Skip script/witness */
 	int wally_err;
 
+	tal_wally_start();
 	wally_err = wally_psbt_add_input_at(psbt, insert_at, flags, input);
 	assert(wally_err == WALLY_OK);
-	tal_gather_wally(psbt);
+	tal_wally_end(psbt);
 	return &psbt->inputs[insert_at];
 }
 
@@ -127,6 +132,7 @@ struct wally_psbt_input *psbt_append_input(struct wally_psbt *psbt,
 	const u32 flags = WALLY_PSBT_FLAG_NON_FINAL; /* Skip script/witness */
 	int wally_err;
 
+	tal_wally_start();
 	if (chainparams->is_elements) {
 		if (wally_tx_elements_input_init_alloc(txid->shad.sha.u.u8,
 						       sizeof(txid->shad.sha.u.u8),
@@ -149,7 +155,7 @@ struct wally_psbt_input *psbt_append_input(struct wally_psbt *psbt,
 	wally_err = wally_psbt_add_input_at(psbt, input_num, flags, tx_in);
 	assert(wally_err == WALLY_OK);
 	wally_tx_input_free(tx_in);
-	tal_gather_wally(psbt);
+	tal_wally_end(psbt);
 
 	if (input_wscript) {
 		/* Add the prev output's data into the PSBT struct */
@@ -157,11 +163,12 @@ struct wally_psbt_input *psbt_append_input(struct wally_psbt *psbt,
 	}
 
 	if (redeemscript) {
+		tal_wally_start();
 		wally_err = wally_psbt_input_set_redeem_script(&psbt->inputs[input_num],
 							       redeemscript,
 							       tal_bytelen(redeemscript));
 		assert(wally_err == WALLY_OK);
-		tal_gather_wally(psbt);
+		tal_wally_end(psbt);
 	}
 
 	return &psbt->inputs[input_num];
@@ -178,9 +185,12 @@ struct wally_psbt_output *psbt_add_output(struct wally_psbt *psbt,
 					  struct wally_tx_output *output,
 					  size_t insert_at)
 {
-	int wally_err = wally_psbt_add_output_at(psbt, insert_at, 0, output);
+	int wally_err;
+
+	tal_wally_start();
+	wally_err = wally_psbt_add_output_at(psbt, insert_at, 0, output);
 	assert(wally_err == WALLY_OK);
-	tal_gather_wally(psbt);
+	tal_wally_end(psbt);
 	return &psbt->outputs[insert_at];
 }
 
@@ -237,12 +247,13 @@ void psbt_input_add_pubkey(struct wally_psbt *psbt, size_t in,
 	/* we serialize the compressed version of the key, wally likes this */
 	pubkey_to_der(pk_der, pubkey);
 
+	tal_wally_start();
 	wally_err = wally_psbt_input_add_keypath_item(&psbt->inputs[in],
 						      pk_der, sizeof(pk_der),
 						      fingerprint, sizeof(fingerprint),
 						      empty_path, ARRAY_SIZE(empty_path));
 	assert(wally_err == WALLY_OK);
-	tal_gather_wally(psbt);
+	tal_wally_end(psbt);
 }
 
 bool psbt_input_set_signature(struct wally_psbt *psbt, size_t in,
@@ -250,19 +261,20 @@ bool psbt_input_set_signature(struct wally_psbt *psbt, size_t in,
 			      const struct bitcoin_signature *sig)
 {
 	u8 pk_der[PUBKEY_CMPR_LEN];
+	bool ok;
 
 	assert(in < psbt->num_inputs);
 
 	/* we serialize the compressed version of the key, wally likes this */
 	pubkey_to_der(pk_der, pubkey);
+	tal_wally_start();
 	wally_psbt_input_set_sighash(&psbt->inputs[in], sig->sighash_type);
-	if (wally_psbt_input_add_signature(&psbt->inputs[in],
+	ok = wally_psbt_input_add_signature(&psbt->inputs[in],
 					    pk_der, sizeof(pk_der),
 					    sig->s.data,
-					   sizeof(sig->s.data)) != WALLY_OK)
-		return false;
-	tal_gather_wally(psbt);
-	return true;
+					    sizeof(sig->s.data)) == WALLY_OK;
+	tal_wally_end(psbt);
+	return ok;
 }
 
 void psbt_input_set_wit_utxo(struct wally_psbt *psbt, size_t in,
@@ -273,6 +285,7 @@ void psbt_input_set_wit_utxo(struct wally_psbt *psbt, size_t in,
 
 	assert(in < psbt->num_inputs);
 	assert(tal_bytelen(scriptPubkey) > 0);
+	tal_wally_start();
 	if (is_elements(chainparams)) {
 		u8 value[9];
 		wally_err =
@@ -298,23 +311,26 @@ void psbt_input_set_wit_utxo(struct wally_psbt *psbt, size_t in,
 	wally_err = wally_psbt_input_set_witness_utxo(&psbt->inputs[in], tx_out);
 	wally_tx_output_free(tx_out);
 	assert(wally_err == WALLY_OK);
-	tal_gather_wally(psbt);
+	tal_wally_end(psbt);
 }
 
 void psbt_input_set_witscript(struct wally_psbt *psbt, size_t in, const u8 *wscript)
 {
 	int wally_err;
 
+	tal_wally_start();
 	wally_err = wally_psbt_input_set_witness_script(&psbt->inputs[in],
 							wscript,
 							tal_bytelen(wscript));
 	assert(wally_err == WALLY_OK);
-	tal_gather_wally(psbt);
+	tal_wally_end(psbt);
 }
 
 void psbt_elements_input_set_asset(struct wally_psbt *psbt, size_t in,
 				   struct amount_asset *asset)
 {
+	tal_wally_start();
+
 	if (asset->value > 0)
 		if (wally_psbt_input_set_value(&psbt->inputs[in],
 					       asset->value) != WALLY_OK)
@@ -325,7 +341,7 @@ void psbt_elements_input_set_asset(struct wally_psbt *psbt, size_t in,
 				       asset->asset + 1,
 				       ELEMENTS_ASSET_LEN - 1) != WALLY_OK)
 		abort();
-	tal_gather_wally(psbt);
+	tal_wally_end(psbt);
 }
 
 void psbt_elements_normalize_fees(struct wally_psbt *psbt)
@@ -370,7 +386,6 @@ void psbt_elements_normalize_fees(struct wally_psbt *psbt)
 	/* We need to add a fee output */
 	if (fee_output_idx == psbt->num_outputs) {
 		psbt_append_output(psbt, NULL, total_in);
-		tal_gather_wally(psbt);
 	} else {
 		u64 sats = total_in.satoshis; /* Raw: wally API */
 		struct wally_tx_output *out = &psbt->tx->outputs[fee_output_idx];
@@ -489,12 +504,13 @@ void psbt_input_add_unknown(const tal_t *ctx,
 			    const void *value,
 			    size_t value_len)
 {
+	tal_wally_start();
 	if (wally_map_add(&in->unknowns,
 			  cast_const(unsigned char *, key), tal_bytelen(key),
 			  (unsigned char *) memcheck(value, value_len), value_len)
 			!= WALLY_OK)
 		abort();
-	tal_gather_wally(ctx);
+	tal_wally_end(ctx);
 }
 
 void *psbt_get_unknown(const struct wally_map *map,
@@ -533,12 +549,13 @@ void psbt_output_add_unknown(const tal_t *ctx,
 			     const void *value,
 			     size_t value_len)
 {
+	tal_wally_start();
 	if (wally_map_add(&out->unknowns,
 			  cast_const(unsigned char *, key), tal_bytelen(key),
 			  (unsigned char *) memcheck(value, value_len), value_len)
 			!= WALLY_OK)
 		abort();
-	tal_gather_wally(ctx);
+	tal_wally_end(ctx);
 }
 
 /* Use the destructor to free the wally_tx */
@@ -547,28 +564,17 @@ static void wally_tx_destroy(struct wally_tx *wtx)
 	wally_tx_free(wtx);
 }
 
-struct wally_tx *psbt_finalize(const tal_t *ctx,
-			       struct wally_psbt *psbt, bool finalize_in_place)
+bool psbt_finalize(struct wally_psbt *psbt)
 {
-	struct wally_psbt *tmppsbt;
-	struct wally_tx *wtx;
+	bool ok;
 
-	/* We want the 'finalized' tx since that includes any signature
-	 * data, not the global tx. But 'finalizing' a tx destroys some fields
-	 * so we 'clone' it first and then finalize it */
-	if (!finalize_in_place) {
-		if (wally_psbt_clone_alloc(psbt, 0, &tmppsbt) != WALLY_OK)
-			return NULL;
-		/* Explicitly freed below */
-		tal_steal(NULL, tmppsbt);
-	} else
-		tmppsbt = cast_const(struct wally_psbt *, psbt);
+	tal_wally_start();
 
 	/* Wally doesn't know how to finalize P2WSH; this happens with
 	 * option_anchor_outputs, and finalizing is trivial. */
 	/* FIXME: miniscript! miniscript! miniscript! */
-	for (size_t i = 0; i < tmppsbt->num_inputs; i++) {
-		struct wally_psbt_input *input = &tmppsbt->inputs[i];
+	for (size_t i = 0; i < psbt->num_inputs; i++) {
+		struct wally_psbt_input *input = &psbt->inputs[i];
 		struct wally_tx_witness_stack *stack;
 
 		if (!is_anchor_witness_script(input->witness_script,
@@ -600,31 +606,29 @@ struct wally_tx *psbt_finalize(const tal_t *ctx,
 					   input->witness_script,
 					   input->witness_script_len);
 		input->final_witness = stack;
-		tal_gather_wally(tmppsbt);
 	}
 
-	if (wally_psbt_finalize(tmppsbt) != WALLY_OK) {
-		if (!finalize_in_place)
-			wally_psbt_free(tmppsbt);
-		return NULL;
+	ok = (wally_psbt_finalize(psbt) == WALLY_OK);
+	tal_wally_end(psbt);
+
+	return ok && psbt_is_finalized(psbt);
 }
 
-	/* wally_psbt_finalize() may or may not have done allocations! */
-	if (tal_first(wally_tal_ctx))
-		tal_gather_wally(tmppsbt);
+struct wally_tx *psbt_final_tx(const tal_t *ctx, const struct wally_psbt *psbt)
+{
+	struct wally_tx *wtx;
+
+	if (!psbt_is_finalized(psbt))
+		return NULL;
 
-	if (psbt_is_finalized(tmppsbt)
-		&& wally_psbt_extract(tmppsbt, &wtx) == WALLY_OK) {
-		tal_gather_wally(tal_steal(ctx, wtx));
+	tal_wally_start();
+	if (wally_psbt_extract(psbt, &wtx) == WALLY_OK)
 		tal_add_destructor(wtx, wally_tx_destroy);
-		if (!finalize_in_place)
-			wally_psbt_free(tmppsbt);
-		return wtx;
-	}
+	else
+		wtx = NULL;
 
-	if (!finalize_in_place)
-		wally_psbt_free(tmppsbt);
-	return NULL;
+	tal_wally_end(tal_steal(ctx, wtx));
+	return wtx;
 }
 
 struct wally_psbt *psbt_from_b64(const tal_t *ctx,
@@ -634,12 +638,12 @@ struct wally_psbt *psbt_from_b64(const tal_t *ctx,
 	struct wally_psbt *psbt;
 	char *str = tal_strndup(tmpctx, b64, b64len);
 
-	if (wally_psbt_from_base64(str, &psbt) != WALLY_OK)
-		return NULL;
-
-	/* We promised it would be owned by ctx: libwally uses a dummy owner */
-	tal_gather_wally(tal_steal(ctx, psbt));
+	tal_wally_start();
+	if (wally_psbt_from_base64(str, &psbt) == WALLY_OK)
 		tal_add_destructor(psbt, psbt_destroy);
+	else
+		psbt = NULL;
+	tal_wally_end(tal_steal(ctx, psbt));
 	return psbt;
 }
 
@@ -648,10 +652,12 @@ char *psbt_to_b64(const tal_t *ctx, const struct wally_psbt *psbt)
 	char *serialized_psbt;
 	int ret;
 
+	tal_wally_start();
 	ret = wally_psbt_to_base64(psbt, 0, &serialized_psbt);
 	assert(ret == WALLY_OK);
+	tal_wally_end(tal_steal(ctx, serialized_psbt));
 
-	return tal_steal(ctx, serialized_psbt);
+	return serialized_psbt;
 }
 REGISTER_TYPE_TO_STRING(wally_psbt, psbt_to_b64);
 
@@ -682,12 +688,13 @@ struct wally_psbt *psbt_from_bytes(const tal_t *ctx, const u8 *bytes,
 {
 	struct wally_psbt *psbt;
 
-	if (wally_psbt_from_bytes(bytes, byte_len, &psbt) != WALLY_OK)
-		return NULL;
-
-	/* We promised it would be owned by ctx: libwally uses a dummy owner */
-	tal_gather_wally(tal_steal(ctx, psbt));
+	tal_wally_start();
+	if (wally_psbt_from_bytes(bytes, byte_len, &psbt) == WALLY_OK)
 		tal_add_destructor(psbt, psbt_destroy);
+	else
+		psbt = NULL;
+	tal_wally_end(tal_steal(ctx, psbt));
+
 	return psbt;
 }
 
@@ -742,7 +749,7 @@ void psbt_txid(const tal_t *ctx,
 	/* You can *almost* take txid of global tx.  But @niftynei thought
 	 * about this far more than me and pointed out that P2SH
 	 * inputs would not be represented, so here we go. */
-
+	tal_wally_start();
 	wally_tx_clone_alloc(psbt->tx, 0, &tx);
 
 	for (size_t i = 0; i < tx->num_inputs; i++) {
@@ -761,7 +768,7 @@ void psbt_txid(const tal_t *ctx,
 			wally_tx_set_input_script(tx, i, script, tal_bytelen(script));
 		}
 	}
-	tal_gather_wally(tal_steal(ctx, tx));
+	tal_wally_end(tal_steal(ctx, tx));
 
 	wally_txid(tx, txid);
 	if (wtx)

bitcoin/psbt.h

@@ -79,8 +79,22 @@ void psbt_txid(const tal_t *ctx,
  */
 void psbt_elements_normalize_fees(struct wally_psbt *psbt);
 
-struct wally_tx *psbt_finalize(const tal_t *ctx,
-			       struct wally_psbt *psbt, bool finalize_in_place);
+/**
+ * psbt_finalize - finalize this psbt.
+ *
+ * Returns false if we can't, otherwise returns true and psbt_is_finalized()
+ * is true.
+ */
+bool psbt_finalize(struct wally_psbt *psbt);
+
+/**
+ * psbt_final_tx - extract transaction from finalized psbt.
+ * @ctx: context to tallocate return
+ * @psbt: psbt to extract.
+ *
+ * If @psbt isn't final, or we can't extract tx, returns NULL.
+ */
+struct wally_tx *psbt_final_tx(const tal_t *ctx, const struct wally_psbt *psbt);
 
 /* psbt_make_key - Create a new, proprietary c-lightning key
  *

bitcoin/tx.c

@@ -34,6 +34,7 @@ struct wally_tx_output *wally_tx_output(const tal_t *ctx,
 	struct wally_tx_output *output;
 	int ret;
 
+	tal_wally_start();
 	if (chainparams->is_elements) {
 		u8 value[9];
 		ret = wally_tx_confidential_value_from_satoshi(satoshis, value,
@@ -42,8 +43,10 @@ struct wally_tx_output *wally_tx_output(const tal_t *ctx,
 		ret = wally_tx_elements_output_init_alloc(
 		    script, tal_bytelen(script), chainparams->fee_asset_tag, 33,
 		    value, sizeof(value), NULL, 0, NULL, 0, NULL, 0, &output);
-		if (ret != WALLY_OK)
-			return NULL;
+		if (ret != WALLY_OK) {
+			output = NULL;
+			goto done;
+		}
 
 		/* Cheat a bit by also setting the numeric satoshi value,
 		 * otherwise we end up converting a number of times */
@@ -51,11 +54,14 @@ struct wally_tx_output *wally_tx_output(const tal_t *ctx,
 	} else {
 		ret = wally_tx_output_init_alloc(satoshis, script,
 						 tal_bytelen(script), &output);
-		if (ret != WALLY_OK)
-			return NULL;
+		if (ret != WALLY_OK) {
+			output = NULL;
+			goto done;
+		}
 	}
 
-	tal_gather_wally(tal_steal(ctx, output));
+done:
+	tal_wally_end(tal_steal(ctx, output));
 	return output;
 }
 
@@ -74,17 +80,19 @@ int bitcoin_tx_add_output(struct bitcoin_tx *tx, const u8 *script,
 
 	output = wally_tx_output(NULL, script, amount);
 	assert(output);
+	tal_wally_start();
 	ret = wally_tx_add_output(tx->wtx, output);
 	assert(ret == WALLY_OK);
-	tal_gather_wally(tx->wtx);
+	tal_wally_end(tx->wtx);
 
 	psbt_out = psbt_add_output(tx->psbt, output, i);
 	if (wscript) {
+		tal_wally_start();
 		ret = wally_psbt_output_set_witness_script(psbt_out,
 							   wscript,
 							   tal_bytelen(wscript));
 		assert(ret == WALLY_OK);
-		tal_gather_wally(tx->psbt);
+		tal_wally_end(tx->psbt);
 	}
 
 	wally_tx_output_free(output);
@@ -210,6 +218,8 @@ int bitcoin_tx_add_input(struct bitcoin_tx *tx, const struct bitcoin_txid *txid,
 	assert(scriptPubkey);
 	psbt_input_set_wit_utxo(tx->psbt, input_num,
 				scriptPubkey, amount);
+
+	tal_wally_start();
 	wally_err = wally_tx_add_input(tx->wtx,
 				       &tx->psbt->tx->inputs[input_num]);
 	assert(wally_err == WALLY_OK);
@@ -217,8 +227,7 @@ int bitcoin_tx_add_input(struct bitcoin_tx *tx, const struct bitcoin_txid *txid,
 	/* scriptsig isn't actually stored in psbt input, so add that now */
 	wally_tx_set_input_script(tx->wtx, input_num,
 				  scriptSig, tal_bytelen(scriptSig));
-	if (scriptSig)
-		tal_gather_wally(tx->wtx);
+	tal_wally_end(tx->wtx);
 
 	if (is_elements(chainparams)) {
 		struct amount_asset asset;
@@ -329,6 +338,7 @@ void bitcoin_tx_input_set_witness(struct bitcoin_tx *tx, int innum,
 	struct wally_tx_witness_stack *stack = NULL;
 	size_t stack_size = tal_count(witness);
 
+	tal_wally_start();
 	/* Free any lingering witness */
 	if (witness) {
 		wally_tx_witness_stack_init_alloc(stack_size, &stack);
@@ -336,14 +346,17 @@ void bitcoin_tx_input_set_witness(struct bitcoin_tx *tx, int innum,
 			wally_tx_witness_stack_add(stack, witness[i],
 						   tal_bytelen(witness[i]));
 	}
+	tal_wally_end(tmpctx);
+
+	tal_wally_start();
 	wally_tx_set_input_witness(tx->wtx, innum, stack);
-	if (stack)
-		tal_gather_wally(tx->wtx);
+	tal_wally_end(tx->wtx);
 
 	/* Also add to the psbt */
 	if (stack) {
+		tal_wally_start();
 		wally_psbt_input_set_final_witness(&tx->psbt->inputs[innum], stack);
-		tal_gather_wally(tx->psbt);
+		tal_wally_end(tx->psbt);
 	} else {
 		/* FIXME: libwally-psbt doesn't allow 'unsetting' of witness via
 		 * the set method at the moment, so we do it manually*/
@@ -353,8 +366,6 @@ void bitcoin_tx_input_set_witness(struct bitcoin_tx *tx, int innum,
 		in->final_witness = NULL;
 	}
 
-	if (stack)
-		wally_tx_witness_stack_free(stack);
 	if (taken(witness))
 		tal_free(witness);
 }
@@ -362,14 +373,17 @@ void bitcoin_tx_input_set_witness(struct bitcoin_tx *tx, int innum,
 void bitcoin_tx_input_set_script(struct bitcoin_tx *tx, int innum, u8 *script)
 {
 	struct wally_psbt_input *in;
+
+	tal_wally_start();
 	wally_tx_set_input_script(tx->wtx, innum, script, tal_bytelen(script));
-	tal_gather_wally(tx->wtx);
+	tal_wally_end(tx->wtx);
 
 	/* Also add to the psbt */
 	assert(innum < tx->psbt->num_inputs);
 	in = &tx->psbt->inputs[innum];
+	tal_wally_start();
 	wally_psbt_input_set_final_scriptsig(in, script, tal_bytelen(script));
-	tal_gather_wally(tx->psbt);
+	tal_wally_end(tx->psbt);
 }
 
 const u8 *bitcoin_tx_input_get_witness(const tal_t *ctx,
@@ -490,10 +504,11 @@ struct bitcoin_tx *bitcoin_tx(const tal_t *ctx,
 	if (chainparams->is_elements)
 		output_count += 1;
 
+	tal_wally_start();
 	wally_tx_init_alloc(WALLY_TX_VERSION_2, nlocktime, input_count, output_count,
 			    &tx->wtx);
-	tal_gather_wally(tal_steal(tx, tx->wtx));
 	tal_add_destructor(tx, bitcoin_tx_destroy);
+	tal_wally_end(tal_steal(tx, tx->wtx));
 
 	tx->chainparams = chainparams;
 	tx->psbt = new_psbt(tx, tx->wtx);
@@ -514,11 +529,16 @@ struct bitcoin_tx *bitcoin_tx_with_psbt(const tal_t *ctx, struct wally_psbt *psb
 					   psbt->tx->num_outputs,
 					   psbt->tx->locktime);
 	wally_tx_free(tx->wtx);
-	tx->wtx = psbt_finalize(tx, psbt, false);
+
+	psbt_finalize(psbt);
+	tx->wtx = psbt_final_tx(tx, psbt);
 	if (!tx->wtx) {
+		tal_wally_start();
 		if (wally_tx_clone_alloc(psbt->tx, 0, &tx->wtx) != WALLY_OK)
-			return NULL;
-		tal_gather_wally(tal_steal(tx, tx->wtx));
+			tx->wtx = NULL;
+		tal_wally_end(tal_steal(tx, tx->wtx));
+		if (!tx->wtx)
+			return tal_free(tx);
 	}
 
 	tal_free(tx->psbt);
@@ -530,29 +550,29 @@ struct bitcoin_tx *bitcoin_tx_with_psbt(const tal_t *ctx, struct wally_psbt *psb
 struct bitcoin_tx *pull_bitcoin_tx(const tal_t *ctx, const u8 **cursor,
 				   size_t *max)
 {
-	size_t wsize;
 	int flags = WALLY_TX_FLAG_USE_WITNESS;
 	struct bitcoin_tx *tx = tal(ctx, struct bitcoin_tx);
 
 	if (chainparams->is_elements)
 		flags |= WALLY_TX_FLAG_USE_ELEMENTS;
 
+	tal_wally_start();
 	if (wally_tx_from_bytes(*cursor, *max, flags, &tx->wtx) != WALLY_OK) {
 		fromwire_fail(cursor, max);
-		return tal_free(tx);
+		tx = tal_free(tx);
 	}
+	tal_wally_end(tx);
 
-	tal_gather_wally(tx);
+	if (tx) {
+		size_t wsize;
 		tal_add_destructor(tx, bitcoin_tx_destroy);
-
-	wally_tx_get_length(tx->wtx, flags, &wsize);
-
 		tx->chainparams = chainparams;
-
 		tx->psbt = new_psbt(tx, tx->wtx);
 
+		wally_tx_get_length(tx->wtx, flags, &wsize);
 		*cursor += wsize;
 		*max -= wsize;
+	}
 	return tx;
 }
 

bitcoin/tx_parts.c

@@ -9,8 +9,7 @@ static void destroy_wally_tx_input(struct wally_tx_input *in)
 	wally_tx_input_free(in);
 }
 
-static struct wally_tx_input *clone_input(const tal_t *ctx,
-					  const struct wally_tx_input *src)
+static struct wally_tx_input *clone_input(const struct wally_tx_input *src)
 {
 	struct wally_tx_input *in;
 	int ret;
@@ -39,7 +38,6 @@ static struct wally_tx_input *clone_input(const tal_t *ctx,
 	}
 	assert(ret == WALLY_OK);
 
-	tal_gather_wally(tal_steal(ctx, in));
 	tal_add_destructor(in, destroy_wally_tx_input);
 	return in;
 }
@@ -49,8 +47,7 @@ static void destroy_wally_tx_output(struct wally_tx_output *out)
 	wally_tx_output_free(out);
 }
 
-static struct wally_tx_output *clone_output(const tal_t *ctx,
-					  const struct wally_tx_output *src)
+static struct wally_tx_output *clone_output(const struct wally_tx_output *src)
 {
 	struct wally_tx_output *out;
 	int ret;
@@ -71,7 +68,6 @@ static struct wally_tx_output *clone_output(const tal_t *ctx,
 	}
 	assert(ret == WALLY_OK);
 
-	tal_gather_wally(tal_steal(ctx, out));
 	tal_add_destructor(out, destroy_wally_tx_output);
 	return out;
 }
@@ -86,17 +82,19 @@ struct tx_parts *tx_parts_from_wally_tx(const tal_t *ctx,
 	txp->inputs = tal_arrz(txp, struct wally_tx_input *, wtx->num_inputs);
 	txp->outputs = tal_arrz(txp, struct wally_tx_output *, wtx->num_outputs);
 
+	tal_wally_start();
 	for (size_t i = 0; i < wtx->num_inputs; i++) {
 		if (input != -1 && input != i)
 			continue;
-		txp->inputs[i] = clone_input(txp->inputs, &wtx->inputs[i]);
+		txp->inputs[i] = clone_input(&wtx->inputs[i]);
 	}
 
 	for (size_t i = 0; i < wtx->num_outputs; i++) {
 		if (output != -1 && output != i)
 			continue;
-		txp->outputs[i] = clone_output(txp->outputs, &wtx->outputs[i]);
+		txp->outputs[i] = clone_output(&wtx->outputs[i]);
 	}
+	tal_wally_end(txp);
 
 	return txp;
 }
@@ -120,6 +118,7 @@ fromwire_wally_tx_witness_stack(const tal_t *ctx,
 	if (num == 0)
 		return NULL;
 
+	tal_wally_start();
 	ret = wally_tx_witness_stack_init_alloc(num, &ws);
 	if (ret != WALLY_OK) {
 		fromwire_fail(cursor, max);
@@ -134,12 +133,14 @@ fromwire_wally_tx_witness_stack(const tal_t *ctx,
 		if (ret != WALLY_OK) {
 			wally_tx_witness_stack_free(ws);
 			fromwire_fail(cursor, max);
-			return NULL;
+			ws = NULL;
+			goto out;
 		}
 	}
 
-	tal_gather_wally(tal_steal(ctx, ws));
 	tal_add_destructor(ws, destroy_wally_tx_witness_stack);
+out:
+	tal_wally_end(tal_steal(ctx, ws));
 	return ws;
 }
 
@@ -181,6 +182,7 @@ static struct wally_tx_input *fromwire_wally_tx_input(const tal_t *ctx,
 		script = tal_free(script);
 	ws = fromwire_wally_tx_witness_stack(tmpctx, cursor, max);
 
+	tal_wally_start();
 	if (is_elements(chainparams)) {
 		u8 *blinding_nonce, *entropy, *issuance_amount,
 			*inflation_keys, *issuance_amount_rangeproof,
@@ -231,11 +233,12 @@ static struct wally_tx_input *fromwire_wally_tx_input(const tal_t *ctx,
 	}
 	if (ret != WALLY_OK) {
 		fromwire_fail(cursor, max);
-		return NULL;
+		in = NULL;
+	} else {
+		tal_add_destructor(in, destroy_wally_tx_input);
 	}
 
-	tal_gather_wally(tal_steal(ctx, in));
-	tal_add_destructor(in, destroy_wally_tx_input);
+	tal_wally_end(tal_steal(ctx, in));
 	return in;
 }
 
@@ -250,6 +253,7 @@ static struct wally_tx_output *fromwire_wally_tx_output(const tal_t *ctx,
 	script = fromwire_tal_arrn(tmpctx,
 				   cursor, max, fromwire_u32(cursor, max));
 
+	tal_wally_start();
 	if (is_elements(chainparams)) {
 		u8 *asset, *value, *nonce, *surjectionproof, *rangeproof;
 
@@ -285,11 +289,12 @@ static struct wally_tx_output *fromwire_wally_tx_output(const tal_t *ctx,
 	}
 	if (ret != WALLY_OK) {
 		fromwire_fail(cursor, max);
-		return NULL;
+		out = NULL;
+	} else {
+		tal_add_destructor(out, destroy_wally_tx_output);
 	}
+	tal_wally_end(tal_steal(ctx, out));
 
-	tal_gather_wally(tal_steal(ctx, out));
-	tal_add_destructor(out, destroy_wally_tx_output);
 	return out;
 }
 

common/daemon.c

@@ -81,19 +81,13 @@ void send_backtrace(const char *why)
 int daemon_poll(struct pollfd *fds, nfds_t nfds, int timeout)
 {
 	const char *t;
-	char *wally_leak;
 
 	t = taken_any();
 	if (t)
 		errx(1, "Outstanding taken pointers: %s", t);
 
-	wally_leak = tal_first(wally_tal_ctx);
-	if (wally_leak) {
-		/* Trigger valgrind to tell us about this! */
-		tal_free(wally_leak);
-		*wally_leak = 0;
-		errx(1, "Outstanding wally allocations");
-	}
+	if (wally_tal_ctx)
+		errx(1, "Outstanding tal_wally_start!");
 
 	clean_tmpctx();
 

common/psbt_open.c

@@ -66,8 +66,11 @@ static const u8 *linearize_input(const tal_t *ctx,
 	struct wally_psbt *psbt = create_psbt(NULL, 1, 0, 0);
 	size_t byte_len;
 
+	tal_wally_start();
 	if (wally_tx_add_input(psbt->tx, tx_in) != WALLY_OK)
 		abort();
+	tal_wally_end(psbt->tx);
+
 	psbt->inputs[0] = *in;
 	psbt->num_inputs++;
 
@@ -104,8 +107,10 @@ static const u8 *linearize_output(const tal_t *ctx,
 	memset(&txid, 0, sizeof(txid));
 	psbt_append_input(psbt, &txid, 0, 0, NULL, NULL, NULL);
 
+	tal_wally_start();
 	if (wally_tx_add_output(psbt->tx, tx_out) != WALLY_OK)
 		abort();
+	tal_wally_end(psbt->tx);
 
 	psbt->outputs[0] = *out;
 	psbt->num_outputs++;

common/setup.c

@@ -1,3 +1,4 @@
+#include <assert.h>
 #include <ccan/ccan/err/err.h>
 #include <common/memleak.h>
 #include <common/setup.h>
@@ -7,6 +8,7 @@
 
 static void *wally_tal(size_t size)
 {
+	assert(wally_tal_ctx);
 	return tal_arr_label(wally_tal_ctx, u8, size, "wally_tal");
 }
 
@@ -33,7 +35,6 @@ void common_setup(const char *argv0)
 		     " available ?");
 
 	/* We set up Wally, the bitcoin wallet lib */
-	wally_tal_ctx = tal_label(NULL, char, "wally_ctx_notleak");
 	wally_init(0);
 	wally_set_operations(&wally_tal_ops);
 	secp256k1_ctx = wally_get_secp_context();

common/test/exp-run-psbt_diff.c

@@ -136,8 +136,10 @@ int main(int argc, const char *argv[])
 
 	/* Create two psbts! */
 	end = create_psbt(tmpctx, 1, 1, 0);
+	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
+	tal_wally_end(tmpctx);
 	diff_count(start, end, 0, 0);
 	diff_count(end, start, 0, 0);
 
@@ -147,22 +149,28 @@ int main(int argc, const char *argv[])
 	diff_count(end, start, 0, 1);
 
 	/* Add another one, before previous */
+	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
+	tal_wally_end(tmpctx);
 	add_in_out_with_serial(end, 5, 2);
 	diff_count(start, end, 1, 0);
 	diff_count(end, start, 0, 1);
 
 	/* Add another, at end */
+	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
+	tal_wally_end(tmpctx);
 	add_in_out_with_serial(end, 15, 3);
 	diff_count(start, end, 1, 0);
 	diff_count(end, start, 0, 1);
 
 	/* Add another, in middle */
+	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
+	tal_wally_end(tmpctx);
 	add_in_out_with_serial(end, 11, 4);
 	diff_count(start, end, 1, 0);
 	diff_count(end, start, 0, 1);
@@ -171,8 +179,10 @@ int main(int argc, const char *argv[])
 	 * (we accomplish this by removing and then
 	 * readding an input/output with the same serial_id
 	 * but different value) */
+	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
+	tal_wally_end(tmpctx);
 	psbt_rm_output(end, 0);
 	psbt_rm_input(end, 0);
 	add_in_out_with_serial(end, 5, 5);

common/utils.c

@@ -16,13 +16,25 @@ bool is_elements(const struct chainparams *chainparams)
 	return chainparams->is_elements;
 }
 
-/* Steal any wally allocations onto this context. */
-void tal_gather_wally(const tal_t *ctx)
+void tal_wally_start(void)
+{
+	if (wally_tal_ctx) {
+		/* This makes valgrind show us backtraces! */
+		*(u8 *)wally_tal_ctx = '\0';
+		abort();
+	}
+
+	wally_tal_ctx = tal_arr(NULL, char, 0);
+}
+
+void tal_wally_end(const tal_t *parent)
 {
 	tal_t *p;
-	assert(tal_first(wally_tal_ctx));
-	while ((p = tal_first(wally_tal_ctx)) != NULL)
-		tal_steal(ctx, p);
+	while ((p = tal_first(wally_tal_ctx)) != NULL) {
+		if (p != parent)
+			tal_steal(parent, p);
+	}
+	wally_tal_ctx = tal_free(wally_tal_ctx);
 }
 
 #if DEVELOPER

common/utils.h

@@ -84,8 +84,11 @@ void setup_tmpctx(void);
 /* Free any children of tmpctx. */
 void clean_tmpctx(void);
 
-/* Steal any wally allocations onto this context. */
-void tal_gather_wally(const tal_t *ctx);
+/* Call this before any libwally function which allocates. */
+void tal_wally_start(void);
+/* Then call this to reparent everything onto this parent (which must
+ * have been tal_steal() if it was allocated by libwally here) */
+void tal_wally_end(const tal_t *parent);
 
 /* Define sha256_eq. */
 STRUCTEQ_DEF(sha256, 0, u);

hsmd/hsmd.c

@@ -1582,6 +1582,7 @@ static void sign_our_inputs(struct utxo **utxos, struct wally_psbt *psbt)
 							scriptpubkey_p2wsh(psbt, wscript),
 							utxo->amount);
 			}
+			tal_wally_start();
 			if (wally_psbt_sign(psbt, privkey.secret.data,
 					    sizeof(privkey.secret.data),
 					    EC_FLAG_GRIND_R) != WALLY_OK)
@@ -1591,10 +1592,9 @@ static void sign_our_inputs(struct utxo **utxos, struct wally_psbt *psbt)
 							     &pubkey),
 					      type_to_string(tmpctx, struct wally_psbt,
 							     psbt));
-
+			tal_wally_end(psbt);
 		}
 	}
-	tal_gather_wally(psbt);
 }
 
 /*~ lightningd asks us to sign a withdrawal; same as above but in theory

lightningd/dual_open_control.c

@@ -614,7 +614,7 @@ openchannel2_sign_hook_cb(struct openchannel2_psbt_payload *payload STEALS)
 	tal_steal(tmpctx, payload);
 
 	/* Finalize it, if not already. It shouldn't work entirely */
-	psbt_finalize(tmpctx, payload->psbt, true);
+	psbt_finalize(payload->psbt);
 
 	if (!psbt_side_finalized(payload->ld->log, payload->psbt, REMOTE))
 		fatal("Plugin must return a 'psbt' with signatures for their inputs"

wallet/reservation.c

@@ -151,10 +151,12 @@ static struct command_result *json_unreserveinputs(struct command *cmd,
 		wally_tx_input_get_txid(&psbt->tx->inputs[i], &txid);
 		utxo_tx = wallet_transaction_get(psbt, cmd->ld->wallet,
 						 &txid);
-		if (utxo_tx)
+		if (utxo_tx) {
+			tal_wally_start();
 			wally_psbt_input_set_utxo(&psbt->inputs[i],
 						  utxo_tx->wtx);
-		else
+			tal_wally_end(psbt);
+		} else
 			log_broken(cmd->ld->log,
 				   "No transaction found for UTXO %s",
 				   type_to_string(tmpctx, struct bitcoin_txid,

wallet/walletrpc.c

@@ -812,7 +812,9 @@ static struct command_result *json_sendpsbt(struct command *cmd,
 
 	sending = tal(cmd, struct sending_psbt);
 	sending->cmd = cmd;
-	sending->wtx = psbt_finalize(sending, psbt, true);
+
+	psbt_finalize(psbt);
+	sending->wtx = psbt_final_tx(sending, psbt);
 	if (!sending->wtx)
 		return command_fail(cmd, LIGHTNINGD,
 				    "PSBT not finalizeable %s",