Add `peer_[dis]connected` to `CustomMessageHandler` trait

There's not much reason not to have these methods on the
`CustomMessageHandler` trait, and they can be quite useful, so we
add them here.

diff --git a/lightning-custom-message/src/lib.rs b/lightning-custom-message/src/lib.rs
index a0a70c9de03ff26dfd499bb7b8cb1dc6eeee8dbe..bbaf82ddff3ac3b95a5ff6779baf6c2d1922fcec 100644 (file)
--- a/lightning-custom-message/src/lib.rs
+++ b/lightning-custom-message/src/lib.rs
@@ -19,7 +19,7 @@
 //!
 //! # use bitcoin::secp256k1::PublicKey;
 //! # use lightning::io;
-//! # use lightning::ln::msgs::{DecodeError, LightningError};
+//! # use lightning::ln::msgs::{DecodeError, LightningError, Init};
 //! # use lightning::ln::features::{InitFeatures, NodeFeatures};
 //! use lightning::ln::peer_handler::CustomMessageHandler;
 //! use lightning::ln::wire::{CustomMessageReader, self};
@@ -67,6 +67,8 @@
 //! #     fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> {
 //! #         unimplemented!()
 //! #     }
+//! #     fn peer_connected(&self, _: &PublicKey, _: &Init, _: bool) -> Result<(), ()> { Ok(()) }
+//! #     fn peer_disconnected(&self, _: &PublicKey) { }
 //! #     fn provided_node_features(&self) -> NodeFeatures {
 //! #         unimplemented!()
 //! #     }
@@ -113,6 +115,8 @@
 //! #     fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> {
 //! #         unimplemented!()
 //! #     }
+//! #     fn peer_connected(&self, _: &PublicKey, _: &Init, _: bool) -> Result<(), ()> { Ok(()) }
+//! #     fn peer_disconnected(&self, _: &PublicKey) { }
 //! #     fn provided_node_features(&self) -> NodeFeatures {
 //! #         unimplemented!()
 //! #     }
@@ -159,6 +163,8 @@
 //! #     fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> {
 //! #         unimplemented!()
 //! #     }
+//! #     fn peer_connected(&self, _: &PublicKey, _: &Init, _: bool) -> Result<(), ()> { Ok(()) }
+//! #     fn peer_disconnected(&self, _: &PublicKey) { }
 //! #     fn provided_node_features(&self) -> NodeFeatures {
 //! #         unimplemented!()
 //! #     }
@@ -288,6 +294,40 @@ macro_rules! composite_custom_message_handler {
                                        .collect()
                        }
 
