Add test utilities and integration tests for onion messages

lightning/src/ln/onion_message.rs

@@ -697,3 +697,106 @@ pub type SimpleArcOnionMessenger<L> = OnionMessenger<InMemorySigner, Arc<KeysMan
 /// Useful for simplifying the parameters of [`SimpleRefChannelManager`] and
 /// [`SimpleRefPeerManager`]. See their docs for more details.
 pub type SimpleRefOnionMessenger<'a, 'b, L> = OnionMessenger<InMemorySigner, &'a KeysManager, &'b L>;
+
+#[cfg(test)]
+mod tests {
+	use chain::keysinterface::{KeysInterface, Recipient};
+	use ln::msgs::OnionMessageHandler;
+	use ln::onion_message::{BlindedRoute, Destination, OnionMessenger};
+	use util::enforcing_trait_impls::EnforcingSigner;
+	use util::events::{MessageSendEvent, MessageSendEventsProvider};
+	use util::test_utils;
+
+	use bitcoin::network::constants::Network;
+	use bitcoin::secp256k1::{PublicKey, Secp256k1};
+
+	use sync::Arc;
+
+	struct MessengerNode {
+		keys_manager: Arc<test_utils::TestKeysInterface>,
+		messenger: OnionMessenger<EnforcingSigner, Arc<test_utils::TestKeysInterface>, Arc<test_utils::TestLogger>>,
+		logger: Arc<test_utils::TestLogger>,
+	}
+
+	impl MessengerNode {
+		fn get_node_pk(&self) -> PublicKey {
+			let secp_ctx = Secp256k1::new();
+			PublicKey::from_secret_key(&secp_ctx, &self.keys_manager.get_node_secret(Recipient::Node).unwrap())
+		}
+	}
+
+	fn create_nodes(num_messengers: u8) -> Vec<MessengerNode> {
+		let mut res = Vec::new();
+		for i in 0..num_messengers {
+			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));
+			res.push(MessengerNode {
+				keys_manager: keys_manager.clone(),
+				messenger: OnionMessenger::new(keys_manager, logger.clone()),
+				logger,
+			});
+		}
+		res
+	}
+
+	fn pass_along_path(path: Vec<&MessengerNode>, expected_path_id: Option<[u8; 32]>) {
+		let mut prev_node = path[0];
+		for (idx, node) in path.iter().enumerate().skip(1) {
+			let events = prev_node.messenger.get_and_clear_pending_msg_events();
+			assert_eq!(events.len(), 1);
+			let onion_msg = match &events[0] {
+				MessageSendEvent::SendOnionMessage { msg, .. } => msg.clone(),
+				_ => panic!("Unexpected event"),
+			};
+			node.messenger.handle_onion_message(&prev_node.get_node_pk(), &onion_msg);
+			if idx == path.len() - 1 {
+				node.logger.assert_log_contains(
+					"lightning::ln::onion_message".to_string(),
+					format!("Received an onion message with path_id: {:02x?}", expected_path_id).to_string(), 1);
+				break
+			}
+			prev_node = node;
+		}
+	}
+
+	#[test]
+	fn one_hop() {
+		let mut nodes = create_nodes(2);
+		let (node1, node2) = (nodes.remove(0), nodes.remove(0));
+
+		node1.messenger.send_onion_message(vec![], Destination::Node(node2.get_node_pk())).unwrap();
+		pass_along_path(vec![&node1, &node2], None);
+	}
+
+	#[test]
+	fn two_unblinded_hops() {
+		let mut nodes = create_nodes(3);
+		let (node1, node2, node3) = (nodes.remove(0), nodes.remove(0), nodes.remove(0));
+
+		node1.messenger.send_onion_message(vec![node2.get_node_pk()], Destination::Node(node3.get_node_pk())).unwrap();
+		pass_along_path(vec![&node1, &node2, &node3], None);
+	}
+
+	#[test]
+	fn two_unblinded_two_blinded() {
+		let mut nodes = create_nodes(5);
+		let (node1, node2, node3, node4, node5) = (nodes.remove(0), nodes.remove(0), nodes.remove(0), nodes.remove(0), nodes.remove(0));
+
+		let blinded_route = BlindedRoute::new(vec![node4.get_node_pk(), node5.get_node_pk()], &node5.keys_manager).unwrap();
+
+		node1.messenger.send_onion_message(vec![node2.get_node_pk(), node3.get_node_pk()], Destination::BlindedRoute(blinded_route)).unwrap();
+		pass_along_path(vec![&node1, &node2, &node3, &node4, &node5], Some([42; 32]));
+	}
+
+	#[test]
+	fn three_blinded_hops() {
+		let mut nodes = create_nodes(4);
+		let (node1, node2, node3, node4) = (nodes.remove(0), nodes.remove(0), nodes.remove(0), nodes.remove(0));
+
+		let blinded_route = BlindedRoute::new(vec![node2.get_node_pk(), node3.get_node_pk(), node4.get_node_pk()], &node4.keys_manager).unwrap();
+
+		node1.messenger.send_onion_message(vec![], Destination::BlindedRoute(blinded_route)).unwrap();
+		pass_along_path(vec![&node1, &node2, &node3, &node4], Some([42; 32]));
+	}
+}