f break BlindedPaymentTlv variants out into structs

Author: valentinewallace

lightning/src/blinded_path/mod.rs

@@ -15,7 +15,7 @@ pub(crate) mod utils;
 
 use bitcoin::secp256k1::{self, PublicKey, Secp256k1, SecretKey};
 
-use crate::blinded_path::payment::BlindedPaymentTlvs;
+use crate::blinded_path::payment::{ForwardTlvs, ReceiveTlvs};
 use crate::sign::EntropySource;
 use crate::ln::msgs::DecodeError;
 use crate::util::ser::{Readable, Writeable, Writer};
@@ -82,16 +82,18 @@ impl BlindedPath {
 	/// Errors if `path` is empty or a node id in `path` is invalid.
 	//  TODO: make all payloads the same size with padding + add dummy hops
 	pub fn new_for_payment<ES: EntropySource, T: secp256k1::Signing + secp256k1::Verification>(
-		path: &[(PublicKey, BlindedPaymentTlvs)], entropy_source: &ES, secp_ctx: &Secp256k1<T>
+		intermediate_nodes: &[(PublicKey, ForwardTlvs)], payee_node_id: PublicKey,
+		payee_tlvs: ReceiveTlvs, entropy_source: &ES, secp_ctx: &Secp256k1<T>
 	) -> Result<Self, ()> {
-		if path.len() < 1 { return Err(()) }
 		let blinding_secret_bytes = entropy_source.get_secure_random_bytes();
 		let blinding_secret = SecretKey::from_slice(&blinding_secret_bytes[..]).expect("RNG is busted");
 
 		Ok(BlindedPath {
-			introduction_node_id: path[0].0,
+			introduction_node_id: intermediate_nodes.first().map_or(payee_node_id, |n| n.0),
 			blinding_point: PublicKey::from_secret_key(secp_ctx, &blinding_secret),
-			blinded_hops: payment::blinded_hops(secp_ctx, path, &blinding_secret).map_err(|_| ())?,
+			blinded_hops: payment::blinded_hops(
+				secp_ctx, intermediate_nodes, payee_node_id, payee_tlvs, &blinding_secret
+			).map_err(|_| ())?,
 		})
 	}
 }

lightning/src/blinded_path/payment.rs

@@ -13,31 +13,43 @@ use crate::ln::msgs::DecodeError;
 use crate::prelude::*;
 use crate::util::ser::{Readable, Writeable, Writer};
 
+/// Data to construct a [`BlindedHop`] for forwarding a payment.
+pub struct ForwardTlvs {
+	/// The short channel id this payment is being forwarded over.
+	short_channel_id: u64,
+	/// Payment parameters for relaying over this channel.
+	payment_relay: PaymentRelay,
+	/// Payment constraints when relaying over this channel.
+	payment_constraints: PaymentConstraints,
+	/// Supported and required features when relaying a payment onion containing this object's
+	/// corresponding [`BlindedHop`].
+	///
+	/// [`BlindedHop`]: crate::blinded_path::BlindedHop
+	features: BlindedHopFeatures,
+}
+
+/// Data to construct a [`BlindedHop`] for receiving a payment.
+pub struct ReceiveTlvs {
+	/// Used to authenticate the sender of a payment to the receiver and tie MPP HTLCs together.
+	payment_secret: PaymentSecret,
+	/// Constraints for the receiver of this payment.
+	payment_constraints: PaymentConstraints,
+}
+
 /// Data to construct a [`BlindedHop`] for sending a payment over.
 ///
 /// [`BlindedHop`]: crate::blinded_path::BlindedHop
