Merge pull request #2973 from valentinewallace/2024-03-om-mailbox

Support intercepting onion messages for offline peers

Author: arik-so

lightning/src/events/mod.rs

@@ -1141,6 +1141,30 @@ pub enum Event {
 	///
 	/// [`ChannelHandshakeConfig::negotiate_anchors_zero_fee_htlc_tx`]: crate::util::config::ChannelHandshakeConfig::negotiate_anchors_zero_fee_htlc_tx
 	BumpTransaction(BumpTransactionEvent),
+	/// We received an onion message that is intended to be forwarded to a peer
+	/// that is currently offline. This event will only be generated if the
+	/// `OnionMessenger` was initialized with
+	/// [`OnionMessenger::new_with_offline_peer_interception`], see its docs.
+	///
+	/// [`OnionMessenger::new_with_offline_peer_interception`]: crate::onion_message::messenger::OnionMessenger::new_with_offline_peer_interception
+	OnionMessageIntercepted {
+		/// The node id of the offline peer.
+		peer_node_id: PublicKey,
+		/// The onion message intended to be forwarded to `peer_node_id`.
+		message: msgs::OnionMessage,
+	},
+	/// Indicates that an onion message supporting peer has come online and it may
+	/// be time to forward any onion messages that were previously intercepted for
+	/// them. This event will only be generated if the `OnionMessenger` was
+	/// initialized with
+	/// [`OnionMessenger::new_with_offline_peer_interception`], see its docs.
+	///
+	/// [`OnionMessenger::new_with_offline_peer_interception`]: crate::onion_message::messenger::OnionMessenger::new_with_offline_peer_interception
+	OnionMessagePeerConnected {
+		/// The node id of the peer we just connected to, who advertises support for
+		/// onion messages.
+		peer_node_id: PublicKey,
+	}
 }
 
 impl Writeable for Event {
@@ -1404,6 +1428,19 @@ impl Writeable for Event {
 				35u8.write(writer)?;
 				// Never write ConnectionNeeded events as buffered onion messages aren't serialized.
 			},
+			&Event::OnionMessageIntercepted { ref peer_node_id, ref message } => {
+				37u8.write(writer)?;
+				write_tlv_fields!(writer, {
+					(0, peer_node_id, required),
+					(2, message, required),
+				});
+			},
+			&Event::OnionMessagePeerConnected { ref peer_node_id } => {
+				39u8.write(writer)?;
+				write_tlv_fields!(writer, {
+					(0, peer_node_id, required),
+				});
+			}
 			// Note that, going forward, all new events must only write data inside of
 			// `write_tlv_fields`. Versions 0.0.101+ will ignore odd-numbered events that write
 			// data via `write_tlv_fields`.
@@ -1814,6 +1851,29 @@ impl MaybeReadable for Event {
 			},
 			// Note that we do not write a length-prefixed TLV for ConnectionNeeded events.
 			35u8 => Ok(None),
+			37u8 => {
+				let mut f = || {
+					_init_and_read_len_prefixed_tlv_fields!(reader, {
+						(0, peer_node_id, required),
+						(2, message, required),
+					});
+					Ok(Some(Event::OnionMessageIntercepted {
+						peer_node_id: peer_node_id.0.unwrap(), message: message.0.unwrap()
+					}))
+				};
+				f()
+			},
+			39u8 => {
+				let mut f = || {
+					_init_and_read_len_prefixed_tlv_fields!(reader, {
+						(0, peer_node_id, required),
+					});
+					Ok(Some(Event::OnionMessagePeerConnected {
+						peer_node_id: peer_node_id.0.unwrap()
+					}))
+				};
+				f()
+			},
 			// Versions prior to 0.0.100 did not ignore odd types, instead returning InvalidValue.
 			// Version 0.0.100 failed to properly ignore odd types, possibly resulting in corrupt
 			// reads.

lightning/src/onion_message/functional_tests.rs

@@ -59,6 +59,17 @@ struct MessengerNode {
 	>>
 }
 
+impl Drop for MessengerNode {
+	fn drop(&mut self) {
+		#[cfg(feature = "std")] {
+			if std::thread::panicking() {
+				return;
+			}
+		}
+		assert!(release_events(self).is_empty());
+	}
+}
+
 struct TestOffersMessageHandler {}
 
 impl OffersMessageHandler for TestOffersMessageHandler {
@@ -162,22 +173,41 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
 }
 
 fn create_nodes(num_messengers: u8) -> Vec<MessengerNode> {
-	let secrets = (1..=num_messengers)
+	let cfgs = (1..=num_messengers)
 		.into_iter()
-		.map(|i| SecretKey::from_slice(&[i; 32]).unwrap())
+		.map(|_| MessengerCfg::new())
 		.collect();
-	create_nodes_using_secrets(secrets)
+	create_nodes_using_cfgs(cfgs)
 }
 
