From b32cc981a4fb38abb9d5dc762df4b174db1dacbf Mon Sep 17 00:00:00 2001
From: Lucas Soriano del Pino <lucas_soriano@fastmail.com>
Date: Mon, 17 Jul 2023 13:48:42 +1000
Subject: [PATCH] fix: Clear certain unprocessed messages after peer disconnect

The DLC message handler implements the `OnionMessageHandler` and
`OnionMessageProvider` to be able to hook onto the `peer_disconnected`
event. This allows for clearing certain `rust-dlc` messages precisely
when a disconnect happens.

We drop the `SubChannelConfirm` and `SubChannelCloseConfirm` messages
to avoid bugs such as the one reported here[1].

The rest of the implementations of `OnionMessageHandler` and
`OnionMessageProvider` is copied from the `IgnoringMessageHandler`.

[1]: https://github.com/get10101/10101/issues/792.

Co-authored-by: Richard Holzeis <richard@holzeis.me>
---
 dlc-messages/Cargo.toml             |  1 +
 dlc-messages/src/message_handler.rs | 72 +++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/dlc-messages/Cargo.toml b/dlc-messages/Cargo.toml
index 4d03d597..631b1ce4 100644
--- a/dlc-messages/Cargo.toml
+++ b/dlc-messages/Cargo.toml
@@ -16,6 +16,7 @@ dlc = {version = "0.4.0", path = "../dlc"}
 lightning = {version = "0.0.114"}
 secp256k1-zkp = {version = "0.7.0", features = ["bitcoin_hashes", "rand", "rand-std"]}
 serde = {version = "1.0", features = ["derive"], optional = true}
+log = "0.4.14"
 
 [dev-dependencies]
 bitcoin = {version = "0.29.2"}
diff --git a/dlc-messages/src/message_handler.rs b/dlc-messages/src/message_handler.rs
index 34228bb1..8fc79570 100644
--- a/dlc-messages/src/message_handler.rs
+++ b/dlc-messages/src/message_handler.rs
@@ -38,6 +38,46 @@ impl Default for MessageHandler {
     }
 }
 
+impl lightning::util::events::OnionMessageProvider for MessageHandler {
+    fn next_onion_message_for_peer(
+        &self,
+        _peer_node_id: PublicKey,
+    ) -> Option<lightning::ln::msgs::OnionMessage> {
+        None
+    }
+}
+
+impl lightning::ln::msgs::OnionMessageHandler for MessageHandler {
+    fn handle_onion_message(
+        &self,
+        _their_node_id: &PublicKey,
+        _msg: &lightning::ln::msgs::OnionMessage,
+    ) {
+    }
+    fn peer_connected(
+        &self,
+        _their_node_id: &PublicKey,
+        _init: &lightning::ln::msgs::Init,
+        _inbound: bool,
+    ) -> Result<(), ()> {
+        Ok(())
+    }
+
+    fn peer_disconnected(&self, their_node_id: &PublicKey) {
+        self.clear_pending_messages_sent_by_node(their_node_id);
+    }
+
+    fn provided_node_features(&self) -> lightning::ln::features::NodeFeatures {
+        lightning::ln::features::NodeFeatures::empty()
+    }
+    fn provided_init_features(
+        &self,
+        _their_node_id: &PublicKey,
+    ) -> lightning::ln::features::InitFeatures {
+        lightning::ln::features::InitFeatures::empty()
+    }
+}
+
 impl MessageHandler {
     /// Creates a new instance of a [`MessageHandler`]
     pub fn new() -> Self {
@@ -56,6 +96,38 @@ impl MessageHandler {
         ret
     }
 
+    /// Drops certain pending messages sent by the remote node.
+    pub fn clear_pending_messages_sent_by_node(&self, disconnected_node_id: &PublicKey) {
+        let messages = &mut *self.msg_received.lock().unwrap();
+
+        messages.retain(|(node_id, message)| {
+            match message {
+                Message::SubChannel(SubChannelMessage::Confirm(message))
+                    if node_id == disconnected_node_id =>
+                {
+                    log::warn!(
+                        "Dropping SubChannelConfirm message for channel {:?} \
+                         after peer {node_id} disconnected",
+                        message.channel_id
+                    );
+                    false
+                }
+                Message::SubChannel(SubChannelMessage::CloseConfirm(message))
+                    if node_id == disconnected_node_id =>
+                {
+                    log::warn!(
+                        "Dropping SubChannelCloseConfirm message for channel {:?} \
+                         after peer {node_id} disconnected",
+                        message.channel_id
+                    );
+                    false
+                }
+                // Keep any other message
+                _ => true,
+            }
+        });
+    }
+
     /// Send a message to the peer with given node id. Not that the message is not
     /// sent right away, but only when the LDK
     /// [`lightning::ln::peer_handler::PeerManager::process_events`] is next called.