f: Introduce Padding a message tlvs field instead of WithPadding

- WithPadding approach has a crucial flaw: it doesn't enforce that
  the tlv of the tlvs struct is not already utilizing the `(1,...)`
  position, which should be used for padding.
- Introduce padding as a field in the `ForwardTlvs` and `ReceiveTlvs`
  structs to ensure the above condition is met.
- Calculate the padding length prior to padding the TLVs and use that
  data to update the padding field later.

Co-authored-by: Jeffrey Czyz <[email protected]>

Author: shaavan

lightning/src/blinded_path/message.rs

@@ -17,7 +17,7 @@ use bitcoin::secp256k1::{self, PublicKey, Secp256k1, SecretKey};
 use crate::prelude::*;
 
 use crate::blinded_path::{BlindedHop, BlindedPath, IntroductionNode, NextMessageHop, NodeIdLookUp};
-use crate::blinded_path::utils::{self, WithPadding};
+use crate::blinded_path::utils::{self, Padding};
 use crate::io;
 use crate::io::Cursor;
 use crate::ln::channelmanager::PaymentId;
@@ -46,6 +46,8 @@ pub struct ForwardNode {
 /// route, they are encoded into [`BlindedHop::encrypted_payload`].
 #[derive(Clone)]
 pub(crate) struct ForwardTlvs {
+	/// The padding data used to make all packets of a Blinded Path of same size
+	pub(crate) padding: Option<Padding>,
 	/// The next hop in the onion message's path.
 	pub(crate) next_hop: NextMessageHop,
 	/// Senders to a blinded path use this value to concatenate the route they find to the
@@ -56,6 +58,8 @@ pub(crate) struct ForwardTlvs {
 /// Similar to [`ForwardTlvs`], but these TLVs are for the final node.
 #[derive(Clone)]
 pub(crate) struct ReceiveTlvs {
+	/// The padding data used to make all packets of a Blinded Path of same size
+	pub(crate) padding: Option<Padding>,
 	/// If `context` is `Some`, it is used to identify the blinded path that this onion message is
 	/// sending to. This is useful for receivers to check that said blinded path is being used in
 	/// the right context.
@@ -69,6 +73,7 @@ impl Writeable for ForwardTlvs {
 			NextMessageHop::ShortChannelId(scid) => (None, Some(scid)),
 		};
 		encode_tlv_stream!(writer, {
+			(1, self.padding, option),
 			(2, short_channel_id, option),
 			(4, next_node_id, option),
 			(8, self.next_blinding_override, option)
@@ -80,6 +85,7 @@ impl Writeable for ForwardTlvs {
 impl Writeable for ReceiveTlvs {
 	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
 		encode_tlv_stream!(writer, {
+			(1, self.padding, option),
 			(65537, self.context, option),
 		});
 		Ok(())
@@ -184,15 +190,15 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 			Some(scid) => NextMessageHop::ShortChannelId(scid),
 			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) })));
+		.map(|next_hop| ControlTlvs::Forward(ForwardTlvs { padding: None, next_hop, next_blinding_override: None }))
+		.chain(core::iter::once(ControlTlvs::Receive(ReceiveTlvs{ padding: None, context: Some(context) })));
 
 	let max_length = tlvs.clone()
 		.map(|tlv| tlv.serialized_length())
 		.max()
 		.unwrap_or(0);
 
-	let length_tlvs = tlvs.map(|tlvs| WithPadding { max_length, tlvs });
+	let length_tlvs = tlvs.map(|tlv| tlv.pad_to_length(max_length));
 
 	utils::construct_blinded_hops(secp_ctx, pks, length_tlvs, session_priv)
 }
@@ -216,7 +222,7 @@ where
 	let mut reader = FixedLengthReader::new(&mut s, encrypted_control_tlvs.len() as u64);
 	match ChaChaPolyReadAdapter::read(&mut reader, rho) {
 		Ok(ChaChaPolyReadAdapter {
-			readable: ControlTlvs::Forward(ForwardTlvs { next_hop, next_blinding_override })
+			readable: ControlTlvs::Forward(ForwardTlvs {padding: _, next_hop, next_blinding_override })
 		}) => {
 			let next_node_id = match next_hop {
 				NextMessageHop::NodeId(pubkey) => pubkey,

lightning/src/blinded_path/payment.rs

@@ -14,7 +14,7 @@
 use bitcoin::secp256k1::{self, PublicKey, Secp256k1, SecretKey};
 
 use crate::blinded_path::{BlindedHop, BlindedPath, IntroductionNode, NodeIdLookUp};
-use crate::blinded_path::utils::{self, WithPadding};
+use crate::blinded_path::utils;
 use crate::crypto::streams::ChaChaPolyReadAdapter;
 use crate::io;
 use crate::io::Cursor;
@@ -279,14 +279,7 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 	let tlvs = intermediate_nodes.iter().map(|node| BlindedPaymentTlvsRef::Forward(&node.tlvs))
 		.chain(core::iter::once(BlindedPaymentTlvsRef::Receive(&payee_tlvs)));
 
-	let max_length = tlvs.clone()
-		.map(|tlv| tlv.serialized_length())
-		.max()
-		.unwrap_or(0);
-
-	let length_tlvs = tlvs.map(|tlvs| WithPadding { max_length, tlvs });
-
-	utils::construct_blinded_hops(secp_ctx, pks, length_tlvs, session_priv)
+	utils::construct_blinded_hops(secp_ctx, pks, tlvs, session_priv)
 }
 
 // Advance the blinded onion payment path by one hop, so make the second hop into the new

lightning/src/blinded_path/utils.rs

@@ -137,7 +137,8 @@ fn encrypt_payload<P: Writeable>(payload: P, encrypted_tlvs_rho: [u8; 32]) -> Ve
 
 /// Represents optional padding for encrypted payloads.
 /// Padding is used to ensure payloads have a consistent length.
-pub(crate) struct Padding {
+#[derive(Clone, Debug)]
+pub struct Padding {
 	length: usize,
 }
 
@@ -177,27 +178,3 @@ impl Writeable for Padding {
 		Ok(())
 	}
 }
-
-
-/// A wrapper struct that stores the largest packet size for a [`BlindedPath`].
-/// This helps us calculate the appropriate padding size for the tlvs when writing them.
-pub(super) struct WithPadding<T: Writeable> {
-	/// Length of the packet with the largest size in the [`BlindedPath`].
-	pub(super) max_length: usize,
-	/// The current packet's TLVs.
-	pub(super) tlvs: T,
-}
-
-impl<T: Writeable> Writeable for WithPadding<T> {
-	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
-		let length = self.max_length.checked_sub(self.tlvs.serialized_length());
-		debug_assert!(length.is_some(), "Size of this packet should not be larger than the size of largest packet.");
-		let padding = Some(Padding::new(length.unwrap()));
-
-		encode_tlv_stream!(writer, {
-			(1, padding, option)
-		});
-
-		self.tlvs.write(writer)
-	}
-}

lightning/src/onion_message/messenger.rs

@@ -950,7 +950,7 @@ where
 		(control_tlvs_ss, custom_handler.deref(), logger.deref())
 	) {
 		Ok((Payload::Receive::<ParsedOnionMessageContents<<<CMH as Deref>::Target as CustomOnionMessageHandler>::CustomMessage>> {
-			message, control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { context }), reply_path,
+			message, control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { padding: _, context }), reply_path,
 		}, None)) => {
 			match (&message, &context) {
 				(_, None) => {
@@ -969,7 +969,7 @@ where
 			}
 		},
 		Ok((Payload::Forward(ForwardControlTlvs::Unblinded(ForwardTlvs {
-			next_hop, next_blinding_override
+			padding: _, next_hop, next_blinding_override
 		})), Some((next_hop_hmac, new_packet_bytes)))) => {
 			// TODO: we need to check whether `next_hop` is our node, in which case this is a dummy
 			// blinded hop and this onion message is destined for us. In this situation, we should keep
@@ -1772,6 +1772,7 @@ fn packet_payloads_and_keys<T: OnionMessageContents, S: secp256k1::Signing + sec
 				if let Some(ss) = prev_control_tlvs_ss.take() {
 					payloads.push((Payload::Forward(ForwardControlTlvs::Unblinded(
 						ForwardTlvs {
+							padding: None,
 							next_hop: NextMessageHop::NodeId(unblinded_pk_opt.unwrap()),
 							next_blinding_override: None,
 						}
@@ -1782,6 +1783,7 @@ fn packet_payloads_and_keys<T: OnionMessageContents, S: secp256k1::Signing + sec
 			} 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 {
+						padding: None,
 						next_hop: NextMessageHop::NodeId(intro_node_id),
 						next_blinding_override: Some(blinding_pt),
 					})), control_tlvs_ss));
@@ -1817,7 +1819,7 @@ fn packet_payloads_and_keys<T: OnionMessageContents, S: secp256k1::Signing + sec
 		}, prev_control_tlvs_ss.unwrap()));
 	} else {
 		payloads.push((Payload::Receive {
-			control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { context: None }),
+			control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { padding: None, context: None }),
 			reply_path: reply_path.take(),
 			message,
 		}, prev_control_tlvs_ss.unwrap()));

