Support (de)serializing payment_data in onion TLVs and track them

This is the first step in Base AMP support, just tracking the
relevant data in internal datastructures.

Author: TheBlueMatt

lightning/src/ln/channelmanager.rs

@@ -72,7 +72,9 @@ enum PendingForwardReceiveHTLCInfo {
 		onion_packet: msgs::OnionPacket,
 		short_channel_id: u64, // This should be NonZero<u64> eventually
 	},
-	Receive {},
+	Receive {
+		payment_data: Option<msgs::FinalOnionHopData>,
+	},
 }
 
 #[derive(Clone)] // See Channel::revoke_and_ack for why, tl;dr: Rust bug
@@ -117,6 +119,12 @@ pub(super) struct HTLCPreviousHopData {
 	incoming_packet_shared_secret: [u8; 32],
 }
 
+struct ClaimableHTLC {
+	src: HTLCPreviousHopData,
+	value: u64,
+	payment_data: Option<msgs::FinalOnionHopData>,
+}
+
 /// Tracks the inbound corresponding to an outbound HTLC
 #[derive(Clone, PartialEq)]
 pub(super) enum HTLCSource {
@@ -274,12 +282,11 @@ pub(super) struct ChannelHolder<ChanSigner: ChannelKeys> {
 	/// guarantees are made about the existence of a channel with the short id here, nor the short
 	/// ids in the PendingHTLCInfo!
 	pub(super) forward_htlcs: HashMap<u64, Vec<HTLCForwardInfo>>,
-	/// payment_hash -> Vec<(amount_received, htlc_source)> for tracking things that were to us and
-	/// can be failed/claimed by the user
+	/// Tracks things that were to us and can be failed/claimed by the user
 	/// Note that while this is held in the same mutex as the channels themselves, no consistency
 	/// guarantees are made about the channels given here actually existing anymore by the time you
 	/// go to read them!
-	pub(super) claimable_htlcs: HashMap<PaymentHash, Vec<(u64, HTLCPreviousHopData)>>,
+	claimable_htlcs: HashMap<PaymentHash, Vec<ClaimableHTLC>>,
 	/// Messages to send to peers - pushed to in the same lock that they are generated in (except
 	/// for broadcast messages, where ordering isn't as strict).
 	pub(super) pending_msg_events: Vec<events::MessageSendEvent>,
@@ -288,7 +295,7 @@ pub(super) struct MutChannelHolder<'a, ChanSigner: ChannelKeys + 'a> {
 	pub(super) by_id: &'a mut HashMap<[u8; 32], Channel<ChanSigner>>,
 	pub(super) short_to_id: &'a mut HashMap<u64, [u8; 32]>,
 	pub(super) forward_htlcs: &'a mut HashMap<u64, Vec<HTLCForwardInfo>>,
-	pub(super) claimable_htlcs: &'a mut HashMap<PaymentHash, Vec<(u64, HTLCPreviousHopData)>>,
+	claimable_htlcs: &'a mut HashMap<PaymentHash, Vec<ClaimableHTLC>>,
 	pub(super) pending_msg_events: &'a mut Vec<events::MessageSendEvent>,
 }
 impl<ChanSigner: ChannelKeys> ChannelHolder<ChanSigner> {
@@ -993,13 +1000,19 @@ impl<ChanSigner: ChannelKeys> ChannelManager<ChanSigner> {
 					return_err!("Upstream node set CLTV to the wrong value", 18, &byte_utils::be32_to_array(msg.cltv_expiry));
 				}
 
+				let payment_data = match next_hop_data.format {
+					msgs::OnionHopDataFormat::Legacy { .. } => None,
+					msgs::OnionHopDataFormat::NonFinalNode { .. } => return_err!("Got non final data with an HMAC of 0", 0x4000 | 22, &[0;0]),
+					msgs::OnionHopDataFormat::FinalNode { payment_data } => payment_data,
+				};
+
 				// Note that we could obviously respond immediately with an update_fulfill_htlc
 				// message, however that would leak that we are the recipient of this payment, so
 				// instead we stay symmetric with the forwarding case, only responding (after a
 				// delay) once they've send us a commitment_signed!
 
 				PendingHTLCStatus::Forward(PendingHTLCInfo {
-					type_data: PendingForwardReceiveHTLCInfo::Receive {},
+					type_data: PendingForwardReceiveHTLCInfo::Receive { payment_data },
 					payment_hash: msg.payment_hash.clone(),
 					incoming_shared_secret: shared_secret,
 					amt_to_forward: next_hop_data.amt_to_forward,
@@ -1548,17 +1561,18 @@ impl<ChanSigner: ChannelKeys> ChannelManager<ChanSigner> {
 					for forward_info in pending_forwards.drain(..) {
 						match forward_info {
 							HTLCForwardInfo::AddHTLC { prev_short_channel_id, prev_htlc_id, forward_info: PendingHTLCInfo {
-									type_data: PendingForwardReceiveHTLCInfo::Receive { },
+									type_data: PendingForwardReceiveHTLCInfo::Receive { payment_data },
 									incoming_shared_secret, payment_hash, amt_to_forward, .. }, } => {
 								let prev_hop_data = HTLCPreviousHopData {
 									short_channel_id: prev_short_channel_id,
 									htlc_id: prev_htlc_id,
 									incoming_packet_shared_secret: incoming_shared_secret,
 								};
-								match channel_state.claimable_htlcs.entry(payment_hash) {
-									hash_map::Entry::Occupied(mut entry) => entry.get_mut().push((amt_to_forward, prev_hop_data)),
-									hash_map::Entry::Vacant(entry) => { entry.insert(vec![(amt_to_forward, prev_hop_data)]); },
-								};
+								channel_state.claimable_htlcs.entry(payment_hash).or_insert(Vec::new()).push(ClaimableHTLC {
+									src: prev_hop_data,
+									value: amt_to_forward,
+									payment_data,
+								});
 								new_events.push(events::Event::PaymentReceived {
 									payment_hash: payment_hash,
 									amt: amt_to_forward,
@@ -1631,11 +1645,11 @@ impl<ChanSigner: ChannelKeys> ChannelManager<ChanSigner> {
 		let mut channel_state = Some(self.channel_state.lock().unwrap());
 		let removed_source = channel_state.as_mut().unwrap().claimable_htlcs.remove(payment_hash);
 		if let Some(mut sources) = removed_source {
-			for (recvd_value, htlc_with_hash) in sources.drain(..) {
+			for htlc in sources.drain(..) {
 				if channel_state.is_none() { channel_state = Some(self.channel_state.lock().unwrap()); }
 				self.fail_htlc_backwards_internal(channel_state.take().unwrap(),
-						HTLCSource::PreviousHopData(htlc_with_hash), payment_hash,
-						HTLCFailReason::Reason { failure_code: 0x4000 | 15, data: byte_utils::be64_to_array(recvd_value).to_vec() });
+						HTLCSource::PreviousHopData(htlc.src), payment_hash,
+						HTLCFailReason::Reason { failure_code: 0x4000 | 15, data: byte_utils::be64_to_array(htlc.value).to_vec() });
 			}
 			true
 		} else { false }
@@ -1759,17 +1773,17 @@ impl<ChanSigner: ChannelKeys> ChannelManager<ChanSigner> {
 		let mut channel_state = Some(self.channel_state.lock().unwrap());
 		let removed_source = channel_state.as_mut().unwrap().claimable_htlcs.remove(&payment_hash);
 		if let Some(mut sources) = removed_source {
-			for (received_amount, htlc_with_hash) in sources.drain(..) {
+			for htlc in sources.drain(..) {
 				if channel_state.is_none() { channel_state = Some(self.channel_state.lock().unwrap()); }
-				if received_amount < expected_amount || received_amount > expected_amount * 2 {
-					let mut htlc_msat_data = byte_utils::be64_to_array(received_amount).to_vec();
+				if htlc.value < expected_amount || htlc.value > expected_amount * 2 {
+					let mut htlc_msat_data = byte_utils::be64_to_array(htlc.value).to_vec();
 					let mut height_data = byte_utils::be32_to_array(self.latest_block_height.load(Ordering::Acquire) as u32).to_vec();
 					htlc_msat_data.append(&mut height_data);
 					self.fail_htlc_backwards_internal(channel_state.take().unwrap(),
-									 HTLCSource::PreviousHopData(htlc_with_hash), &payment_hash,
+									 HTLCSource::PreviousHopData(htlc.src), &payment_hash,
 									 HTLCFailReason::Reason { failure_code: 0x4000|15, data: htlc_msat_data });
 				} else {
-					self.claim_funds_internal(channel_state.take().unwrap(), HTLCSource::PreviousHopData(htlc_with_hash), payment_preimage);
+					self.claim_funds_internal(channel_state.take().unwrap(), HTLCSource::PreviousHopData(htlc.src), payment_preimage);
 				}
 			}
 			true
@@ -3079,8 +3093,9 @@ impl Writeable for PendingHTLCInfo {
 				onion_packet.write(writer)?;
 				short_channel_id.write(writer)?;
 			},
-			&PendingForwardReceiveHTLCInfo::Receive { } => {
+			&PendingForwardReceiveHTLCInfo::Receive { ref payment_data } => {
 				1u8.write(writer)?;
+				payment_data.write(writer)?;
 			},
 		}
 		self.incoming_shared_secret.write(writer)?;
@@ -3099,7 +3114,9 @@ impl<R: ::std::io::Read> Readable<R> for PendingHTLCInfo {
 					onion_packet: Readable::read(reader)?,
 					short_channel_id: Readable::read(reader)?,
 				},
-				1u8 => PendingForwardReceiveHTLCInfo::Receive { },
+				1u8 => PendingForwardReceiveHTLCInfo::Receive {
+					payment_data: Readable::read(reader)?,
+				},
 				_ => return Err(DecodeError::InvalidValue),
 			},
 			incoming_shared_secret: Readable::read(reader)?,
@@ -3304,9 +3321,10 @@ impl<ChanSigner: ChannelKeys + Writeable> Writeable for ChannelManager<ChanSigne
 		for (payment_hash, previous_hops) in channel_state.claimable_htlcs.iter() {
 			payment_hash.write(writer)?;
 			(previous_hops.len() as u64).write(writer)?;
-			for &(recvd_amt, ref previous_hop) in previous_hops.iter() {
-				recvd_amt.write(writer)?;
-				previous_hop.write(writer)?;
+			for htlc in previous_hops.iter() {
+				htlc.src.write(writer)?;
+				htlc.value.write(writer)?;
+				htlc.payment_data.write(writer)?;
 			}
 		}
 
@@ -3448,7 +3466,11 @@ impl<'a, R : ::std::io::Read, ChanSigner: ChannelKeys + Readable<R>> ReadableArg
 			let previous_hops_len: u64 = Readable::read(reader)?;
 			let mut previous_hops = Vec::with_capacity(cmp::min(previous_hops_len as usize, 2));
 			for _ in 0..previous_hops_len {
-				previous_hops.push((Readable::read(reader)?, Readable::read(reader)?));
+				previous_hops.push(ClaimableHTLC {
+					src: Readable::read(reader)?,
+					value: Readable::read(reader)?,
+					payment_data: Readable::read(reader)?,
+				});
 			}
 			claimable_htlcs.insert(payment_hash, previous_hops);
 		}

lightning/src/ln/msgs.rs

@@ -660,6 +660,11 @@ pub trait RoutingMessageHandler : Send + Sync {
 mod fuzzy_internal_msgs {
 	// These types aren't intended to be pub, but are exposed for direct fuzzing (as we deserialize
 	// them from untrusted input):
+	#[derive(Clone)]
+	pub(crate) struct FinalOnionHopData {
+		pub(crate) payment_secret: [u8; 32],
+		pub(crate) total_msat: u64,
+	}
 
 	pub(crate) enum OnionHopDataFormat {
 		Legacy { // aka Realm-0
@@ -668,7 +673,9 @@ mod fuzzy_internal_msgs {
 		NonFinalNode {
 			short_channel_id: u64,
 		},
-		FinalNode,
+		FinalNode {
+			payment_data: Option<FinalOnionHopData>,
+		},
 	}
 
 	pub struct OnionHopData {
@@ -1013,6 +1020,11 @@ impl_writeable!(UpdateAddHTLC, 32+8+8+32+4+1366, {
 	onion_routing_packet
 });
 
+impl_writeable!(FinalOnionHopData, 32+8, {
+	payment_secret,
+	total_msat
+});
+
 impl Writeable for OnionHopData {
 	fn write<W: Writer>(&self, w: &mut W) -> Result<(), ::std::io::Error> {
 		w.size_hint(33);
@@ -1030,11 +1042,19 @@ impl Writeable for OnionHopData {
 					(6, short_channel_id)
 				});
 			},
-			OnionHopDataFormat::FinalNode => {
-				encode_varint_length_prefixed_tlv!(w, {
-					(2, self.amt_to_forward),
-					(4, self.outgoing_cltv_value)
-				});
+			OnionHopDataFormat::FinalNode { ref payment_data } => {
+				if let Some(final_data) = payment_data {
+					encode_varint_length_prefixed_tlv!(w, {
+						(2, self.amt_to_forward),
+						(4, self.outgoing_cltv_value),
+						(8, final_data)
+					});
+				} else {
+					encode_varint_length_prefixed_tlv!(w, {
+						(2, self.amt_to_forward),
+						(4, self.outgoing_cltv_value)
+					});
+				}
 			},
 		}
 		match self.format {
@@ -1060,19 +1080,24 @@ impl<R: Read> Readable<R> for OnionHopData {
 			let mut amt: u64 = 0;
 			let mut cltv_value: u32 = 0;
 			let mut short_id: Option<u64> = None;
+			let mut payment_data: Option<FinalOnionHopData> = None;
 			decode_tlv!(&mut rd, {
 				(2, amt),
 				(4, cltv_value)
 			}, {
-				(6, short_id)
+				(6, short_id),
+				(8, payment_data)
 			});
 			rd.eat_remaining().map_err(|_| DecodeError::ShortRead)?;
 			let format = if let Some(short_channel_id) = short_id {
+				if payment_data.is_some() { return Err(DecodeError::InvalidValue); }
 				OnionHopDataFormat::NonFinalNode {
 					short_channel_id,
 				}
 			} else {
-				OnionHopDataFormat::FinalNode
+				OnionHopDataFormat::FinalNode {
+					payment_data
+				}
 			};
 			(format, amt, cltv_value)
 		} else {
@@ -2078,15 +2103,17 @@ mod tests {
 	#[test]
 	fn encoding_final_onion_hop_data() {
 		let mut msg = msgs::OnionHopData {
-			format: OnionHopDataFormat::FinalNode,
+			format: OnionHopDataFormat::FinalNode {
+				payment_data: None,
+			},
 			amt_to_forward: 0x0badf00d01020304,
 			outgoing_cltv_value: 0xffffffff,
 		};
 		let encoded_value = msg.encode();
 		let target_value = hex::decode("1002080badf00d010203040404ffffffff").unwrap();
 		assert_eq!(encoded_value, target_value);
 		msg = Readable::read(&mut Cursor::new(&target_value[..])).unwrap();
-		if let OnionHopDataFormat::FinalNode = msg.format { } else { panic!(); }
+		if let OnionHopDataFormat::FinalNode { payment_data: None } = msg.format { } else { panic!(); }
 		assert_eq!(msg.amt_to_forward, 0x0badf00d01020304);
 		assert_eq!(msg.outgoing_cltv_value, 0xffffffff);
 	}

lightning/src/ln/onion_utils.rs

@@ -123,7 +123,9 @@ pub(super) fn build_onion_payloads(route: &Route, starting_htlc_offset: u32) ->
 		res.insert(0, msgs::OnionHopData {
 			format: if hop.node_features.supports_variable_length_onion() {
 				if idx == 0 {
-					msgs::OnionHopDataFormat::FinalNode
+					msgs::OnionHopDataFormat::FinalNode {
+						payment_data: None,
+					}
 				} else {
 					msgs::OnionHopDataFormat::NonFinalNode {
 						short_channel_id: last_short_channel_id,