Use a single iterator to construct a BlindedPath

Instead of using separate iterators for pubkeys and TLVs, use an
iterator that yields a tuple of them. This allows for holding a mutable
reference to the TLVs such that they can be padded. With two iterators,
the pubkey iterator would have a reference to the ForwardNode slice when
constructing a payment path. However, this would prevent holding mutable
references in the TLVs iterator.

Author: jkczyz

lightning/src/blinded_path/message.rs

@@ -386,6 +386,8 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 		.map(|next_hop| ControlTlvs::Forward(ForwardTlvs { next_hop, next_blinding_override: None }))
 		.chain(core::iter::once(ControlTlvs::Receive(ReceiveTlvs{ context: Some(context) })));
 
-	utils::construct_blinded_hops(secp_ctx, pks, tlvs, session_priv)
+	let path = pks.zip(tlvs);
+
+	utils::construct_blinded_hops(secp_ctx, path, session_priv)
 }
 

lightning/src/blinded_path/payment.rs

@@ -466,7 +466,10 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 		.chain(core::iter::once(payee_node_id));
 	let tlvs = intermediate_nodes.iter().map(|node| BlindedPaymentTlvsRef::Forward(&node.tlvs))
 		.chain(core::iter::once(BlindedPaymentTlvsRef::Receive(&payee_tlvs)));
-	utils::construct_blinded_hops(secp_ctx, pks, tlvs, session_priv)
+
+	let path = pks.zip(tlvs);
+
+	utils::construct_blinded_hops(secp_ctx, path, session_priv)
 }
 
 /// `None` if underflow occurs.

lightning/src/blinded_path/utils.rs

@@ -118,13 +118,14 @@ where
 }
 
 #[inline]
-pub(super) fn construct_keys_callback_for_blinded_path<'a, T, I, F>(
+pub(super) fn construct_keys_callback_for_blinded_path<'a, T, I, F, H>(
 	secp_ctx: &Secp256k1<T>, unblinded_path: I, session_priv: &SecretKey, mut callback: F,
 ) -> Result<(), secp256k1::Error>
 where
 	T: secp256k1::Signing + secp256k1::Verification,
-	I: Iterator<Item=PublicKey>,
-	F: FnMut(PublicKey, SharedSecret, PublicKey, [u8; 32], Option<PublicKey>, Option<Vec<u8>>),
+	H: Borrow<PublicKey>,
+	I: Iterator<Item=H>,
+	F: FnMut(PublicKey, SharedSecret, PublicKey, [u8; 32], Option<H>, Option<Vec<u8>>),
 {
 	build_keys_helper!(session_priv, secp_ctx, callback);
 
@@ -134,23 +135,32 @@ where
 	Ok(())
 }
 
-// Panics if `unblinded_tlvs` length is less than `unblinded_pks` length
-pub(crate) fn construct_blinded_hops<'a, T, I1, I2>(
-	secp_ctx: &Secp256k1<T>, unblinded_pks: I1, mut unblinded_tlvs: I2, session_priv: &SecretKey
+struct PublicKeyWithTlvs<W: Writeable> {
+	 pubkey: PublicKey,
+	 tlvs: W,
+}
+
+impl<W: Writeable> Borrow<PublicKey> for PublicKeyWithTlvs<W> {
+	fn borrow(&self) -> &PublicKey {
+		&self.pubkey
+	}
+}
+
+pub(crate) fn construct_blinded_hops<'a, T, I, W>(
+	secp_ctx: &Secp256k1<T>, unblinded_path: I, session_priv: &SecretKey,
 ) -> Result<Vec<BlindedHop>, secp256k1::Error>
 where
 	T: secp256k1::Signing + secp256k1::Verification,
-	I1: Iterator<Item=PublicKey>,
-	I2: Iterator,
-	I2::Item: Writeable
+	I: Iterator<Item=(PublicKey, W)>,
+	W: Writeable
 {
-	let mut blinded_hops = Vec::with_capacity(unblinded_pks.size_hint().0);
+	let mut blinded_hops = Vec::with_capacity(unblinded_path.size_hint().0);
 	construct_keys_callback_for_blinded_path(
-		secp_ctx, unblinded_pks, session_priv,
-		|blinded_node_id, _, _, encrypted_payload_rho, _, _| {
+		secp_ctx, unblinded_path.map(|(pubkey, tlvs)| PublicKeyWithTlvs { pubkey, tlvs }), session_priv,
+		|blinded_node_id, _, _, encrypted_payload_rho, unblinded_hop_data, _| {
 			blinded_hops.push(BlindedHop {
 				blinded_node_id,
-				encrypted_payload: encrypt_payload(unblinded_tlvs.next().unwrap(), encrypted_payload_rho),
+				encrypted_payload: encrypt_payload(unblinded_hop_data.unwrap().tlvs, encrypted_payload_rho),
 			});
 		})?;
 	Ok(blinded_hops)

lightning/src/ln/blinded_payment_tests.rs

@@ -1392,19 +1392,17 @@ fn route_blinding_spec_test_vector() {
 	let blinding_override = PublicKey::from_secret_key(&secp_ctx, &dave_eve_session_priv);
 	assert_eq!(blinding_override, pubkey_from_hex("031b84c5567b126440995d3ed5aaba0565d71e1834604819ff9c17f5e9d5dd078f"));
 	// Can't use the public API here as the encrypted payloads contain unknown TLVs.
+	let path = [(dave_node_id, WithoutLength(&dave_unblinded_tlvs)), (eve_node_id, WithoutLength(&eve_unblinded_tlvs))];
 	let mut dave_eve_blinded_hops = blinded_path::utils::construct_blinded_hops(
-		&secp_ctx, [dave_node_id, eve_node_id].iter(),
-		&mut [WithoutLength(&dave_unblinded_tlvs), WithoutLength(&eve_unblinded_tlvs)].iter(),
-		&dave_eve_session_priv
+		&secp_ctx, path.into_iter(), &dave_eve_session_priv
 	).unwrap();
 
 	// Concatenate an additional Bob -> Carol blinded path to the Eve -> Dave blinded path.
 	let bob_carol_session_priv = secret_from_hex("0202020202020202020202020202020202020202020202020202020202020202");
 	let bob_blinding_point = PublicKey::from_secret_key(&secp_ctx, &bob_carol_session_priv);
+	let path = [(bob_node_id, WithoutLength(&bob_unblinded_tlvs)), (carol_node_id, WithoutLength(&carol_unblinded_tlvs))];
 	let bob_carol_blinded_hops = blinded_path::utils::construct_blinded_hops(
-		&secp_ctx, [bob_node_id, carol_node_id].iter(),
-		&mut [WithoutLength(&bob_unblinded_tlvs), WithoutLength(&carol_unblinded_tlvs)].iter(),
-		&bob_carol_session_priv
+		&secp_ctx, path.into_iter(), &bob_carol_session_priv
 	).unwrap();
 
 	let mut blinded_hops = bob_carol_blinded_hops;