From d224f980ed208445d2425f697eb4563aaa432fa3 Mon Sep 17 00:00:00 2001
From: Valentine Wallace <vwallace@protonmail.com>
Date: Wed, 9 Aug 2023 14:07:58 -0700
Subject: [PATCH] Simplify onion message blinded hop construction

Also adds a util for general blinded hop creation to be reused for blinded
payment paths.
---
 lightning/src/blinded_path/message.rs | 35 ++++++---------------------
 lightning/src/blinded_path/utils.rs   | 26 ++++++++++++++++++--
 lightning/src/onion_message/packet.rs | 12 ++++++++-
 3 files changed, 42 insertions(+), 31 deletions(-)

diff --git a/lightning/src/blinded_path/message.rs b/lightning/src/blinded_path/message.rs
index ca0cb4766..d2e81444e 100644
--- a/lightning/src/blinded_path/message.rs
+++ b/lightning/src/blinded_path/message.rs
@@ -57,35 +57,14 @@ impl Writeable for ReceiveTlvs {
 pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 	secp_ctx: &Secp256k1<T>, unblinded_path: &[PublicKey], session_priv: &SecretKey
 ) -> Result<Vec<BlindedHop>, secp256k1::Error> {
-	let mut blinded_hops = Vec::with_capacity(unblinded_path.len());
+	let blinded_tlvs = unblinded_path.iter()
+		.skip(1) // The first node's TLVs contains the next node's pubkey
+		.map(|pk| {
+			ControlTlvs::Forward(ForwardTlvs { next_node_id: *pk, next_blinding_override: None })
+		})
+		.chain(core::iter::once(ControlTlvs::Receive(ReceiveTlvs { path_id: None })));
 
-	let mut prev_ss_and_blinded_node_id = None;
-	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 };
-		blinded_hops.push(BlindedHop {
-			blinded_node_id: final_blinded_node_id,
-			encrypted_payload: utils::encrypt_payload(final_payload, final_ss),
-		});
-	} else { debug_assert!(false) }
-
-	Ok(blinded_hops)
+	utils::construct_blinded_hops(secp_ctx, unblinded_path.iter(), blinded_tlvs, session_priv)
 }
 
 // Advance the blinded onion message path by one hop, so make the second hop into the new
diff --git a/lightning/src/blinded_path/utils.rs b/lightning/src/blinded_path/utils.rs
index 1ac651945..9b1ce50b1 100644
--- a/lightning/src/blinded_path/utils.rs
+++ b/lightning/src/blinded_path/utils.rs
@@ -15,7 +15,7 @@ use bitcoin::hashes::sha256::Hash as Sha256;
 use bitcoin::secp256k1::{self, PublicKey, Secp256k1, SecretKey, Scalar};
 use bitcoin::secp256k1::ecdh::SharedSecret;
 
-use super::BlindedPath;
+use super::{BlindedHop, BlindedPath};
 use crate::ln::msgs::DecodeError;
 use crate::ln::onion_utils;
 use crate::onion_message::Destination;
@@ -105,8 +105,30 @@ where
 	Ok(())
 }
 
+// Panics if `unblinded_tlvs` length is less than `unblinded_pks` length
+pub(super) fn construct_blinded_hops<'a, T, I1, I2>(
+	secp_ctx: &Secp256k1<T>, unblinded_pks: I1, mut unblinded_tlvs: I2, session_priv: &SecretKey
+) -> Result<Vec<BlindedHop>, secp256k1::Error>
+where
+	T: secp256k1::Signing + secp256k1::Verification,
+	I1: ExactSizeIterator<Item=&'a PublicKey>,
+	I2: Iterator,
+	I2::Item: Writeable
+{
+	let mut blinded_hops = Vec::with_capacity(unblinded_pks.len());
+	construct_keys_callback(
+		secp_ctx, unblinded_pks, None, session_priv,
+		|blinded_node_id, _, _, encrypted_payload_rho, _, _| {
+			blinded_hops.push(BlindedHop {
+				blinded_node_id,
+				encrypted_payload: encrypt_payload(unblinded_tlvs.next().unwrap(), encrypted_payload_rho),
+			});
+		})?;
+	Ok(blinded_hops)
+}
+
 /// Encrypt TLV payload to be used as a [`crate::blinded_path::BlindedHop::encrypted_payload`].
-pub(super) fn encrypt_payload<P: Writeable>(payload: P, encrypted_tlvs_ss: [u8; 32]) -> Vec<u8> {
+fn encrypt_payload<P: Writeable>(payload: P, encrypted_tlvs_ss: [u8; 32]) -> Vec<u8> {
 	let mut writer = VecWriter(Vec::new());
 	let write_adapter = ChaChaPolyWriteAdapter::new(encrypted_tlvs_ss, &payload);
 	write_adapter.write(&mut writer).expect("In-memory writes cannot fail");
diff --git a/lightning/src/onion_message/packet.rs b/lightning/src/onion_message/packet.rs
index 8896941c8..9eb2e21c2 100644
--- a/lightning/src/onion_message/packet.rs
+++ b/lightning/src/onion_message/packet.rs
@@ -265,7 +265,8 @@ ReadableArgs<(SharedSecret, &H, &L)> for Payload<<H as CustomOnionMessageHandler
 
 /// When reading a packet off the wire, we don't know a priori whether the packet is to be forwarded
 /// or received. Thus we read a ControlTlvs rather than reading a ForwardControlTlvs or
-/// ReceiveControlTlvs directly.
+/// ReceiveControlTlvs directly. Also useful on the encoding side to keep forward and receive TLVs
+/// in the same iterator.
 pub(crate) enum ControlTlvs {
 	/// This onion message is intended to be forwarded.
 	Forward(ForwardTlvs),
@@ -304,3 +305,12 @@ impl Readable for ControlTlvs {
 		Ok(payload_fmt)
 	}
 }
+
+impl Writeable for ControlTlvs {
+	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
+		match self {
+			Self::Forward(tlvs) => tlvs.write(w),
+			Self::Receive(tlvs) => tlvs.write(w),
+		}
+	}
+}
-- 
2.39.5