-fn create_nodes_using_secrets(secrets: Vec<SecretKey>) -> Vec<MessengerNode> {
+struct MessengerCfg {
+	secret_override: Option<SecretKey>,
+	intercept_offline_peer_oms: bool,
+}
+impl MessengerCfg {
+	fn new() -> Self {
+		Self { secret_override: None, intercept_offline_peer_oms: false }
+	}
+	fn with_node_secret(mut self, secret: SecretKey) -> Self {
+		self.secret_override = Some(secret);
+		self
+	}
+	fn with_offline_peer_interception(mut self) -> Self {
+		self.intercept_offline_peer_oms = true;
+		self
+	}
+}
+
+fn create_nodes_using_cfgs(cfgs: Vec<MessengerCfg>) -> Vec<MessengerNode> {
 	let gossip_logger = Arc::new(test_utils::TestLogger::with_id("gossip".to_string()));
 	let network_graph = Arc::new(NetworkGraph::new(Network::Testnet, gossip_logger.clone()));
 	let gossip_sync = Arc::new(
 		P2PGossipSync::new(network_graph.clone(), None, gossip_logger)
 	);
 
 	let mut nodes = Vec::new();
-	for (i, secret_key) in secrets.into_iter().enumerate() {
+	for (i, cfg) in cfgs.into_iter().enumerate() {
+		let secret_key = cfg.secret_override.unwrap_or(SecretKey::from_slice(&[(i + 1) as u8; 32]).unwrap());
 		let logger = Arc::new(test_utils::TestLogger::with_id(format!("node {}", i)));
 		let seed = [i as u8; 32];
 		let entropy_source = Arc::new(test_utils::TestKeysInterface::new(&seed, Network::Testnet));
@@ -189,14 +219,24 @@ fn create_nodes_using_secrets(secrets: Vec<SecretKey>) -> Vec<MessengerNode> {
 		);
 		let offers_message_handler = Arc::new(TestOffersMessageHandler {});
 		let custom_message_handler = Arc::new(TestCustomMessageHandler::new());
+		let messenger = if cfg.intercept_offline_peer_oms {
+			OnionMessenger::new_with_offline_peer_interception(
+				entropy_source.clone(), node_signer.clone(), logger.clone(),
+				node_id_lookup, message_router, offers_message_handler,
+				custom_message_handler.clone()
+			)
+		} else {
+			OnionMessenger::new(
+				entropy_source.clone(), node_signer.clone(), logger.clone(),
+				node_id_lookup, message_router, offers_message_handler,
+				custom_message_handler.clone()
+			)
+		};
 		nodes.push(MessengerNode {
 			privkey: secret_key,
 			node_id: node_signer.get_node_id(Recipient::Node).unwrap(),
-			entropy_source: entropy_source.clone(),
-			messenger: OnionMessenger::new(
-				entropy_source, node_signer, logger.clone(), node_id_lookup, message_router,
-				offers_message_handler, custom_message_handler.clone()
-			),
+			entropy_source,
+			messenger,
 			custom_message_handler,
 			gossip_sync: gossip_sync.clone(),
 		});
@@ -369,11 +409,10 @@ fn we_are_intro_node() {
 
 #[test]
 fn invalid_blinded_path_error() {
-	// Make sure we error as expected if a provided blinded path has 0 or 1 hops.
+	// Make sure we error as expected if a provided blinded path has 0 hops.
 	let nodes = create_nodes(3);
 	let test_msg = TestCustomMessage::Response;
 
-	// 0 hops
 	let secp_ctx = Secp256k1::new();
 	let mut blinded_path = BlindedPath::new_for_message(&[nodes[1].node_id, nodes[2].node_id], &*nodes[2].entropy_source, &secp_ctx).unwrap();
 	blinded_path.blinded_hops.clear();
@@ -549,18 +588,83 @@ fn drops_buffered_messages_waiting_for_peer_connection() {
 	assert!(nodes[0].messenger.next_onion_message_for_peer(nodes[1].node_id).is_none());
 }
 
+#[test]
+fn intercept_offline_peer_oms() {
+	// Ensure that if OnionMessenger is initialized with
+	// new_with_offline_peer_interception, we will intercept OMs for offline
+	// peers, generate the right events, and forward OMs when they are re-injected
+	// by the user.
+	let node_cfgs = vec![MessengerCfg::new(), MessengerCfg::new().with_offline_peer_interception(), MessengerCfg::new()];
+	let mut nodes = create_nodes_using_cfgs(node_cfgs);
+
+	let peer_conn_evs = release_events(&nodes[1]);
+	assert_eq!(peer_conn_evs.len(), 2);
+	for (i, ev) in peer_conn_evs.iter().enumerate() {
+		match ev {
+			Event::OnionMessagePeerConnected { peer_node_id } => {
+				let node_idx = if i == 0 { 0 } else { 2 };
+				assert_eq!(peer_node_id, &nodes[node_idx].node_id);
+			},
+			_ => panic!()
+		}
+	}
+
+	let message = TestCustomMessage::Response;
+	let secp_ctx = Secp256k1::new();
+	let blinded_path = BlindedPath::new_for_message(
+		&[nodes[1].node_id, nodes[2].node_id], &*nodes[2].entropy_source, &secp_ctx
+	).unwrap();
+	let destination = Destination::BlindedPath(blinded_path);
+
+	// Disconnect the peers to ensure we intercept the OM.
+	disconnect_peers(&nodes[1], &nodes[2]);
+	nodes[0].messenger.send_onion_message(message, destination, None).unwrap();
+	let mut final_node_vec = nodes.split_off(2);
+	pass_along_path(&nodes);
+
+	let mut events = release_events(&nodes[1]);
+	assert_eq!(events.len(), 1);
+	let onion_message = match events.remove(0) {
+		Event::OnionMessageIntercepted { peer_node_id, message } => {
+			assert_eq!(peer_node_id, final_node_vec[0].node_id);
+			message
+		},
+		_ => panic!()
+	};
+
+	// Ensure that we'll refuse to forward the re-injected OM until after the
+	// outbound peer comes back online.
+	let err = nodes[1].messenger.forward_onion_message(onion_message.clone(), &final_node_vec[0].node_id).unwrap_err();
+	assert_eq!(err, SendError::InvalidFirstHop(final_node_vec[0].node_id));
+
+	connect_peers(&nodes[1], &final_node_vec[0]);
+	let peer_conn_ev = release_events(&nodes[1]);
+	assert_eq!(peer_conn_ev.len(), 1);
+	match peer_conn_ev[0] {
+		Event::OnionMessagePeerConnected { peer_node_id } => {
+			assert_eq!(peer_node_id, final_node_vec[0].node_id);
+		},
+		_ => panic!()
+	}
+
+	nodes[1].messenger.forward_onion_message(onion_message, &final_node_vec[0].node_id).unwrap();
+	final_node_vec[0].custom_message_handler.expect_message(TestCustomMessage::Response);
+	pass_along_path(&vec![nodes.remove(1), final_node_vec.remove(0)]);
+}
+
 #[test]
 fn spec_test_vector() {
-	let secret_keys = [
+	let node_cfgs = [
 		"4141414141414141414141414141414141414141414141414141414141414141", // Alice
 		"4242424242424242424242424242424242424242424242424242424242424242", // Bob
 		"4343434343434343434343434343434343434343434343434343434343434343", // Carol
 		"4444444444444444444444444444444444444444444444444444444444444444", // Dave
 	]
 		.iter()
-		.map(|secret| SecretKey::from_slice(&<Vec<u8>>::from_hex(secret).unwrap()).unwrap())
+		.map(|secret_hex| SecretKey::from_slice(&<Vec<u8>>::from_hex(secret_hex).unwrap()).unwrap())
+		.map(|secret| MessengerCfg::new().with_node_secret(secret))
 		.collect();
-	let nodes = create_nodes_using_secrets(secret_keys);
+	let nodes = create_nodes_using_cfgs(node_cfgs);
 
 	// Hardcode the sender->Alice onion message, because it includes an unknown TLV of type 1, which
 	// LDK doesn't support constructing.