OM pathfinding: ensure first_hops overrides network graph

Author: valentinewallace

lightning/src/routing/onion_message.rs

@@ -40,13 +40,15 @@ pub fn find_path<L: Deref, GL: Deref>(
 
 	// Add our start and first-hops to `frontier`.
 	let start = NodeId::from_pubkey(&our_node_pubkey);
+	let mut valid_first_hops = HashSet::new();
 	let mut frontier = BinaryHeap::new();
 	frontier.push(PathBuildingHop { cost: 0, node_id: start, parent_node_id: start });
 	if let Some(first_hops) = first_hops {
 		for hop in first_hops {
 			if !hop.counterparty.features.supports_onion_messages() { continue; }
 			let node_id = NodeId::from_pubkey(&hop.counterparty.node_id);
 			frontier.push(PathBuildingHop { cost: 1, node_id, parent_node_id: start });
+			valid_first_hops.insert(node_id);
 		}
 	}
 
@@ -60,7 +62,7 @@ pub fn find_path<L: Deref, GL: Deref>(
 			return Ok(reverse_path(visited, our_node_id, dest_node_id, logger)?)
 		}
 		if let Some(node_info) = network_nodes.get(&node_id) {
-			if node_id == our_node_id {
+			if valid_first_hops.contains(&node_id) || node_id == our_node_id {
 			} else if let Some(node_ann) = &node_info.announcement_info {
 				if !node_ann.features.supports_onion_messages() || node_ann.features.requires_unknown_bits()
 				{ continue; }
@@ -146,7 +148,7 @@ fn reverse_path<L: Deref>(
 #[cfg(test)]
 mod tests {
 	use ln::features::{InitFeatures, NodeFeatures};
-	use routing::test_utils::{add_or_update_node, build_graph_with_features, build_line_graph, get_nodes};
+	use routing::test_utils::{add_or_update_node, build_graph_with_features, build_line_graph, get_channel_details, get_nodes};
 
 	use sync::Arc;
 
@@ -219,6 +221,13 @@ mod tests {
 		// If all nodes require some features we don't understand, route should fail
 		let err = super::find_path(&our_id, &node_pks[2], &network_graph, None, Arc::clone(&logger)).unwrap_err();
 		assert_eq!(err, super::Error::PathNotFound);
+
+		// If we specify a channel to node7, that overrides our local channel view and that gets used
+		let our_chans = vec![get_channel_details(Some(42), node_pks[7].clone(), features, 250_000_000)];
+		let path = super::find_path(&our_id, &node_pks[2], &network_graph, Some(&our_chans.iter().collect::<Vec<_>>()), Arc::clone(&logger)).unwrap();
+		assert_eq!(path.len(), 2);
+		assert_eq!(path[0], node_pks[7]);
+		assert_eq!(path[1], node_pks[2]);
 	}
 
 	#[test]
@@ -236,6 +245,12 @@ mod tests {
 		assert_eq!(path[0], node_pks[1]);
 		assert_eq!(path[1], node_pks[2]);
 		assert_eq!(path[2], node_pks[0]);
+
+		// If we specify a channel to node1, that overrides our local channel view and that gets used
+		let our_chans = vec![get_channel_details(Some(42), node_pks[0].clone(), features, 250_000_000)];
+		let path = super::find_path(&our_id, &node_pks[0], &network_graph, Some(&our_chans.iter().collect::<Vec<_>>()), Arc::clone(&logger)).unwrap();
+		assert_eq!(path.len(), 1);
+		assert_eq!(path[0], node_pks[0]);
 	}
 }