Introduce custom TLV support for OnionMessage

Author: shaavan

lightning/src/blinded_path/message.rs

@@ -22,6 +22,8 @@ pub(crate) struct ForwardTlvs {
 	/// Senders to a blinded path use this value to concatenate the route they find to the
 	/// introduction node with the blinded path.
 	pub(crate) next_blinding_override: Option<PublicKey>,
+	/// Custom Tlvs
+	pub(crate) custom_tlvs: Vec<(u64, Vec<u8>)>,
 }
 
 /// Similar to [`ForwardTlvs`], but these TLVs are for the final node.
@@ -30,6 +32,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 {
@@ -60,9 +64,9 @@ pub(super) fn blinded_hops<T: secp256k1::Signing + secp256k1::Verification>(
 	let blinded_tlvs = unblinded_path.iter()
 		.skip(1) // The first node's TLVs contains the next node's pubkey
 		.map(|pk| {
-			ControlTlvs::Forward(ForwardTlvs { next_node_id: *pk, next_blinding_override: None })
+			ControlTlvs::Forward(ForwardTlvs { next_node_id: *pk, next_blinding_override: None, custom_tlvs: Vec::new() })
 		})
-		.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)
 }
@@ -80,7 +84,7 @@ pub(crate) fn advance_path_by_one<NS: Deref, T: secp256k1::Signing + secp256k1::
 	match ChaChaPolyReadAdapter::read(&mut reader, rho) {
 		Ok(ChaChaPolyReadAdapter { readable: ControlTlvs::Forward(ForwardTlvs {
 			mut next_node_id, next_blinding_override,
-		})}) => {
+			custom_tlvs})}) => {
 			let mut new_blinding_point = match next_blinding_override {
 				Some(blinding_point) => blinding_point,
 				None => {

lightning/src/onion_message/messenger.rs

@@ -627,13 +627,13 @@ 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
-		})), Some((next_hop_hmac, new_packet_bytes)))) => {
+			next_node_id, next_blinding_override,
+			custom_tlvs})), 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
 			// unwrapping the onion layers to get to the final payload. Since we don't have the option
@@ -1109,6 +1109,7 @@ fn packet_payloads_and_keys<T: OnionMessageContents, S: secp256k1::Signing + sec
 						ForwardTlvs {
 							next_node_id: unblinded_pk_opt.unwrap(),
 							next_blinding_override: None,
+							custom_tlvs: Vec::new(),
 						}
 					)), ss));
 				}
@@ -1119,6 +1120,7 @@ fn packet_payloads_and_keys<T: OnionMessageContents, S: secp256k1::Signing + sec
 					payloads.push((Payload::Forward(ForwardControlTlvs::Unblinded(ForwardTlvs {
 						next_node_id: intro_node_id,
 						next_blinding_override: Some(blinding_pt),
+						custom_tlvs: Vec::new()
 					})), control_tlvs_ss));
 				}
 			}
@@ -1152,7 +1154,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;
@@ -299,10 +309,12 @@ impl Readable for ControlTlvs {
 			ControlTlvs::Forward(ForwardTlvs {
 				next_node_id: next_node_id.unwrap(),
 				next_blinding_override,
+				custom_tlvs,
 			})
 		} 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.