-pub enum BlindedPaymentTlvs {
+pub(crate) enum BlindedPaymentTlvs {
 	/// This blinded payment data is to be forwarded.
-	Forward {
-		/// The short channel id this payment is being forwarded over.
-		short_channel_id: u64,
-		/// Payment parameters for relaying over this channel.
-		payment_relay: PaymentRelay,
-		/// Payment constraints when relaying over this channel.
-		payment_constraints: PaymentConstraints,
-		/// Supported and required features when relaying a payment onion containing this object's
-		/// corresponding [`BlindedHop`].
-		///
-		/// [`BlindedHop`]: crate::blinded_path::BlindedHop
-		features: BlindedHopFeatures,
-	},
+	Forward(ForwardTlvs),
 	/// This blinded payment data is to be received.
-	Receive {
-		/// Used to authenticate the sender of a payment to the receiver and tie MPP HTLCs together.
-		payment_secret: PaymentSecret,
-		/// Constraints for the receiver of this payment.
-		payment_constraints: PaymentConstraints,
-	},
+	Receive(ReceiveTlvs),
+}
+
+// Used to include forward and receive TLVs in the same iterator for encoding.
+enum BlindedPaymentTlvsRef<'a> {
+	Forward(&'a ForwardTlvs),
+	Receive(&'a ReceiveTlvs),
 }
 
 /// Parameters for relaying over a given [`BlindedHop`].
@@ -72,24 +84,24 @@ pub struct PaymentConstraints {
 	pub htlc_minimum_msat: u64,
 }
 
-impl Writeable for BlindedPaymentTlvs {
+impl_writeable_tlv_based!(ForwardTlvs, {
+	(2, short_channel_id, required),
+	(10, payment_relay, required),
+	(12, payment_constraints, required),
+	(14, features, required),
+});
+
+impl_writeable_tlv_based!(ReceiveTlvs, {
+	(6, payment_secret, required),
+	(12, payment_constraints, required),
+});
+
+impl<'a> Writeable for BlindedPaymentTlvsRef<'a> {
 	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
 		// TODO: write padding
 		match self {
-			Self::Forward { short_channel_id, payment_relay, payment_constraints, features } => {
-				encode_tlv_stream!(w, {
-					(2, short_channel_id, required),
-					(10, payment_relay, required),
-					(12, payment_constraints, required),
-					(14, features, required),
-				});
-			},
-			Self::Receive { payment_secret, payment_constraints } => {
-				encode_tlv_stream!(w, {
-					(6, payment_secret, required),
-					(12, payment_constraints, required),
-				});
-			}
+			Self::Forward(tlvs) => tlvs.write(w)?,
+			Self::Receive(tlvs) => tlvs.write(w)?,
 		}
 		Ok(())
 	}
@@ -107,28 +119,32 @@ impl Readable for BlindedPaymentTlvs {
 		});
 		if let Some(short_channel_id) = scid {
 			if payment_secret.is_some() { return Err(DecodeError::InvalidValue) }
-			Ok(BlindedPaymentTlvs::Forward {
+			Ok(BlindedPaymentTlvs::Forward(ForwardTlvs {
 				short_channel_id,
 				payment_relay: payment_relay.ok_or(DecodeError::InvalidValue)?,
 				payment_constraints: payment_constraints.0.unwrap(),
 				features: features.ok_or(DecodeError::InvalidValue)?,
-			})
+			}))
 		} else {
 			if payment_relay.is_some() | features.is_some() { return Err(DecodeError::InvalidValue) }
-			Ok(BlindedPaymentTlvs::Receive {
+			Ok(BlindedPaymentTlvs::Receive(ReceiveTlvs {
 				payment_secret: payment_secret.ok_or(DecodeError::InvalidValue)?,
 				payment_constraints: payment_constraints.0.unwrap(),
-			})
+			}))
 		}
 	}
 }
 
-/// Construct blinded payment hops for the given `unblinded_path`.
+/// Construct blinded payment hops for the given `intermediate_nodes` and payee info.
 pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
-	secp_ctx: &Secp256k1<T>, unblinded_path: &[(PublicKey, BlindedPaymentTlvs)], session_priv: &SecretKey
+	secp_ctx: &Secp256k1<T>, intermediate_nodes: &[(PublicKey, ForwardTlvs)],
+	payee_node_id: PublicKey, payee_tlvs: ReceiveTlvs, session_priv: &SecretKey
 ) -> Result<Vec<BlindedHop>, secp256k1::Error> {
-	utils::construct_blinded_hops(secp_ctx, unblinded_path.iter().map(|(pk, _)| pk),
-		unblinded_path.iter().map(|(_, tlvs)| tlvs), session_priv)
+	let pks = intermediate_nodes.iter().map(|(pk, _)| pk)
+		.chain(core::iter::once(&payee_node_id));
+	let tlvs = intermediate_nodes.iter().map(|(_, tlvs)| BlindedPaymentTlvsRef::Forward(tlvs))
+		.chain(core::iter::once(BlindedPaymentTlvsRef::Receive(&payee_tlvs)));
+	utils::construct_blinded_hops(secp_ctx, pks, tlvs, session_priv)
 }
 
 impl_writeable_msg!(PaymentRelay, {

lightning/src/blinded_path/utils.rs

@@ -111,11 +111,11 @@ pub(super) fn construct_blinded_hops<'a, T, I1, I2>(
 ) -> Result<Vec<BlindedHop>, secp256k1::Error>
 where
 	T: secp256k1::Signing + secp256k1::Verification,
-	I1: ExactSizeIterator<Item=&'a PublicKey>,
+	I1: Iterator<Item=&'a PublicKey>,
 	I2: Iterator,
 	I2::Item: Writeable
 {
-	let mut blinded_hops = Vec::with_capacity(unblinded_pks.len());
+	let mut blinded_hops = Vec::with_capacity(unblinded_pks.size_hint().0);
 	let mut curr_hop_idx = 0;
 	construct_keys_callback(
 		secp_ctx, unblinded_pks, None, session_priv,