Merge pull request #2776 from jkczyz/2023-12-direct-connect-follow-ups

Folllow-ups to #2723
lightningdevkit · Dec 8, 2023 · 4b53654 · 4b53654
2 parents 2d26679 + 1c4d328
commit 4b53654
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 25 deletions.
diff --git a/fuzz/src/onion_message.rs b/fuzz/src/onion_message.rs
@@ -79,7 +79,7 @@ impl MessageRouter for TestMessageRouter {
 		Ok(OnionMessagePath {
 			intermediate_nodes: vec![],
 			destination,
-			addresses: None,
+			first_node_addresses: None,
 		})
 	}
 }

diff --git a/lightning/src/onion_message/functional_tests.rs b/lightning/src/onion_message/functional_tests.rs
@@ -51,7 +51,8 @@ impl MessageRouter for TestMessageRouter {
 		Ok(OnionMessagePath {
 			intermediate_nodes: vec![],
 			destination,
-			addresses: Some(vec![SocketAddress::TcpIpV4 { addr: [127, 0, 0, 1], port: 1000 }]),
+			first_node_addresses:
+				Some(vec![SocketAddress::TcpIpV4 { addr: [127, 0, 0, 1], port: 1000 }]),
 		})
 	}
 }
@@ -227,7 +228,7 @@ fn one_unblinded_hop() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::Node(nodes[1].node_id),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
 	nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response);
@@ -242,7 +243,7 @@ fn two_unblinded_hops() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![nodes[1].node_id],
 		destination: Destination::Node(nodes[2].node_id),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
 	nodes[2].custom_message_handler.expect_message(TestCustomMessage::Response);
@@ -259,7 +260,7 @@ fn one_blinded_hop() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::BlindedPath(blinded_path),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
 	nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response);
@@ -276,7 +277,7 @@ fn two_unblinded_two_blinded() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![nodes[1].node_id, nodes[2].node_id],
 		destination: Destination::BlindedPath(blinded_path),
-		addresses: None,
+		first_node_addresses: None,
 	};
 
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
@@ -294,7 +295,7 @@ fn three_blinded_hops() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::BlindedPath(blinded_path),
-		addresses: None,
+		first_node_addresses: None,
 	};
 
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
@@ -313,7 +314,7 @@ fn too_big_packet_error() {
 	let path = OnionMessagePath {
 		intermediate_nodes: hops,
 		destination: Destination::Node(hop_node_id),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	let err = nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap_err();
 	assert_eq!(err, SendError::TooBigPacket);
@@ -331,7 +332,7 @@ fn we_are_intro_node() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::BlindedPath(blinded_path),
-		addresses: None,
+		first_node_addresses: None,
 	};
 
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg.clone(), None).unwrap();
@@ -343,7 +344,7 @@ fn we_are_intro_node() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::BlindedPath(blinded_path),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
 	nodes[1].custom_message_handler.expect_message(TestCustomMessage::Response);
@@ -364,7 +365,7 @@ fn invalid_blinded_path_error() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::BlindedPath(blinded_path),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	let err = nodes[0].messenger.send_onion_message_using_path(path, test_msg.clone(), None).unwrap_err();
 	assert_eq!(err, SendError::TooFewBlindedHops);
@@ -380,7 +381,7 @@ fn reply_path() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![nodes[1].node_id, nodes[2].node_id],
 		destination: Destination::Node(nodes[3].node_id),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	let reply_path = BlindedPath::new_for_message(&[nodes[2].node_id, nodes[1].node_id, nodes[0].node_id], &*nodes[0].entropy_source, &secp_ctx).unwrap();
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg.clone(), Some(reply_path)).unwrap();
@@ -396,7 +397,7 @@ fn reply_path() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::BlindedPath(blinded_path),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	let reply_path = BlindedPath::new_for_message(&[nodes[2].node_id, nodes[1].node_id, nodes[0].node_id], &*nodes[0].entropy_source, &secp_ctx).unwrap();
 
@@ -431,7 +432,7 @@ fn invalid_custom_message_type() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::Node(nodes[1].node_id),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	let err = nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap_err();
 	assert_eq!(err, SendError::InvalidMessage);
