om mailbox peer limit and test

Author: joostjager

src/event.rs

@@ -7,7 +7,7 @@
 
 use crate::om_mailbox::OnionMessageMailbox;
 use crate::payment::asynchronous::static_invoice_store::StaticInvoiceStore;
-use crate::types::{CustomTlvRecord, DynStore, PaymentStore, Sweeper, Wallet};
+use crate::types::{CustomTlvRecord, DynStore, OnionMessenger, PaymentStore, Sweeper, Wallet};
 use crate::{
 	hex_utils, BumpTransactionEventHandler, ChannelManager, Error, Graph, PeerInfo, PeerStore,
 	UserChannelId,
@@ -460,6 +460,7 @@ where
 	logger: L,
 	config: Arc<Config>,
 	static_invoice_store: Option<StaticInvoiceStore>,
+	onion_messenger: Arc<OnionMessenger>,
 	om_mailbox: OnionMessageMailbox,
 }
 
@@ -474,8 +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>, om_mailbox: OnionMessageMailbox,
-		runtime: Arc<Runtime>, logger: L, config: Arc<Config>,
+		static_invoice_store: Option<StaticInvoiceStore>, onion_messenger: Arc<OnionMessenger>,
+		om_mailbox: OnionMessageMailbox, runtime: Arc<Runtime>, logger: L, config: Arc<Config>,
 	) -> Self {
 		Self {
 			event_queue,
@@ -492,6 +493,7 @@ where
 			runtime,
 			config,
 			static_invoice_store,
+			onion_messenger,
 			om_mailbox,
 		}
 	}
@@ -1498,7 +1500,11 @@ where
 				self.om_mailbox.onion_message_intercepted(peer_node_id, message);
 			},
 			LdkEvent::OnionMessagePeerConnected { peer_node_id } => {
-				self.om_mailbox.onion_message_peer_connected(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

@@ -507,7 +507,7 @@ impl Node {
 		} else {
 			None
 		};
-		let om_mailbox = OnionMessageMailbox::new(Arc::clone(&self.onion_messenger));
+		let om_mailbox = OnionMessageMailbox::new();
 
 		let event_handler = Arc::new(EventHandler::new(
 			Arc::clone(&self.event_queue),
@@ -521,6 +521,7 @@ impl Node {
 			Arc::clone(&self.payment_store),
 			Arc::clone(&self.peer_store),
 			static_invoice_store,
+			Arc::clone(&self.onion_messenger),
 			om_mailbox,
 			Arc::clone(&self.runtime),
 			Arc::clone(&self.logger),

src/om_mailbox.rs

@@ -1,20 +1,20 @@
 use std::{
 	collections::{hash_map::Entry, HashMap, VecDeque},
-	sync::{Arc, Mutex},
+	sync::Mutex,
 };
 
-use crate::types::OnionMessenger;
+use lightning::ln::msgs::OnionMessage;
 
 pub(crate) struct OnionMessageMailbox {
 	map: Mutex<HashMap<bitcoin::secp256k1::PublicKey, VecDeque<lightning::ln::msgs::OnionMessage>>>,
-	onion_messenger: Arc<OnionMessenger>,
 }
 
 impl OnionMessageMailbox {
 	const MAX_MESSAGES_PER_PEER: usize = 100;
+	const MAX_PEERS: usize = 100;
 
-	pub fn new(onion_messenger: Arc<OnionMessenger>) -> Self {
-		Self { map: Mutex::new(HashMap::new()), onion_messenger }
+	pub fn new() -> Self {
+		Self { map: Mutex::new(HashMap::new()) }
 	}
 
 	pub(crate) fn onion_message_intercepted(
@@ -28,18 +28,74 @@ impl OnionMessageMailbox {
 			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) {
+	pub(crate) fn onion_message_peer_connected(
+		&self, peer_node_id: bitcoin::secp256k1::PublicKey,
+	) -> Vec<OnionMessage> {
 		let mut map = self.map.lock().unwrap();
 
-		let queue = match map.entry(peer_node_id) {
-			Entry::Occupied(entry) => entry.into_mut(),
-			Entry::Vacant(_) => return,
+		match map.entry(peer_node_id) {
+			Entry::Occupied(mut entry) => {
+				let queue = std::mem::take(entry.get_mut());
+				queue.into()
+			},
+			Entry::Vacant(_) => Vec::new(),
+		}
+	}
+}
+
+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());
 
-		while let Some(message) = queue.pop_front() {
-			let _ = self.onion_messenger.forward_onion_message(message, &peer_node_id);
-		}
+		let messages = mailbox.onion_message_peer_connected(peer_node_id);
+		assert_eq!(messages.len(), 1);
+		assert_eq!(messages[0], message);
+
+		let messages = mailbox.onion_message_peer_connected(peer_node_id);
+		assert_eq!(messages.len(), 0);
 	}
 }