+                       fn peer_connected(
+                               &self, their_node_id: &$crate::bitcoin::secp256k1::PublicKey,
+                               init_msg: &$crate::lightning::ln::msgs::Init, inbound: bool
+                       ) -> Result<(), ()> {
+                               let mut should_disconnect = false;
+                               let mut handlers_connected = 0;
+                               $(
+                                       if !should_disconnect {
+                                               if self.$field.peer_connected(their_node_id, init_msg, inbound).is_err() {
+                                                       should_disconnect = true;
+                                               } else {
+                                                       handlers_connected += 1;
+                                               }
+                                       }
+                               )*
+                               if should_disconnect {
+                                       let mut i = 0;
+                                       $(
+                                               if i < handlers_connected {
+                                                       self.$field.peer_disconnected(their_node_id);
+                                               }
+                                               #[allow(unused_assignments)]
+                                               i += 1;
+                                       )*
+                                       Err(())
+                               } else { Ok(()) }
+                       }
+
+                       fn peer_disconnected(&self, their_node_id: &$crate::bitcoin::secp256k1::PublicKey) {
+                               $(
+                                       self.$field.peer_disconnected(their_node_id);
+                               )*
+                       }
+
                        fn provided_node_features(&self) -> $crate::lightning::ln::features::NodeFeatures {
                                $crate::lightning::ln::features::NodeFeatures::empty()
                                        $(

diff --git a/lightning/src/ln/peer_handler.rs b/lightning/src/ln/peer_handler.rs
index 8e91023b1b132f0665ef49e6fa5c6f132419fe59..77087f38cff2ef50a329046b385aefe24db1bb6d 100644 (file)
--- a/lightning/src/ln/peer_handler.rs
+++ b/lightning/src/ln/peer_handler.rs
@@ -68,6 +68,16 @@ pub trait CustomMessageHandler: wire::CustomMessageReader {
        /// connection to the node exists, then the message is simply not sent.
        fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)>;
 
+       /// Indicates a connection to the peer failed/an existing connection was lost.
+       fn peer_disconnected(&self, their_node_id: &PublicKey);
+
+       /// Handle a peer reconnecting, possibly generating `channel_reestablish` message(s).
+       ///
+       /// May return an `Err(())` if the features the peer supports are not sufficient to communicate
+       /// with us. Implementors should be somewhat conservative about doing so, however, as other
+       /// message handlers may still wish to communicate with this peer.
+       fn peer_connected(&self, their_node_id: &PublicKey, msg: &msgs::Init, inbound: bool) -> Result<(), ()>;
+
        /// Gets the node feature flags which this handler itself supports. All available handlers are
        /// queried similarly and their feature flags are OR'd together to form the [`NodeFeatures`]
        /// which are broadcasted in our [`NodeAnnouncement`] message.
@@ -170,6 +180,10 @@ impl CustomMessageHandler for IgnoringMessageHandler {
 
        fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> { Vec::new() }
 
+       fn peer_disconnected(&self, _: &PublicKey) {}
+
+       fn peer_connected(&self, _: &PublicKey, _: &msgs::Init, _: bool) -> Result<(), ()> { Ok(()) }
+
        fn provided_node_features(&self) -> NodeFeatures { NodeFeatures::empty() }
 
        fn provided_init_features(&self, _their_node_id: &PublicKey) -> InitFeatures {
@@ -1549,6 +1563,10 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
                                log_debug!(self.logger, "Onion Message Handler decided we couldn't communicate with peer {}", log_pubkey!(their_node_id));
                                return Err(PeerHandleError { }.into());
                        }
+                       if let Err(()) = self.message_handler.custom_message_handler.peer_connected(&their_node_id, &msg, peer_lock.inbound_connection) {
+                               log_debug!(self.logger, "Custom Message Handler decided we couldn't communicate with peer {}", log_pubkey!(their_node_id));
+                               return Err(PeerHandleError { }.into());
+                       }
 
                        peer_lock.their_features = Some(msg.features);
                        return Ok(None);
@@ -2235,6 +2253,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
                        log_trace!(self.logger, "Disconnecting peer with id {} due to {}", node_id, reason);
                        self.message_handler.chan_handler.peer_disconnected(&node_id);
                        self.message_handler.onion_message_handler.peer_disconnected(&node_id);
+                       self.message_handler.custom_message_handler.peer_disconnected(&node_id);
                }
                descriptor.disconnect_socket();
        }
@@ -2257,6 +2276,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
                                        if !peer.handshake_complete() { return; }
                                        self.message_handler.chan_handler.peer_disconnected(&node_id);
                                        self.message_handler.onion_message_handler.peer_disconnected(&node_id);
+                                       self.message_handler.custom_message_handler.peer_disconnected(&node_id);
                                }
                        }
                };
@@ -2551,6 +2571,9 @@ mod tests {
 
                fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> { Vec::new() }
 
+               fn peer_disconnected(&self, _: &PublicKey) {}
+               fn peer_connected(&self, _: &PublicKey, _: &msgs::Init, _: bool) -> Result<(), ()> { Ok(()) }
+
                fn provided_node_features(&self) -> NodeFeatures { NodeFeatures::empty() }
 
                fn provided_init_features(&self, _: &PublicKey) -> InitFeatures {