From 83f0dbc0021335dce183450d7dce7e9f284ff0b6 Mon Sep 17 00:00:00 2001
From: Alec Chen <alecchendev@gmail.com>
Date: Wed, 16 Aug 2023 14:02:21 -0500
Subject: [PATCH] Address custom HTLC TLV fixups

Don't collect iterators to compare, minorly simplify encoding the
keysend TLV, combine the _encode_tlv_stream variants to check that the
ordering of TLVs is correct including custom TLVs.
---
 lightning/src/ln/msgs.rs             |  6 ++----
 lightning/src/ln/outbound_payment.rs |  6 +++---
 lightning/src/util/ser_macros.rs     | 26 +++++++++++---------------
 3 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/lightning/src/ln/msgs.rs b/lightning/src/ln/msgs.rs
index b6ec022d..515dcbba 100644
--- a/lightning/src/ln/msgs.rs
+++ b/lightning/src/ln/msgs.rs
@@ -1987,10 +1987,8 @@ impl Writeable for OutboundOnionPayload {
 				// We need to update [`ln::outbound_payment::RecipientOnionFields::with_custom_tlvs`]
 				// to reject any reserved types in the experimental range if new ones are ever
 				// standardized.
-				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().chain(preimage.iter()).collect();
+				let keysend_tlv = keysend_preimage.map(|preimage| (5482373484, preimage.encode()));
+				let mut custom_tlvs: Vec<&(u64, Vec<u8>)> = custom_tlvs.iter().chain(keysend_tlv.iter()).collect();
 				custom_tlvs.sort_unstable_by_key(|(typ, _)| *typ);
 				_encode_varint_length_prefixed_tlv!(w, {
 					(2, HighZeroBytesDroppedBigSize(*amt_msat), required),
diff --git a/lightning/src/ln/outbound_payment.rs b/lightning/src/ln/outbound_payment.rs
index 9c8d6645..bb974dc2 100644
--- a/lightning/src/ln/outbound_payment.rs
+++ b/lightning/src/ln/outbound_payment.rs
@@ -515,9 +515,9 @@ impl RecipientOnionFields {
 		let tlvs = &mut self.custom_tlvs;
 		let further_tlvs = &mut further_htlc_fields.custom_tlvs;
 
-		let even_tlvs: Vec<&(u64, Vec<u8>)> = tlvs.iter().filter(|(typ, _)| *typ % 2 == 0).collect();
-		let further_even_tlvs: Vec<&(u64, Vec<u8>)> = further_tlvs.iter().filter(|(typ, _)| *typ % 2 == 0).collect();
-		if even_tlvs != further_even_tlvs { return Err(()) }
+		let even_tlvs = tlvs.iter().filter(|(typ, _)| *typ % 2 == 0);
+		let further_even_tlvs = further_tlvs.iter().filter(|(typ, _)| *typ % 2 == 0);
+		if even_tlvs.ne(further_even_tlvs) { return Err(()) }
 
 		tlvs.retain(|tlv| further_tlvs.iter().any(|further_tlv| tlv == further_tlv));
 		further_tlvs.retain(|further_tlv| tlvs.iter().any(|tlv| tlv == further_tlv));
diff --git a/lightning/src/util/ser_macros.rs b/lightning/src/util/ser_macros.rs
index 41185b17..a68dae39 100644
--- a/lightning/src/util/ser_macros.rs
+++ b/lightning/src/util/ser_macros.rs
@@ -143,6 +143,9 @@ macro_rules! encode_tlv_stream {
 #[macro_export]
 macro_rules! _encode_tlv_stream {
 	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),* $(,)*}) => { {
+		$crate::_encode_tlv_stream!($stream, { $(($type, $field, $fieldty)),* }, &[])
+	} };
+	($stream: expr, {$(($type: expr, $field: expr, $fieldty: tt)),* $(,)*}, $extra_tlvs: expr) => { {
 		#[allow(unused_imports)]
 		use $crate::{
 			ln::msgs::DecodeError,
@@ -154,6 +157,10 @@ macro_rules! _encode_tlv_stream {
 		$(
 			$crate::_encode_tlv!($stream, $type, $field, $fieldty);
 		)*
+		for tlv in $extra_tlvs {
+			let (typ, value): &(u64, Vec<u8>) = tlv;
+			$crate::_encode_tlv!($stream, *typ, *value, required_vec);
+		}
 
 		#[allow(unused_mut, unused_variables, unused_assignments)]
 		#[cfg(debug_assertions)]
@@ -162,18 +169,8 @@ 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, required_vec);
-		}
-
-		#[cfg(debug_assertions)] {
-			let mut last_seen: Option<u64> = None;
-			for tlv in $tlvs {
-				let (typ, _): &&(u64, Vec<u8>) = tlv;
+			for tlv in $extra_tlvs {
+				let (typ, _): &(u64, Vec<u8>) = tlv;
 				$crate::_check_encoded_tlv_order!(last_seen, *typ, required_vec);
 			}
 		}
@@ -246,14 +243,13 @@ macro_rules! _encode_varint_length_prefixed_tlv {
 				$crate::_get_varint_length_prefixed_tlv_length!(len, $type, $field, $fieldty);
 			)*
 			for tlv in $extra_tlvs {
-				let (typ, value): &&(u64, Vec<u8>) = tlv;
+				let (typ, value): &(u64, Vec<u8>) = tlv;
 				$crate::_get_varint_length_prefixed_tlv_length!(len, *typ, *value, required_vec);
 			}
 			len.0
 		};
 		BigSize(len as u64).write($stream)?;
-		$crate::_encode_tlv_stream!($stream, { $(($type, $field, $fieldty)),* });
-		$crate::_encode_tlv_stream!($stream, $extra_tlvs);
+		$crate::_encode_tlv_stream!($stream, { $(($type, $field, $fieldty)),* }, $extra_tlvs);
 	} };
 }
 
-- 
2.30.2