psbt: have the unknown map 'add' be a 'set' instead

bitcoin/psbt.c

@@ -509,19 +509,49 @@ u8 *psbt_make_key(const tal_t *ctx, u8 key_subtype, const u8 *key_data)
 	return key;
 }
 
-void psbt_input_add_unknown(const tal_t *ctx,
+static bool wally_map_set_unknown(const tal_t *ctx,
+				  struct wally_map *map,
+				  const u8 *key,
+				  const void *value,
+				  size_t value_len)
+{
+	size_t exists_at;
+	struct wally_map_item *item;
+
+	assert(value_len != 0);
+	if (wally_map_find(map,
+			   cast_const(unsigned char *, key), tal_bytelen(key),
+			   &exists_at) != WALLY_OK)
+		return false;
+
+	/* If not exists, add */
+	if (exists_at == 0) {
+		bool ok;
+		tal_wally_start();
+		ok = wally_map_add(map, cast_const(unsigned char *, key), tal_bytelen(key),
+			      (unsigned char *) memcheck(value, value_len), value_len)
+			== WALLY_OK;
+		tal_wally_end(ctx);
+		return ok;
+	}
+
+	/* Already in map, update entry */
+	item = &map->items[exists_at - 1];
+	tal_resize(&item->value, value_len);
+	memcpy(item->value, memcheck(value, value_len), value_len);
+	item->value_len = value_len;
+
+	return true;
+}
+
+void psbt_input_set_unknown(const tal_t *ctx,
 			    struct wally_psbt_input *in,
 			    const u8 *key,
 			    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)
+	if (!wally_map_set_unknown(ctx, &in->unknowns, key, value, value_len))
 		abort();
-	tal_wally_end(ctx);
 }
 
 void *psbt_get_unknown(const struct wally_map *map,
@@ -554,19 +584,14 @@ void *psbt_get_lightning(const struct wally_map *map,
 }
 
 
-void psbt_output_add_unknown(const tal_t *ctx,
+void psbt_output_set_unknown(const tal_t *ctx,
 			     struct wally_psbt_output *out,
 			     const u8 *key,
 			     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)
+	if (!wally_map_set_unknown(ctx, &out->unknowns, key, value, value_len))
 		abort();
-	tal_wally_end(ctx);
 }
 
 /* Use the destructor to free the wally_tx */

bitcoin/psbt.h

@@ -174,14 +174,14 @@ void psbt_elements_input_init_witness(struct wally_psbt *psbt, size_t in,
 				      const u8 *nonce);
 bool psbt_input_set_redeemscript(struct wally_psbt *psbt, size_t in,
 				 const u8 *redeemscript);
-/* psbt_input_add_unknown - Add the given Key-Value to the psbt's input keymap
+/* psbt_input_set_unknown - Set the given Key-Value in the psbt's input keymap
  * @ctx - tal context for allocations
- * @in - psbt input to add key-value to
+ * @in - psbt input to set key-value on
  * @key - key for key-value pair
- * @value - value to add
+ * @value - value to set
  * @value_len - length of {@value}
  */
-void psbt_input_add_unknown(const tal_t *ctx,
+void psbt_input_set_unknown(const tal_t *ctx,
 			    struct wally_psbt_input *in,
 			    const u8 *key,
 			    const void *value,
@@ -208,15 +208,15 @@ void *psbt_get_lightning(const struct wally_map *map,
 			 const u8 proprietary_type,
 			 size_t *val_len);
 
-/* psbt_output_add_unknown - Add the given Key-Value to the psbt's output keymap
+/* psbt_output_set_unknown - Set the given Key-Value in the psbt's output keymap
  *
  * @ctx - tal context for allocations
- * @out - psbt output to add key-value to
+ * @out - psbt output to set key-value on
  * @key - key for key-value pair
- * @value - value to add
+ * @value - value to set
  * @value_len - length of {@value}
  */
-void psbt_output_add_unknown(const tal_t *ctx,
+void psbt_output_set_unknown(const tal_t *ctx,
 			     struct wally_psbt_output *out,
 			     const u8 *key, const void *value,
 			     size_t value_len);

common/psbt_open.c

@@ -398,24 +398,24 @@ u8 *psbt_changeset_get_next(const tal_t *ctx, struct channel_id *cid,
 	return NULL;
 }
 
-void psbt_input_add_serial_id(const tal_t *ctx,
+void psbt_input_set_serial_id(const tal_t *ctx,
 			      struct wally_psbt_input *input,
 			      u16 serial_id)
 {
 	u8 *key = psbt_make_key(tmpctx, PSBT_TYPE_SERIAL_ID, NULL);
 	beint16_t bev = cpu_to_be16(serial_id);
 
-	psbt_input_add_unknown(ctx, input, key, &bev, sizeof(bev));
+	psbt_input_set_unknown(ctx, input, key, &bev, sizeof(bev));
 }
 
 
-void psbt_output_add_serial_id(const tal_t *ctx,
-			      struct wally_psbt_output *output,
+void psbt_output_set_serial_id(const tal_t *ctx,
+			       struct wally_psbt_output *output,
 			       u16 serial_id)
 {
 	u8 *key = psbt_make_key(tmpctx, PSBT_TYPE_SERIAL_ID, NULL);
 	beint16_t bev = cpu_to_be16(serial_id);
-	psbt_output_add_unknown(ctx, output, key, &bev, sizeof(bev));
+	psbt_output_set_unknown(ctx, output, key, &bev, sizeof(bev));
 }
 
 int psbt_find_serial_input(struct wally_psbt *psbt, u16 serial_id)

common/psbt_open.h

@@ -85,22 +85,22 @@ struct psbt_changeset *psbt_get_changeset(const tal_t *ctx,
 u8 *psbt_changeset_get_next(const tal_t *ctx, struct channel_id *cid,
 			    struct psbt_changeset *set);
 
-/* psbt_input_add_serial_id - Adds a serial id to given input
+/* psbt_input_set_serial_id - Sets a serial id on given input
  *
  * @ctx - tal context for allocations
- * @input - to add serial_id to
- * @serial_id - to add
+ * @input - to set serial_id on
+ * @serial_id - to set
  */
-void psbt_input_add_serial_id(const tal_t *ctx,
+void psbt_input_set_serial_id(const tal_t *ctx,
 			      struct wally_psbt_input *input,
-			      u16 serial_id);
-/* psbt_output_add_serial_id - Adds a serial id to given output
+			       u16 serial_id);
+/* psbt_output_set_serial_id - Sets a serial id on given output
  *
  * @ctx - tal context for allocations
- * @output - to add serial_id to
- * @serial_id - to add
+ * @output - to set serial_id on
+ * @serial_id - to set
  */
-void psbt_output_add_serial_id(const tal_t *ctx,
+void psbt_output_set_serial_id(const tal_t *ctx,
 			       struct wally_psbt_output *output,
 			       u16 serial_id);
 

common/test/exp-run-psbt_diff.c

@@ -117,7 +117,7 @@ static void add_in_out_with_serial(struct wally_psbt *psbt,
 			       NULL, NULL, NULL);
 	if (!in)
 		abort();
-	psbt_input_add_serial_id(psbt, in, serial_id);
+	psbt_input_set_serial_id(psbt, in, serial_id);
 
 	script = tal_arr(tmpctx, u8, 20);
 	memset(script, default_value, 20);
@@ -125,7 +125,25 @@ static void add_in_out_with_serial(struct wally_psbt *psbt,
 	out = psbt_append_output(psbt, script, sat);
 	if (!out)
 		abort();
-	psbt_output_add_serial_id(psbt, out, serial_id);
+	psbt_output_set_serial_id(psbt, out, serial_id);
+}
+
+/* Try changing up the serial ids */
+static void change_serials(void)
+{
+	struct wally_psbt *psbt;
+
+	psbt = create_psbt(tmpctx, 1, 1, 0);
+	add_in_out_with_serial(psbt, 10, 1);
+
+	psbt_output_set_serial_id(psbt, &psbt->outputs[0], 2);
+	assert(psbt_find_serial_output(psbt, 2) == 0);
+	assert(psbt_find_serial_output(psbt, 10) == -1);
+
+	psbt_input_set_serial_id(psbt, &psbt->inputs[0], 4);
+	assert(psbt_find_serial_input(psbt, 4) == 0);
+	assert(psbt_find_serial_input(psbt, 10) == -1);
+
 }
 
 int main(int argc, const char *argv[])
@@ -195,8 +213,8 @@ int main(int argc, const char *argv[])
 	/* Add some extra unknown info to a PSBT */
 	u8 *key = psbt_make_key(tmpctx, 0x05, NULL);
 	char *val = tal_fmt(tmpctx, "hello");
-	psbt_input_add_unknown(end, &end->inputs[1], key, val, tal_bytelen(val));
-	psbt_input_add_unknown(start, &start->inputs[1], key, val, tal_bytelen(val));
+	psbt_input_set_unknown(end, &end->inputs[1], key, val, tal_bytelen(val));
+	psbt_input_set_unknown(start, &start->inputs[1], key, val, tal_bytelen(val));
 
 	/* Swap locations */
 	struct wally_map_item tmp;
@@ -208,6 +226,8 @@ int main(int argc, const char *argv[])
 	diff_count(start, end, 1, 1);
 	diff_count(end, start, 1, 1);
 
+	change_serials();
+
 	/* No memory leaks please */
 	common_shutdown();
 	return 0;

lightningd/dual_open_control.c

@@ -365,7 +365,7 @@ static void psbt_add_serials(struct wally_psbt *psbt, enum tx_role role)
 			continue;
 
 		serial_id = psbt_new_input_serial(psbt, role);
-		psbt_input_add_serial_id(psbt, &psbt->inputs[i], serial_id);
+		psbt_input_set_serial_id(psbt, &psbt->inputs[i], serial_id);
 	}
 	for (size_t i = 0; i < psbt->num_outputs; i++) {
 		/* Skip ones that already have a serial id */
@@ -373,7 +373,7 @@ static void psbt_add_serials(struct wally_psbt *psbt, enum tx_role role)
 			continue;
 
 		serial_id = psbt_new_output_serial(psbt, role);
-		psbt_output_add_serial_id(psbt, &psbt->outputs[i], serial_id);
+		psbt_output_set_serial_id(psbt, &psbt->outputs[i], serial_id);
 	}
 }
 

openingd/dualopend.c

@@ -688,7 +688,7 @@ static bool run_tx_interactive(struct state *state, struct wally_psbt **orig_psb
 				psbt_elements_input_set_asset(psbt, outnum, &asset);
 			}
 
-			psbt_input_add_serial_id(psbt, in, serial_id);
+			psbt_input_set_serial_id(psbt, in, serial_id);
 
 			/* FIXME: what's in the tlv? */
 			break;
@@ -749,7 +749,7 @@ static bool run_tx_interactive(struct state *state, struct wally_psbt **orig_psb
 					    "Duplicate serial_id rcvd. %u", serial_id);
 			amt = amount_sat(value);
 			out = psbt_append_output(psbt, scriptpubkey, amt);
-			psbt_output_add_serial_id(psbt, out, serial_id);
+			psbt_output_set_serial_id(psbt, out, serial_id);
 			break;
 		}
 		case WIRE_TX_REMOVE_OUTPUT: {
@@ -1391,7 +1391,7 @@ static u8 *opener_start(struct state *state, u8 *msg)
 					 total);
 	/* Add a serial_id for this output */
 	serial_id = psbt_new_input_serial(psbt, TX_INITIATOR);
-	psbt_output_add_serial_id(psbt, funding_out, serial_id);
+	psbt_output_set_serial_id(psbt, funding_out, serial_id);
 
 	/* Add all of our inputs/outputs to the changeset */
 	init_changeset(state, psbt);