Introduce padding for Onion Messages

Author: shaavan

lightning/src/blinded_path/message.rs

@@ -65,7 +65,6 @@ impl Writeable for ForwardTlvs {
 			NextMessageHop::NodeId(pubkey) => (Some(pubkey), None),
 			NextMessageHop::ShortChannelId(scid) => (None, Some(scid)),
 		};
-		// TODO: write padding
 		encode_tlv_stream!(writer, {
 			(2, short_channel_id, option),
 			(4, next_node_id, option),
@@ -77,7 +76,6 @@ 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, {
 			(65537, self.context, option),
 		});
@@ -138,17 +136,25 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 ) -> Result<Vec<BlindedHop>, secp256k1::Error> {
 	let pks = intermediate_nodes.iter().map(|node| &node.node_id)
 		.chain(core::iter::once(&recipient_node_id));
-	let tlvs = pks.clone()
+	let tlvs: Vec<ControlTlvs> = pks.clone()
 		.skip(1) // The first node's TLVs contains the next node's pubkey
 		.zip(intermediate_nodes.iter().map(|node| node.short_channel_id))
 		.map(|(pubkey, scid)| match scid {
 			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) })));
+		.chain(core::iter::once(ControlTlvs::Receive(ReceiveTlvs{ context: Some(context) })))
+		.collect();
 
-	utils::construct_blinded_hops(secp_ctx, pks, tlvs, session_priv)
+	let max_length = tlvs.iter()
+		.max_by_key(|c| c.serialized_length())
+		.map(|c| c.serialized_length())
+		.unwrap_or(0);
+
+	let length_tlvs = tlvs.into_iter().map(move |tlv| (max_length, tlv));
+
+	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

lightning/src/blinded_path/utils.rs

@@ -20,7 +20,7 @@ use crate::ln::msgs::DecodeError;
 use crate::ln::onion_utils;
 use crate::onion_message::messenger::Destination;
 use crate::crypto::streams::ChaChaPolyWriteAdapter;
-use crate::util::ser::{Readable, Writeable};
+use crate::util::ser::{Readable, Writeable, Writer};
 
 use crate::io;
 
@@ -135,17 +135,49 @@ fn encrypt_payload<P: Writeable>(payload: P, encrypted_tlvs_rho: [u8; 32]) -> Ve
 	write_adapter.encode()
 }
 
-/// Blinded path encrypted payloads may be padded to ensure they are equal length.
-///
-/// Reads padding to the end, ignoring what's read.
-pub(crate) struct Padding {}
+/// Represents optional padding for encrypted payloads.
+/// Padding is used to ensure payloads have a consistent length.
+pub(crate) struct Padding {
+	padding: Option<Vec<u8>>
+}
+
+impl Padding {
+	/// Creates a new [`Padding`] instance with a specified size.
+	/// Use this method when defining the padding size before writing
+	/// an encrypted payload.
+	pub fn new(length: usize) -> Self {
+		Self {
+			padding: Some(vec![0; length]),
+		}
+	}
+
+	/// Creates a [`Padding`] instance with no size.
+	/// This method is used when returning the [`Padding`] after reading.
+	/// It essentially discards any existing padding data.
+	pub fn none() -> Self {
+		Self {
+			padding: None,
+		}
+	}
+}
+
 impl Readable for Padding {
+	/// Reads and discards the padding data.
 	#[inline]
 	fn read<R: io::Read>(reader: &mut R) -> Result<Self, DecodeError> {
 		loop {
 			let mut buf = [0; 8192];
 			if reader.read(&mut buf[..])? == 0 { break; }
 		}
-		Ok(Self {})
+		Ok(Self::none())
+	}
+}
+
+impl Writeable for Padding {
+	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
+		match &self.padding {
+			Some(padding) => padding.write(writer),
+			None => return Ok(()),
+		}
 	}
 }

lightning/src/onion_message/packet.rs

@@ -359,3 +359,16 @@ impl Writeable for ControlTlvs {
 		}
 	}
 }
+
+impl Writeable for (usize, ControlTlvs) {
+	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
+		let length = self.0 - self.1.serialized_length();
+		let padding = Some(Padding::new(length));
+
+		encode_tlv_stream!(writer, {
+			(1, padding, option)
+		});
+
+		self.1.write(writer)
+	}
+}