Add test of an initial message other than Init

This test fails without the previous commit.

Author: TheBlueMatt

lightning/src/ln/peer_handler.rs

@@ -2171,6 +2171,7 @@ fn is_gossip_msg(type_id: u16) -> bool {
 #[cfg(test)]
 mod tests {
 	use crate::chain::keysinterface::{NodeSigner, Recipient};
+	use crate::ln::peer_channel_encryptor::PeerChannelEncryptor;
 	use crate::ln::peer_handler::{PeerManager, MessageHandler, SocketDescriptor, IgnoringMessageHandler, filter_addresses};
 	use crate::ln::{msgs, wire};
 	use crate::ln::msgs::NetAddress;
@@ -2279,19 +2280,15 @@ mod tests {
 		// Simple test which builds a network of PeerManager, connects and brings them to NoiseState::Finished and
 		// push a DisconnectPeer event to remove the node flagged by id
 		let cfgs = create_peermgr_cfgs(2);
-		let chan_handler = test_utils::TestChannelMessageHandler::new();
-		let mut peers = create_network(2, &cfgs);
+		let peers = create_network(2, &cfgs);
 		establish_connection(&peers[0], &peers[1]);
 		assert_eq!(peers[0].peers.read().unwrap().len(), 1);
 
 		let their_id = peers[1].node_signer.get_node_id(Recipient::Node).unwrap();
-
-		chan_handler.pending_events.lock().unwrap().push(events::MessageSendEvent::HandleError {
+		cfgs[0].chan_handler.pending_events.lock().unwrap().push(events::MessageSendEvent::HandleError {
 			node_id: their_id,
 			action: msgs::ErrorAction::DisconnectPeer { msg: None },
 		});
-		assert_eq!(chan_handler.pending_events.lock().unwrap().len(), 1);
-		peers[0].message_handler.chan_handler = &chan_handler;
 
 		peers[0].process_events();
 		assert_eq!(peers[0].peers.read().unwrap().len(), 0);
@@ -2325,6 +2322,35 @@ mod tests {
 		assert_eq!(peers[1].read_event(&mut fd_b, &a_data).unwrap(), false);
 	}
 
+	#[test]
+	fn test_non_init_first_msg() {
+		// Simple test of the first message received over a connection being something other than
+		// Init. This results in an immediate disconnection, which previously included a spurious
+		// peer_disconnected event handed to event handlers (which would panic in
+		// `TestChannelMessageHandler` here).
+		let cfgs = create_peermgr_cfgs(2);
+		let peers = create_network(2, &cfgs);
+
+		let mut fd_dup = FileDescriptor { fd: 3, outbound_data: Arc::new(Mutex::new(Vec::new())) };
+		let addr_dup = NetAddress::IPv4{addr: [127, 0, 0, 1], port: 1003};
+		let id_a = cfgs[0].node_signer.get_node_id(Recipient::Node).unwrap();
+		peers[0].new_inbound_connection(fd_dup.clone(), Some(addr_dup.clone())).unwrap();
+
+		let mut dup_encryptor = PeerChannelEncryptor::new_outbound(id_a, SecretKey::from_slice(&[42; 32]).unwrap());
+		let initial_data = dup_encryptor.get_act_one(&peers[1].secp_ctx);
+		assert_eq!(peers[0].read_event(&mut fd_dup, &initial_data).unwrap(), false);
+		peers[0].process_events();
+
+		let a_data = fd_dup.outbound_data.lock().unwrap().split_off(0);
+		let (act_three, _) =
+			dup_encryptor.process_act_two(&a_data[..], &&cfgs[1].node_signer).unwrap();
+		assert_eq!(peers[0].read_event(&mut fd_dup, &act_three).unwrap(), false);
+
+		let not_init_msg = msgs::Ping { ponglen: 4, byteslen: 0 };
+		let msg_bytes = dup_encryptor.encrypt_message(&not_init_msg);
+		assert!(peers[0].read_event(&mut fd_dup, &msg_bytes).is_err());
+	}
+
 	#[test]
 	fn test_disconnect_all_peer() {
 		// Simple test which builds a network of PeerManager, connects and brings them to NoiseState::Finished and

lightning/src/util/test_utils.rs

@@ -330,13 +330,15 @@ impl chaininterface::BroadcasterInterface for TestBroadcaster {
 pub struct TestChannelMessageHandler {
 	pub pending_events: Mutex<Vec<events::MessageSendEvent>>,
 	expected_recv_msgs: Mutex<Option<Vec<wire::Message<()>>>>,
+	connected_peers: Mutex<HashSet<PublicKey>>,
 }
 
 impl TestChannelMessageHandler {
 	pub fn new() -> Self {
 		TestChannelMessageHandler {
 			pending_events: Mutex::new(Vec::new()),
 			expected_recv_msgs: Mutex::new(None),
+			connected_peers: Mutex::new(HashSet::new()),
 		}
 	}
 
@@ -421,8 +423,11 @@ impl msgs::ChannelMessageHandler for TestChannelMessageHandler {
 	fn handle_channel_reestablish(&self, _their_node_id: &PublicKey, msg: &msgs::ChannelReestablish) {
 		self.received_msg(wire::Message::ChannelReestablish(msg.clone()));
 	}
-	fn peer_disconnected(&self, _their_node_id: &PublicKey, _no_connection_possible: bool) {}
-	fn peer_connected(&self, _their_node_id: &PublicKey, _msg: &msgs::Init) -> Result<(), ()> {
+	fn peer_disconnected(&self, their_node_id: &PublicKey, _no_connection_possible: bool) {
+		assert!(self.connected_peers.lock().unwrap().remove(their_node_id));
+	}
+	fn peer_connected(&self, their_node_id: &PublicKey, _msg: &msgs::Init) -> 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(())