sim-ln/refactor: Improve activity dest node lookup via Graph

chuksys · chuksys · commit 9f3a4f2ed54c · 2025-05-09T16:14:27.000+01:00
Switches from an async closure to direct NetworkGraph usage for activity destination
node lookup. This improves efficiency and maintainability, and enables destination
aliases by validating against the graph and erroring on duplicates.
diff --git a/sim-cli/src/parsing.rs b/sim-cli/src/parsing.rs
@@ -15,7 +15,6 @@ use simln_lib::{
 use simln_lib::{ShortChannelID, SimulationError};
 use std::collections::HashMap;
 use std::fs;
-use std::ops::AsyncFn;
 use std::path::PathBuf;
 use std::sync::Arc;
 use tokio::sync::Mutex;
@@ -179,6 +178,15 @@ pub struct ActivityParser {
     pub amount_msat: Amount,
 }
 
+struct ActivityValidationParams {
+    pk_node_map: HashMap<PublicKey, NodeInfo>,
+    alias_node_map: HashMap<String, NodeInfo>,
+    graph_nodes_by_pk: HashMap<PublicKey, NodeInfo>,
+    // Store graph nodes' information keyed by their alias.
+    // An alias can be mapped to multiple nodes because it is not a unique identifier.
+    graph_nodes_by_alias: HashMap<String, Vec<NodeInfo>>,
+}
+
 impl TryFrom<&Cli> for SimulationCfg {
     type Error = anyhow::Error;
 
@@ -378,15 +386,20 @@ fn add_node_to_maps(nodes: &HashMap<PublicKey, NodeInfo>) -> Result<NodeMapping,
 /// have been configured.
 async fn validate_activities(
     activity: Vec<ActivityParser>,
-    pk_node_map: HashMap<PublicKey, NodeInfo>,
-    alias_node_map: HashMap<String, NodeInfo>,
-    get_node_info: impl AsyncFn(&PublicKey) -> Result<NodeInfo, LightningError>,
+    activity_validation_params: ActivityValidationParams,
 ) -> Result<Vec<ActivityDefinition>, LightningError> {
     let mut validated_activities = vec![];
 
+    let ActivityValidationParams {
+        pk_node_map,
+        alias_node_map,
+        graph_nodes_by_pk,
+        graph_nodes_by_alias,
+    } = activity_validation_params;
+
     // Make all the activities identifiable by PK internally
     for act in activity.into_iter() {
-        // We can only map aliases to nodes we control, so if either the source or destination alias
+        // We can only map source aliases to nodes we control, so if the source alias
         // is not in alias_node_map, we fail
         let source = if let Some(source) = match &act.source {
             NodeId::PublicKey(pk) => pk_node_map.get(pk),
@@ -402,8 +415,21 @@ async fn validate_activities(
 
         let destination = match &act.destination {
             NodeId::Alias(a) => {
-                if let Some(info) = alias_node_map.get(a) {
-                    info.clone()
+                if let Some(node_info) = alias_node_map.get(a) {
+                    node_info.clone()
+                } else if let Some(node_infos) = graph_nodes_by_alias.get(a) {
+                    if node_infos.len() > 1 {
+                        let pks: Vec<PublicKey> = node_infos
+                            .iter()
+                            .map(|node_info| node_info.pubkey)
+                            .collect();
+                        return Err(LightningError::ValidationError(format!(
+                            "Multiple nodes in the graph have the same destination alias - {}. 
+                            Use one of these public keys as the destination instead - {:?}",
+                            a, pks
+                        )));
+                    }
+                    node_infos[0].clone()
                 } else {
                     return Err(LightningError::ValidationError(format!(
                         "unknown activity destination: {}.",
@@ -412,10 +438,15 @@ async fn validate_activities(
                 }
             },
             NodeId::PublicKey(pk) => {
-                if let Some(info) = pk_node_map.get(pk) {
-                    info.clone()
+                if let Some(node_info) = pk_node_map.get(pk) {
+                    node_info.clone()
+                } else if let Some(node_info) = graph_nodes_by_pk.get(pk) {
+                    node_info.clone()
                 } else {
-                    get_node_info(pk).await?
+                    return Err(LightningError::ValidationError(format!(
+                        "unknown activity destination: {}.",
+                        act.destination
+                    )));
                 }
             },
         };
@@ -507,18 +538,35 @@ pub async fn get_validated_activities(
 ) -> Result<Vec<ActivityDefinition>, LightningError> {
     // We need to be able to look up destination nodes in the graph, because we allow defined activities to send to
     // nodes that we do not control. To do this, we can just grab the first node in our map and perform the lookup.
-    let get_node = async |pk: &PublicKey| -> Result<NodeInfo, LightningError> {
-        if let Some(c) = clients.values().next() {
-            return c.lock().await.get_node_info(pk).await;
-        }
-        Err(LightningError::GetNodeInfoError(
-            "no nodes for query".to_string(),
-        ))
-    };
+    let graph = match clients.values().next() {
+        Some(client) => client
+            .lock()
+            .await
+            .get_graph()
+            .await
+            .map_err(|e| LightningError::GetGraphError(format!("Error getting graph {:?}", e))),
+        None => Err(LightningError::GetGraphError("Graph is empty".to_string())),
+    }?;
+    let mut graph_nodes_by_alias: HashMap<String, Vec<NodeInfo>> = HashMap::new();
+
+    for node in &graph.nodes_by_pk {
+        graph_nodes_by_alias
+            .entry(node.1.alias.clone())
+            .or_default()
+            .push(node.1.clone());
+    }
+
     let NodeMapping {
         pk_node_map,
         alias_node_map,
     } = add_node_to_maps(&nodes_info)?;
 
-    validate_activities(activity.to_vec(), pk_node_map, alias_node_map, get_node).await
+    let activity_validation_params = ActivityValidationParams {
+        pk_node_map,
+        alias_node_map,
+        graph_nodes_by_pk: graph.nodes_by_pk,
+        graph_nodes_by_alias,
+    };
+
+    validate_activities(activity.to_vec(), activity_validation_params).await
 }
diff --git a/simln-lib/src/cln.rs b/simln-lib/src/cln.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use async_trait::async_trait;
 use bitcoin::secp256k1::PublicKey;
 use bitcoin::Network;
@@ -17,7 +19,8 @@ use tonic::transport::{Certificate, Channel, ClientTlsConfig, Identity};
 use triggered::Listener;
 
 use crate::{
-    serializers, LightningError, LightningNode, NodeId, NodeInfo, PaymentOutcome, PaymentResult,
+    serializers, Graph, LightningError, LightningNode, NodeId, NodeInfo, PaymentOutcome,
+    PaymentResult,
 };
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -263,6 +266,34 @@ impl LightningNode for ClnNode {
         node_channels.extend(self.node_channels(false).await?);
         Ok(node_channels)
     }
+
+    async fn get_graph(&mut self) -> Result<Graph, LightningError> {
+        let nodes: Vec<cln_grpc::pb::ListnodesNodes> = self
+            .client
+            .list_nodes(ListnodesRequest { id: None })
+            .await
+            .map_err(|err| LightningError::GetNodeInfoError(err.to_string()))?
+            .into_inner()
+            .nodes;
+
+        let mut nodes_by_pk: HashMap<PublicKey, NodeInfo> = HashMap::new();
+
+        for node in nodes {
+            nodes_by_pk.insert(
+                PublicKey::from_slice(&node.nodeid).expect("Public Key not valid"),
+                NodeInfo {
+                    pubkey: PublicKey::from_slice(&node.nodeid).expect("Public Key not valid"),
+                    alias: node.clone().alias.unwrap_or(String::new()),
+                    features: node
+                        .features
+                        .clone()
+                        .map_or(NodeFeatures::empty(), NodeFeatures::from_be_bytes),
+                },
+            );
+        }
+
+        Ok(Graph { nodes_by_pk })
+    }
 }
 
 async fn reader(filename: &str) -> Result<Vec<u8>, Error> {
diff --git a/simln-lib/src/eclair.rs b/simln-lib/src/eclair.rs
@@ -1,5 +1,6 @@
 use crate::{
-    serializers, LightningError, LightningNode, NodeId, NodeInfo, PaymentOutcome, PaymentResult,
+    serializers, Graph, LightningError, LightningNode, NodeId, NodeInfo, PaymentOutcome,
+    PaymentResult,
 };
 use async_trait::async_trait;
 use bitcoin::secp256k1::PublicKey;
@@ -243,6 +244,29 @@ impl LightningNode for EclairNode {
 
         Ok(capacities_msat)
     }
+
+    async fn get_graph(&mut self) -> Result<Graph, LightningError> {
+        let nodes: NodesResponse = self
+            .client
+            .request("nodes", None)
+            .await
+            .map_err(|err| LightningError::GetNodeInfoError(err.to_string()))?;
+
+        let mut nodes_by_pk: HashMap<PublicKey, NodeInfo> = HashMap::new();
+
+        for node in nodes {
+            nodes_by_pk.insert(
+                PublicKey::from_str(&node.node_id).expect("Public Key not valid"),
+                NodeInfo {
+                    pubkey: PublicKey::from_str(&node.node_id).expect("Public Key not valid"),
+                    alias: node.alias.clone(),
+                    features: parse_json_to_node_features(&node.features),
+                },
+            );
+        }
+
+        Ok(Graph { nodes_by_pk })
+    }
 }
 
 #[derive(Debug, Deserialize)]
@@ -288,7 +312,16 @@ struct NodeResponse {
     announcement: Announcement,
 }
 
+#[derive(Debug, Deserialize)]
+struct NodeInGraph {
+    #[serde(rename = "nodeId")]
+    node_id: String,
+    alias: String,
+    features: Value,
+}
+
 type ChannelsResponse = Vec<Channel>;
+type NodesResponse = Vec<NodeInGraph>;
 
 #[derive(Debug, Deserialize)]
 struct Channel {
diff --git a/simln-lib/src/lib.rs b/simln-lib/src/lib.rs
@@ -258,6 +258,9 @@ pub enum LightningError {
     /// Error that occurred while listing channels.
     #[error("List channels error: {0}")]
     ListChannelsError(String),
+    /// Error that occurred while getting graph.
+    #[error("Get graph error: {0}")]
+    GetGraphError(String),
 }
 
 /// Information about a Lightning Network node.
@@ -286,6 +289,33 @@ impl Display for NodeInfo {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct ChannelInfo {
+    pub channel_id: ShortChannelID,
+    pub capacity_msat: u64,
+}
+
+#[derive(Debug, Clone)]
+/// Graph represents the network graph of the simulated network and is useful for efficient lookups.
+pub struct Graph {
+    // Store nodes' information keyed by their public key.
+    pub nodes_by_pk: HashMap<PublicKey, NodeInfo>,
+}
+
+impl Graph {
+    pub fn new() -> Self {
+        Graph {
+            nodes_by_pk: HashMap::new(),
+        }
+    }
+}
+
+impl Default for Graph {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 /// LightningNode represents the functionality that is required to execute events on a lightning node.
 #[async_trait]
 pub trait LightningNode: Send {
@@ -310,6 +340,8 @@ pub trait LightningNode: Send {
     /// Lists all channels, at present only returns a vector of channel capacities in msat because no further
     /// information is required.
     async fn list_channels(&mut self) -> Result<Vec<u64>, LightningError>;
+    /// Get the network graph from the point of view of a given node.
+    async fn get_graph(&mut self) -> Result<Graph, LightningError>;
 }
 
 /// Represents an error that occurs when generating a destination for a payment.
diff --git a/simln-lib/src/lnd.rs b/simln-lib/src/lnd.rs
@@ -2,7 +2,8 @@ use std::collections::HashSet;
 use std::{collections::HashMap, str::FromStr};
 
 use crate::{
-    serializers, LightningError, LightningNode, NodeId, NodeInfo, PaymentOutcome, PaymentResult,
+    serializers, Graph, LightningError, LightningNode, NodeId, NodeInfo, PaymentOutcome,
+    PaymentResult,
 };
 use async_trait::async_trait;
 use bitcoin::hashes::{sha256, Hash};
@@ -12,7 +13,9 @@ use lightning::ln::features::NodeFeatures;
 use lightning::ln::{PaymentHash, PaymentPreimage};
 use serde::{Deserialize, Serialize};
 use tonic_lnd::lnrpc::{payment::PaymentStatus, GetInfoRequest, GetInfoResponse};
-use tonic_lnd::lnrpc::{ListChannelsRequest, NodeInfoRequest, PaymentFailureReason};
+use tonic_lnd::lnrpc::{
+    ChannelGraphRequest, ListChannelsRequest, NodeInfoRequest, PaymentFailureReason,
+};
 use tonic_lnd::routerrpc::TrackPaymentRequest;
 use tonic_lnd::tonic::Code::Unavailable;
 use tonic_lnd::tonic::Status;
@@ -275,6 +278,34 @@ impl LightningNode for LndNode {
             .map(|channel| 1000 * channel.capacity as u64)
             .collect())
     }
+
+    async fn get_graph(&mut self) -> Result<Graph, LightningError> {
+        let nodes = self
+            .client
+            .lightning()
+            .describe_graph(ChannelGraphRequest {
+                include_unannounced: false,
+            })
+            .await
+            .map_err(|err| LightningError::GetNodeInfoError(err.to_string()))?
+            .into_inner()
+            .nodes;
+
+        let mut nodes_by_pk: HashMap<PublicKey, NodeInfo> = HashMap::new();
+
+        for node in nodes {
+            nodes_by_pk.insert(
+                PublicKey::from_str(&node.pub_key).expect("Public Key not valid"),
+                NodeInfo {
+                    pubkey: PublicKey::from_str(&node.pub_key).expect("Public Key not valid"),
+                    alias: node.alias.clone(),
+                    features: parse_node_features(node.features.keys().cloned().collect()),
+                },
+            );
+        }
+
+        Ok(Graph { nodes_by_pk })
+    }
 }
 
 fn string_to_payment_hash(hash: &str) -> Result<PaymentHash, LightningError> {
diff --git a/simln-lib/src/sim_node.rs b/simln-lib/src/sim_node.rs
diff --git a/simln-lib/src/test_utils.rs b/simln-lib/src/test_utils.rs
