De/serialize custom TLVs on OnionHopData::FinalNode

When serialized, the TLVs in `OnionHopData`, unlike a normal TLV
stream, are prefixed with the length of the stream. To allow a user to
add arbitrary custom TLVs, we aren't able to communicate to our
serialization macros exactly which fields to expect, so this commit
adds new macro variants to allow appending an extra set of bytes (and
modifying the prefixed length accordingly).

Because the keysend preimage TLV has a type number in the custom type
range, and a user's TLVs may have type numbers above and/or below
keysend's type number, and because TLV streams must be serialized in
increasing order by type number, this commit also ensures the keysend
TLV is properly sorted/serialized amongst the custom TLVs.

Author: alecchendev

lightning/src/ln/channelmanager.rs

@@ -2458,7 +2458,7 @@ where
 					msg: "Got non final data with an HMAC of 0",
 				});
 			},
-			msgs::OnionHopDataFormat::FinalNode { payment_data, keysend_preimage, payment_metadata } => {
+			msgs::OnionHopDataFormat::FinalNode { payment_data, keysend_preimage, payment_metadata, custom_tlvs } => {
 				if let Some(payment_preimage) = keysend_preimage {
 					// We need to check that the sender knows the keysend preimage before processing this
 					// payment further. Otherwise, an intermediary routing hop forwarding non-keysend-HTLC X

lightning/src/ln/msgs.rs

@@ -43,7 +43,7 @@ use crate::io_extras::read_to_end;
 
 use crate::events::{MessageSendEventsProvider, OnionMessageProvider};
 use crate::util::logger;
-use crate::util::ser::{LengthReadable, Readable, ReadableArgs, Writeable, Writer, WithoutLength, FixedLengthReader, HighZeroBytesDroppedBigSize, Hostname, TransactionU16LenLimited};
+use crate::util::ser::{LengthReadable, Readable, ReadableArgs, Writeable, Writer, WithoutLength, FixedLengthReader, HighZeroBytesDroppedBigSize, Hostname, TransactionU16LenLimited, BigSize};
 
 use crate::ln::{PaymentPreimage, PaymentHash, PaymentSecret};
 
@@ -1429,6 +1429,8 @@ mod fuzzy_internal_msgs {
 			payment_data: Option<FinalOnionHopData>,
 			payment_metadata: Option<Vec<u8>>,
 			keysend_preimage: Option<PaymentPreimage>,
+			/// Serialization will fail if this is not a valid TLV stream
+			custom_tlvs: Option<Vec<(u64, Vec<u8>)>>,
 		},
 	}
 
@@ -1954,14 +1956,19 @@ impl Writeable for OnionHopData {
 					(6, short_channel_id, required)
 				});
 			},
-			OnionHopDataFormat::FinalNode { ref payment_data, ref payment_metadata, ref keysend_preimage } => {
+			OnionHopDataFormat::FinalNode { ref payment_data, ref payment_metadata, ref keysend_preimage, ref custom_tlvs } => {
+				let preimage = if let Some(ref preimage) = keysend_preimage {
+					Some((5482373484, preimage.encode()))
+				} else { None };
+				let mut custom_tlvs: Vec<&(u64, Vec<u8>)> = custom_tlvs.iter().flatten().chain(preimage.iter()).collect();
+				custom_tlvs.sort_by_key(|(typ, _)| *typ);
+
 				_encode_varint_length_prefixed_tlv!(w, {
 					(2, HighZeroBytesDroppedBigSize(self.amt_to_forward), required),
 					(4, HighZeroBytesDroppedBigSize(self.outgoing_cltv_value), required),
 					(8, payment_data, option),
-					(16, payment_metadata.as_ref().map(|m| WithoutLength(m)), option),
-					(5482373484, keysend_preimage, option)
-				});
+					(16, payment_metadata.as_ref().map(|m| WithoutLength(m)), option)
+				}, custom_tlvs.iter());
 			},
 		}
 		Ok(())
