Introduce custom TLV support for OnionMessage

Author: shaavan

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/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.