lightning/src/onion_message/packet.rs

@@ -314,6 +314,21 @@ pub(crate) enum ControlTlvs {
 	Receive(ReceiveTlvs),
 }
 
+impl ControlTlvs {
+	pub(crate) fn pad_to_length(mut self, length: usize) -> Self {
+		let pad_length = length.checked_sub(self.serialized_length());
+		debug_assert!(pad_length.is_some(), "Size of this packet should not be larger than the size of largest packet.");
+		let padding = Some(Padding::new(pad_length.unwrap()));
+
+		match &mut self {
+			ControlTlvs::Forward(tlvs) => tlvs.padding = padding,
+			ControlTlvs::Receive(tlvs) => tlvs.padding = padding,
+		}
+
+		self
+	}
+}
+
 impl Readable for ControlTlvs {
 	fn read<R: Read>(r: &mut R) -> Result<Self, DecodeError> {
 		_init_and_read_tlv_stream!(r, {
@@ -337,11 +352,13 @@ impl Readable for ControlTlvs {
 
 		let payload_fmt = if valid_fwd_fmt {
 			ControlTlvs::Forward(ForwardTlvs {
+				padding: None,
 				next_hop: next_hop.unwrap(),
 				next_blinding_override,
 			})
 		} else if valid_recv_fmt {
 			ControlTlvs::Receive(ReceiveTlvs {
+				padding: None,
 				context,
 			})
 		} else {