Refactor BlindedPath construction utils

lightning/src/blinded_path/message.rs

@@ -374,18 +374,20 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 	secp_ctx: &Secp256k1<T>, intermediate_nodes: &[ForwardNode], recipient_node_id: PublicKey,
 	context: MessageContext, session_priv: &SecretKey
 ) -> Result<Vec<BlindedHop>, secp256k1::Error> {
-	let pks = intermediate_nodes.iter().map(|node| &node.node_id)
-		.chain(core::iter::once(&recipient_node_id));
+	let pks = intermediate_nodes.iter().map(|node| node.node_id)
+		.chain(core::iter::once(recipient_node_id));
 	let tlvs = pks.clone()
 		.skip(1) // The first node's TLVs contains the next node's pubkey
 		.zip(intermediate_nodes.iter().map(|node| node.short_channel_id))
 		.map(|(pubkey, scid)| match scid {
 			Some(scid) => NextMessageHop::ShortChannelId(scid),
-			None => NextMessageHop::NodeId(*pubkey),
+			None => NextMessageHop::NodeId(pubkey),
 		})
 		.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

@@ -462,11 +462,14 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 	secp_ctx: &Secp256k1<T>, intermediate_nodes: &[ForwardNode],
 	payee_node_id: PublicKey, payee_tlvs: ReceiveTlvs, session_priv: &SecretKey
 ) -> Result<Vec<BlindedHop>, secp256k1::Error> {
-	let pks = intermediate_nodes.iter().map(|node| &node.node_id)
-		.chain(core::iter::once(&payee_node_id));
+	let pks = intermediate_nodes.iter().map(|node| node.node_id)
+		.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

@@ -25,42 +25,38 @@ use crate::util::ser::{Readable, Writeable};
 
 use crate::io;
 
+use core::borrow::Borrow;
+
 #[allow(unused_imports)]
 use crate::prelude::*;
 
 // TODO: DRY with onion_utils::construct_onion_keys_callback
-#[inline]
-pub(crate) fn construct_keys_callback<'a, T, I, F>(
-	secp_ctx: &Secp256k1<T>, unblinded_path: I, destination: Option<Destination>,
-	session_priv: &SecretKey, mut callback: F
-) -> Result<(), secp256k1::Error>
-where
-	T: secp256k1::Signing + secp256k1::Verification,
-	I: Iterator<Item=&'a PublicKey>,
-	F: FnMut(PublicKey, SharedSecret, PublicKey, [u8; 32], Option<PublicKey>, Option<Vec<u8>>),
+macro_rules! build_keys_helper {
+	($session_priv: ident, $secp_ctx: ident, $callback: ident) =>
 {
-	let mut msg_blinding_point_priv = session_priv.clone();
-	let mut msg_blinding_point = PublicKey::from_secret_key(secp_ctx, &msg_blinding_point_priv);
+	let mut msg_blinding_point_priv = $session_priv.clone();
+	let mut msg_blinding_point = PublicKey::from_secret_key($secp_ctx, &msg_blinding_point_priv);
 	let mut onion_packet_pubkey_priv = msg_blinding_point_priv.clone();
 	let mut onion_packet_pubkey = msg_blinding_point.clone();
 
 	macro_rules! build_keys {
-		($pk: expr, $blinded: expr, $encrypted_payload: expr) => {{
-			let encrypted_data_ss = SharedSecret::new(&$pk, &msg_blinding_point_priv);
+		($hop: expr, $blinded: expr, $encrypted_payload: expr) => {{
+			let pk = *$hop.borrow();
+			let encrypted_data_ss = SharedSecret::new(&pk, &msg_blinding_point_priv);
 
-			let blinded_hop_pk = if $blinded { $pk } else {
+			let blinded_hop_pk = if $blinded { pk } else {
 				let hop_pk_blinding_factor = {
 					let mut hmac = HmacEngine::<Sha256>::new(b"blinded_node_id");
 					hmac.input(encrypted_data_ss.as_ref());
 					Hmac::from_engine(hmac).to_byte_array()
 				};
-				$pk.mul_tweak(secp_ctx, &Scalar::from_be_bytes(hop_pk_blinding_factor).unwrap())?
+				pk.mul_tweak($secp_ctx, &Scalar::from_be_bytes(hop_pk_blinding_factor).unwrap())?
 			};
 			let onion_packet_ss = SharedSecret::new(&blinded_hop_pk, &onion_packet_pubkey_priv);
 
 			let rho = onion_utils::gen_rho_from_shared_secret(encrypted_data_ss.as_ref());
-			let unblinded_pk_opt = if $blinded { None } else { Some($pk) };
-			callback(blinded_hop_pk, onion_packet_ss, onion_packet_pubkey, rho, unblinded_pk_opt, $encrypted_payload);
+			let unblinded_hop_opt = if $blinded { None } else { Some($hop) };
+			$callback(blinded_hop_pk, onion_packet_ss, onion_packet_pubkey, rho, unblinded_hop_opt, $encrypted_payload);
 			(encrypted_data_ss, onion_packet_ss)
 		}}
 	}
@@ -77,7 +73,7 @@ where
 			};
 
 			msg_blinding_point_priv = msg_blinding_point_priv.mul_tweak(&Scalar::from_be_bytes(msg_blinding_point_blinding_factor).unwrap())?;
-			msg_blinding_point = PublicKey::from_secret_key(secp_ctx, &msg_blinding_point_priv);
+			msg_blinding_point = PublicKey::from_secret_key($secp_ctx, &msg_blinding_point_priv);
 
 			let onion_packet_pubkey_blinding_factor = {
 				let mut sha = Sha256::engine();
@@ -86,45 +82,83 @@ where
 				Sha256::from_engine(sha).to_byte_array()
 			};
 			onion_packet_pubkey_priv = onion_packet_pubkey_priv.mul_tweak(&Scalar::from_be_bytes(onion_packet_pubkey_blinding_factor).unwrap())?;
-			onion_packet_pubkey = PublicKey::from_secret_key(secp_ctx, &onion_packet_pubkey_priv);
+			onion_packet_pubkey = PublicKey::from_secret_key($secp_ctx, &onion_packet_pubkey_priv);
 		};
 	}
+}}
+
+#[inline]
+pub(crate) fn construct_keys_for_onion_message<'a, T, I, F>(
+	secp_ctx: &Secp256k1<T>, unblinded_path: I, destination: Destination, 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>>),
+{
+	build_keys_helper!(session_priv, secp_ctx, callback);
 
 	for pk in unblinded_path {
-		build_keys_in_loop!(*pk, false, None);
+		build_keys_in_loop!(pk, false, None);
 	}
-	if let Some(dest) = destination {
-		match dest {
-			Destination::Node(pk) => {
-				build_keys!(pk, false, None);
-			},
-			Destination::BlindedPath(BlindedMessagePath(BlindedPath { blinded_hops, .. })) => {
-				for hop in blinded_hops {
-					build_keys_in_loop!(hop.blinded_node_id, true, Some(hop.encrypted_payload));
-				}
-			},
-		}
+	match destination {
+		Destination::Node(pk) => {
+			build_keys!(pk, false, None);
+		},
+		Destination::BlindedPath(BlindedMessagePath(BlindedPath { blinded_hops, .. })) => {
+			for hop in blinded_hops {
+				build_keys_in_loop!(hop.blinded_node_id, true, Some(hop.encrypted_payload));
+			}
+		},
+	}
+	Ok(())
+}
+
+#[inline]
+pub(super) fn construct_keys_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,
+	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);
+
+	for pk in unblinded_path {
+		build_keys_in_loop!(pk, false, None);
 	}
 	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=&'a 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);
