Simplify onion message blinded hop construction

Also adds a util for general blinded hop creation to be reused for blinded
payment paths.

diff --git a/lightning/src/blinded_path/message.rs b/lightning/src/blinded_path/message.rs
index ca0cb476655915bcea24aa4cc7db9dc2ec37f168..d2e81444ef6307a2cac41082c444f18568af7a85 100644 (file)
--- 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 1ac6519452c1eb10af3757240d41a046d381efaa..9b1ce50b10ebefa61943680694c5ce5852d20d43 100644 (file)
--- 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 8896941c8753789518944a260e15d3fa3112e6cd..9eb2e21c2657d944d6d591c9c6f542f49df6ce4f 100644 (file)
--- 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),
+               }
+       }
+}