@@ -1976,14 +1983,24 @@ impl Readable for OnionHopData {
 		let mut payment_data: Option<FinalOnionHopData> = None;
 		let mut payment_metadata: Option<WithoutLength<Vec<u8>>> = None;
 		let mut keysend_preimage: Option<PaymentPreimage> = None;
-		read_tlv_fields!(r, {
+		let mut custom_tlvs = Vec::new();
+
+		let tlv_len = BigSize::read(r)?;
+		let rd = FixedLengthReader::new(r, tlv_len.0);
+		decode_tlv_stream_with_custom_tlv_decode!(rd, {
 			(2, amt, required),
 			(4, cltv_value, required),
 			(6, short_id, option),
 			(8, payment_data, option),
 			(16, payment_metadata, option),
 			// See https://github.com/lightning/blips/blob/master/blip-0003.md
 			(5482373484, keysend_preimage, 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 format = if let Some(short_channel_id) = short_id {
@@ -2002,6 +2019,7 @@ impl Readable for OnionHopData {
 				payment_data,
 				payment_metadata: payment_metadata.map(|w| w.0),
 				keysend_preimage,
+				custom_tlvs: if custom_tlvs.is_empty() { None } else { Some(custom_tlvs) },
 			}
 		};
 
@@ -3543,6 +3561,7 @@ mod tests {
 				payment_data: None,
 				payment_metadata: None,
 				keysend_preimage: None,
+				custom_tlvs: None,
 			},
 			amt_to_forward: 0x0badf00d01020304,
 			outgoing_cltv_value: 0xffffffff,
@@ -3567,6 +3586,7 @@ mod tests {
 				}),
 				payment_metadata: None,
 				keysend_preimage: None,
+				custom_tlvs: None,
 			},
 			amt_to_forward: 0x0badf00d01020304,
 			outgoing_cltv_value: 0xffffffff,
@@ -3582,13 +3602,81 @@ mod tests {
 			}),
 			payment_metadata: None,
 			keysend_preimage: None,
+			custom_tlvs: None,
 		} = msg.format {
 			assert_eq!(payment_secret, expected_payment_secret);
 		} else { panic!(); }
 		assert_eq!(msg.amt_to_forward, 0x0badf00d01020304);
 		assert_eq!(msg.outgoing_cltv_value, 0xffffffff);
 	}
 
+	#[test]
+	fn encoding_final_onion_hop_data_with_bad_custom_tlvs() {
+		// If custom TLVs have type number within the range reserved for protocol, treat them as if
+		// they're unknown
+		let bad_type_range_tlvs = vec![
+			((1 << 16) - 4, vec![42]),
+			((1 << 16) - 2, vec![42; 32]),
+		];
+		let mut msg = msgs::OnionHopData {
+			format: OnionHopDataFormat::FinalNode {
+				payment_data: None,
+				payment_metadata: None,
+				keysend_preimage: None,
+				custom_tlvs: Some(bad_type_range_tlvs),
+			},
+			amt_to_forward: 0x0badf00d01020304,
+			outgoing_cltv_value: 0xffffffff,
+		};
+		let encoded_value = msg.encode();
+		assert!(msgs::OnionHopData::read(&mut Cursor::new(&encoded_value[..])).is_err());
+		let good_type_range_tlvs = vec![
+			((1 << 16) - 3, vec![42]),
+			((1 << 16) - 1, vec![42; 32]),
+		];
+		if let OnionHopDataFormat::FinalNode { ref mut custom_tlvs, .. } = msg.format {
+			*custom_tlvs = Some(good_type_range_tlvs.clone());
+		}
+		let encoded_value = msg.encode();
+		msg = Readable::read(&mut Cursor::new(&encoded_value[..])).unwrap();
+		match msg.format {
+			OnionHopDataFormat::FinalNode { custom_tlvs, .. } => assert!(custom_tlvs.is_none()),
+			_ => panic!(),
+		}
+	}
+
+	#[test]
+	fn encoding_final_onion_hop_data_with_custom_tlvs() {
+		let expected_custom_tlvs = vec![
+			(5482373483, vec![0x12, 0x34]),
+			(5482373487, vec![0x42u8; 8]),
+		];
+		let mut msg = msgs::OnionHopData {
+			format: OnionHopDataFormat::FinalNode {
+				payment_data: None,
+				payment_metadata: None,
+				keysend_preimage: None,
+				custom_tlvs: Some(expected_custom_tlvs.clone()),
+			},
+			amt_to_forward: 0x0badf00d01020304,
+			outgoing_cltv_value: 0xffffffff,
+		};
+		let encoded_value = msg.encode();
+		let target_value = hex::decode("2e02080badf00d010203040404ffffffffff0000000146c6616b021234ff0000000146c6616f084242424242424242").unwrap();
+		assert_eq!(encoded_value, target_value);
+		msg = Readable::read(&mut Cursor::new(&target_value[..])).unwrap();
+		if let OnionHopDataFormat::FinalNode {
+			payment_data: None,
+			payment_metadata: None,
+			keysend_preimage: None,
+			custom_tlvs,
+		} = msg.format {
+			assert_eq!(custom_tlvs.unwrap(), expected_custom_tlvs);
+		} else { panic!(); }
+		assert_eq!(msg.amt_to_forward, 0x0badf00d01020304);
+		assert_eq!(msg.outgoing_cltv_value, 0xffffffff);
+	}
+
 	#[test]
 	fn query_channel_range_end_blocknum() {
 		let tests: Vec<(u32, u32, u32)> = vec![

lightning/src/ln/onion_utils.rs

@@ -172,6 +172,7 @@ pub(super) fn build_onion_payloads(path: &Path, total_msat: u64, mut recipient_o
 					} else { None },
 					payment_metadata: recipient_onion.payment_metadata.take(),
 					keysend_preimage: *keysend_preimage,
+					custom_tlvs: recipient_onion.custom_tlvs.take(),
 				}
 			} else {
 				msgs::OnionHopDataFormat::NonFinalNode {

lightning/src/util/ser_macros.rs

@@ -117,6 +117,16 @@ macro_rules! _check_encoded_tlv_order {
 /// [`Writer`]: crate::util::ser::Writer
 #[macro_export]
 macro_rules! encode_tlv_stream {
+	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),* $(,)*}) => {
+		$crate::_encode_tlv_stream!($stream, {$(($type, $field, $fieldty)),*})
+	}
+}
+
+/// Implementation of [`encode_tlv_stream`].
+/// This is exported for use by other exported macros, do not use directly.
+#[doc(hidden)]
+#[macro_export]
+macro_rules! _encode_tlv_stream {
 	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),* $(,)*}) => { {
 		#[allow(unused_imports)]
 		use $crate::{
@@ -138,7 +148,21 @@ macro_rules! encode_tlv_stream {
 				$crate::_check_encoded_tlv_order!(last_seen, $type, $fieldty);
 			)*
 		}
-	} }
+	} };
+	($stream: expr, $tlvs: expr) => { {
+		for tlv in $tlvs {
+			let (typ, value): &&(u64, Vec<u8>) = tlv;
+			$crate::_encode_tlv!($stream, *typ, *value, vec_type);
+		}
+
+		#[cfg(debug_assertions)] {
+			let mut last_seen: Option<u64> = None;
+			for tlv in $tlvs {
+				let (typ, _): &&(u64, Vec<u8>) = tlv;
+				$crate::_check_encoded_tlv_order!(last_seen, *typ, vec_type);
+			}
+		}
+	} };
 }
 
 /// Adds the length of the serialized field to a [`LengthCalculatingWriter`].
