Support sending and receiving reply paths

valentinewallace · valentinewallace · commit 9e6c7dc6c85e · 2022-08-22T14:40:22.000-04:00
diff --git a/lightning/src/onion_message/blinded_route.rs b/lightning/src/onion_message/blinded_route.rs
@@ -13,8 +13,9 @@ use bitcoin::secp256k1::{self, PublicKey, Secp256k1, SecretKey};
 
 use chain::keysinterface::{KeysInterface, Sign};
 use super::utils;
+use ln::msgs::DecodeError;
 use util::chacha20poly1305rfc::ChaChaPolyWriteAdapter;
-use util::ser::{VecWriter, Writeable, Writer};
+use util::ser::{FixedLengthReader, LengthReadable, Readable, ReadTrackingReader, VecReadWrapper, VecWriter, Writeable, Writer};
 
 use core::iter::FromIterator;
 use io;
@@ -113,6 +114,41 @@ fn encrypt_payload<P: Writeable>(payload: P, encrypted_tlvs_ss: [u8; 32]) -> Vec
 	writer.0
 }
 
+impl Writeable for BlindedRoute {
+	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
+		self.introduction_node_id.write(w)?;
+		self.blinding_point.write(w)?;
+		(self.blinded_hops.len() as u8).write(w)?;
+		for hop in &self.blinded_hops {
+			hop.write(w)?;
+		}
+		Ok(())
+	}
+}
+
+impl Readable for BlindedRoute {
+	fn read<R: io::Read>(r: &mut R) -> Result<Self, DecodeError> {
+		let introduction_node_id = Readable::read(r)?;
+		let blinding_point = Readable::read(r)?;
+		let num_hops: u8 = Readable::read(r)?;
+		if num_hops == 0 { return Err(DecodeError::InvalidValue) }
+		let mut blinded_hops: Vec<BlindedHop> = Vec::with_capacity(num_hops.into());
+		for i in 0..num_hops {
+			blinded_hops.push(Readable::read(r)?);
+		}
+		Ok(BlindedRoute {
+			introduction_node_id,
+			blinding_point,
+			blinded_hops,
+		})
+	}
+}
+
+impl_writeable!(BlindedHop, {
+	blinded_node_id,
+	encrypted_payload
+});
+
 /// TLVs to encode in an intermediate onion message packet's hop data. When provided in a blinded
 /// route, they are encoded into [`BlindedHop::encrypted_payload`].
 pub(crate) struct ForwardTlvs {
diff --git a/lightning/src/onion_message/functional_tests.rs b/lightning/src/onion_message/functional_tests.rs
@@ -72,15 +72,15 @@ fn pass_along_path(path: &Vec<MessengerNode>, expected_path_id: Option<[u8; 32]>
 fn one_hop() {
 	let nodes = create_nodes(2);
 
-	nodes[0].messenger.send_onion_message(&[], Destination::Node(nodes[1].get_node_pk())).unwrap();
+	nodes[0].messenger.send_onion_message(&[], Destination::Node(nodes[1].get_node_pk()), None).unwrap();
 	pass_along_path(&nodes, None);
 }
 
 #[test]
 fn two_unblinded_hops() {
 	let nodes = create_nodes(3);
 
-	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk()], Destination::Node(nodes[2].get_node_pk())).unwrap();
+	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk()], Destination::Node(nodes[2].get_node_pk()), None).unwrap();
 	pass_along_path(&nodes, None);
 }
 
@@ -91,7 +91,7 @@ fn two_unblinded_two_blinded() {
 	let secp_ctx = Secp256k1::new();
 	let blinded_route = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[3].get_node_pk(), nodes[4].get_node_pk()], &*nodes[4].keys_manager, &secp_ctx).unwrap();
 
-	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], Destination::BlindedRoute(blinded_route)).unwrap();
+	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], Destination::BlindedRoute(blinded_route), None).unwrap();
 	pass_along_path(&nodes, None);
 }
 
