From 5a0d6d05ef16074f7d80f191d58b6c865850c342 Mon Sep 17 00:00:00 2001
From: Valentine Wallace <vwallace@protonmail.com>
Date: Sun, 7 Aug 2022 23:46:19 -0400
Subject: [PATCH] Implement some rate limiting for onion messages.

In this commit, we add business logic for checking if a peer's outbound buffer
has room for onion messages, and if so pulls them from an implementer of a new
trait, OnionMessageProvider.

Makes sure channel messages are prioritized over OMs.

The onion_message module remains private until further rate limiting is added.
---
 lightning/src/ln/peer_handler.rs         | 45 +++++++++++++++++++++++-
 lightning/src/onion_message/messenger.rs | 10 ++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/lightning/src/ln/peer_handler.rs b/lightning/src/ln/peer_handler.rs
index ebb7f749a43..0bcdce29eda 100644
--- a/lightning/src/ln/peer_handler.rs
+++ b/lightning/src/ln/peer_handler.rs
@@ -306,6 +306,10 @@ enum InitSyncTracker{
 /// forwarding gossip messages to peers altogether.
 const FORWARD_INIT_SYNC_BUFFER_LIMIT_RATIO: usize = 2;
 
+/// The ratio between buffer sizes at which we stop sending initial sync messages vs when we pause
+/// forwarding onion messages to peers altogether.
+const OM_BUFFER_LIMIT_RATIO: usize = 2;
+
 /// When the outbound buffer has this many messages, we'll stop reading bytes from the peer until
 /// we have fewer than this many messages in the outbound buffer again.
 /// We also use this as the target number of outbound gossip messages to keep in the write buffer,
@@ -315,6 +319,10 @@ const OUTBOUND_BUFFER_LIMIT_READ_PAUSE: usize = 10;
 /// the peer.
 const OUTBOUND_BUFFER_LIMIT_DROP_GOSSIP: usize = OUTBOUND_BUFFER_LIMIT_READ_PAUSE * FORWARD_INIT_SYNC_BUFFER_LIMIT_RATIO;
 
+/// When the outbound buffer has this many messages, we won't poll for new onion messages for this
+/// peer.
+const OUTBOUND_BUFFER_LIMIT_PAUSE_OMS: usize = 16;
+
 /// If we've sent a ping, and are still awaiting a response, we may need to churn our way through
 /// the socket receive buffer before receiving the ping.
 ///
@@ -417,6 +425,14 @@ impl Peer {
 		}
 		true
 	}
+
+	/// Returns the number of onion messages we can fit in this peer's buffer.
+	fn onion_message_buffer_slots_available(&self) -> usize {
+		cmp::min(
+			OUTBOUND_BUFFER_LIMIT_PAUSE_OMS.saturating_sub(self.pending_outbound_buffer.len()),
+			(BUFFER_DRAIN_MSGS_PER_TICK * OM_BUFFER_LIMIT_RATIO).saturating_sub(self.msgs_sent_since_pong))
+	}
+
 }
 
 /// SimpleArcPeerManager is useful when you need a PeerManager with a static lifetime, e.g.
@@ -824,8 +840,12 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 	/// ready to call `[write_buffer_space_avail`] again if a write call generated here isn't
 	/// sufficient!
 	///
+	/// If any bytes are written, [`process_events`] should be called afterwards.
+	//  TODO: why?
+	///
 	/// [`send_data`]: SocketDescriptor::send_data
 	/// [`write_buffer_space_avail`]: PeerManager::write_buffer_space_avail
+	/// [`process_events`]: PeerManager::process_events
 	pub fn write_buffer_space_avail(&self, descriptor: &mut Descriptor) -> Result<(), PeerHandleError> {
 		let peers = self.peers.read().unwrap();
 		match peers.get(descriptor) {
@@ -1668,7 +1688,30 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 			}
 
 			for (descriptor, peer_mutex) in peers.iter() {
-				self.do_attempt_write_data(&mut (*descriptor).clone(), &mut *peer_mutex.lock().unwrap());
+				let peer_node_id = {
+					let mut peer = peer_mutex.lock().unwrap();
+					self.do_attempt_write_data(&mut (*descriptor).clone(), &mut *peer);
+					peer.their_node_id
+				};
+
+				// Only see if we have room for onion messages after we've written all channel messages, to
+				// ensure the latter take priority.
+				if let Some(peer_node_id) = peer_node_id {
+					loop {
+						let om_buffer_slots_avail = {
+							let peer = peer_mutex.lock().unwrap();
+							peer.onion_message_buffer_slots_available()
+						};
+						if om_buffer_slots_avail == 0 { break; }
+						let onion_msgs = self.message_handler.onion_message_handler.next_onion_messages_for_peer(
+							peer_node_id, om_buffer_slots_avail);
+						if onion_msgs.len() == 0 { break; }
+						for msg in onion_msgs {
+							self.enqueue_message(&mut *get_peer_for_forwarding!(&peer_node_id), &msg);
+						}
+						self.do_attempt_write_data(&mut (*descriptor).clone(), &mut *peer_mutex.lock().unwrap());
+					}
+				}
 			}
 		}
 		if !peers_to_disconnect.is_empty() {
diff --git a/lightning/src/onion_message/messenger.rs b/lightning/src/onion_message/messenger.rs
index d7c935179c9..946297faee5 100644
--- a/lightning/src/onion_message/messenger.rs
+++ b/lightning/src/onion_message/messenger.rs
@@ -292,6 +292,16 @@ impl<Signer: Sign, K: Deref, L: Deref> OnionMessageProvider for OnionMessenger<S
 	      L::Target: Logger,
 {
 	fn next_onion_messages_for_peer(&self, peer_node_id: PublicKey, max_messages: usize) -> Vec<msgs::OnionMessage> {
+		let mut pending_msgs = self.pending_messages.lock().unwrap();
+		if let Some(msgs) = pending_msgs.get_mut(&peer_node_id) {
+			if max_messages >= msgs.len() {
+				let mut peer_pending_msgs = Vec::new();
+				mem::swap(msgs, &mut peer_pending_msgs);
+				return peer_pending_msgs
+			} else {
+				return msgs.split_off(max_messages)
+			}
+		}
 		Vec::new()
 	}
 }