Introduce custom TLV support for OnionMessage

lightning/src/blinded_path/message.rs

@@ -30,6 +30,8 @@ pub(crate) struct ReceiveTlvs {
 	/// sending to. This is useful for receivers to check that said blinded path is being used in
 	/// the right context.
 	pub(crate) path_id: Option<[u8; 32]>,
+	/// Custom Tlvs
+	pub(crate) custom_tlvs: Vec<(u64, Vec<u8>)>,
 }
 
 impl Writeable for ForwardTlvs {
@@ -47,6 +49,7 @@ impl Writeable for ReceiveTlvs {
 	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), io::Error> {
 		// TODO: write padding
 		encode_tlv_stream!(writer, {
+			(1, self.custom_tlvs, optional_vec),
 			(6, self.path_id, option),
 		});
 		Ok(())
@@ -62,7 +65,7 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 		.map(|pk| {
 			ControlTlvs::Forward(ForwardTlvs { next_node_id: *pk, next_blinding_override: None })
 		})
-		.chain(core::iter::once(ControlTlvs::Receive(ReceiveTlvs { path_id: None })));
+		.chain(core::iter::once(ControlTlvs::Receive(ReceiveTlvs { path_id: None, custom_tlvs: Vec::new() })));
 
 	utils::construct_blinded_hops(secp_ctx, unblinded_path.iter(), blinded_tlvs, session_priv)
 }

lightning/src/blinded_path/payment.rs

@@ -13,9 +13,10 @@ use crate::ln::features::BlindedHopFeatures;
 use crate::ln::msgs::DecodeError;
 use crate::offers::invoice::BlindedPayInfo;
 use crate::prelude::*;
-use crate::util::ser::{Readable, Writeable, Writer};
+use crate::util::ser::{BigSize, FixedLengthReader, Readable, Writeable, Writer};
 
 use core::convert::TryFrom;
+use crate::io::Read;
 
 /// An intermediate node, its outbound channel, and relay parameters.
 #[derive(Clone, Debug)]
@@ -53,6 +54,8 @@ pub struct ReceiveTlvs {
 	pub payment_secret: PaymentSecret,
 	/// Constraints for the receiver of this payment.
 	pub payment_constraints: PaymentConstraints,
+	/// Custom Tlvs
+	pub custom_tlvs: Vec<(u64, Vec<u8>)>,
 }
 
 /// Data to construct a [`BlindedHop`] for sending a payment over.
@@ -121,6 +124,7 @@ impl Writeable for ForwardTlvs {
 impl Writeable for ReceiveTlvs {
 	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
 		encode_tlv_stream!(w, {
+			(1, self.custom_tlvs, optional_vec),
 			(12, self.payment_constraints, required),
 			(65536, self.payment_secret, required)
 		});
@@ -141,18 +145,29 @@ impl<'a> Writeable for BlindedPaymentTlvsRef<'a> {
 
 impl Readable for BlindedPaymentTlvs {
 	fn read<R: io::Read>(r: &mut R) -> Result<Self, DecodeError> {
-		_init_and_read_tlv_stream!(r, {
+		let mut custom_tlvs = Vec::new();
+
+		let tlv_len = BigSize::read(r)?;
+		let rd = FixedLengthReader::new(r, tlv_len.0);
+		_init_and_read_tlv_stream_with_custom_tlv_decode!(rd, {
 			(1, _padding, option),
 			(2, scid, option),
 			(10, payment_relay, option),
 			(12, payment_constraints, required),
 			(14, features, option),
 			(65536, payment_secret, option),
+		}, |msg_type: u64, msg_reader: &mut FixedLengthReader<_>| -> Result<bool, DecodeError> {
+			if msg_type < 1 << 16 { return Ok(false) }
+			let mut value = Vec::new();
+			msg_reader.read_to_end(&mut value)?;
+			custom_tlvs.push((msg_type, value));
+			Ok(true)
 		});
 		let _padding: Option<utils::Padding> = _padding;
 
 		if let Some(short_channel_id) = scid {
 			if payment_secret.is_some() { return Err(DecodeError::InvalidValue) }
+			if !custom_tlvs.is_empty() { return Err(DecodeError::InvalidValue) }
 			Ok(BlindedPaymentTlvs::Forward(ForwardTlvs {
 				short_channel_id,
 				payment_relay: payment_relay.ok_or(DecodeError::InvalidValue)?,
@@ -164,6 +179,7 @@ impl Readable for BlindedPaymentTlvs {
 			Ok(BlindedPaymentTlvs::Receive(ReceiveTlvs {
 				payment_secret: payment_secret.ok_or(DecodeError::InvalidValue)?,
 				payment_constraints: payment_constraints.0.unwrap(),
+				custom_tlvs,
 			}))
 		}
 	}
@@ -325,6 +341,7 @@ mod tests {
 				max_cltv_expiry: 0,
 				htlc_minimum_msat: 1,
 			},
+			custom_tlvs: Vec::new(),
 		};
 		let htlc_maximum_msat = 100_000;
 		let blinded_payinfo = super::compute_payinfo(&intermediate_nodes[..], &recv_tlvs, htlc_maximum_msat).unwrap();
@@ -343,6 +360,7 @@ mod tests {
 				max_cltv_expiry: 0,
 				htlc_minimum_msat: 1,
 			},
+			custom_tlvs: Vec::new(),
 		};
 		let blinded_payinfo = super::compute_payinfo(&[], &recv_tlvs, 4242).unwrap();
 		assert_eq!(blinded_payinfo.fee_base_msat, 0);
@@ -396,6 +414,7 @@ mod tests {
 				max_cltv_expiry: 0,
 				htlc_minimum_msat: 3,
 			},
+			custom_tlvs: Vec::new(),
 		};
 		let htlc_maximum_msat = 100_000;
 		let blinded_payinfo = super::compute_payinfo(&intermediate_nodes[..], &recv_tlvs, htlc_maximum_msat).unwrap();
@@ -446,6 +465,7 @@ mod tests {
 				max_cltv_expiry: 0,
 				htlc_minimum_msat: 1,
 			},
+			custom_tlvs: Vec::new(),
 		};
 		let htlc_minimum_msat = 3798;
 		assert!(super::compute_payinfo(&intermediate_nodes[..], &recv_tlvs, htlc_minimum_msat - 1).is_err());
@@ -500,6 +520,7 @@ mod tests {
 				max_cltv_expiry: 0,
 				htlc_minimum_msat: 1,
 			},
+			custom_tlvs: Vec::new(),
 		};
 
 		let blinded_payinfo = super::compute_payinfo(&intermediate_nodes[..], &recv_tlvs, 10_000).unwrap();

lightning/src/ln/blinded_payment_tests.rs

@@ -56,6 +56,7 @@ pub fn get_blinded_route_parameters(
 			max_cltv_expiry: u32::max_value(),
 			htlc_minimum_msat: channel_upds.last().unwrap().htlc_minimum_msat,
 		},
+		custom_tlvs: Vec::new(),
 	};
 	let mut secp_ctx = Secp256k1::new();
 	let blinded_path = BlindedPath::new_for_payment(
@@ -89,6 +90,7 @@ fn do_one_hop_blinded_path(success: bool) {
 			max_cltv_expiry: u32::max_value(),
 			htlc_minimum_msat: chan_upd.htlc_minimum_msat,
 		},
+		custom_tlvs: Vec::new(),
 	};
 	let mut secp_ctx = Secp256k1::new();
 	let blinded_path = BlindedPath::one_hop_for_payment(
@@ -131,6 +133,7 @@ fn mpp_to_one_hop_blinded_path() {
 			max_cltv_expiry: u32::max_value(),
 			htlc_minimum_msat: chan_upd_1_3.htlc_minimum_msat,
 		},
+		custom_tlvs: Vec::new(),
 	};
 	let blinded_path = BlindedPath::one_hop_for_payment(
 		nodes[3].node.get_our_node_id(), payee_tlvs, &chanmon_cfgs[3].keys_manager, &secp_ctx

lightning/src/ln/channelmanager.rs

@@ -7939,6 +7939,7 @@ where
 				max_cltv_expiry,
 				htlc_minimum_msat: 1,
 			},
+			custom_tlvs: Vec::new(),
 		};
 		self.router.create_blinded_payment_paths(
 			payee_node_id, first_hops, payee_tlvs, amount_msats, entropy_source, secp_ctx

lightning/src/ln/msgs.rs

@@ -2397,8 +2397,8 @@ impl<NS: Deref> ReadableArgs<(Option<PublicKey>, &NS)> for InboundOnionPayload w
 					})
 				},
 				ChaChaPolyReadAdapter { readable: BlindedPaymentTlvs::Receive(ReceiveTlvs {
-					payment_secret, payment_constraints
-				})} => {
+					payment_secret, payment_constraints,
+					custom_tlvs: _ })} => {
 					if total_msat.unwrap_or(0) > MAX_VALUE_MSAT { return Err(DecodeError::InvalidValue) }
 					Ok(Self::BlindedReceive {
 						amt_msat: amt.ok_or(DecodeError::InvalidValue)?,

lightning/src/onion_message/messenger.rs

@@ -627,12 +627,12 @@ 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 { path_id }), reply_path,
+			message, control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id, custom_tlvs: _ }), reply_path,
 		}, None)) => {
 			Ok(PeeledOnion::Receive(message, path_id, reply_path))
 		},
 		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
