Add onion mailbox for async receivers

Author: joostjager

src/builder.rs

@@ -27,6 +27,7 @@ use crate::liquidity::{
 };
 use crate::logger::{log_error, LdkLogger, LogLevel, LogWriter, Logger};
 use crate::message_handler::NodeCustomMessageHandler;
+use crate::om_mailbox::OnionMessageMailbox;
 use crate::peer_store::PeerStore;
 use crate::runtime::Runtime;
 use crate::tx_broadcaster::TransactionBroadcaster;
@@ -1378,6 +1379,10 @@ fn build_with_store_internal(
 			100;
 	}
 
+	if config.async_payment_services_enabled {
+		user_config.accept_forwards_to_priv_channels = true;
+	}
+
 	let message_router =
 		Arc::new(MessageRouter::new(Arc::clone(&network_graph), Arc::clone(&keys_manager)));
 
@@ -1448,17 +1453,31 @@ fn build_with_store_internal(
 	}
 
 	// Initialize the PeerManager
-	let onion_messenger: Arc<OnionMessenger> = Arc::new(OnionMessenger::new(
-		Arc::clone(&keys_manager),
-		Arc::clone(&keys_manager),
-		Arc::clone(&logger),
-		Arc::clone(&channel_manager),
-		message_router,
-		Arc::clone(&channel_manager),
-		Arc::clone(&channel_manager),
-		IgnoringMessageHandler {},
-		IgnoringMessageHandler {},
-	));
+	let onion_messenger: Arc<OnionMessenger> = if config.async_payment_services_enabled {
+		Arc::new(OnionMessenger::new_with_offline_peer_interception(
+			Arc::clone(&keys_manager),
+			Arc::clone(&keys_manager),
+			Arc::clone(&logger),
+			Arc::clone(&channel_manager),
+			message_router,
+			Arc::clone(&channel_manager),
+			Arc::clone(&channel_manager),
+			IgnoringMessageHandler {},
+			IgnoringMessageHandler {},
+		))
+	} else {
+		Arc::new(OnionMessenger::new(
+			Arc::clone(&keys_manager),
+			Arc::clone(&keys_manager),
+			Arc::clone(&logger),
+			Arc::clone(&channel_manager),
+			message_router,
+			Arc::clone(&channel_manager),
+			Arc::clone(&channel_manager),
+			IgnoringMessageHandler {},
+			IgnoringMessageHandler {},
+		))
+	};
 	let ephemeral_bytes: [u8; 32] = keys_manager.get_secure_random_bytes();
 
 	// Initialize the GossipSource
@@ -1645,6 +1664,8 @@ fn build_with_store_internal(
 		},
 	};
 
+	let om_mailbox = Arc::new(OnionMessageMailbox::new());
+
 	let (stop_sender, _) = tokio::sync::watch::channel(());
 	let (background_processor_stop_sender, _) = tokio::sync::watch::channel(());
 	let is_running = Arc::new(RwLock::new(false));
@@ -1677,6 +1698,7 @@ fn build_with_store_internal(
 		is_running,
 		is_listening,
 		node_metrics,
+		om_mailbox,
 	})
 }
 

src/config.rs

@@ -330,6 +330,8 @@ pub(crate) fn default_user_config(config: &Config) -> UserConfig {
 		user_config.channel_handshake_limits.force_announced_channel_preference = true;
 	}
 
+	user_config.enable_htlc_hold = true;
+
 	user_config
 }
 

src/event.rs

@@ -5,7 +5,8 @@
 // http://opensource.org/licenses/MIT>, at your option. You may not use this file except in
 // accordance with one or both of these licenses.
 
-use crate::types::{CustomTlvRecord, DynStore, PaymentStore, Sweeper, Wallet};
+use crate::om_mailbox::OnionMessageMailbox;
+use crate::types::{CustomTlvRecord, DynStore, OnionMessenger, PaymentStore, Sweeper, Wallet};
 use crate::{
 	hex_utils, BumpTransactionEventHandler, ChannelManager, Error, Graph, PeerInfo, PeerStore,
 	UserChannelId,
@@ -459,6 +460,8 @@ where
 	logger: L,
 	config: Arc<Config>,
 	static_invoice_store: Option<StaticInvoiceStore>,
+	onion_messenger: Arc<OnionMessenger>,
+	om_mailbox: Arc<OnionMessageMailbox>,
 }
 
 impl<L: Deref + Clone + Sync + Send + 'static> EventHandler<L>
