Add BlindedRoute Control Tlvs Padding

ViktorT-11 · ViktorT-11 · commit 10d48515f16f · 2022-10-14T21:24:33.000+02:00
diff --git a/lightning/src/onion_message/blinded_route.rs b/lightning/src/onion_message/blinded_route.rs
@@ -13,9 +13,11 @@ use bitcoin::secp256k1::{self, PublicKey, Secp256k1, SecretKey};
 
 use chain::keysinterface::KeysInterface;
 use super::utils;
+use ::get_control_tlv_length;
 use ln::msgs::DecodeError;
 use util::chacha20poly1305rfc::ChaChaPolyWriteAdapter;
 use util::ser::{Readable, VecWriter, Writeable, Writer};
+use super::packet::{ControlTlvs, Padding};
 
 use io;
 use prelude::*;
@@ -54,10 +56,11 @@ impl BlindedRoute {
 	/// will be the destination node.
 	///
 	/// Errors if less than two hops are provided or if `node_pk`(s) are invalid.
-	//  TODO: make all payloads the same size with padding + add dummy hops
-	pub fn new<K: KeysInterface, T: secp256k1::Signing + secp256k1::Verification>
-		(node_pks: &[PublicKey], keys_manager: &K, secp_ctx: &Secp256k1<T>) -> Result<Self, ()>
-	{
+	//  TODO: Add dummy hops
+	pub fn new<K: KeysInterface, T: secp256k1::Signing + secp256k1::Verification> (
+		node_pks: &[PublicKey], keys_manager: &K, secp_ctx: &Secp256k1<T>,
+		include_next_blinding_override_padding: bool
+	) -> Result<Self, ()> {
 		if node_pks.len() < 2 { return Err(()) }
 		let blinding_secret_bytes = keys_manager.get_secure_random_bytes();
 		let blinding_secret = SecretKey::from_slice(&blinding_secret_bytes[..]).expect("RNG is busted");
@@ -66,16 +69,18 @@ impl BlindedRoute {
 		Ok(BlindedRoute {
 			introduction_node_id,
 			blinding_point: PublicKey::from_secret_key(secp_ctx, &blinding_secret),
-			blinded_hops: blinded_hops(secp_ctx, node_pks, &blinding_secret).map_err(|_| ())?,
+			blinded_hops: blinded_hops(secp_ctx, node_pks, &blinding_secret, include_next_blinding_override_padding).map_err(|_| ())?,
 		})
 	}
 }
 
 /// Construct blinded hops for the given `unblinded_path`.
 fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
-	secp_ctx: &Secp256k1<T>, unblinded_path: &[PublicKey], session_priv: &SecretKey
+	secp_ctx: &Secp256k1<T>, unblinded_path: &[PublicKey], session_priv: &SecretKey,
+	include_next_blinding_override_padding: bool
 ) -> Result<Vec<BlindedHop>, secp256k1::Error> {
 	let mut blinded_hops = Vec::with_capacity(unblinded_path.len());
+	let max_length = get_control_tlv_length!(true, include_next_blinding_override_padding);
 
 	let mut prev_ss_and_blinded_node_id = None;
 	utils::construct_keys_callback(secp_ctx, unblinded_path, None, session_priv, |blinded_node_id, _, _, encrypted_payload_ss, unblinded_pk, _| {
@@ -84,6 +89,7 @@ fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 				let payload = ForwardTlvs {
 					next_node_id: pk,
 					next_blinding_override: None,
+					total_length: max_length,
 				};
 				blinded_hops.push(BlindedHop {
 					blinded_node_id: prev_blinded_node_id,
@@ -95,7 +101,7 @@ fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 	})?;
 
 	if let Some((final_ss, final_blinded_node_id)) = prev_ss_and_blinded_node_id {
-		let final_payload = ReceiveTlvs { path_id: None };
+		let final_payload = ReceiveTlvs { path_id: None, total_length: max_length, };
 		blinded_hops.push(BlindedHop {
 			blinded_node_id: final_blinded_node_id,
 			encrypted_payload: encrypt_payload(final_payload, final_ss),
@@ -150,37 +156,45 @@ impl_writeable!(BlindedHop, {
 
 /// 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, Copy)]
 pub(crate) struct ForwardTlvs {
 	/// The node id of the next hop in the onion message's path.
 	pub(super) next_node_id: PublicKey,
 	/// Senders to a blinded route use this value to concatenate the route they find to the
 	/// introduction node with the blinded route.
 	pub(super) next_blinding_override: Option<PublicKey>,
+	/// The length the tlv should have when it's serialized, with padding included if needed.
+	/// Used to ensure that all control tlvs in a blinded route have the same length.
+	pub(super) total_length: u16,
 }
 
 /// Similar to [`ForwardTlvs`], but these TLVs are for the final node.
+#[derive(Clone, Copy)]
 pub(crate) struct ReceiveTlvs {
 	/// If `path_id` is `Some`, it is used to identify the blinded route that this onion message is
 	/// sending to. This is useful for receivers to check that said blinded route is being used in
 	/// the right context.
 	pub(super) path_id: Option<[u8; 32]>,
+	/// The length the tlv should have when it's serialized, with padding included if needed.
+	/// Used to ensure that all control tlvs in a blinded route have the same length.
+	pub(super) total_length: u16,
 }
 
 impl Writeable for ForwardTlvs {
 	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
-		// TODO: write padding
 		encode_tlv_stream!(writer, {
+			(1, Padding::new_from_tlv(ControlTlvs::Forward(*self)), option),
 			(4, self.next_node_id, required),
-			(8, self.next_blinding_override, option)
+			(8, self.next_blinding_override, option),
 		});
 		Ok(())
 	}
 }
 
 impl Writeable for ReceiveTlvs {
 	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
-		// TODO: write padding
 		encode_tlv_stream!(writer, {
+			(1, Padding::new_from_tlv(ControlTlvs::Receive(*self)), option),
 			(6, self.path_id, option),
 		});
 		Ok(())
diff --git a/lightning/src/onion_message/functional_tests.rs b/lightning/src/onion_message/functional_tests.rs
@@ -98,7 +98,7 @@ fn two_unblinded_two_blinded() {
 	let nodes = create_nodes(5);
 
 	let secp_ctx = Secp256k1::new();
-	let blinded_route = BlindedRoute::new(&[nodes[3].get_node_pk(), nodes[4].get_node_pk()], &*nodes[4].keys_manager, &secp_ctx).unwrap();
+	let blinded_route = BlindedRoute::new(&[nodes[3].get_node_pk(), nodes[4].get_node_pk()], &*nodes[4].keys_manager, &secp_ctx, true).unwrap();
 
 	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], Destination::BlindedRoute(blinded_route), None).unwrap();
 	pass_along_path(&nodes, None);
@@ -109,7 +109,7 @@ fn three_blinded_hops() {
 	let nodes = create_nodes(4);
 
 	let secp_ctx = Secp256k1::new();
-	let blinded_route = BlindedRoute::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk(), nodes[3].get_node_pk()], &*nodes[3].keys_manager, &secp_ctx).unwrap();
+	let blinded_route = BlindedRoute::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk(), nodes[3].get_node_pk()], &*nodes[3].keys_manager, &secp_ctx, true).unwrap();
 
 	nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), None).unwrap();
 	pass_along_path(&nodes, None);
@@ -133,13 +133,13 @@ fn invalid_blinded_route_error() {
 
 	// 0 hops
 	let secp_ctx = Secp256k1::new();
-	let mut blinded_route = BlindedRoute::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], &*nodes[2].keys_manager, &secp_ctx).unwrap();
+	let mut blinded_route = BlindedRoute::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], &*nodes[2].keys_manager, &secp_ctx, true).unwrap();
 	blinded_route.blinded_hops.clear();
 	let err = nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), None).unwrap_err();
 	assert_eq!(err, SendError::TooFewBlindedHops);
 
 	// 1 hop
