Update: Convert TestMessageRouter to enum

Author: shaavan

lightning/src/ln/channelmanager.rs

@@ -16386,7 +16386,7 @@ pub mod bench {
 		let scorer = RwLock::new(test_utils::TestScorer::new());
 		let entropy = test_utils::TestKeysInterface::new(&[0u8; 32], network);
 		let router = test_utils::TestRouter::new(Arc::new(NetworkGraph::new(network, &logger_a)), &logger_a, &scorer);
-		let message_router = test_utils::TestMessageRouter::new(Arc::new(NetworkGraph::new(network, &logger_a)), &entropy);
+		let message_router = test_utils::TestMessageRouter::new_default(Arc::new(NetworkGraph::new(network, &logger_a)), &entropy);
 
 		let mut config: UserConfig = Default::default();
 		config.channel_config.max_dust_htlc_exposure = MaxDustHTLCExposure::FeeRateMultiplier(5_000_000 / 253);

lightning/src/ln/functional_test_utils.rs

@@ -722,7 +722,7 @@ impl<'a, 'b, 'c> Drop for Node<'a, 'b, 'c> {
 					signer_provider: self.keys_manager,
 					fee_estimator: &test_utils::TestFeeEstimator::new(253),
 					router: &test_utils::TestRouter::new(Arc::clone(&network_graph), &self.logger, &scorer),
-					message_router: &test_utils::TestMessageRouter::new(network_graph, self.keys_manager),
+					message_router: &test_utils::TestMessageRouter::new_default(network_graph, self.keys_manager),
 					chain_monitor: self.chain_monitor,
 					tx_broadcaster: &broadcaster,
 					logger: &self.logger,
@@ -3310,7 +3310,7 @@ pub fn create_node_cfgs_with_persisters<'a>(node_count: usize, chanmon_cfgs: &'a
 			tx_broadcaster: &chanmon_cfgs[i].tx_broadcaster,
 			fee_estimator: &chanmon_cfgs[i].fee_estimator,
 			router: test_utils::TestRouter::new(network_graph.clone(), &chanmon_cfgs[i].logger, &chanmon_cfgs[i].scorer),
-			message_router: test_utils::TestMessageRouter::new(network_graph.clone(), &chanmon_cfgs[i].keys_manager),
+			message_router: test_utils::TestMessageRouter::new_default(network_graph.clone(), &chanmon_cfgs[i].keys_manager),
 			chain_monitor,
 			keys_manager: &chanmon_cfgs[i].keys_manager,
 			node_seed: seed,
@@ -3322,6 +3322,34 @@ pub fn create_node_cfgs_with_persisters<'a>(node_count: usize, chanmon_cfgs: &'a
 	nodes
 }
 
+pub fn create_node_cfgs_with_node_id_message_router<'a>(node_count: usize, chanmon_cfgs: &'a Vec<TestChanMonCfg>) -> Vec<NodeCfg<'a>> {
+	let persisters: Vec<_> = chanmon_cfgs.iter().map(|c| &c.persister).collect();
+
+	let mut nodes = Vec::new();
+
+	for i in 0..node_count {
+		let chain_monitor = test_utils::TestChainMonitor::new(Some(&chanmon_cfgs[i].chain_source), &chanmon_cfgs[i].tx_broadcaster, &chanmon_cfgs[i].logger, &chanmon_cfgs[i].fee_estimator, persisters[i], &chanmon_cfgs[i].keys_manager);
+		let network_graph = Arc::new(NetworkGraph::new(Network::Testnet, &chanmon_cfgs[i].logger));
+		let seed = [i as u8; 32];
+		nodes.push(NodeCfg {
+			chain_source: &chanmon_cfgs[i].chain_source,
+			logger: &chanmon_cfgs[i].logger,
+			tx_broadcaster: &chanmon_cfgs[i].tx_broadcaster,
+			fee_estimator: &chanmon_cfgs[i].fee_estimator,
+			router: test_utils::TestRouter::new(network_graph.clone(), &chanmon_cfgs[i].logger, &chanmon_cfgs[i].scorer),
+			message_router: test_utils::TestMessageRouter::new_node_id_router(network_graph.clone(), &chanmon_cfgs[i].keys_manager),
+			chain_monitor,
+			keys_manager: &chanmon_cfgs[i].keys_manager,
+			node_seed: seed,
+			network_graph,
+			override_init_features: Rc::new(RefCell::new(None)),
+		});
+	}
+
+	nodes
+}
+
+
 pub fn test_default_channel_config() -> UserConfig {
 	let mut default_config = UserConfig::default();
 	// Set cltv_expiry_delta slightly lower to keep the final CLTV values inside one byte in our

lightning/src/ln/functional_tests.rs

@@ -5987,7 +5987,7 @@ pub fn test_key_derivation_params() {
 	let network_graph = Arc::new(NetworkGraph::new(Network::Testnet, &chanmon_cfgs[0].logger));
 	let scorer = RwLock::new(test_utils::TestScorer::new());
 	let router = test_utils::TestRouter::new(network_graph.clone(), &chanmon_cfgs[0].logger, &scorer);
-	let message_router = test_utils::TestMessageRouter::new(network_graph.clone(), &keys_manager);
+	let message_router = test_utils::TestMessageRouter::new_default(network_graph.clone(), &keys_manager);
 	let node = NodeCfg { chain_source: &chanmon_cfgs[0].chain_source, logger: &chanmon_cfgs[0].logger, tx_broadcaster: &chanmon_cfgs[0].tx_broadcaster, fee_estimator: &chanmon_cfgs[0].fee_estimator, router, message_router, chain_monitor, keys_manager: &keys_manager, network_graph, node_seed: seed, override_init_features: alloc::rc::Rc::new(core::cell::RefCell::new(None)) };
 	let mut node_cfgs = create_node_cfgs(3, &chanmon_cfgs);
 	node_cfgs.remove(0);

lightning/src/util/test_utils.rs

@@ -33,7 +33,7 @@ use crate::ln::types::ChannelId;
 use crate::ln::{msgs, wire};
 use crate::offers::invoice::UnsignedBolt12Invoice;
 use crate::onion_message::messenger::{
-	DefaultMessageRouter, Destination, MessageRouter, OnionMessagePath,
+	DefaultMessageRouter, Destination, MessageRouter, NodeIdMessageRouter, OnionMessagePath
 };
 use crate::routing::gossip::{EffectiveCapacity, NetworkGraph, NodeId, RoutingFees};
 use crate::routing::router::{
@@ -311,41 +311,66 @@ impl<'a> Drop for TestRouter<'a> {
 	}
 }
 
-pub struct TestMessageRouter<'a> {
-	inner: DefaultMessageRouter<
-		Arc<NetworkGraph<&'a TestLogger>>,
-		&'a TestLogger,
-		&'a TestKeysInterface,
-	>,
+pub enum TestMessageRouter<'a> {
+	Default {
+		inner: DefaultMessageRouter<
+			Arc<NetworkGraph<&'a TestLogger>>,
+			&'a TestLogger,
+			&'a TestKeysInterface,
+		>,
+	},
+
+	NodeId {
+		inner: NodeIdMessageRouter<
+			Arc<NetworkGraph<&'a TestLogger>>,
+			&'a TestLogger,
+			&'a TestKeysInterface,
+		>
+	}
 }
 
 impl<'a> TestMessageRouter<'a> {
-	pub fn new(
+	pub fn new_default(
 		network_graph: Arc<NetworkGraph<&'a TestLogger>>, entropy_source: &'a TestKeysInterface,
 	) -> Self {
-		Self { inner: DefaultMessageRouter::new(network_graph, entropy_source) }
+		Self::Default { inner: DefaultMessageRouter::new(network_graph, entropy_source) }
+	}
+
+	pub fn new_node_id_router(
+		network_graph: Arc<NetworkGraph<&'a TestLogger>>, entropy_source: &'a TestKeysInterface,
+	) -> Self {
+		Self::NodeId { inner: NodeIdMessageRouter::new(network_graph, entropy_source) }
 	}
 }
 
 impl<'a> MessageRouter for TestMessageRouter<'a> {
 	fn find_path(
 		&self, sender: PublicKey, peers: Vec<PublicKey>, destination: Destination,
 	) -> Result<OnionMessagePath, ()> {
-		self.inner.find_path(sender, peers, destination)
+		match self {
+			Self::Default { inner } => inner.find_path(sender, peers, destination),
+			Self::NodeId { inner } => inner.find_path(sender, peers, destination),
+		}
 	}
 
 	fn create_blinded_paths<T: secp256k1::Signing + secp256k1::Verification>(
 		&self, recipient: PublicKey, context: MessageContext, peers: Vec<MessageForwardNode>,
 		secp_ctx: &Secp256k1<T>,
 	) -> Result<Vec<BlindedMessagePath>, ()> {
-		self.inner.create_blinded_paths(recipient, context, peers, secp_ctx)
+		match self {
+			Self::Default { inner } => inner.create_blinded_paths(recipient, context, peers, secp_ctx),
+			Self::NodeId { inner } => inner.create_blinded_paths(recipient, context, peers, secp_ctx),
+		}
 	}
 
 	fn create_compact_blinded_paths<T: secp256k1::Signing + secp256k1::Verification>(
 		&self, recipient: PublicKey, context: MessageContext, peers: Vec<MessageForwardNode>,
 		secp_ctx: &Secp256k1<T>,
 	) -> Result<Vec<BlindedMessagePath>, ()> {
-		self.inner.create_compact_blinded_paths(recipient, context, peers, secp_ctx)
+		match self {
+			Self::Default { inner } => inner.create_compact_blinded_paths(recipient, context, peers, secp_ctx),
+			Self::NodeId { inner } => inner.create_compact_blinded_paths(recipient, context, peers, secp_ctx),
+		}
 	}
 }