Skip to content

Commit

Permalink
add peer_(dis)connected to custom message handler
Browse files Browse the repository at this point in the history
  • Loading branch information
@johncantrell97
johncantrell97 committed Jun 6, 2024
1 parent 2701bc5 commit 1c063e4
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 3 deletions.
36 changes: 35 additions & 1 deletion lightning-custom-message/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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!()
//! # }
Expand Down Expand Up @@ -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!()
//! # }
Expand Down Expand Up @@ -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!()
//! # }
Expand Down Expand Up @@ -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()
$(
Expand Down
30 changes: 28 additions & 2 deletions lightning/src/ln/peer_handler.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 was lost.
fn peer_disconnected(&self, their_node_id: &PublicKey);

/// Handle a peer reconnecting.
///
/// 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.
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -2430,6 +2449,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
log_trace!(WithContext::from(&self.logger, Some(node_id), None, None), "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();
}
Expand All @@ -2452,6 +2472,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);
}
}
};
Expand Down Expand Up @@ -2680,7 +2701,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;
Expand Down Expand Up @@ -2747,6 +2768,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 {
Expand Down

0 comments on commit 1c063e4

Please sign in to comment.