@@ -1152,7 +1152,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 { path_id: None, }),
+			control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id: None, custom_tlvs: Vec::new() }),
 			reply_path: reply_path.take(),
 			message,
 		}, prev_control_tlvs_ss.unwrap()));

lightning/src/onion_message/packet.rs

@@ -282,12 +282,22 @@ pub(crate) enum ControlTlvs {
 
 impl Readable for ControlTlvs {
 	fn read<R: Read>(r: &mut R) -> Result<Self, DecodeError> {
-		_init_and_read_tlv_stream!(r, {
+		let mut custom_tlvs = Vec::new();
+
+		let tlv_len = BigSize::read(r)?;
+		let rd = FixedLengthReader::new(r, tlv_len.0);
+		_init_and_read_tlv_stream_with_custom_tlv_decode!(rd, {
 			(1, _padding, option),
 			(2, _short_channel_id, option),
 			(4, next_node_id, option),
 			(6, path_id, option),
 			(8, next_blinding_override, option),
+		}, |msg_type: u64, msg_reader: &mut FixedLengthReader<_>| -> Result<bool, DecodeError> {
+			if msg_type < 1 << 16 { return Ok(false) }
+			let mut value = Vec::new();
+			msg_reader.read_to_end(&mut value)?;
+			custom_tlvs.push((msg_type, value));
+			Ok(true)
 		});
 		let _padding: Option<Padding> = _padding;
 		let _short_channel_id: Option<u64> = _short_channel_id;
@@ -296,13 +306,15 @@ impl Readable for ControlTlvs {
 		let valid_recv_fmt = next_node_id.is_none() && next_blinding_override.is_none();
 
 		let payload_fmt = if valid_fwd_fmt {
+			if !custom_tlvs.is_empty() { return Err(DecodeError::InvalidValue) }
 			ControlTlvs::Forward(ForwardTlvs {
 				next_node_id: next_node_id.unwrap(),
 				next_blinding_override,
 			})
 		} else if valid_recv_fmt {
 			ControlTlvs::Receive(ReceiveTlvs {
 				path_id,
+				custom_tlvs,
 			})
 		} else {
 			return Err(DecodeError::InvalidValue)

lightning/src/util/ser_macros.rs

@@ -472,6 +472,7 @@ macro_rules! decode_tlv_stream {
 ///
 /// [`FixedLengthReader`]: crate::util::ser::FixedLengthReader
 /// [`DecodeError`]: crate::ln::msgs::DecodeError
+#[macro_export]
 macro_rules! decode_tlv_stream_with_custom_tlv_decode {
 	($stream: expr, {$(($type: expr, $field: ident, $fieldty: tt)),* $(,)*}
 	 $(, $decode_custom_tlv: expr)?) => { {
@@ -825,6 +826,24 @@ macro_rules! _init_and_read_tlv_stream {
 	}
 }
 
+/// Equivalent to running [`_init_tlv_field_var`] then [`decode_tlv_stream_with_custom_tlv_decode`].
+///
+/// If any unused values are read, their type MUST be specified or else `rustc` will read them as an
+/// `i64`.
+
+macro_rules! _init_and_read_tlv_stream_with_custom_tlv_decode {
+    ($reader: ident, {$(($type: expr, $field: ident, $fieldty: tt)),* $(,)*}
+     $(, $decode_custom_tlv: expr)?) => {
+        $(
+            $crate::_init_tlv_field_var!($field, $fieldty);
+        )*
+
+        $crate::decode_tlv_stream_with_custom_tlv_decode!(
+            $reader, {$(($type, $field, $fieldty)),*} $(, $decode_custom_tlv)?
+        );
+    }
+}
+
 /// Implements [`Readable`]/[`Writeable`] for a struct storing it as a set of TLVs
 /// If `$fieldty` is `required`, then `$field` is a required field that is not an [`Option`] nor a [`Vec`].
 /// If `$fieldty` is `(default_value, $default)`, then `$field` will be set to `$default` if not present.