@@ -102,7 +102,7 @@ fn three_blinded_hops() {
 	let secp_ctx = Secp256k1::new();
 	let blinded_route = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[1].get_node_pk(), nodes[2].get_node_pk(), nodes[3].get_node_pk()], &*nodes[3].keys_manager, &secp_ctx).unwrap();
 
-	nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route)).unwrap();
+	nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), None).unwrap();
 	pass_along_path(&nodes, None);
 }
 
@@ -116,7 +116,7 @@ fn too_big_packet_error() {
 	let hop_node_id = PublicKey::from_secret_key(&secp_ctx, &hop_secret);
 
 	let hops = [hop_node_id; 400];
-	let err = nodes[0].messenger.send_onion_message(&hops, Destination::Node(hop_node_id)).unwrap_err();
+	let err = nodes[0].messenger.send_onion_message(&hops, Destination::Node(hop_node_id), None).unwrap_err();
 	assert_eq!(err, SendError::TooBigPacket);
 }
 
@@ -129,13 +129,57 @@ fn invalid_blinded_route_error() {
 	let secp_ctx = Secp256k1::new();
 	let mut blinded_route = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], &*nodes[2].keys_manager, &secp_ctx).unwrap();
 	blinded_route.blinded_hops.clear();
-	let err = nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route)).unwrap_err();
+	let err = nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), None).unwrap_err();
 	assert_eq!(err, SendError::TooFewBlindedHops);
 
 	// 1 hop
 	let mut blinded_route = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], &*nodes[2].keys_manager, &secp_ctx).unwrap();
 	blinded_route.blinded_hops.remove(0);
 	assert_eq!(blinded_route.blinded_hops.len(), 1);
-	let err = nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route)).unwrap_err();
+	let err = nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), None).unwrap_err();
 	assert_eq!(err, SendError::TooFewBlindedHops);
 }
+
+#[test]
+fn two_hop_reply_path() {
+	let mut nodes = create_nodes(3);
+	let secp_ctx = Secp256k1::new();
+
+	// Destination::Node
+	let reply_path = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx).unwrap();
+	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk()], Destination::Node(nodes[2].get_node_pk()), Some(reply_path)).unwrap();
+	pass_along_path(&nodes, None);
+	nodes.reverse();
+	pass_along_path(&nodes, None);
+
+	// Destination::BlindedRoute
+	let blinded_route = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], &*nodes[2].keys_manager, &secp_ctx).unwrap();
+	let reply_path = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx).unwrap();
+
+	nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), Some(reply_path)).unwrap();
+	pass_along_path(&nodes, None);
+	nodes.reverse();
+	pass_along_path(&nodes, None);
+}
+
+#[test]
+fn three_hop_reply_path() {
+	let mut nodes = create_nodes(4);
+	let secp_ctx = Secp256k1::new();
+
+	// Destination::Node
+	let reply_path = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[2].get_node_pk(), nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx).unwrap();
+	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], Destination::Node(nodes[3].get_node_pk()), Some(reply_path)).unwrap();
+	pass_along_path(&nodes, None);
+	nodes.reverse();
+	pass_along_path(&nodes, None);
+
+	// Destination::BlindedRoute
+	let blinded_route = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[1].get_node_pk(), nodes[2].get_node_pk(), nodes[3].get_node_pk()], &*nodes[3].keys_manager, &secp_ctx).unwrap();
+	let reply_path = BlindedRoute::new::<EnforcingSigner, _, _>(&[nodes[2].get_node_pk(), nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx).unwrap();
+
+	nodes[0].messenger.send_onion_message(&[], Destination::BlindedRoute(blinded_route), Some(reply_path)).unwrap();
+	pass_along_path(&nodes, None);
+	nodes.reverse();
+	pass_along_path(&nodes, None);
+}
diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs
@@ -142,7 +142,7 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
 
 	/// Send an empty onion message to `destination`, routing it through `intermediate_nodes`.
 	/// See [`OnionMessenger`] for example usage.
