PSBT: Add audi_psbt routine

A routine that audit’s and asserts PSBT memory to confirm it has a sane memory allocation hierarchy.

Changelog-None

bitcoin/psbt.c

@@ -102,6 +102,79 @@ struct wally_psbt *combine_psbt(const tal_t *ctx,
 	return combined_psbt;
 }
 
+static bool parent_or_grandparent(const tal_t *goal, const tal_t *child)
+{
+	const tal_t *parent = tal_parent(child);
+	if (!parent)
+		return false;
+	return parent == goal || parent_or_grandparent(goal, parent);
+}
+
+#define NULL_OR_MATCH_P(item, parent) \
+	((item) == NULL || parent_or_grandparent((parent), (item)))
+
+static void audit_map(const tal_t *ctx, const struct wally_map *map)
+{
+	assert(NULL_OR_MATCH_P(map->items, ctx));
+	for (size_t i = 0; i < map->num_items; i++) {
+		assert(NULL_OR_MATCH_P(map->items[i].key, ctx));
+		assert(NULL_OR_MATCH_P(map->items[i].value, ctx));
+		assert(!map->items[i].key
+			|| tal_bytelen(map->items[i].key)
+				== map->items[i].key_len);
+		assert(!map->items[i].value
+			|| tal_bytelen(map->items[i].value)
+				== map->items[i].value_len);
+	}
+}
+
+void audit_psbt(const tal_t *ctx, const struct wally_psbt *psbt)
+{
+	assert(psbt);
+	assert(NULL_OR_MATCH_P(psbt->tx, ctx));
+	assert(NULL_OR_MATCH_P(psbt->inputs, ctx));
+	assert(NULL_OR_MATCH_P(psbt->outputs, ctx));
+	audit_map(ctx, &psbt->unknowns);
+	audit_map(ctx, &psbt->global_xpubs);
+#ifndef WALLY_ABI_NO_ELEMENTS
+	audit_map(ctx, &psbt->global_scalars);
+#endif
+	for (size_t i = 0; i < psbt->num_inputs; i++) {
+		assert(NULL_OR_MATCH_P(psbt->inputs[i].utxo, ctx));
+		assert(NULL_OR_MATCH_P(psbt->inputs[i].witness_utxo, ctx));
+		assert(NULL_OR_MATCH_P(psbt->inputs[i].final_witness, ctx));
+		audit_map(ctx, &psbt->inputs[i].keypaths);
+		audit_map(ctx, &psbt->inputs[i].signatures);
+		audit_map(ctx, &psbt->inputs[i].unknowns);
+		audit_map(ctx, &psbt->inputs[i].preimages);
+		audit_map(ctx, &psbt->inputs[i].psbt_fields);
+		audit_map(ctx, &psbt->inputs[i].taproot_leaf_signatures);
+		audit_map(ctx, &psbt->inputs[i].taproot_leaf_scripts);
+		audit_map(ctx, &psbt->inputs[i].taproot_leaf_hashes);
+		audit_map(ctx, &psbt->inputs[i].taproot_leaf_paths);
+		/* DTODO: Investigate if taproot wally maps have child maps */
+#ifndef WALLY_ABI_NO_ELEMENTS
+		assert(NULL_OR_MATCH_P(psbt->inputs[i].pegin_tx, ctx));
+		assert(NULL_OR_MATCH_P(psbt->inputs[i].pegin_witness, ctx));
+		audit_map(ctx, &psbt->inputs[i].pset_fields);
+#endif
+	}
+	for (size_t i = 0; i < psbt->num_outputs; i++) {
+		assert(NULL_OR_MATCH_P(psbt->outputs[i].script, ctx));
+		assert(psbt->outputs[i].script_len
+			== tal_bytelen(psbt->outputs[i].script));
+		audit_map(ctx, &psbt->outputs[i].keypaths);
+		audit_map(ctx, &psbt->outputs[i].unknowns);
+		audit_map(ctx, &psbt->outputs[i].psbt_fields);
+		audit_map(ctx, &psbt->outputs[i].taproot_tree);
+		audit_map(ctx, &psbt->outputs[i].taproot_leaf_hashes);
+		audit_map(ctx, &psbt->outputs[i].taproot_leaf_paths);
+#ifndef WALLY_ABI_NO_ELEMENTS
+		audit_map(ctx, &psbt->outputs[i].pset_fields);
+#endif
+	}
+}
+
 bool psbt_is_finalized(const struct wally_psbt *psbt)
 {
 	size_t is_finalized;

bitcoin/psbt.h

@@ -61,6 +61,17 @@ struct wally_psbt *combine_psbt(const tal_t *ctx,
 				const struct wally_psbt *psbt0,
 				const struct wally_psbt *psbt1);
 
+/**
+ * audit_psbt - Audit the memory structure of the PSBT.
+ *
+ * This checks all known memory allocations in the PSBT and asserts that they
+ * are all allocated with 'ctx' being it's parent.
+ *
+ * ctx - the ctx all memory *should* be attached to
+ * psbt - the PSBT to audit.
+ * */
+void audit_psbt(const tal_t *ctx, const struct wally_psbt *psbt);
+
 /**
  * psbt_is_finalized - Check if tx is ready to be extracted
  *

bitcoin/test/run-psbt-from-tx.c

@@ -114,6 +114,7 @@ int main(int argc, char *argv[])
 	const struct wally_map_item *final_scriptsig = wally_map_get_integer(&tx2->psbt->inputs[0].psbt_fields, /* PSBT_IN_FINAL_SCRIPTSIG */ 0x07);
 	assert(final_scriptsig->value_len > 0);
 	assert(tx2->psbt->inputs[0].final_witness != NULL);