@@ -472,7 +475,8 @@ where
 		output_sweeper: Arc<Sweeper>, network_graph: Arc<Graph>,
 		liquidity_source: Option<Arc<LiquiditySource<Arc<Logger>>>>,
 		payment_store: Arc<PaymentStore>, peer_store: Arc<PeerStore<L>>,
-		static_invoice_store: Option<StaticInvoiceStore>, runtime: Arc<Runtime>, logger: L,
+		static_invoice_store: Option<StaticInvoiceStore>, onion_messenger: Arc<OnionMessenger>,
+		om_mailbox: Arc<OnionMessageMailbox>, runtime: Arc<Runtime>, logger: L,
 		config: Arc<Config>,
 	) -> Self {
 		Self {
@@ -490,6 +494,8 @@ where
 			runtime,
 			config,
 			static_invoice_store,
+			onion_messenger,
+			om_mailbox,
 		}
 	}
 
@@ -1491,11 +1497,15 @@ where
 
 				self.bump_tx_event_handler.handle_event(&bte).await;
 			},
-			LdkEvent::OnionMessageIntercepted { .. } => {
-				debug_assert!(false, "We currently don't support onion message interception, so this event should never be emitted.");
+			LdkEvent::OnionMessageIntercepted { peer_node_id, message } => {
+				self.om_mailbox.onion_message_intercepted(peer_node_id, message);
 			},
-			LdkEvent::OnionMessagePeerConnected { .. } => {
-				debug_assert!(false, "We currently don't support onion message interception, so this event should never be emitted.");
+			LdkEvent::OnionMessagePeerConnected { peer_node_id } => {
+				let messages = self.om_mailbox.onion_message_peer_connected(peer_node_id);
+
+				for message in messages {
+					let _ = self.onion_messenger.forward_onion_message(message, &peer_node_id);
+				}
 			},
 
 			LdkEvent::PersistStaticInvoice {

src/lib.rs

@@ -92,6 +92,7 @@ pub mod io;
 pub mod liquidity;
 pub mod logger;
 mod message_handler;
+mod om_mailbox;
 pub mod payment;
 mod peer_store;
 mod runtime;
@@ -171,6 +172,8 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex, RwLock};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 
+use crate::om_mailbox::OnionMessageMailbox;
+
 #[cfg(feature = "uniffi")]
 uniffi::include_scaffolding!("ldk_node");
 
@@ -205,6 +208,7 @@ pub struct Node {
 	is_running: Arc<RwLock<bool>>,
 	is_listening: Arc<AtomicBool>,
 	node_metrics: Arc<RwLock<NodeMetrics>>,
+	om_mailbox: Arc<OnionMessageMailbox>,
 }
 
 impl Node {
@@ -517,6 +521,8 @@ impl Node {
 			Arc::clone(&self.payment_store),
 			Arc::clone(&self.peer_store),
 			static_invoice_store,
+			Arc::clone(&self.onion_messenger),
+			Arc::clone(&self.om_mailbox),
 			Arc::clone(&self.runtime),
 			Arc::clone(&self.logger),
 			Arc::clone(&self.config),
@@ -1491,6 +1497,11 @@ impl Node {
 				Error::PersistenceFailed
 			})
 	}
+
+	#[allow(missing_docs)]
+	pub fn om_mailbox_is_empty(&self) -> bool {
+		self.om_mailbox.is_empty()
+	}
 }
 
 impl Drop for Node {

src/om_mailbox.rs

@@ -0,0 +1,107 @@
+use std::{
+	collections::{HashMap, VecDeque},
+	sync::Mutex,
+};
+
+use lightning::ln::msgs::OnionMessage;
+
+pub(crate) struct OnionMessageMailbox {
+	map: Mutex<HashMap<bitcoin::secp256k1::PublicKey, VecDeque<lightning::ln::msgs::OnionMessage>>>,
+}
+
+impl OnionMessageMailbox {
+	const MAX_MESSAGES_PER_PEER: usize = 100;
+	const MAX_PEERS: usize = 100;
+
+	pub fn new() -> Self {
+		Self { map: Mutex::new(HashMap::new()) }
+	}
+
+	pub(crate) fn onion_message_intercepted(
+		&self, peer_node_id: bitcoin::secp256k1::PublicKey,
+		message: lightning::ln::msgs::OnionMessage,
+	) {
+		let mut map = self.map.lock().unwrap();
+
+		let queue = map.entry(peer_node_id).or_insert_with(VecDeque::new);
+		if queue.len() >= Self::MAX_MESSAGES_PER_PEER {
+			queue.pop_front();
+		}
+		queue.push_back(message);
+
+		// Enforce a peers limit. If exceeded, evict the peer with the longest queue.
+		if map.len() > Self::MAX_PEERS {
+			let peer_to_remove = map
+				.iter()
+				.max_by_key(|(_, queue)| queue.len())
+				.map(|(peer, _)| peer.clone())
+				.unwrap();
+
+			map.remove(&peer_to_remove);
+		}
+	}
+
+	pub(crate) fn onion_message_peer_connected(
+		&self, peer_node_id: bitcoin::secp256k1::PublicKey,
+	) -> Vec<OnionMessage> {
+		let mut map = self.map.lock().unwrap();
+
+		if let Some(queue) = map.remove(&peer_node_id) {
+			queue.into()
+		} else {
+			Vec::new()
+		}
+	}
+
+	pub(crate) fn is_empty(&self) -> bool {
+		let map = self.map.lock().unwrap();
+		map.is_empty()
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use bitcoin::{
+		key::Secp256k1,
+		secp256k1::{PublicKey, SecretKey},
+	};
+	use lightning::onion_message;
+
+	use crate::om_mailbox::OnionMessageMailbox;
+
+	#[test]
+	fn onion_message_mailbox() {
+		let mailbox = OnionMessageMailbox::new();
+
+		let secp = Secp256k1::new();
+		let sk_bytes = [12; 32];
+		let sk = SecretKey::from_slice(&sk_bytes).unwrap();
+		let peer_node_id = PublicKey::from_secret_key(&secp, &sk);
+
+		let blinding_sk = SecretKey::from_slice(&[13; 32]).unwrap();
+		let blinding_point = PublicKey::from_secret_key(&secp, &blinding_sk);
+
+		let message_sk = SecretKey::from_slice(&[13; 32]).unwrap();
+		let message_point = PublicKey::from_secret_key(&secp, &message_sk);
+
+		let message = lightning::ln::msgs::OnionMessage {
+			blinding_point,
+			onion_routing_packet: onion_message::packet::Packet {
+				version: 0,
+				public_key: message_point,
+				hop_data: vec![1, 2, 3],
+				hmac: [0; 32],
+			},
+		};
+		mailbox.onion_message_intercepted(peer_node_id, message.clone());
+
+		let messages = mailbox.onion_message_peer_connected(peer_node_id);
+		assert_eq!(messages.len(), 1);
+		assert_eq!(messages[0], message);
+
+		assert!(mailbox.is_empty());
+
+		let messages = mailbox.onion_message_peer_connected(peer_node_id);
+		assert_eq!(messages.len(), 0);
+	}
+}

tests/common/mod.rs

@@ -589,14 +589,21 @@ pub(crate) fn bump_fee_and_broadcast<E: ElectrumApi>(
 pub fn open_channel(
 	node_a: &TestNode, node_b: &TestNode, funding_amount_sat: u64, should_announce: bool,
 	electrsd: &ElectrsD,
+) -> OutPoint {
+	open_channel_push_amt(node_a, node_b, funding_amount_sat, None, should_announce, electrsd)
+}
+
+pub fn open_channel_push_amt(
+	node_a: &TestNode, node_b: &TestNode, funding_amount_sat: u64, push_amount_msat: Option<u64>,
+	should_announce: bool, electrsd: &ElectrsD,
 ) -> OutPoint {
 	if should_announce {
 		node_a
 			.open_announced_channel(
 				node_b.node_id(),
 				node_b.listening_addresses().unwrap().first().unwrap().clone(),
 				funding_amount_sat,
-				None,
+				push_amount_msat,
 				None,
 			)
 			.unwrap();
@@ -606,7 +613,7 @@ pub fn open_channel(
 				node_b.node_id(),
 				node_b.listening_addresses().unwrap().first().unwrap().clone(),
 				funding_amount_sat,
-				None,
+				push_amount_msat,
 				None,
 			)
 			.unwrap();