fixup! Limit OnionMessenger outbound buffer size

Author: valentinewallace

lightning/src/onion_message/functional_tests.rs

@@ -178,5 +178,5 @@ fn peer_buffer_full() {
 		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);
+	assert_eq!(err, SendError::BufferFull);
 }

lightning/src/onion_message/messenger.rs

@@ -123,8 +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,
+	/// Our next-hop peer's buffer was full or our total outbound buffer was full.
+	BufferFull,
 }
 
 impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
@@ -172,10 +172,10 @@ 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(&introduction_node_id, &pending_per_peer_msgs) { return Err(SendError::BufferFull) }
 		match pending_per_peer_msgs.entry(introduction_node_id) {
 			hash_map::Entry::Vacant(_) => Err(SendError::InvalidFirstHop),
 			hash_map::Entry::Occupied(mut e) => {
-				if peer_buffer_full(e.get()) { return Err(SendError::PeerBufferFull) }
 				e.get_mut().push_back(msgs::OnionMessage { blinding_point, onion_routing_packet });
 				Ok(())
 			}
@@ -196,13 +196,24 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessenger<Signer, K, L>
 	}
 }
 
-fn peer_buffer_full(peer_buf: &VecDeque<msgs::OnionMessage>) -> bool {
-	const MAX_BUFFER_SIZE: usize = 1024 * 256;
-	let mut buffered_bytes = 0;
-	for om in peer_buf {
-		buffered_bytes += om.serialized_length();
-		if buffered_bytes >= MAX_BUFFER_SIZE {
-			return true
+fn outbound_buffer_full(peer_node_id: &PublicKey, buffer: &HashMap<PublicKey, VecDeque<msgs::OnionMessage>>) -> bool {
+	const MAX_TOTAL_BUFFER_SIZE: usize = (2 << 20) * 128;
+	const MAX_PER_PEER_BUFFER_SIZE: usize = (2 << 10) * 256;
+	let mut total_buffered_bytes = 0;
+	let mut peer_buffered_bytes = 0;
+	for (pk, peer_buf) in buffer {
+		for om in peer_buf {
+			let om_len = om.serialized_length();
+			if pk == peer_node_id {
+				peer_buffered_bytes += om_len;
+			}
+			total_buffered_bytes += om_len;
+
+			if total_buffered_bytes >= MAX_TOTAL_BUFFER_SIZE ||
+				peer_buffered_bytes >= MAX_PER_PEER_BUFFER_SIZE
+			{
+				return true
+			}
 		}
 	}
 	false
@@ -294,6 +305,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(&next_node_id, &pending_per_peer_msgs) {
+					log_trace!(self.logger, "Dropping forwarded onion message to peer {:?}: outbound buffer full", next_node_id);
+					return
+				}
 
 				#[cfg(fuzzing)]
 				pending_per_peer_msgs.entry(next_node_id).or_insert_with(|| VecDeque::new());