-	let mut blinded_route = BlindedRoute::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], &*nodes[2].keys_manager, &secp_ctx).unwrap();
+	let mut blinded_route = BlindedRoute::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], &*nodes[2].keys_manager, &secp_ctx, true).unwrap();
 	blinded_route.blinded_hops.remove(0);
 	assert_eq!(blinded_route.blinded_hops.len(), 1);
 	let err = nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), None).unwrap_err();
@@ -152,7 +152,7 @@ fn reply_path() {
 	let secp_ctx = Secp256k1::new();
 
 	// Destination::Node
-	let reply_path = BlindedRoute::new(&[nodes[2].get_node_pk(), nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx).unwrap();
+	let reply_path = BlindedRoute::new(&[nodes[2].get_node_pk(), nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx, false).unwrap();
 	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], Destination::Node(nodes[3].get_node_pk()), Some(reply_path)).unwrap();
 	pass_along_path(&nodes, None);
 	// Make sure the last node successfully decoded the reply path.
@@ -161,8 +161,8 @@ fn reply_path() {
 		format!("Received an onion message with path_id: None and reply_path").to_string(), 1);
 
 	// Destination::BlindedRoute
-	let blinded_route = BlindedRoute::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk(), nodes[3].get_node_pk()], &*nodes[3].keys_manager, &secp_ctx).unwrap();
-	let reply_path = BlindedRoute::new(&[nodes[2].get_node_pk(), nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx).unwrap();
+	let blinded_route = BlindedRoute::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk(), nodes[3].get_node_pk()], &*nodes[3].keys_manager, &secp_ctx, true).unwrap();
+	let reply_path = BlindedRoute::new(&[nodes[2].get_node_pk(), nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx, false).unwrap();
 
 	nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), Some(reply_path)).unwrap();
 	pass_along_path(&nodes, None);
diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs
@@ -22,6 +22,7 @@ use ln::onion_utils;
 use super::blinded_route::{BlindedRoute, ForwardTlvs, ReceiveTlvs};
 use super::packet::{BIG_PACKET_HOP_DATA_LEN, ForwardControlTlvs, Packet, Payload, ReceiveControlTlvs, SMALL_PACKET_HOP_DATA_LEN};
 use super::utils;
+use ::get_control_tlv_length;
 use util::events::OnionMessageProvider;
 use util::logger::Logger;
 use util::ser::Writeable;
@@ -256,14 +257,14 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessageHandler for OnionMessenger<Si
 			msg.onion_routing_packet.hmac, control_tlvs_ss)
 		{
 			Ok((Payload::Receive {
-				control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id }), reply_path,
+				control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id, .. }), reply_path,
 			}, None)) => {
 				log_info!(self.logger,
 					"Received an onion message with path_id: {:02x?} and {}reply_path",
 						path_id, if reply_path.is_some() { "" } else { "no " });
 			},
 			Ok((Payload::Forward(ForwardControlTlvs::Unblinded(ForwardTlvs {
-				next_node_id, next_blinding_override
+				next_node_id, next_blinding_override, ..
 			})), Some((next_hop_hmac, new_packet_bytes)))) => {
 				// TODO: we need to check whether `next_node_id` 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
@@ -418,6 +419,7 @@ fn packet_payloads_and_keys<T: secp256k1::Signing + secp256k1::Verification>(
 					ForwardTlvs {
 						next_node_id: unblinded_pk_opt.unwrap(),
 						next_blinding_override: None,
+						total_length: get_control_tlv_length!(true, false),
 					}
 				)), ss));
 			}
