Introduce Option<Padding> for ControlTlvs variants.

- 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;
+use crate::blinded_path::utils::{self, Padding};
 use crate::io;
 use crate::io::Cursor;
 use crate::ln::channelmanager::PaymentId;
@@ -44,7 +44,10 @@ pub struct ForwardNode {
 
 /// TLVs to encode in an intermediate onion message packet's hop data. When provided in a blinded
 /// 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
@@ -53,7 +56,10 @@ 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.
@@ -66,8 +72,8 @@ impl Writeable for ForwardTlvs {
 			NextMessageHop::NodeId(pubkey) => (Some(pubkey), None),
 			NextMessageHop::ShortChannelId(scid) => (None, Some(scid)),
 		};
-		// TODO: write padding
 		encode_tlv_stream!(writer, {
+			(1, self.padding, option),
 			(2, short_channel_id, option),
 			(4, next_node_id, option),
 			(8, self.next_blinding_override, option)
@@ -78,8 +84,8 @@ impl Writeable for ForwardTlvs {
 
 impl Writeable for ReceiveTlvs {
 	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
-		// TODO: write padding
 		encode_tlv_stream!(writer, {
+			(1, self.padding, option),
 			(65537, self.context, option),
 		});
 		Ok(())
@@ -184,10 +190,17 @@ 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) })));
 
-	utils::construct_blinded_hops(secp_ctx, pks, tlvs, session_priv)
+	let max_length = tlvs.clone()
+		.map(|tlv| tlv.serialized_length())
+		.max()
+		.unwrap_or(0);
+
+	let length_tlvs = tlvs.map(|tlv| tlv.pad_to_length(max_length));
+
+	utils::construct_blinded_hops(secp_ctx, pks, length_tlvs, session_priv)
 }
 
 // Advance the blinded onion message path by one hop, so make the second hop into the new
@@ -209,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

@@ -86,6 +86,7 @@ pub(crate) enum BlindedPaymentTlvs {
 }
 
 // Used to include forward and receive TLVs in the same iterator for encoding.
+#[derive(Clone)]
 enum BlindedPaymentTlvsRef<'a> {
 	Forward(&'a ForwardTlvs),
 	Receive(&'a ReceiveTlvs),
@@ -226,7 +227,6 @@ impl Writeable for ReceiveTlvs {
 
 impl<'a> Writeable for BlindedPaymentTlvsRef<'a> {
 	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
-		// TODO: write padding
 		match self {
 			Self::Forward(tlvs) => tlvs.write(w)?,
 			Self::Receive(tlvs) => tlvs.write(w)?,
@@ -278,6 +278,7 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 		.chain(core::iter::once(&payee_node_id));
 	let tlvs = intermediate_nodes.iter().map(|node| BlindedPaymentTlvsRef::Forward(&node.tlvs))
 		.chain(core::iter::once(BlindedPaymentTlvsRef::Receive(&payee_tlvs)));
+
 	utils::construct_blinded_hops(secp_ctx, pks, tlvs, session_priv)
 }
 

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,
 }
 

lightning/src/ln/blinded_payment_tests.rs

@@ -289,7 +289,7 @@ fn do_forward_checks_failure(check: ForwardCheckFail, intro_fails: bool) {
 	let mut route_params = get_blinded_route_parameters(amt_msat, payment_secret, 1, 1_0000_0000,
 		nodes.iter().skip(1).map(|n| n.node.get_our_node_id()).collect(),
 		&[&chan_upd_1_2, &chan_upd_2_3], &chanmon_cfgs[3].keys_manager);
-	route_params.payment_params.max_path_length = 18;
+	route_params.payment_params.max_path_length = 17;
 
 	let route = get_route(&nodes[0], &route_params).unwrap();
 	node_cfgs[0].router.expect_find_route(route_params.clone(), Ok(route.clone()));

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

@@ -306,13 +306,29 @@ for Payload<ParsedOnionMessageContents<<H as CustomOnionMessageHandler>::CustomM
 /// or received. Thus we read a `ControlTlvs` rather than reading a [`ForwardTlvs`] or
 /// [`ReceiveTlvs`] directly. Also useful on the encoding side to keep forward and receive TLVs in
 /// the same iterator.
+#[derive(Clone)]
 pub(crate) enum ControlTlvs {
 	/// This onion message is intended to be forwarded.
 	Forward(ForwardTlvs),
 	/// This onion message is intended to be received.
 	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, {
@@ -336,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 {