Test onion message replies

Author: jkczyz

lightning/src/onion_message/functional_tests.rs

@@ -20,7 +20,9 @@ use bitcoin::network::constants::Network;
 use bitcoin::secp256k1::{PublicKey, Secp256k1};
 
 use crate::io;
-use crate::sync::Arc;
+use crate::sync::{Arc, Mutex};
+
+use crate::prelude::*;
 
 struct MessengerNode {
 	keys_manager: Arc<test_utils::TestKeysInterface>,
@@ -32,7 +34,7 @@ struct MessengerNode {
 		Arc<TestOffersMessageHandler>,
 		Arc<TestCustomMessageHandler>
 	>,
-	logger: Arc<test_utils::TestLogger>,
+	custom_handler: Arc<TestCustomMessageHandler>,
 }
 
 impl MessengerNode {
@@ -45,7 +47,7 @@ struct TestMessageRouter {}
 
 impl MessageRouter for TestMessageRouter {
 	fn find_route(&self, _sender: &PublicKey, _destination: &Destination) -> Vec<PublicKey> {
-		todo!()
+		vec![]
 	}
 }
 
@@ -63,7 +65,7 @@ impl ResponseErrorHandler for TestOffersMessageHandler {
 	}
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug, PartialEq)]
 enum TestCustomMessage {
 	Request,
 	Response,
@@ -92,11 +94,38 @@ impl Writeable for TestCustomMessage {
 	}
 }
 
-struct TestCustomMessageHandler {}
+struct TestCustomMessageHandler {
+	expected_messages: Mutex<VecDeque<TestCustomMessage>>,
+}
+
+impl TestCustomMessageHandler {
+	fn new() -> Self {
+		Self { expected_messages: Mutex::new(VecDeque::new()) }
+	}
+
+	fn expect_message(&self, message: TestCustomMessage) {
+		self.expected_messages.lock().unwrap().push_back(message);
+	}
+}
+
+impl Drop for TestCustomMessageHandler {
+	fn drop(&mut self) {
+		#[cfg(feature = "std")] {
+			if !std::thread::panicking() {
+				assert!(self.expected_messages.lock().unwrap().is_empty());
+			}
+		}
+	}
+}
 
 impl CustomOnionMessageHandler for TestCustomMessageHandler {
 	type CustomMessage = TestCustomMessage;
 	fn handle_custom_message(&self, msg: Self::CustomMessage) -> Option<Self::CustomMessage> {
+		match self.expected_messages.lock().unwrap().pop_front() {
+			Some(expected_msg) => assert_eq!(expected_msg, msg),
+			None => panic!("Unexpected message: {:?}", msg),
+		}
+
 		match msg {
 			TestCustomMessage::Request => Some(TestCustomMessage::Response),
 			TestCustomMessage::Response => None,
@@ -133,14 +162,14 @@ fn create_nodes(num_messengers: u8) -> Vec<MessengerNode> {
 		let logger = Arc::new(test_utils::TestLogger::with_id(format!("node {}", i)));
 		let seed = [i as u8; 32];
 		let keys_manager = Arc::new(test_utils::TestKeysInterface::new(&seed, Network::Testnet));
+		let custom_handler = Arc::new(TestCustomMessageHandler::new());
 		nodes.push(MessengerNode {
 			keys_manager: keys_manager.clone(),
 			messenger: OnionMessenger::new(
-				keys_manager.clone(), keys_manager.clone(), logger.clone(),
-				Arc::new(TestMessageRouter {}), Arc::new(TestOffersMessageHandler {}),
-				Arc::new(TestCustomMessageHandler {})
+				Arc::clone(&keys_manager), keys_manager, logger, Arc::new(TestMessageRouter {}),
+				Arc::new(TestOffersMessageHandler {}), Arc::clone(&custom_handler)
 			),
-			logger,
+			custom_handler,
 		});
 	}
 	for idx in 0..num_messengers - 1 {
@@ -156,20 +185,14 @@ fn create_nodes(num_messengers: u8) -> Vec<MessengerNode> {
 
 fn pass_along_path(path: &Vec<MessengerNode>, expected_path_id: Option<[u8; 32]>) {
 	let mut prev_node = &path[0];
-	let num_nodes = path.len();
-	for (idx, node) in path.into_iter().skip(1).enumerate() {
+	for node in path.into_iter().skip(1) {
 		let events = prev_node.messenger.release_pending_msgs();
 		let onion_msg =  {
 			let msgs = events.get(&node.get_node_pk()).unwrap();
 			assert_eq!(msgs.len(), 1);
 			msgs[0].clone()
 		};
 		node.messenger.handle_onion_message(&prev_node.get_node_pk(), &onion_msg);
-		if idx == num_nodes - 1 {
-			node.logger.assert_log_contains(
-				"lightning::onion_message::messenger".to_string(),
-				format!("Received an onion message with path_id: {:02x?}", expected_path_id).to_string(), 1);
-		}
 		prev_node = node;
 	}
 }
@@ -180,6 +203,7 @@ fn one_hop() {
 	let test_msg = OnionMessageContents::Custom(TestCustomMessage::Response);
 
 	nodes[0].messenger.send_onion_message(&[], Destination::Node(nodes[1].get_node_pk()), test_msg, None).unwrap();
+	nodes[1].custom_handler.expect_message(TestCustomMessage::Response);
 	pass_along_path(&nodes, None);
 }
 
@@ -189,6 +213,7 @@ fn two_unblinded_hops() {
 	let test_msg = OnionMessageContents::Custom(TestCustomMessage::Response);
 
 	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk()], Destination::Node(nodes[2].get_node_pk()), test_msg, None).unwrap();
+	nodes[2].custom_handler.expect_message(TestCustomMessage::Response);
 	pass_along_path(&nodes, None);
 }
 
@@ -201,6 +226,7 @@ fn two_unblinded_two_blinded() {
 	let blinded_path = BlindedPath::new(&[nodes[3].get_node_pk(), nodes[4].get_node_pk()], &*nodes[4].keys_manager, &secp_ctx).unwrap();
 
 	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], Destination::BlindedPath(blinded_path), test_msg, None).unwrap();
+	nodes[4].custom_handler.expect_message(TestCustomMessage::Response);
 	pass_along_path(&nodes, None);
 }
 
@@ -213,6 +239,7 @@ fn three_blinded_hops() {
 	let blinded_path = BlindedPath::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk(), nodes[3].get_node_pk()], &*nodes[3].keys_manager, &secp_ctx).unwrap();
 
 	nodes[0].messenger.send_onion_message(&[], Destination::BlindedPath(blinded_path), test_msg, None).unwrap();
+	nodes[3].custom_handler.expect_message(TestCustomMessage::Response);
 	pass_along_path(&nodes, None);
 }
 
