Limit OnionMessenger outbound buffer size

Drop OMs if they push us over the max OnionMessenger outbound buffer size

Author: valentinewallace

lightning/src/onion_message/functional_tests.rs

@@ -170,3 +170,13 @@ fn reply_path() {
 		"lightning::onion_message::messenger".to_string(),
 		format!("Received an onion message with path_id: None and reply_path").to_string(), 2);
 }
+
+#[test]
+fn peer_buffer_full() {
+	let nodes = create_nodes(2);
+	for _ in 0..188 {
+		nodes[0].messenger.send_onion_message(&[], Destination::Node(nodes[1].get_node_pk()), None).unwrap();
+	}
+	let err = nodes[0].messenger.send_onion_message(&[], Destination::Node(nodes[1].get_node_pk()), None).unwrap_err();
+	assert_eq!(err, SendError::PeerBufferFull);
+}

lightning/src/onion_message/messenger.rs

@@ -23,6 +23,7 @@ use super::packet::{BIG_PACKET_HOP_DATA_LEN, ForwardControlTlvs, Packet, Payload
 use super::utils;
 use util::events::OnionMessageProvider;
 use util::logger::Logger;
+use util::ser::Writeable;
 
 use core::ops::Deref;
 use sync::{Arc, Mutex};
@@ -122,6 +123,8 @@ pub enum SendError {
 	TooFewBlindedHops,
 	/// Our next-hop peer was offline or does not support onion message forwarding.
 	InvalidFirstHop,
+	/// Our next-hop peer's buffer was full.
+	PeerBufferFull,
 }
 
 impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
@@ -169,6 +172,7 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
 			packet_payloads, packet_keys, prng_seed).map_err(|()| SendError::TooBigPacket)?;
 
 		let mut pending_per_peer_msgs = self.pending_messages.lock().unwrap();
+		if outbound_buffer_full(&pending_per_peer_msgs) { return Err(SendError::PeerBufferFull) }
 		match pending_per_peer_msgs.entry(introduction_node_id) {
 			hash_map::Entry::Vacant(_) => Err(SendError::InvalidFirstHop),
 			hash_map::Entry::Occupied(mut e) => {
@@ -177,7 +181,6 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
 			}
 		}
 	}
-
 	#[cfg(test)]
 	pub(super) fn release_pending_msgs(&self) -> HashMap<PublicKey, VecDeque<msgs::OnionMessage>> {
 		let mut pending_msgs = self.pending_messages.lock().unwrap();
@@ -192,6 +195,20 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
 	}
 }
 
+fn outbound_buffer_full(buffer: &HashMap<PublicKey, VecDeque<msgs::OnionMessage>>) -> bool {
+	const MAX_BUFFER_SIZE: usize = 262144; // 256KiB
+	let mut buffered_bytes = 0;
+	for peer_buf in buffer.values() {
+		for om in peer_buf {
+			buffered_bytes += om.serialized_length();
+			if buffered_bytes >= MAX_BUFFER_SIZE {
+				return true
+			}
+		}
+	}
+	false
+}
+
 impl<Signer: Sign, K: Deref, L: Deref> OnionMessageHandler for OnionMessenger<Signer, K, L>
 	where K::Target: KeysInterface<Signer = Signer>,
 	      L::Target: Logger,
@@ -256,6 +273,10 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessageHandler for OnionMessenger<Si
 				};
 
 				let mut pending_per_peer_msgs = self.pending_messages.lock().unwrap();
+				if outbound_buffer_full(&pending_per_peer_msgs) {
+					log_trace!(self.logger, "Dropping forwarded onion message to peer {:?}: outbound buffer full", next_node_id);
+					return
+				}
 				match pending_per_peer_msgs.entry(next_node_id) {
 					hash_map::Entry::Vacant(_) => {
 						log_trace!(self.logger, "Dropping forwarded onion message to disconnected peer {:?}", next_node_id);