-	construct_keys_callback(
-		secp_ctx, unblinded_pks, None, session_priv,
-		|blinded_node_id, _, _, encrypted_payload_rho, _, _| {
+	let mut blinded_hops = Vec::with_capacity(unblinded_path.size_hint().0);
+	construct_keys_for_blinded_path(
+		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;

lightning/src/onion_message/messenger.rs

@@ -925,7 +925,7 @@ where
 		}
 	};
 	let (packet_payloads, packet_keys) = packet_payloads_and_keys(
-		&secp_ctx, &intermediate_nodes, destination, contents, reply_path, &blinding_secret
+		&secp_ctx, intermediate_nodes, destination, contents, reply_path, &blinding_secret
 	)?;
 
 	let prng_seed = entropy_source.get_secure_random_bytes();
@@ -1784,7 +1784,7 @@ pub type SimpleRefOnionMessenger<
 /// Construct onion packet payloads and keys for sending an onion message along the given
 /// `unblinded_path` to the given `destination`.
 fn packet_payloads_and_keys<T: OnionMessageContents, S: secp256k1::Signing + secp256k1::Verification>(
-	secp_ctx: &Secp256k1<S>, unblinded_path: &[PublicKey], destination: Destination, message: T,
+	secp_ctx: &Secp256k1<S>, unblinded_path: Vec<PublicKey>, destination: Destination, message: T,
 	mut reply_path: Option<BlindedMessagePath>, session_priv: &SecretKey
 ) -> Result<(Vec<(Payload<T>, [u8; 32])>, Vec<onion_utils::OnionKeys>), SendError> {
 	let num_hops = unblinded_path.len() + destination.num_hops();
@@ -1809,7 +1809,8 @@ fn packet_payloads_and_keys<T: OnionMessageContents, S: secp256k1::Signing + sec
 	let mut blinded_path_idx = 0;
 	let mut prev_control_tlvs_ss = None;
 	let mut final_control_tlvs = None;
-	utils::construct_keys_callback(secp_ctx, unblinded_path.iter(), Some(destination), session_priv,
+	utils::construct_keys_for_onion_message(
+		secp_ctx, unblinded_path.into_iter(), destination, session_priv,
 		|_, onion_packet_ss, ephemeral_pubkey, control_tlvs_ss, unblinded_pk_opt, enc_payload_opt| {
 			if num_unblinded_hops != 0 && unblinded_path_idx < num_unblinded_hops {
 				if let Some(ss) = prev_control_tlvs_ss.take() {