@@ -239,11 +266,13 @@ fn we_are_intro_node() {
 	let blinded_path = BlindedPath::new(&[nodes[0].get_node_pk(), nodes[1].get_node_pk(), nodes[2].get_node_pk()], &*nodes[2].keys_manager, &secp_ctx).unwrap();
 
 	nodes[0].messenger.send_onion_message(&[], Destination::BlindedPath(blinded_path), OnionMessageContents::Custom(test_msg.clone()), None).unwrap();
+	nodes[2].custom_handler.expect_message(TestCustomMessage::Response);
 	pass_along_path(&nodes, None);
 
 	// Try with a two-hop blinded path where we are the introduction node.
 	let blinded_path = BlindedPath::new(&[nodes[0].get_node_pk(), nodes[1].get_node_pk()], &*nodes[1].keys_manager, &secp_ctx).unwrap();
 	nodes[0].messenger.send_onion_message(&[], Destination::BlindedPath(blinded_path), OnionMessageContents::Custom(test_msg), None).unwrap();
+	nodes[1].custom_handler.expect_message(TestCustomMessage::Response);
 	nodes.remove(2);
 	pass_along_path(&nodes, None);
 }
@@ -271,28 +300,31 @@ fn invalid_blinded_path_error() {
 
 #[test]
 fn reply_path() {
-	let nodes = create_nodes(4);
-	let test_msg = TestCustomMessage::Response;
+	let mut nodes = create_nodes(4);
+	let test_msg = TestCustomMessage::Request;
 	let secp_ctx = Secp256k1::new();
 
 	// Destination::Node
 	let reply_path = BlindedPath::new(&[nodes[2].get_node_pk(), nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx).unwrap();
 	nodes[0].messenger.send_onion_message(&[nodes[1].get_node_pk(), nodes[2].get_node_pk()], Destination::Node(nodes[3].get_node_pk()), OnionMessageContents::Custom(test_msg.clone()), Some(reply_path)).unwrap();
+	nodes[3].custom_handler.expect_message(TestCustomMessage::Request);
 	pass_along_path(&nodes, None);
 	// Make sure the last node successfully decoded the reply path.
-	nodes[3].logger.assert_log_contains(
-		"lightning::onion_message::messenger".to_string(),
-		format!("Received an onion message with path_id None and a reply_path").to_string(), 1);
+	nodes[0].custom_handler.expect_message(TestCustomMessage::Response);
+	nodes.reverse();
+	pass_along_path(&nodes, None);
 
 	// Destination::BlindedPath
 	let blinded_path = BlindedPath::new(&[nodes[1].get_node_pk(), nodes[2].get_node_pk(), nodes[3].get_node_pk()], &*nodes[3].keys_manager, &secp_ctx).unwrap();
 	let reply_path = BlindedPath::new(&[nodes[2].get_node_pk(), nodes[1].get_node_pk(), nodes[0].get_node_pk()], &*nodes[0].keys_manager, &secp_ctx).unwrap();
 
 	nodes[0].messenger.send_onion_message(&[], Destination::BlindedPath(blinded_path), OnionMessageContents::Custom(test_msg), Some(reply_path)).unwrap();
+	nodes[3].custom_handler.expect_message(TestCustomMessage::Request);
+	pass_along_path(&nodes, None);
+	// Make sure the last node successfully decoded the reply path.
+	nodes[0].custom_handler.expect_message(TestCustomMessage::Response);
+	nodes.reverse();
 	pass_along_path(&nodes, None);
-	nodes[3].logger.assert_log_contains(
-		"lightning::onion_message::messenger".to_string(),
-		format!("Received an onion message with path_id None and a reply_path").to_string(), 2);
 }
 
 #[test]