@@ -192,18 +216,27 @@ macro_rules! _get_varint_length_prefixed_tlv_length {
 #[macro_export]
 macro_rules! _encode_varint_length_prefixed_tlv {
 	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),*}) => { {
+		_encode_varint_length_prefixed_tlv!($stream, {$(($type, $field, $fieldty)),*}, &[])
+	} };
+	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),*}, $extra_tlvs: expr) => { {
 		use $crate::util::ser::BigSize;
+		use alloc::vec::Vec;
 		let len = {
 			#[allow(unused_mut)]
 			let mut len = $crate::util::ser::LengthCalculatingWriter(0);
 			$(
 				$crate::_get_varint_length_prefixed_tlv_length!(len, $type, $field, $fieldty);
 			)*
+			for tlv in $extra_tlvs {
+				let (typ, value): &&(u64, Vec<u8>) = tlv;
+				$crate::_get_varint_length_prefixed_tlv_length!(len, *typ, *value, vec_type);
+			}
 			len.0
 		};
 		BigSize(len as u64).write($stream)?;
-		$crate::encode_tlv_stream!($stream, { $(($type, $field, $fieldty)),* });
-	} }
+		$crate::_encode_tlv_stream!($stream, { $(($type, $field, $fieldty)),* });
+		$crate::_encode_tlv_stream!($stream, $extra_tlvs);
+	} };
 }
 
 /// Errors if there are missing required TLV types between the last seen type and the type currently being processed.
@@ -766,7 +799,8 @@ macro_rules! _init_and_read_tlv_fields {
 ///
 /// For example,
 /// ```
-/// # use lightning::impl_writeable_tlv_based;
+/// # use lightning::{impl_writeable_tlv_based, _encode_varint_length_prefixed_tlv};
+/// # extern crate alloc;
 /// struct LightningMessage {
 /// 	tlv_integer: u32,
 /// 	tlv_default_integer: u32,