add peer_(dis)connected to custom message handler

Author: johncantrell97

lightning-custom-message/src/lib.rs

@@ -20,7 +20,7 @@
 //!
 //! # use bitcoin::secp256k1::PublicKey;
 //! # use lightning::io;
-//! # use lightning::ln::msgs::{DecodeError, LightningError};
+//! # use lightning::ln::msgs::{DecodeError, Init, LightningError};
 //! # use lightning::ln::features::{InitFeatures, NodeFeatures};
 //! use lightning::ln::peer_handler::CustomMessageHandler;
 //! use lightning::ln::wire::{CustomMessageReader, self};
@@ -68,6 +68,12 @@
 //! #     fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> {
 //! #         unimplemented!()
 //! #     }
+//! #     fn peer_disconnected(&self, _their_node_id: &PublicKey) {
+//! #         unimplemented!()
+//! #     }
+//! #     fn peer_connected(&self, _their_node_id: &PublicKey, _msg: &Init, _inbound: bool) -> Result<(), ()> {
+//! #         unimplemented!()
+//! #     }
 //! #     fn provided_node_features(&self) -> NodeFeatures {
 //! #         unimplemented!()
 //! #     }
@@ -114,6 +120,12 @@
 //! #     fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> {
 //! #         unimplemented!()
 //! #     }
+//! #     fn peer_disconnected(&self, _their_node_id: &PublicKey) {
+//! #         unimplemented!()
+//! #     }
+//! #     fn peer_connected(&self, _their_node_id: &PublicKey, _msg: &Init, _inbound: bool) -> Result<(), ()> {
+//! #         unimplemented!()
+//! #     }
 //! #     fn provided_node_features(&self) -> NodeFeatures {
 //! #         unimplemented!()
 //! #     }
@@ -160,6 +172,12 @@
 //! #     fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> {
 //! #         unimplemented!()
 //! #     }
+//! #     fn peer_disconnected(&self, _their_node_id: &PublicKey) {
+//! #         unimplemented!()
+//! #     }
+//! #     fn peer_connected(&self, _their_node_id: &PublicKey, _msg: &Init, _inbound: bool) -> Result<(), ()> {
+//! #         unimplemented!()
+//! #     }
 //! #     fn provided_node_features(&self) -> NodeFeatures {
 //! #         unimplemented!()
 //! #     }
@@ -287,6 +305,22 @@ macro_rules! composite_custom_message_handler {
 					.collect()
 			}
 
+			fn peer_disconnected(&self, their_node_id: &$crate::bitcoin::secp256k1::PublicKey) {
+				$(
+					self.$field.peer_disconnected(their_node_id);
+				)*
+			}
+
+			fn peer_connected(&self, their_node_id: &$crate::bitcoin::secp256k1::PublicKey, msg: &$crate::lightning::ln::msgs::Init, inbound: bool) -> Result<(), ()> {
+				let mut result = Ok(());
+				$(
+						if let Err(e) = self.$field.peer_connected(their_node_id, msg, inbound) {
+								result = Err(e);
+						}
+				)*
+				result
+			}
+
 			fn provided_node_features(&self) -> $crate::lightning::ln::features::NodeFeatures {
 				$crate::lightning::ln::features::NodeFeatures::empty()
 					$(

lightning/src/ln/peer_handler.rs

@@ -23,7 +23,7 @@ use crate::events::{MessageSendEvent, MessageSendEventsProvider};
 use crate::ln::types::ChannelId;
 use crate::ln::features::{InitFeatures, NodeFeatures};
 use crate::ln::msgs;
-use crate::ln::msgs::{ChannelMessageHandler, LightningError, SocketAddress, OnionMessageHandler, RoutingMessageHandler};
+use crate::ln::msgs::{ChannelMessageHandler, Init, LightningError, SocketAddress, OnionMessageHandler, RoutingMessageHandler};
 use crate::util::ser::{VecWriter, Writeable, Writer};
 use crate::ln::peer_channel_encryptor::{PeerChannelEncryptor, NextNoiseStep, MessageBuf, MSG_BUF_ALLOC_SIZE};
 use crate::ln::wire;
@@ -79,6 +79,17 @@ 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)>;
 
+	// Connection loss/reestablish:
+	/// 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: &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.
@@ -190,6 +201,10 @@ impl CustomMessageHandler for IgnoringMessageHandler {
 
 	fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> { Vec::new() }
 
+	fn peer_disconnected(&self, _their_node_id: &PublicKey) {}
+
+	fn peer_connected(&self, _their_node_id: &PublicKey, _msg: &Init, _inbound: bool) -> Result<(), ()> { Ok(()) }
+
 	fn provided_node_features(&self) -> NodeFeatures { NodeFeatures::empty() }
 
 	fn provided_init_features(&self, _their_node_id: &PublicKey) -> InitFeatures {
@@ -1680,6 +1695,10 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 				log_debug!(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!(logger, "Custom Message Handler decided we couldn't communicate with peer {}", log_pubkey!(their_node_id));
+				return Err(PeerHandleError { }.into());
+			}
 
 			peer_lock.awaiting_pong_timer_tick_intervals = 0;
 			peer_lock.their_features = Some(msg.features);
@@ -2680,7 +2699,7 @@ mod tests {
 	use crate::ln::peer_channel_encryptor::PeerChannelEncryptor;
 	use crate::ln::peer_handler::{CustomMessageHandler, PeerManager, MessageHandler, SocketDescriptor, IgnoringMessageHandler, filter_addresses, ErroringMessageHandler, MAX_BUFFER_DRAIN_TICK_INTERVALS_PER_PEER};
 	use crate::ln::{msgs, wire};
-	use crate::ln::msgs::{LightningError, SocketAddress};
+	use crate::ln::msgs::{Init, LightningError, SocketAddress};
 	use crate::util::test_utils;
 
 	use bitcoin::Network;
@@ -2747,6 +2766,11 @@ mod tests {
 
 		fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> { Vec::new() }
 
+
+		fn peer_disconnected(&self, _their_node_id: &PublicKey) {}
+
+		fn peer_connected(&self, _their_node_id: &PublicKey, _msg: &Init, _inbound: bool) -> Result<(), ()> { Ok(()) }
+
 		fn provided_node_features(&self) -> NodeFeatures { NodeFeatures::empty() }
 
 		fn provided_init_features(&self, _: &PublicKey) -> InitFeatures {