@@ -428,6 +430,7 @@ fn packet_payloads_and_keys<T: secp256k1::Signing + secp256k1::Verification>(
 				payloads.push((Payload::Forward(ForwardControlTlvs::Unblinded(ForwardTlvs {
 					next_node_id: intro_node_id,
 					next_blinding_override: Some(blinding_pt),
+					total_length: get_control_tlv_length!(true, true),
 				})), control_tlvs_ss));
 			}
 			if let Some(encrypted_payload) = enc_payload_opt {
@@ -460,7 +463,7 @@ fn packet_payloads_and_keys<T: secp256k1::Signing + secp256k1::Verification>(
 
 	if let Some(control_tlvs_ss) = prev_control_tlvs_ss {
 		payloads.push((Payload::Receive {
-			control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id: None, }),
+			control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id: None, total_length: get_control_tlv_length!(false), }),
 			reply_path: reply_path.take(),
 		}, control_tlvs_ss));
 	}
diff --git a/lightning/src/onion_message/packet.rs b/lightning/src/onion_message/packet.rs
@@ -14,6 +14,7 @@ use bitcoin::secp256k1::ecdh::SharedSecret;
 
 use ln::msgs::DecodeError;
 use ln::onion_utils;
+use ::get_control_tlv_length;
 use super::blinded_route::{BlindedRoute, ForwardTlvs, ReceiveTlvs};
 use util::chacha20poly1305rfc::{ChaChaPolyReadAdapter, ChaChaPolyWriteAdapter};
 use util::ser::{BigSize, FixedLengthReader, LengthRead, LengthReadable, LengthReadableArgs, Readable, ReadableArgs, Writeable, Writer};