+	audit_psbt(tx2->psbt, tx2->psbt);
 
 	common_shutdown();
 	return 0;

channeld/channeld.c

@@ -3677,6 +3677,8 @@ static void resume_splice_negotiation(struct peer *peer,
 							      new_output_index);
 		wire_sync_write(MASTER_FD, take(msg));
 	}
+
+	audit_psbt(current_psbt, current_psbt);
 }
 
 static struct inflight *inflights_new(struct peer *peer)
@@ -3995,10 +3997,15 @@ static void splice_initiator(struct peer *peer, const u8 *inmsg)
 	outmsg = towire_channeld_splice_confirmed_init(NULL, psbt);
 	wire_sync_write(MASTER_FD, take(outmsg));
 
+	audit_psbt(psbt, psbt);
+
 	/* We reset current_psbt to empty as now it represends the difference
 	 * what we've sent our peer so far */
 	tal_free(peer->splicing->current_psbt);
 	peer->splicing->current_psbt = create_psbt(peer->splicing, 0, 0, 0);
+
+	audit_psbt(peer->splicing->current_psbt, peer->splicing->current_psbt);
+	assert(tal_parent(peer->splicing->current_psbt) != tmpctx);
 }
 
 /* This occurs when the user has marked they are done making changes to the
@@ -4102,6 +4109,8 @@ static void splice_initiator_user_finalized(struct peer *peer)
 	new_inflight->force_sign_first = peer->splicing->force_sign_first;
 	new_inflight->locked_scid = NULL;
 
+	audit_psbt(ictx->current_psbt, ictx->current_psbt);
+
 	/* Switch over to using inflight psbt. This allows us to be reentrant.
 	 * On restart we *will* have inflight psbt but we will not have any
 	 * normal in-memory copy of the psbt: peer->splicing/ictx->current_psbt.
@@ -4144,6 +4153,9 @@ static void splice_initiator_user_finalized(struct peer *peer)
 							 true,
 							 !sign_first);
 	wire_sync_write(MASTER_FD, take(outmsg));
+
+	audit_psbt(new_inflight->psbt, new_inflight->psbt);
+	assert(tal_parent(new_inflight->psbt) != tmpctx);
 }
 
 /* During a splice the user may call splice_update mulitple times adding
@@ -4227,6 +4239,8 @@ static void splice_initiator_user_update(struct peer *peer, const u8 *inmsg)
 							 ictx->current_psbt,
 							 false, false);
 	wire_sync_write(MASTER_FD, take(outmsg));
+	audit_psbt(peer->splicing->current_psbt, peer->splicing->current_psbt);
+	assert(tal_parent(peer->splicing->current_psbt) != tmpctx);
 }
 
 /* This occurs when the user has signed the final version of the PSBT. At this
@@ -4317,7 +4331,13 @@ static void splice_initiator_user_signed(struct peer *peer, const u8 *inmsg)
 	sign_first = do_i_sign_first(peer, inflight->psbt, TX_INITIATOR,
 				     inflight->force_sign_first);
 
+	audit_psbt(inflight->psbt, inflight->psbt);
+	assert(tal_parent(inflight->psbt) != tmpctx);
+
 	resume_splice_negotiation(peer, false, false, true, sign_first);
+
+	audit_psbt(inflight->psbt, inflight->psbt);
+	assert(tal_parent(inflight->psbt) != tmpctx);
 }
 
 /* This occurs once our 'stfu' transition was successful. */

