Merge pull request lightningdevkit#3580 from TheBlueMatt/2025-01-peer-connected-consistency

Ensure `peer_disconnected` is called after a handler refuses a connection

Author: TheBlueMatt

lightning-net-tokio/src/lib.rs

@@ -689,6 +689,7 @@ mod tests {
 		) -> Result<(), ()> {
 			Ok(())
 		}
+		fn peer_disconnected(&self, _their_node_id: PublicKey) {}
 		fn handle_reply_channel_range(
 			&self, _their_node_id: PublicKey, _msg: ReplyChannelRange,
 		) -> Result<(), LightningError> {

lightning/src/ln/msgs.rs

@@ -1579,6 +1579,8 @@ pub trait ChannelMessageHandler : MessageSendEventsProvider {
 	/// 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.
+	///
+	/// [`Self::peer_disconnected`] will not be called if `Err(())` is returned.
 	fn peer_connected(&self, their_node_id: PublicKey, msg: &Init, inbound: bool) -> Result<(), ()>;
 	/// Handle an incoming `channel_reestablish` message from the given peer.
 	fn handle_channel_reestablish(&self, their_node_id: PublicKey, msg: &ChannelReestablish);
@@ -1657,7 +1659,11 @@ pub trait RoutingMessageHandler : MessageSendEventsProvider {
 	/// 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.
+	///
+	/// [`Self::peer_disconnected`] will not be called if `Err(())` is returned.
 	fn peer_connected(&self, their_node_id: PublicKey, init: &Init, inbound: bool) -> Result<(), ()>;
+	/// Indicates a connection to the peer failed/an existing connection was lost.
+	fn peer_disconnected(&self, their_node_id: PublicKey);
 	/// Handles the reply of a query we initiated to learn about channels
 	/// for a given range of blocks. We can expect to receive one or more
 	/// replies to a single query.
@@ -1708,6 +1714,8 @@ pub trait OnionMessageHandler {
 	/// 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.
+	///
+	/// [`Self::peer_disconnected`] will not be called if `Err(())` is returned.
 	fn peer_connected(&self, their_node_id: PublicKey, init: &Init, inbound: bool) -> Result<(), ()>;
 
 	/// Indicates a connection to the peer failed/an existing connection was lost. Allows handlers to

lightning/src/ln/peer_handler.rs

@@ -88,6 +88,8 @@ pub trait CustomMessageHandler: wire::CustomMessageReader {
 	/// 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.
+	///
+	/// [`Self::peer_disconnected`] will not be called if `Err(())` is returned.
 	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
@@ -119,6 +121,7 @@ impl RoutingMessageHandler for IgnoringMessageHandler {
 		Option<(msgs::ChannelAnnouncement, Option<msgs::ChannelUpdate>, Option<msgs::ChannelUpdate>)> { None }
 	fn get_next_node_announcement(&self, _starting_point: Option<&NodeId>) -> Option<msgs::NodeAnnouncement> { None }
 	fn peer_connected(&self, _their_node_id: PublicKey, _init: &msgs::Init, _inbound: bool) -> Result<(), ()> { Ok(()) }
+	fn peer_disconnected(&self, _their_node_id: PublicKey) { }
 	fn handle_reply_channel_range(&self, _their_node_id: PublicKey, _msg: msgs::ReplyChannelRange) -> Result<(), LightningError> { Ok(()) }
 	fn handle_reply_short_channel_ids_end(&self, _their_node_id: PublicKey, _msg: msgs::ReplyShortChannelIdsEnd) -> Result<(), LightningError> { Ok(()) }
 	fn handle_query_channel_range(&self, _their_node_id: PublicKey, _msg: msgs::QueryChannelRange) -> Result<(), LightningError> { Ok(()) }
@@ -1714,14 +1717,20 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 			}
 			if let Err(()) = self.message_handler.chan_handler.peer_connected(their_node_id, &msg, peer_lock.inbound_connection) {
 				log_debug!(logger, "Channel Handler decided we couldn't communicate with peer {}", log_pubkey!(their_node_id));
+				self.message_handler.route_handler.peer_disconnected(their_node_id);
 				return Err(PeerHandleError { }.into());
 			}
 			if let Err(()) = self.message_handler.onion_message_handler.peer_connected(their_node_id, &msg, peer_lock.inbound_connection) {
 				log_debug!(logger, "Onion Message Handler decided we couldn't communicate with peer {}", log_pubkey!(their_node_id));
+				self.message_handler.route_handler.peer_disconnected(their_node_id);
+				self.message_handler.chan_handler.peer_disconnected(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));
+				self.message_handler.route_handler.peer_disconnected(their_node_id);
+				self.message_handler.chan_handler.peer_disconnected(their_node_id);
+				self.message_handler.onion_message_handler.peer_disconnected(their_node_id);
 				return Err(PeerHandleError { }.into());
 			}
 
@@ -2533,6 +2542,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 		debug_assert!(peer.their_node_id.is_some());
 		if let Some((node_id, _)) = peer.their_node_id {
 			log_trace!(WithContext::from(&self.logger, Some(node_id), None, None), "Disconnecting peer with id {} due to {}", node_id, reason);
+			self.message_handler.route_handler.peer_disconnected(node_id);
 			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);
@@ -2557,6 +2567,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 					let removed = self.node_id_to_descriptor.lock().unwrap().remove(&node_id);
 					debug_assert!(removed.is_some(), "descriptor maps should be consistent");
 					if !peer.handshake_complete() { return; }
+					self.message_handler.route_handler.peer_disconnected(node_id);
 					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);
@@ -2856,6 +2867,16 @@ mod tests {
 
 	struct TestCustomMessageHandler {
 		features: InitFeatures,
+		conn_tracker: test_utils::ConnectionTracker,
+	}
+
+	impl TestCustomMessageHandler {
+		fn new(features: InitFeatures) -> Self {
+			Self {
+				features,
+				conn_tracker: test_utils::ConnectionTracker::new(),
+			}
+		}
 	}
 
 	impl wire::CustomMessageReader for TestCustomMessageHandler {
@@ -2872,10 +2893,13 @@ mod tests {
 
 		fn get_and_clear_pending_msg(&self) -> Vec<(PublicKey, Self::CustomMessage)> { Vec::new() }
 
+		fn peer_disconnected(&self, their_node_id: PublicKey) {
+			self.conn_tracker.peer_disconnected(their_node_id);
+		}
 
-		fn peer_disconnected(&self, _their_node_id: PublicKey) {}
-
-		fn peer_connected(&self, _their_node_id: PublicKey, _msg: &Init, _inbound: bool) -> Result<(), ()> { Ok(()) }
+		fn peer_connected(&self, their_node_id: PublicKey, _msg: &Init, _inbound: bool) -> Result<(), ()> {
+			self.conn_tracker.peer_connected(their_node_id)
+		}
 
 		fn provided_node_features(&self) -> NodeFeatures { NodeFeatures::empty() }
 
@@ -2898,7 +2922,7 @@ mod tests {
 					chan_handler: test_utils::TestChannelMessageHandler::new(ChainHash::using_genesis_block(Network::Testnet)),
 					logger: test_utils::TestLogger::with_id(i.to_string()),
 					routing_handler: test_utils::TestRoutingMessageHandler::new(),
-					custom_handler: TestCustomMessageHandler { features },
+					custom_handler: TestCustomMessageHandler::new(features),
 					node_signer: test_utils::TestNodeSigner::new(node_secret),
 				}
 			);
@@ -2921,7 +2945,7 @@ mod tests {
 					chan_handler: test_utils::TestChannelMessageHandler::new(ChainHash::using_genesis_block(Network::Testnet)),
 					logger: test_utils::TestLogger::new(),
 					routing_handler: test_utils::TestRoutingMessageHandler::new(),
-					custom_handler: TestCustomMessageHandler { features },
+					custom_handler: TestCustomMessageHandler::new(features),
 					node_signer: test_utils::TestNodeSigner::new(node_secret),
 				}
 			);
@@ -2941,7 +2965,7 @@ mod tests {
 					chan_handler: test_utils::TestChannelMessageHandler::new(network),
 					logger: test_utils::TestLogger::new(),
 					routing_handler: test_utils::TestRoutingMessageHandler::new(),
-					custom_handler: TestCustomMessageHandler { features },
+					custom_handler: TestCustomMessageHandler::new(features),
 					node_signer: test_utils::TestNodeSigner::new(node_secret),
 				}
 			);
@@ -2965,19 +2989,16 @@ mod tests {
 		peers
 	}
 
-	fn establish_connection<'a>(peer_a: &PeerManager<FileDescriptor, &'a test_utils::TestChannelMessageHandler, &'a test_utils::TestRoutingMessageHandler, IgnoringMessageHandler, &'a test_utils::TestLogger, &'a TestCustomMessageHandler, &'a test_utils::TestNodeSigner>, peer_b: &PeerManager<FileDescriptor, &'a test_utils::TestChannelMessageHandler, &'a test_utils::TestRoutingMessageHandler, IgnoringMessageHandler, &'a test_utils::TestLogger, &'a TestCustomMessageHandler, &'a test_utils::TestNodeSigner>) -> (FileDescriptor, FileDescriptor) {
+	fn try_establish_connection<'a>(peer_a: &PeerManager<FileDescriptor, &'a test_utils::TestChannelMessageHandler, &'a test_utils::TestRoutingMessageHandler, IgnoringMessageHandler, &'a test_utils::TestLogger, &'a TestCustomMessageHandler, &'a test_utils::TestNodeSigner>, peer_b: &PeerManager<FileDescriptor, &'a test_utils::TestChannelMessageHandler, &'a test_utils::TestRoutingMessageHandler, IgnoringMessageHandler, &'a test_utils::TestLogger, &'a TestCustomMessageHandler, &'a test_utils::TestNodeSigner>) -> (FileDescriptor, FileDescriptor, Result<bool, PeerHandleError>, Result<bool, PeerHandleError>) {
+		let addr_a = SocketAddress::TcpIpV4{addr: [127, 0, 0, 1], port: 1000};
+		let addr_b = SocketAddress::TcpIpV4{addr: [127, 0, 0, 1], port: 1001};
+
 		static FD_COUNTER: AtomicUsize = AtomicUsize::new(0);
 		let fd = FD_COUNTER.fetch_add(1, Ordering::Relaxed) as u16;
 
 		let id_a = peer_a.node_signer.get_node_id(Recipient::Node).unwrap();
 		let mut fd_a = FileDescriptor::new(fd);
-		let addr_a = SocketAddress::TcpIpV4{addr: [127, 0, 0, 1], port: 1000};
-
-		let id_b = peer_b.node_signer.get_node_id(Recipient::Node).unwrap();
-		let features_a = peer_a.init_features(id_b);
-		let features_b = peer_b.init_features(id_a);
 		let mut fd_b = FileDescriptor::new(fd);
-		let addr_b = SocketAddress::TcpIpV4{addr: [127, 0, 0, 1], port: 1001};
 
 		let initial_data = peer_b.new_outbound_connection(id_a, fd_b.clone(), Some(addr_a.clone())).unwrap();
 		peer_a.new_inbound_connection(fd_a.clone(), Some(addr_b.clone())).unwrap();
@@ -2989,11 +3010,30 @@ mod tests {
 
 		peer_b.process_events();
 		let b_data = fd_b.outbound_data.lock().unwrap().split_off(0);
-		assert_eq!(peer_a.read_event(&mut fd_a, &b_data).unwrap(), false);
+		let a_refused = peer_a.read_event(&mut fd_a, &b_data);
 
 		peer_a.process_events();
 		let a_data = fd_a.outbound_data.lock().unwrap().split_off(0);
-		assert_eq!(peer_b.read_event(&mut fd_b, &a_data).unwrap(), false);
+		let b_refused = peer_b.read_event(&mut fd_b, &a_data);
+
+		(fd_a, fd_b, a_refused, b_refused)
+	}
+
+
+	fn establish_connection<'a>(peer_a: &PeerManager<FileDescriptor, &'a test_utils::TestChannelMessageHandler, &'a test_utils::TestRoutingMessageHandler, IgnoringMessageHandler, &'a test_utils::TestLogger, &'a TestCustomMessageHandler, &'a test_utils::TestNodeSigner>, peer_b: &PeerManager<FileDescriptor, &'a test_utils::TestChannelMessageHandler, &'a test_utils::TestRoutingMessageHandler, IgnoringMessageHandler, &'a test_utils::TestLogger, &'a TestCustomMessageHandler, &'a test_utils::TestNodeSigner>) -> (FileDescriptor, FileDescriptor) {
+		let addr_a = SocketAddress::TcpIpV4{addr: [127, 0, 0, 1], port: 1000};
+		let addr_b = SocketAddress::TcpIpV4{addr: [127, 0, 0, 1], port: 1001};
+
+		let id_a = peer_a.node_signer.get_node_id(Recipient::Node).unwrap();
+		let id_b = peer_b.node_signer.get_node_id(Recipient::Node).unwrap();
+
+		let features_a = peer_a.init_features(id_b);
+		let features_b = peer_b.init_features(id_a);
+
+		let (fd_a, fd_b, a_refused, b_refused) = try_establish_connection(peer_a, peer_b);
+
+		assert_eq!(a_refused.unwrap(), false);
+		assert_eq!(b_refused.unwrap(), false);
 
 		assert_eq!(peer_a.peer_by_node_id(&id_b).unwrap().counterparty_node_id, id_b);
 		assert_eq!(peer_a.peer_by_node_id(&id_b).unwrap().socket_address, Some(addr_b));
@@ -3246,6 +3286,50 @@ mod tests {
 		assert_eq!(peers[0].peers.read().unwrap().len(), 0);
 	}
 
+	fn do_test_peer_connected_error_disconnects(handler: usize) {
+		// Test that if a message handler fails a connection in `peer_connected` we reliably
+		// produce `peer_disconnected` events for all other message handlers (that saw a
+		// corresponding `peer_connected`).
+		let cfgs = create_peermgr_cfgs(2);
+		let peers = create_network(2, &cfgs);
+
+		match handler & !1 {
+			0 => {
+				peers[handler & 1].message_handler.chan_handler.conn_tracker.fail_connections.store(true, Ordering::Release);
+			}
+			2 => {
+				peers[handler & 1].message_handler.route_handler.conn_tracker.fail_connections.store(true, Ordering::Release);
+			}
+			4 => {
+				peers[handler & 1].message_handler.custom_message_handler.conn_tracker.fail_connections.store(true, Ordering::Release);
+			}
+			_ => panic!(),
+		}
+		let (_sd1, _sd2, a_refused, b_refused) = try_establish_connection(&peers[0], &peers[1]);
+		if handler & 1 == 0 {
+			assert!(a_refused.is_err());
+			assert!(peers[0].list_peers().is_empty());
+		} else {
+			assert!(b_refused.is_err());
+			assert!(peers[1].list_peers().is_empty());
+		}
+		// At least one message handler should have seen the connection.
+		assert!(peers[handler & 1].message_handler.chan_handler.conn_tracker.had_peers.load(Ordering::Acquire) ||
+			peers[handler & 1].message_handler.route_handler.conn_tracker.had_peers.load(Ordering::Acquire) ||
+			peers[handler & 1].message_handler.custom_message_handler.conn_tracker.had_peers.load(Ordering::Acquire));
+		// And both message handlers doing tracking should see the disconnection
+		assert!(peers[handler & 1].message_handler.chan_handler.conn_tracker.connected_peers.lock().unwrap().is_empty());
+		assert!(peers[handler & 1].message_handler.route_handler.conn_tracker.connected_peers.lock().unwrap().is_empty());
+		assert!(peers[handler & 1].message_handler.custom_message_handler.conn_tracker.connected_peers.lock().unwrap().is_empty());
+	}
+
+	#[test]
+	fn test_peer_connected_error_disconnects() {
+		for i in 0..6 {
+			do_test_peer_connected_error_disconnects(i);
+		}
+	}
+
 	#[test]
 	fn test_do_attempt_write_data() {
 		// Create 2 peers with custom TestRoutingMessageHandlers and connect them.

lightning/src/routing/gossip.rs

@@ -706,6 +706,8 @@ where
 		Ok(())
 	}
 
+	fn peer_disconnected(&self, _their_node_id: PublicKey) {}
+
 	fn handle_reply_channel_range(
 		&self, _their_node_id: PublicKey, _msg: ReplyChannelRange,
 	) -> Result<(), LightningError> {

lightning/src/util/test_utils.rs

@@ -887,10 +887,45 @@ impl chaininterface::BroadcasterInterface for TestBroadcaster {
 	}
 }
 
+pub struct ConnectionTracker {
+	pub had_peers: AtomicBool,
+	pub connected_peers: Mutex<Vec<PublicKey>>,
+	pub fail_connections: AtomicBool,
+}
+
+impl ConnectionTracker {
+	pub fn new() -> Self {
+		Self {
+			had_peers: AtomicBool::new(false),
+			connected_peers: Mutex::new(Vec::new()),
+			fail_connections: AtomicBool::new(false),
+		}
+	}
+
+	pub fn peer_connected(&self, their_node_id: PublicKey) -> Result<(), ()> {
+		self.had_peers.store(true, Ordering::Release);
+		let mut connected_peers = self.connected_peers.lock().unwrap();
+		assert!(!connected_peers.contains(&their_node_id));
+		if self.fail_connections.load(Ordering::Acquire) {
+			Err(())
+		} else {
+			connected_peers.push(their_node_id);
+			Ok(())
+		}
+	}
+
+	pub fn peer_disconnected(&self, their_node_id: PublicKey) {
+		assert!(self.had_peers.load(Ordering::Acquire));
+		let mut connected_peers = self.connected_peers.lock().unwrap();
+		assert!(connected_peers.contains(&their_node_id));
+		connected_peers.retain(|id| *id != their_node_id);
+	}
+}
+
 pub struct TestChannelMessageHandler {
 	pub pending_events: Mutex<Vec<events::MessageSendEvent>>,
 	expected_recv_msgs: Mutex<Option<Vec<wire::Message<()>>>>,
-	connected_peers: Mutex<HashSet<PublicKey>>,
+	pub conn_tracker: ConnectionTracker,
 	chain_hash: ChainHash,
 }
 
@@ -905,7 +940,7 @@ impl TestChannelMessageHandler {
 		TestChannelMessageHandler {
 			pending_events: Mutex::new(Vec::new()),
 			expected_recv_msgs: Mutex::new(None),
-			connected_peers: Mutex::new(new_hash_set()),
+			conn_tracker: ConnectionTracker::new(),
 			chain_hash,
 		}
 	}
@@ -1017,15 +1052,14 @@ impl msgs::ChannelMessageHandler for TestChannelMessageHandler {
 		self.received_msg(wire::Message::ChannelReestablish(msg.clone()));
 	}
 	fn peer_disconnected(&self, their_node_id: PublicKey) {
-		assert!(self.connected_peers.lock().unwrap().remove(&their_node_id));
+		self.conn_tracker.peer_disconnected(their_node_id)
 	}
 	fn peer_connected(
 		&self, their_node_id: PublicKey, _msg: &msgs::Init, _inbound: bool,
 	) -> Result<(), ()> {
-		assert!(self.connected_peers.lock().unwrap().insert(their_node_id.clone()));
 		// Don't bother with `received_msg` for Init as its auto-generated and we don't want to
 		// bother re-generating the expected Init message in all tests.
-		Ok(())
+		self.conn_tracker.peer_connected(their_node_id)
 	}
 	fn handle_error(&self, _their_node_id: PublicKey, msg: &msgs::ErrorMessage) {
 		self.received_msg(wire::Message::Error(msg.clone()));
@@ -1155,6 +1189,7 @@ pub struct TestRoutingMessageHandler {
 	pub pending_events: Mutex<Vec<events::MessageSendEvent>>,
 	pub request_full_sync: AtomicBool,
 	pub announcement_available_for_sync: AtomicBool,
+	pub conn_tracker: ConnectionTracker,
 }
 
 impl TestRoutingMessageHandler {
@@ -1166,6 +1201,7 @@ impl TestRoutingMessageHandler {
 			pending_events,
 			request_full_sync: AtomicBool::new(false),
 			announcement_available_for_sync: AtomicBool::new(false),
+			conn_tracker: ConnectionTracker::new(),
 		}
 	}
 }
@@ -1240,7 +1276,12 @@ impl msgs::RoutingMessageHandler for TestRoutingMessageHandler {
 				timestamp_range: u32::max_value(),
 			},
 		});
-		Ok(())
+
+		self.conn_tracker.peer_connected(their_node_id)
+	}
+
+	fn peer_disconnected(&self, their_node_id: PublicKey) {
+		self.conn_tracker.peer_disconnected(their_node_id);
 	}
 
 	fn handle_reply_channel_range(