-	pub fn send_onion_message(&self, intermediate_nodes: &[PublicKey], destination: Destination) -> Result<(), SendError> {
+	pub fn send_onion_message(&self, intermediate_nodes: &[PublicKey], destination: Destination, reply_path: Option<BlindedRoute>) -> Result<(), SendError> {
 		if let Destination::BlindedRoute(BlindedRoute { ref blinded_hops, .. }) = destination {
 			if blinded_hops.len() < 2 {
 				return Err(SendError::TooFewBlindedHops);
@@ -160,7 +160,7 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
 			}
 		};
 		let (packet_payloads, packet_keys) = packet_payloads_and_keys(
-			&self.secp_ctx, intermediate_nodes, destination, &blinding_secret)
+			&self.secp_ctx, intermediate_nodes, destination, reply_path, &blinding_secret)
 			.map_err(|e| SendError::Secp256k1(e))?;
 
 		let prng_seed = self.keys_manager.get_secure_random_bytes();
@@ -209,9 +209,17 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
 			msg.onion_routing_packet.hmac, control_tlvs_ss)
 		{
 			Ok((Payload::Receive {
-				control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id })
+				control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id }), reply_path,
 			}, None)) => {
-				log_info!(self.logger, "Received an onion message with path_id: {:02x?}", path_id);
+				log_info!(self.logger,
+					"Received an onion message with path_id: {:02x?} and {}reply_path",
+						path_id, if reply_path.is_some() { "" } else { "no " });
+				#[cfg(test)]
+				{
+					if let Some(reply_path) = reply_path {
+						self.send_onion_message(&[], Destination::BlindedRoute(reply_path), None).unwrap();
+					}
+				}
 			},
 			Ok((Payload::Forward(ForwardControlTlvs::Unblinded(ForwardTlvs {
 				next_node_id, next_blinding_override
@@ -296,7 +304,8 @@ pub type SimpleRefOnionMessenger<'a, 'b, L> = OnionMessenger<InMemorySigner, &'a
 /// Construct onion packet payloads and keys for sending an onion message along the given
 /// `unblinded_path` to the given `destination`.
 fn packet_payloads_and_keys<T: secp256k1::Signing + secp256k1::Verification>(
-	secp_ctx: &Secp256k1<T>, unblinded_path: &[PublicKey], destination: Destination, session_priv: &SecretKey
+	secp_ctx: &Secp256k1<T>, unblinded_path: &[PublicKey], destination: Destination, mut reply_path:
+	Option<BlindedRoute>, session_priv: &SecretKey
 ) -> Result<(Vec<(Payload, [u8; 32])>, Vec<onion_utils::OnionKeys>), secp256k1::Error> {
 	let num_hops = unblinded_path.len() + destination.num_hops();
 	let mut payloads = Vec::with_capacity(num_hops);
@@ -341,6 +350,7 @@ fn packet_payloads_and_keys<T: secp256k1::Signing + secp256k1::Verification>(
 		} else if let Some(encrypted_payload) = enc_payload_opt {
 			payloads.push((Payload::Receive {
 				control_tlvs: ReceiveControlTlvs::Blinded(encrypted_payload),
+				reply_path: reply_path.take(),
 			}, control_tlvs_ss));
 		}
 
@@ -358,7 +368,8 @@ fn packet_payloads_and_keys<T: secp256k1::Signing + secp256k1::Verification>(
 
 	if let Some(control_tlvs_ss) = prev_control_tlvs_ss {
 		payloads.push((Payload::Receive {
-			control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id: None, })
+			control_tlvs: ReceiveControlTlvs::Unblinded(ReceiveTlvs { path_id: None, }),
+			reply_path: reply_path.take(),
 		}, control_tlvs_ss));
 	}
 
diff --git a/lightning/src/onion_message/packet.rs b/lightning/src/onion_message/packet.rs
@@ -14,7 +14,7 @@ use bitcoin::secp256k1::ecdh::SharedSecret;
 
 use ln::msgs::DecodeError;
 use ln::onion_utils;
-use super::blinded_route::{ForwardTlvs, ReceiveTlvs};
+use super::blinded_route::{BlindedRoute, ForwardTlvs, ReceiveTlvs};
 use util::chacha20poly1305rfc::{ChaChaPolyReadAdapter, ChaChaPolyWriteAdapter};
 use util::ser::{BigSize, FixedLengthReader, LengthRead, LengthReadable, LengthReadableArgs, Readable, ReadableArgs, Writeable, Writer};
 
@@ -98,8 +98,8 @@ pub(super) enum Payload {
 	/// This payload is for the final hop.
 	Receive {
 		control_tlvs: ReceiveControlTlvs,
+		reply_path: Option<BlindedRoute>,
 		// Coming soon:
-		// reply_path: Option<BlindedRoute>,
 		// message: Message,
 	}
 }
@@ -135,21 +135,31 @@ pub(super) enum ReceiveControlTlvs {
 impl Writeable for (Payload, [u8; 32]) {
 	fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
 		match &self.0 {
-			Payload::Forward(ForwardControlTlvs::Blinded(encrypted_bytes)) |
-			Payload::Receive { control_tlvs: ReceiveControlTlvs::Blinded(encrypted_bytes)} => {
+			Payload::Forward(ForwardControlTlvs::Blinded(encrypted_bytes)) => {
 				encode_varint_length_prefixed_tlv!(w, {
 					(4, encrypted_bytes, vec_type)
 				})
 			},
+			Payload::Receive {
+				control_tlvs: ReceiveControlTlvs::Blinded(encrypted_bytes), reply_path
+			} => {
+				encode_varint_length_prefixed_tlv!(w, {
+					(2, reply_path, option),
+					(4, encrypted_bytes, vec_type)
+				})
+			},
 			Payload::Forward(ForwardControlTlvs::Unblinded(control_tlvs)) => {
 				let write_adapter = ChaChaPolyWriteAdapter::new(self.1, &control_tlvs);
 				encode_varint_length_prefixed_tlv!(w, {
 					(4, write_adapter, required)
 				})
 			},
-			Payload::Receive { control_tlvs: ReceiveControlTlvs::Unblinded(control_tlvs)} => {
+			Payload::Receive {
+				control_tlvs: ReceiveControlTlvs::Unblinded(control_tlvs), reply_path,
+			} => {
 				let write_adapter = ChaChaPolyWriteAdapter::new(self.1, &control_tlvs);
 				encode_varint_length_prefixed_tlv!(w, {
+					(2, reply_path, option),
 					(4, write_adapter, required)
 				})
 			},
@@ -163,12 +173,11 @@ impl ReadableArgs<SharedSecret> for Payload {
 	fn read<R: Read>(mut r: &mut R, encrypted_tlvs_ss: SharedSecret) -> Result<Self, DecodeError> {
 		let v: BigSize = Readable::read(r)?;
 		let mut rd = FixedLengthReader::new(r, v.0);
-		// TODO: support reply paths
-		let mut _reply_path_bytes: Option<Vec<u8>> = Some(Vec::new());
+		let mut reply_path: Option<BlindedRoute> = None;
 		let mut read_adapter: Option<ChaChaPolyReadAdapter<ControlTlvs>> = None;
 		let rho = onion_utils::gen_rho_from_shared_secret(&encrypted_tlvs_ss.secret_bytes());
 		decode_tlv_stream!(&mut rd, {
-			(2, _reply_path_bytes, vec_type),
+			(2, reply_path, option),
 			(4, read_adapter, (option: LengthReadableArgs, rho))
 		});
 		rd.eat_remaining().map_err(|_| DecodeError::ShortRead)?;
@@ -179,7 +188,7 @@ impl ReadableArgs<SharedSecret> for Payload {
 				Ok(Payload::Forward(ForwardControlTlvs::Unblinded(tlvs)))
 			},
 			Some(ChaChaPolyReadAdapter { readable: ControlTlvs::Receive(tlvs)}) => {
-				Ok(Payload::Receive { control_tlvs: ReceiveControlTlvs::Unblinded(tlvs)})
+				Ok(Payload::Receive { control_tlvs: ReceiveControlTlvs::Unblinded(tlvs), reply_path })
 			},
 		}
 	}