common/test/run-psbt_diff.c

@@ -162,7 +162,7 @@ int main(int argc, const char *argv[])
 	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
-	tal_wally_end(tmpctx);
+	tal_wally_end_onto(tmpctx, start, struct wally_psbt);
 	diff_count(start, end, 0, 0);
 	diff_count(end, start, 0, 0);
 
@@ -175,7 +175,7 @@ int main(int argc, const char *argv[])
 	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
-	tal_wally_end(tmpctx);
+	tal_wally_end_onto(tmpctx, start, struct wally_psbt);
 	add_in_out_with_serial(end, 5, 2);
 	diff_count(start, end, 1, 0);
 	diff_count(end, start, 0, 1);
@@ -184,7 +184,7 @@ int main(int argc, const char *argv[])
 	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
-	tal_wally_end(tmpctx);
+	tal_wally_end_onto(tmpctx, start, struct wally_psbt);
 	add_in_out_with_serial(end, 15, 3);
 	diff_count(start, end, 1, 0);
 	diff_count(end, start, 0, 1);
@@ -193,7 +193,7 @@ int main(int argc, const char *argv[])
 	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
-	tal_wally_end(tmpctx);
+	tal_wally_end_onto(tmpctx, start, struct wally_psbt);
 	add_in_out_with_serial(end, 11, 4);
 	diff_count(start, end, 1, 0);
 	diff_count(end, start, 0, 1);
@@ -205,7 +205,7 @@ int main(int argc, const char *argv[])
 	tal_wally_start();
 	if (wally_psbt_clone_alloc(end, flags, &start) != WALLY_OK)
 		abort();
-	tal_wally_end(tmpctx);
+	tal_wally_end_onto(tmpctx, start, struct wally_psbt);
 	psbt_rm_output(end, 0);
 	psbt_rm_input(end, 0);
 	add_in_out_with_serial(end, 5, 5);
@@ -219,6 +219,7 @@ int main(int argc, const char *argv[])
 	psbt_input_set_unknown(start, &start->inputs[1], key, val, tal_bytelen(val));
 
 	/* Swap locations */
+	assert(end->inputs[1].unknowns.num_items > 1);
 	struct wally_map_item tmp;
 	tmp = end->inputs[1].unknowns.items[0];
 	end->inputs[1].unknowns.items[0] = end->inputs[1].unknowns.items[1];
@@ -232,6 +233,9 @@ int main(int argc, const char *argv[])
 
 	check_psbt_comparison();
 
+	audit_psbt(start, start);
+	audit_psbt(end, end);
+
 	/* No memory leaks please */
 	common_shutdown();
 	return 0;