Merge pull request #274 from elnosh/custom-records

add custom records to htlc

Author: carlaKC

sim-cli/src/main.rs

@@ -3,7 +3,10 @@ use std::sync::Arc;
 use clap::Parser;
 use log::LevelFilter;
 use sim_cli::parsing::{create_simulation, create_simulation_with_network, parse_sim_params, Cli};
-use simln_lib::{latency_interceptor::LatencyIntercepor, sim_node::Interceptor};
+use simln_lib::{
+    latency_interceptor::LatencyIntercepor,
+    sim_node::{CustomRecords, Interceptor},
+};
 use simple_logger::SimpleLogger;
 use tokio_util::task::TaskTracker;
 
@@ -38,7 +41,14 @@ async fn main() -> anyhow::Result<()> {
         } else {
             vec![]
         };
-        create_simulation_with_network(&cli, &sim_params, tasks.clone(), interceptors).await?
+        create_simulation_with_network(
+            &cli,
+            &sim_params,
+            tasks.clone(),
+            interceptors,
+            CustomRecords::default(),
+        )
+        .await?
     };
     let sim2 = sim.clone();
 

sim-cli/src/parsing.rs

@@ -5,8 +5,8 @@ use log::LevelFilter;
 use serde::{Deserialize, Serialize};
 use simln_lib::clock::SimulationClock;
 use simln_lib::sim_node::{
-    ln_node_from_graph, populate_network_graph, ChannelPolicy, Interceptor, SimGraph,
-    SimulatedChannel,
+    ln_node_from_graph, populate_network_graph, ChannelPolicy, CustomRecords, Interceptor,
+    SimGraph, SimulatedChannel,
 };
 use simln_lib::{
     cln, cln::ClnNode, eclair, eclair::EclairNode, lnd, lnd::LndNode, serializers,
@@ -229,6 +229,7 @@ pub async fn create_simulation_with_network(
     sim_params: &SimParams,
     tasks: TaskTracker,
     interceptors: Vec<Arc<dyn Interceptor>>,
+    custom_records: CustomRecords,
 ) -> Result<(Simulation<SimulationClock>, Vec<ActivityDefinition>), anyhow::Error> {
     let cfg: SimulationCfg = SimulationCfg::try_from(cli)?;
     let SimParams {
@@ -259,6 +260,7 @@ pub async fn create_simulation_with_network(
             channels.clone(),
             tasks.clone(),
             interceptors,
+            custom_records,
             (shutdown_trigger.clone(), shutdown_listener.clone()),
         )
         .map_err(|e| SimulationError::SimulatedNetworkError(format!("{:?}", e)))?,

simln-lib/src/sim_node.rs

@@ -824,7 +824,7 @@ async fn handle_intercepted_htlc(
     shutdown_listener: Listener,
 ) -> Result<Result<CustomRecords, ForwardingError>, CriticalError> {
     if interceptors.is_empty() {
-        return Ok(Ok(HashMap::new()));
+        return Ok(Ok(request.incoming_custom_records));
     }
 
     let mut attached_custom_records: CustomRecords = HashMap::new();
@@ -925,6 +925,9 @@ pub struct SimGraph {
     /// trigger a shutdown signal to other interceptors.
     interceptors: Vec<Arc<dyn Interceptor>>,
 
+    /// Custom records that will be added to the first outgoing HTLC in a payment.
+    default_custom_records: CustomRecords,
+
     /// Shutdown signal that can be used to trigger a shutdown if a critical error occurs. Listener
     /// can be used to listen for shutdown signals coming from upstream.
     shutdown_signal: (Trigger, Listener),
@@ -936,6 +939,7 @@ impl SimGraph {
         graph_channels: Vec<SimulatedChannel>,
         tasks: TaskTracker,
         interceptors: Vec<Arc<dyn Interceptor>>,
+        default_custom_records: CustomRecords,
         shutdown_signal: (Trigger, Listener),
     ) -> Result<Self, SimulationError> {
         let mut nodes: HashMap<PublicKey, Vec<u64>> = HashMap::new();
@@ -971,6 +975,7 @@ impl SimGraph {
             channels: Arc::new(Mutex::new(channels)),
             tasks,
             interceptors,
+            default_custom_records,
             shutdown_signal,
         })
     }
@@ -1091,15 +1096,16 @@ impl SimNetwork for SimGraph {
             },
         };
 
-        self.tasks.spawn(propagate_payment(
-            self.channels.clone(),
+        self.tasks.spawn(propagate_payment(PropagatePaymentRequest {
+            nodes: Arc::clone(&self.channels),
             source,
-            path.clone(),
+            route: path.clone(),
             payment_hash,
             sender,
-            self.interceptors.clone(),
-            self.shutdown_signal.clone(),
-        ));
+            interceptors: self.interceptors.clone(),
+            custom_records: self.default_custom_records.clone(),
+            shutdown_signal: self.shutdown_signal.clone(),
+        }));
     }
 
     /// lookup_node fetches a node's information and channel capacities.
@@ -1149,13 +1155,14 @@ async fn add_htlcs(
     route: Path,
     payment_hash: PaymentHash,
     interceptors: Vec<Arc<dyn Interceptor>>,
+    custom_records: CustomRecords,
     shutdown_listener: Listener,
 ) -> Result<Result<(), (Option<usize>, ForwardingError)>, CriticalError> {
     let mut outgoing_node = source;
     let mut outgoing_amount = route.fee_msat() + route.final_value_msat();
     let mut outgoing_cltv = route.hops.iter().map(|hop| hop.cltv_expiry_delta).sum();
 
-    let mut incoming_custom_records = HashMap::new();
+    let mut incoming_custom_records = custom_records;
 
     // Tracks the hop index that we need to remove htlcs from on payment completion (both success and failure).
     // Given a payment from A to C, over the route A -- B -- C, this index has the following meanings:
@@ -1237,7 +1244,7 @@ async fn add_htlcs(
             forwarding_node: hop.pubkey,
             payment_hash,
             incoming_htlc: incoming_htlc.clone(),
-            incoming_custom_records: incoming_custom_records.clone(),
+            incoming_custom_records,
             outgoing_channel_id: next_scid,
             incoming_amount_msat: outgoing_amount,
             outgoing_amount_msat: outgoing_amount - hop.fee_msat,
@@ -1338,43 +1345,47 @@ async fn remove_htlcs(
     Ok(())
 }
 
-/// Finds a payment path from the source to destination nodes provided, and propagates the appropriate htlcs through
-/// the simulated network, notifying the sender channel provided of the payment outcome. If a critical error occurs,
-/// ie a breakdown of our state machine, it will still notify the payment outcome and will use the shutdown trigger
-/// to signal that we should exit.
-async fn propagate_payment(
+struct PropagatePaymentRequest {
     nodes: Arc<Mutex<HashMap<ShortChannelID, SimulatedChannel>>>,
     source: PublicKey,
     route: Path,
     payment_hash: PaymentHash,
     sender: Sender<Result<PaymentResult, LightningError>>,
     interceptors: Vec<Arc<dyn Interceptor>>,
+    custom_records: CustomRecords,
     shutdown_signal: (Trigger, Listener),
-) {
+}
+
+/// Finds a payment path from the source to destination nodes provided, and propagates the appropriate htlcs through
+/// the simulated network, notifying the sender channel provided of the payment outcome. If a critical error occurs,
+/// ie a breakdown of our state machine, it will still notify the payment outcome and will use the shutdown trigger
+/// to signal that we should exit.
+async fn propagate_payment(request: PropagatePaymentRequest) {
     let notify_result = match add_htlcs(
-        nodes.clone(),
-        source,
-        route.clone(),
-        payment_hash,
-        interceptors.clone(),
-        shutdown_signal.1,
+        request.nodes.clone(),
+        request.source,
+        request.route.clone(),
+        request.payment_hash,
+        request.interceptors.clone(),
+        request.custom_records,
+        request.shutdown_signal.1,
     )
     .await
     {
         Ok(Ok(_)) => {
             // If we successfully added the htlc, go ahead and remove all the htlcs in the route with successful resolution.
             if let Err(e) = remove_htlcs(
-                nodes,
-                route.hops.len() - 1,
-                source,
-                route,
-                payment_hash,
+                request.nodes,
+                request.route.hops.len() - 1,
+                request.source,
+                request.route,
+                request.payment_hash,
                 true,
-                interceptors,
+                request.interceptors,
             )
             .await
             {
-                shutdown_signal.0.trigger();
+                request.shutdown_signal.0.trigger();
                 log::error!("Could not remove htlcs from channel: {e}.");
             }
             PaymentResult {
@@ -1387,35 +1398,35 @@ async fn propagate_payment(
             // state. It's possible that we failed with the very first add, and then we don't need to clean anything up.
             if let Some(resolution_idx) = fail_idx {
                 if remove_htlcs(
-                    nodes,
+                    request.nodes,
                     resolution_idx,
-                    source,
-                    route,
-                    payment_hash,
+                    request.source,
+                    request.route,
+                    request.payment_hash,
                     false,
-                    interceptors,
+                    request.interceptors,
                 )
                 .await
                 .is_err()
                 {
-                    shutdown_signal.0.trigger();
+                    request.shutdown_signal.0.trigger();
                 }
             }
 
             log::debug!(
                 "Forwarding failure for simulated payment {}: {fwd_err}",
-                hex::encode(payment_hash.0)
+                hex::encode(request.payment_hash.0)
             );
             PaymentResult {
                 htlc_count: 0,
                 payment_outcome: PaymentOutcome::Unknown,
             }
         },
         Err(critical_err) => {
-            shutdown_signal.0.trigger();
+            request.shutdown_signal.0.trigger();
             log::debug!(
                 "Critical error in simulated payment {}: {critical_err}",
-                hex::encode(payment_hash.0)
+                hex::encode(request.payment_hash.0)
             );
             PaymentResult {
                 htlc_count: 0,
@@ -1424,7 +1435,7 @@ async fn propagate_payment(
         },
     };
 
-    if let Err(e) = sender.send(Ok(notify_result)) {
+    if let Err(e) = request.sender.send(Ok(notify_result)) {
         log::error!("Could not notify payment result: {:?}.", e);
     }
 }
@@ -1962,7 +1973,11 @@ mod tests {
         /// Alice (100) --- (0) Bob (100) --- (0) Carol (100) --- (0) Dave
         ///
         /// The nodes pubkeys in this chain of channels are provided in-order for easy access.
-        async fn new(capacity: u64, interceptors: Vec<Arc<dyn Interceptor>>) -> Self {
+        async fn new(
+            capacity: u64,
+            interceptors: Vec<Arc<dyn Interceptor>>,
+            custom_records: CustomRecords,
+        ) -> Self {
             let shutdown_signal = triggered::trigger();
             let channels = create_simulated_channels(3, capacity);
             let routing_graph = Arc::new(
@@ -1989,6 +2004,7 @@ mod tests {
                     channels.clone(),
                     TaskTracker::new(),
                     interceptors,
+                    custom_records,
                     shutdown_signal,
                 )
                 .expect("could not create test graph"),
@@ -2069,7 +2085,8 @@ mod tests {
     #[tokio::test]
     async fn test_successful_dispatch() {
         let chan_capacity = 500_000_000;
-        let mut test_kit = DispatchPaymentTestKit::new(chan_capacity, vec![]).await;
+        let mut test_kit =
+            DispatchPaymentTestKit::new(chan_capacity, vec![], CustomRecords::default()).await;
 
         // Send a payment that should succeed from Alice -> Dave.
         let mut amt = 20_000;
@@ -2134,7 +2151,8 @@ mod tests {
     #[tokio::test]
     async fn test_successful_multi_hop() {
         let chan_capacity = 500_000_000;
-        let mut test_kit = DispatchPaymentTestKit::new(chan_capacity, vec![]).await;
+        let mut test_kit =
+            DispatchPaymentTestKit::new(chan_capacity, vec![], CustomRecords::default()).await;
 
         // Send a payment that should succeed from Alice -> Dave.
         let amt = 20_000;
@@ -2164,7 +2182,8 @@ mod tests {
     #[tokio::test]
     async fn test_single_hop_payments() {
         let chan_capacity = 500_000_000;
-        let mut test_kit = DispatchPaymentTestKit::new(chan_capacity, vec![]).await;
+        let mut test_kit =
+            DispatchPaymentTestKit::new(chan_capacity, vec![], CustomRecords::default()).await;
 
         // Send a single hop payment from Alice -> Bob, it will succeed because Alice has all the liquidity.
         let amt = 150_000;
@@ -2196,7 +2215,8 @@ mod tests {
     #[tokio::test]
     async fn test_multi_hop_faiulre() {
         let chan_capacity = 500_000_000;
-        let mut test_kit = DispatchPaymentTestKit::new(chan_capacity, vec![]).await;
+        let mut test_kit =
+            DispatchPaymentTestKit::new(chan_capacity, vec![], CustomRecords::default()).await;
 
         // Drain liquidity between Bob and Carol to force failures on Bob's outgoing linke.
         test_kit
@@ -2242,7 +2262,7 @@ mod tests {
                 channel_id: ShortChannelID(0),
                 index: 0,
             },
-            incoming_custom_records: CustomRecords::new(),
+            incoming_custom_records: CustomRecords::default(),
             outgoing_channel_id: None,
             incoming_amount_msat: 0,
             outgoing_amount_msat: 0,
@@ -2404,7 +2424,8 @@ mod tests {
             .returning(|_| Ok(()));
 
         let mock_1 = Arc::new(mock_interceptor_1);
-        let mut test_kit = DispatchPaymentTestKit::new(500_000_000, vec![mock_1]).await;
+        let mut test_kit =
+            DispatchPaymentTestKit::new(500_000_000, vec![mock_1], CustomRecords::default()).await;
         let (_, result) = test_kit
             .send_test_payment(test_kit.nodes[0], test_kit.nodes[3], 150_000_000)
             .await;
@@ -2425,4 +2446,40 @@ mod tests {
         test_kit.graph.tasks.close();
         test_kit.graph.tasks.wait().await;
     }
+
+    /// Tests custom records set for interceptors in multi-hop payment.
+    #[tokio::test]
+    async fn test_custom_records() {
+        let custom_records = HashMap::from([(1000, vec![1])]);
+
+        let mut mock_interceptor_1 = MockTestInterceptor::new();
+        let custom_records_clone = custom_records.clone();
+        mock_interceptor_1
+            .expect_intercept_htlc()
+            .withf(move |req: &InterceptRequest| {
+                // Check custom records passed to interceptor are the default ones set.
+                req.incoming_custom_records == custom_records_clone
+            })
+            .returning(|_| Ok(Ok(CustomRecords::default()))); // Set empty records for 2nd hop.
+        mock_interceptor_1
+            .expect_notify_resolution()
+            .returning(|_| Ok(()))
+            .times(2);
+
+        mock_interceptor_1
+            .expect_intercept_htlc()
+            .withf(move |req: &InterceptRequest| {
+                // On this 2nd hop, the custom records should be empty.
+                req.incoming_custom_records == CustomRecords::default()
+            })
+            .returning(|_| Ok(Ok(CustomRecords::default())));
+
+        let chan_capacity = 500_000_000;
+        let mock_1: Arc<dyn Interceptor> = Arc::new(mock_interceptor_1);
+        let mut test_kit =
+            DispatchPaymentTestKit::new(chan_capacity, vec![mock_1], custom_records).await;
+        let _ = test_kit
+            .send_test_payment(test_kit.nodes[0], test_kit.nodes[2], 150_000)
+            .await;
+    }
 }