@@ -197,7 +198,7 @@ impl ReadableArgs<SharedSecret> for Payload {
 /// 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.
-pub(super) enum ControlTlvs {
+pub(crate) enum ControlTlvs {
 	/// This onion message is intended to be forwarded.
 	Forward(ForwardTlvs),
 	/// This onion message is intended to be received.
@@ -206,13 +207,13 @@ pub(super) enum ControlTlvs {
 
 impl Readable for ControlTlvs {
 	fn read<R: Read>(mut r: &mut R) -> Result<Self, DecodeError> {
-		let mut _padding: Option<Padding> = None;
+		let mut padding: Option<Padding> = None;
 		let mut _short_channel_id: Option<u64> = None;
 		let mut next_node_id: Option<PublicKey> = None;
 		let mut path_id: Option<[u8; 32]> = None;
 		let mut next_blinding_override: Option<PublicKey> = None;
 		decode_tlv_stream!(&mut r, {
-			(1, _padding, option),
+			(1, padding, option),
 			(2, _short_channel_id, option),
 			(4, next_node_id, option),
 			(6, path_id, option),
@@ -222,13 +223,20 @@ impl Readable for ControlTlvs {
 		let valid_fwd_fmt  = next_node_id.is_some() && path_id.is_none();
 		let valid_recv_fmt = next_node_id.is_none() && next_blinding_override.is_none();
 
+		let mut total_length = get_control_tlv_length!(next_node_id.is_some(), next_blinding_override.is_some(), path_id.is_some(), 0, 0);
+		if padding.is_some() {
+			total_length += padding.unwrap().padding_length + 2 // 2 extra prefix bytes
+		}
+
 		let payload_fmt = if valid_fwd_fmt {
 			ControlTlvs::Forward(ForwardTlvs {
+				total_length,
 				next_node_id: next_node_id.unwrap(),
 				next_blinding_override,
 			})
 		} else if valid_recv_fmt {
 			ControlTlvs::Receive(ReceiveTlvs {
+				total_length,
 				path_id,
 			})
 		} else {
@@ -239,15 +247,57 @@ impl Readable for ControlTlvs {
 	}
 }
 
-/// Reads padding to the end, ignoring what's read.
-pub(crate) struct Padding {}
+pub(crate) struct Padding {
+	padding_length: u16,
+}
+
+/// Reads padding to the end, and validates that the read bytes' content are 0s.
 impl Readable for Padding {
 	#[inline]
 	fn read<R: Read>(reader: &mut R) -> Result<Self, DecodeError> {
+		let mut padding_length = 0;
 		loop {
 			let mut buf = [0; 8192];
-			if reader.read(&mut buf[..])? == 0 { break; }
+			let read_bytes_len = reader.read(&mut buf[..])?;
+			if read_bytes_len == 0 { break; }
+			for n in 0..read_bytes_len {
+				if buf[n] != 0 as u8 { return Err(DecodeError::InvalidValue); }
+			}
+			padding_length += read_bytes_len as u16;
+		}
+		Ok(Self { padding_length })
+	}
+}
+
+impl Writeable for Padding {
+	#[inline]
+	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
+		for _ in 0..self.padding_length {
+			(0 as u8).write(writer)?;
+		}
+		Ok(())
+	}
+}
+
+impl Padding {
+	pub fn new_from_tlv(tlv: ControlTlvs) -> Option<Padding> {
+		let (data_length, tlv_total_length) = match tlv {
+			ControlTlvs::Forward(forward_tlvs) => {
+				let data_length = get_control_tlv_length!(true, forward_tlvs.next_blinding_override.is_some());
+				(data_length, forward_tlvs.total_length as u16)
+			},
+			ControlTlvs::Receive(receive_tlvs) => {
+				let data_length = get_control_tlv_length!(receive_tlvs.path_id.is_some());
+				(data_length, receive_tlvs.total_length as u16)
+			}
+		};
+
+		let extra_bytes_needed = tlv_total_length - data_length;
+		if extra_bytes_needed >= 2 {
+			let padding_length = extra_bytes_needed - 2; // 2 bytes of prefix removed
+			Some(Padding { padding_length })
+		} else {
+			None
 		}
-		Ok(Self {})
 	}
 }
diff --git a/lightning/src/onion_message/utils.rs b/lightning/src/onion_message/utils.rs
@@ -96,3 +96,58 @@ pub(super) fn construct_keys_callback<T: secp256k1::Signing + secp256k1::Verific
 	}
 	Ok(())
 }
+
+#[macro_export]
+macro_rules! get_control_tlv_length {
+	($has_path_id: expr) => {{ // ReceiveControlTlvs
+		get_control_tlv_length!(false, false, $has_path_id, 0, 0)
+	}};
+	($has_next_node_id: expr, $has_next_blinding_override: expr) => {{ // ForwardControlTlvs
+		get_control_tlv_length!($has_next_node_id, $has_next_blinding_override, false, 0, 0)
+	}};
+	($has_next_node_id: expr, $has_next_blinding_override: expr, $has_path_id: expr, $tag_prefix_length: expr, $tag_length: expr) => {{
+		// tag_prefix_length and tag_length refer to custom types in ControlTlvs, not the be
+		// confused with the onion message tag.
+		let mut res = 0;
+
+		macro_rules! add_length {
+			($should_add_len: expr, $prefix_len: expr, $content_len: expr) => {
+				if $should_add_len {
+					res += $prefix_len;
+					res += $content_len;
+				}
+			}
+		}
+
+		add_length!($has_next_node_id, 2, 33);
+		add_length!($has_next_blinding_override, 2, 33);
+		add_length!($has_path_id, 2, 32);
+		add_length!($tag_length > 0, $tag_prefix_length, $tag_length);
+
+		res
+	}}
+
+	/*
+	TODO:
+
+	Also add support for payment_onion ControlTlvs also consisting of:
+
+	payment_relay:
+	2 bytes prefix
+	2 bytes for cltv_expiry_delta
+	4 bytes for fee_proportional_millionths
+	0-4 bytes for fee_base_msat (tu32)
+
+	payment_constraints:
+	2 bytes prefix
+	4 bytes max_cltv_expiry
+	0-8 bytes htlc_minimum_msat (tu64)
+
+	allowed_features:
+	- If IS payment onion AND has NO known allowed_features:
+	  2 bytes prefix only
+	- If IS payment onion AND HAS known allowed_features:
+	  2 bytes prefix
+	  X bytes of allowed_features
+	*/
+}
