From 7c1726b585af71897e15dcb83024d5cab55d4fcd Mon Sep 17 00:00:00 2001
From: Valentine Wallace <vwallace@protonmail.com>
Date: Fri, 16 Jun 2023 14:40:28 -0400
Subject: [PATCH] Update blinded path util to take iterator instead of slice

Useful for blinded payment path construction.
---
 lightning/src/blinded_path/message.rs    | 31 +++++-----
 lightning/src/blinded_path/utils.rs      | 14 +++--
 lightning/src/onion_message/messenger.rs | 76 ++++++++++++------------
 3 files changed, 64 insertions(+), 57 deletions(-)

diff --git a/lightning/src/blinded_path/message.rs b/lightning/src/blinded_path/message.rs
index 2549673b0..ca0cb4766 100644
--- a/lightning/src/blinded_path/message.rs
+++ b/lightning/src/blinded_path/message.rs
@@ -60,21 +60,22 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 	let mut blinded_hops = Vec::with_capacity(unblinded_path.len());
 
 	let mut prev_ss_and_blinded_node_id = None;
-	utils::construct_keys_callback(secp_ctx, unblinded_path, None, session_priv, |blinded_node_id, _, _, encrypted_payload_ss, unblinded_pk, _| {
-		if let Some((prev_ss, prev_blinded_node_id)) = prev_ss_and_blinded_node_id {
-			if let Some(pk) = unblinded_pk {
-				let payload = ForwardTlvs {
-					next_node_id: pk,
-					next_blinding_override: None,
-				};
-				blinded_hops.push(BlindedHop {
-					blinded_node_id: prev_blinded_node_id,
-					encrypted_payload: utils::encrypt_payload(payload, prev_ss),
-				});
-			} else { debug_assert!(false); }
-		}
-		prev_ss_and_blinded_node_id = Some((encrypted_payload_ss, blinded_node_id));
-	})?;
+	utils::construct_keys_callback(secp_ctx, unblinded_path.iter(), None, session_priv,
+		|blinded_node_id, _, _, encrypted_payload_ss, unblinded_pk, _| {
+			if let Some((prev_ss, prev_blinded_node_id)) = prev_ss_and_blinded_node_id {
+				if let Some(pk) = unblinded_pk {
+					let payload = ForwardTlvs {
+						next_node_id: pk,
+						next_blinding_override: None,
+					};
+					blinded_hops.push(BlindedHop {
+						blinded_node_id: prev_blinded_node_id,
+						encrypted_payload: utils::encrypt_payload(payload, prev_ss),
+					});
+				} else { debug_assert!(false); }
+			}
+			prev_ss_and_blinded_node_id = Some((encrypted_payload_ss, blinded_node_id));
+		})?;
 
 	if let Some((final_ss, final_blinded_node_id)) = prev_ss_and_blinded_node_id {
 		let final_payload = ReceiveTlvs { path_id: None };
diff --git a/lightning/src/blinded_path/utils.rs b/lightning/src/blinded_path/utils.rs
index a188274a5..1ac651945 100644
--- a/lightning/src/blinded_path/utils.rs
+++ b/lightning/src/blinded_path/utils.rs
@@ -27,11 +27,15 @@ use crate::prelude::*;
 
 // TODO: DRY with onion_utils::construct_onion_keys_callback
 #[inline]
-pub(crate) fn construct_keys_callback<T: secp256k1::Signing + secp256k1::Verification,
-	FType: FnMut(PublicKey, SharedSecret, PublicKey, [u8; 32], Option<PublicKey>, Option<Vec<u8>>)>(
-	secp_ctx: &Secp256k1<T>, unblinded_path: &[PublicKey], destination: Option<Destination>,
-	session_priv: &SecretKey, mut callback: FType
-) -> Result<(), secp256k1::Error> {
+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>>),
+{
 	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();
diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs
index 09f207e36..1a6e2614b 100644
--- a/lightning/src/onion_message/messenger.rs
+++ b/lightning/src/onion_message/messenger.rs
@@ -650,46 +650,48 @@ fn packet_payloads_and_keys<T: CustomOnionMessageContents, S: secp256k1::Signing
 	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, Some(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() {
-				payloads.push((Payload::Forward(ForwardControlTlvs::Unblinded(
-					ForwardTlvs {
-						next_node_id: unblinded_pk_opt.unwrap(),
-						next_blinding_override: None,
-					}
-				)), ss));
+	utils::construct_keys_callback(secp_ctx, unblinded_path.iter(), Some(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() {
+					payloads.push((Payload::Forward(ForwardControlTlvs::Unblinded(
+						ForwardTlvs {
+							next_node_id: unblinded_pk_opt.unwrap(),
+							next_blinding_override: None,
+						}
+					)), ss));
+				}
+				prev_control_tlvs_ss = Some(control_tlvs_ss);
+				unblinded_path_idx += 1;
+			} else if let Some((intro_node_id, blinding_pt)) = intro_node_id_blinding_pt.take() {
+				if let Some(control_tlvs_ss) = prev_control_tlvs_ss.take() {
+					payloads.push((Payload::Forward(ForwardControlTlvs::Unblinded(ForwardTlvs {
+						next_node_id: intro_node_id,
+						next_blinding_override: Some(blinding_pt),
+					})), control_tlvs_ss));
+				}
 			}
-			prev_control_tlvs_ss = Some(control_tlvs_ss);
-			unblinded_path_idx += 1;
-		} else if let Some((intro_node_id, blinding_pt)) = intro_node_id_blinding_pt.take() {
-			if let Some(control_tlvs_ss) = prev_control_tlvs_ss.take() {
-				payloads.push((Payload::Forward(ForwardControlTlvs::Unblinded(ForwardTlvs {
-					next_node_id: intro_node_id,
-					next_blinding_override: Some(blinding_pt),
-				})), control_tlvs_ss));
+			if blinded_path_idx < num_blinded_hops.saturating_sub(1) && enc_payload_opt.is_some() {
+				payloads.push((Payload::Forward(ForwardControlTlvs::Blinded(enc_payload_opt.unwrap())),
+					control_tlvs_ss));
+				blinded_path_idx += 1;
+			} else if let Some(encrypted_payload) = enc_payload_opt {
+				final_control_tlvs = Some(ReceiveControlTlvs::Blinded(encrypted_payload));
+				prev_control_tlvs_ss = Some(control_tlvs_ss);
 			}
-		}
-		if blinded_path_idx < num_blinded_hops.saturating_sub(1) && enc_payload_opt.is_some() {
-			payloads.push((Payload::Forward(ForwardControlTlvs::Blinded(enc_payload_opt.unwrap())),
-				control_tlvs_ss));
-			blinded_path_idx += 1;
-		} else if let Some(encrypted_payload) = enc_payload_opt {
-			final_control_tlvs = Some(ReceiveControlTlvs::Blinded(encrypted_payload));
-			prev_control_tlvs_ss = Some(control_tlvs_ss);
-		}
 
-		let (rho, mu) = onion_utils::gen_rho_mu_from_shared_secret(onion_packet_ss.as_ref());
-		onion_packet_keys.push(onion_utils::OnionKeys {
-			#[cfg(test)]
-			shared_secret: onion_packet_ss,
-			#[cfg(test)]
-			blinding_factor: [0; 32],
-			ephemeral_pubkey,
-			rho,
-			mu,
-		});
-	})?;
+			let (rho, mu) = onion_utils::gen_rho_mu_from_shared_secret(onion_packet_ss.as_ref());
+			onion_packet_keys.push(onion_utils::OnionKeys {
+				#[cfg(test)]
+				shared_secret: onion_packet_ss,
+				#[cfg(test)]
+				blinding_factor: [0; 32],
+				ephemeral_pubkey,
+				rho,
+				mu,
+			});
+		}
+	)?;
 
 	if let Some(control_tlvs) = final_control_tlvs {
 		payloads.push((Payload::Receive {
-- 
2.39.5