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 a new macro variant 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 introduces a way to sort
TLVs from a byte slice.

Author: alecchendev

lightning/src/ln/channelmanager.rs

@@ -2341,7 +2341,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 payment_data.is_some() && keysend_preimage.is_some() {
 					return Err(ReceiveError {
 						err_code: 0x4000|22,

lightning/src/ln/msgs.rs

@@ -32,6 +32,7 @@ use bitcoin::hash_types::{Txid, BlockHash};
 
 use crate::ln::features::{ChannelFeatures, ChannelTypeFeatures, InitFeatures, NodeFeatures};
 use crate::ln::onion_utils;
+use crate::offers::{TlvStream, TlvRecord};
 use crate::onion_message;
 
 use crate::prelude::*;
@@ -42,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};
 
@@ -1413,6 +1414,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<u8>>,
 		},
 	}
 
@@ -1934,20 +1937,36 @@ 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 mut custom_tlv_stream = custom_tlvs.clone().unwrap_or(Vec::new());
+				encode_tlv_stream!(&mut custom_tlv_stream, {
+					(5482373484, keysend_preimage, option)
+				});
+				let sorted_custom_tlvs = sorted_tlv_stream(&custom_tlv_stream)?;
 				_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)
+				}, &sorted_custom_tlvs);
 			},
 		}
 		Ok(())
 	}
 }
 
+fn sorted_tlv_stream(bytes: &[u8]) -> Result<Vec<u8>, io::Error> {
+	let mut tlv_stream = TlvStream::try_new(bytes)
+		.map_err(|_| io::Error::new(io::ErrorKind::Other, "Could not serialize custom TLV stream"))?
+		.collect::<Vec<TlvRecord>>();
+	tlv_stream.sort_by_key(|tlv| tlv.get_type());
+	let mut sorted_bytes = Vec::with_capacity(bytes.len());
+	for tlv in tlv_stream {
+		tlv.write(&mut sorted_bytes)?;
+	}
+	Ok(sorted_bytes)
+}
+
 impl Readable for OnionHopData {
 	fn read<R: Read>(r: &mut R) -> Result<Self, DecodeError> {
 		let mut amt = HighZeroBytesDroppedBigSize(0u64);
@@ -1956,14 +1975,27 @@ 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) }
+			// Re-serialize into TLV record
+			let mut buffer = Vec::new();
+			msg_reader.read_to_end(&mut buffer)?;
+			BigSize(msg_type).write(&mut custom_tlvs)?;
+			BigSize(buffer.len() as u64).write(&mut custom_tlvs)?;
+			custom_tlvs.append(&mut buffer);
+			Ok(true)
 		});
 
 		let format = if let Some(short_channel_id) = short_id {
@@ -1982,6 +2014,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) },
 			}
 		};
 
@@ -2417,7 +2450,7 @@ mod tests {
 	use crate::ln::features::{ChannelFeatures, ChannelTypeFeatures, InitFeatures, NodeFeatures};
 	use crate::ln::msgs::{self, FinalOnionHopData, OnionErrorPacket, OnionHopDataFormat};
 	use crate::routing::gossip::{NodeAlias, NodeId};
-	use crate::util::ser::{Writeable, Readable, Hostname, TransactionU16LenLimited};
+	use crate::util::ser::{Writeable, Writer, Readable, Hostname, TransactionU16LenLimited};
 
 	use bitcoin::hashes::hex::FromHex;
 	use bitcoin::util::address::Address;
@@ -3513,6 +3546,7 @@ mod tests {
 				payment_data: None,
 				payment_metadata: None,
 				keysend_preimage: None,
+				custom_tlvs: None,
 			},
 			amt_to_forward: 0x0badf00d01020304,
 			outgoing_cltv_value: 0xffffffff,
@@ -3537,6 +3571,7 @@ mod tests {
 				}),
 				payment_metadata: None,
 				keysend_preimage: None,
+				custom_tlvs: None,
 			},
 			amt_to_forward: 0x0badf00d01020304,
 			outgoing_cltv_value: 0xffffffff,
@@ -3552,13 +3587,87 @@ 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() {
+		let bad_tlvs = vec![42; 32];
+		let mut msg = msgs::OnionHopData {
+			format: OnionHopDataFormat::FinalNode {
+				payment_data: None,
+				payment_metadata: None,
+				keysend_preimage: None,
+				custom_tlvs: Some(bad_tlvs),
+			},
+			amt_to_forward: 0x0badf00d01020304,
+			outgoing_cltv_value: 0xffffffff,
+		};
+		assert!(msg.write(&mut Vec::new()).is_err());
+
+		// If custom TLVs have type number within the range reserved for protocol, treat them as if
+		// they're unknown
+		let bad_type_range_tlvs = _get_encoded_tlv_stream!({
+			((1 << 16) - 4, 42, required),
+			((1 << 16) - 2, Some(vec![42; 32]), option),
+		});
+		if let OnionHopDataFormat::FinalNode { ref mut custom_tlvs, .. } = msg.format {
+			*custom_tlvs = Some(bad_type_range_tlvs.clone());
+		}
+		let encoded_value = msg.encode();
+		assert!(msgs::OnionHopData::read(&mut Cursor::new(&encoded_value[..])).is_err());
+		let good_type_range_tlvs = _get_encoded_tlv_stream!({
+			((1 << 16) - 3, 42, required),
+			((1 << 16) - 1, Some(vec![42; 32]), option),
+		});
+		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 = _get_encoded_tlv_stream!({
+			(5482373483, 0x1234, required),
+			(5482373487, Some(vec![0x42u8; 8]), option),
+		});
+		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("3202080badf00d010203040404ffffffffff0000000146c6616b0400001234ff0000000146c6616f0a00084242424242424242").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

@@ -207,19 +207,24 @@ macro_rules! _get_varint_length_prefixed_tlv_length {
 #[doc(hidden)]
 #[macro_export]
 macro_rules! _encode_varint_length_prefixed_tlv {
-	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),*}) => { {
+	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),*}) => {
+		_encode_varint_length_prefixed_tlv!($stream, {$(($type, $field, $fieldty)),*}, &[0u8; 0])
+	};
+	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),*}, $extra: expr) => { {
+		// $extra should be a byte slice holding a valid TLV stream
 		use $crate::util::ser::BigSize;
 		let len = {
 			#[allow(unused_mut)]
 			let mut len = $crate::util::ser::LengthCalculatingWriter(0);
 			$(
 				$crate::_get_varint_length_prefixed_tlv_length!(len, $type, $field, $fieldty);
 			)*
-			len.0
+			len.0 + $extra.len()
 		};
 		BigSize(len as u64).write($stream)?;
 		$crate::encode_tlv_stream!($stream, { $(($type, $field, $fieldty)),* });
-	} }
+		$stream.write_all($extra)?;
+	} };
 }
 
 /// Errors if there are missing required TLV types between the last seen type and the type currently being processed.
@@ -782,7 +787,7 @@ 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};
 /// struct LightningMessage {
 /// 	tlv_integer: u32,
 /// 	tlv_default_integer: u32,
@@ -1063,7 +1068,7 @@ mod tests {
 	use crate::io::{self, Cursor};
 	use crate::prelude::*;
 	use crate::ln::msgs::DecodeError;
-	use crate::util::ser::{Writeable, HighZeroBytesDroppedBigSize, VecWriter};
+	use crate::util::ser::{Writer, Writeable, HighZeroBytesDroppedBigSize, VecWriter};
 	use bitcoin::secp256k1::PublicKey;
 
 	// The BOLT TLV test cases don't include any tests which use our "required-value" logic since