@@ -444,7 +445,7 @@ fn peer_buffer_full() {
 	let path = OnionMessagePath {
 		intermediate_nodes: vec![],
 		destination: Destination::Node(nodes[1].node_id),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	for _ in 0..188 { // Based on MAX_PER_PEER_BUFFER_SIZE in OnionMessenger
 		nodes[0].messenger.send_onion_message_using_path(path.clone(), test_msg.clone(), None).unwrap();
@@ -469,7 +470,7 @@ fn many_hops() {
 	let path = OnionMessagePath {
 		intermediate_nodes,
 		destination: Destination::Node(nodes[num_nodes-1].node_id),
-		addresses: None,
+		first_node_addresses: None,
 	};
 	nodes[0].messenger.send_onion_message_using_path(path, test_msg, None).unwrap();
 	nodes[num_nodes-1].custom_message_handler.expect_message(TestCustomMessage::Response);

diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs
@@ -87,7 +87,7 @@ pub(super) const MAX_TIMER_TICKS: usize = 2;
 /// #         Ok(OnionMessagePath {
 /// #             intermediate_nodes: vec![hop_node_id1, hop_node_id2],
 /// #             destination,
-/// #             addresses: None,
+/// #             first_node_addresses: None,
 /// #         })
 /// #     }
 /// # }
@@ -223,6 +223,13 @@ impl OnionMessageRecipient {
 			*self = OnionMessageRecipient::ConnectedPeer(new_pending_messages);
 		}
 	}
+
+	fn is_connected(&self) -> bool {
+		match self {
+			OnionMessageRecipient::ConnectedPeer(..) => true,
+			OnionMessageRecipient::PendingConnection(..) => false,
+		}
+	}
 }
 
 /// An [`OnionMessage`] for [`OnionMessenger`] to send.
@@ -292,7 +299,9 @@ where
 	) -> Result<OnionMessagePath, ()> {
 		let first_node = destination.first_node();
 		if peers.contains(&first_node) {
-			Ok(OnionMessagePath { intermediate_nodes: vec![], destination, addresses: None })
+			Ok(OnionMessagePath {
+				intermediate_nodes: vec![], destination, first_node_addresses: None
+			})
 		} else {
 			let network_graph = self.network_graph.deref().read_only();
 			let node_announcement = network_graph
@@ -303,8 +312,10 @@ where
 
 			match node_announcement {
 				Some(node_announcement) if node_announcement.features.supports_onion_messages() => {
-					let addresses = Some(node_announcement.addresses.clone());
-					Ok(OnionMessagePath { intermediate_nodes: vec![], destination, addresses })
+					let first_node_addresses = Some(node_announcement.addresses.clone());
+					Ok(OnionMessagePath {
+						intermediate_nodes: vec![], destination, first_node_addresses
+					})
 				},
 				_ => Err(()),
 			}
@@ -325,7 +336,7 @@ pub struct OnionMessagePath {
 	///
 	/// Only needs to be set if a connection to the node is required. [`OnionMessenger`] may use
 	/// this to initiate such a connection.
-	pub addresses: Option<Vec<SocketAddress>>,
+	pub first_node_addresses: Option<Vec<SocketAddress>>,
 }
 
 impl OnionMessagePath {
@@ -459,7 +470,8 @@ pub enum PeeledOnion<T: OnionMessageContents> {
 /// Creates an [`OnionMessage`] with the given `contents` for sending to the destination of
 /// `path`.
 ///
-/// Returns both the node id of the peer to send the message to and the message itself.
+/// Returns the node id of the peer to send the message to, the message itself, and any addresses
+/// need to connect to the first node.
 pub fn create_onion_message<ES: Deref, NS: Deref, T: OnionMessageContents>(
 	entropy_source: &ES, node_signer: &NS, secp_ctx: &Secp256k1<secp256k1::All>,
 	path: OnionMessagePath, contents: T, reply_path: Option<BlindedPath>,
@@ -468,7 +480,7 @@ where
 	ES::Target: EntropySource,
 	NS::Target: NodeSigner,
 {
-	let OnionMessagePath { intermediate_nodes, mut destination, addresses } = path;
+	let OnionMessagePath { intermediate_nodes, mut destination, first_node_addresses } = path;
 	if let Destination::BlindedPath(BlindedPath { ref blinded_hops, .. }) = destination {
 		if blinded_hops.is_empty() {
 			return Err(SendError::TooFewBlindedHops);
@@ -510,7 +522,7 @@ where
 		packet_payloads, packet_keys, prng_seed).map_err(|()| SendError::TooBigPacket)?;
 
 	let message = OnionMessage { blinding_point, onion_routing_packet };
-	Ok((first_node_id, message, addresses))
+	Ok((first_node_id, message, first_node_addresses))
 }
 
 /// Decode one layer of an incoming [`OnionMessage`].
@@ -724,7 +736,11 @@ where
 			},
 			hash_map::Entry::Occupied(mut e) => {
 				e.get_mut().enqueue_message(onion_message);
-				Ok(SendSuccess::Buffered)
+				if e.get().is_connected() {
+					Ok(SendSuccess::Buffered)
+				} else {
+					Ok(SendSuccess::BufferedAwaitingConnection(first_node_id))
+				}
 			},
 		}
 	}