code refactor for trampoline routing

Author: chenyukang

crates/fiber-lib/src/fiber/network.rs

@@ -92,7 +92,7 @@ use crate::fiber::gossip::{GossipConfig, GossipService, SubscribableGossipMessag
 use crate::fiber::payment::SessionRoute;
 use crate::fiber::payment::{
     AttemptStatus, PaymentActor, PaymentActorArguments, PaymentActorMessage, PaymentCustomRecords,
-    PaymentStatus, SendPaymentCommand, SendPaymentWithRouterCommand,
+    PaymentStatus, SendPaymentCommand, SendPaymentWithRouterCommand, TrampolineContext,
 };
 use crate::fiber::serde_utils::EntityHex;
 use crate::fiber::types::{
@@ -2474,19 +2474,21 @@ where
                     ));
                 }
 
-                let remaining_trampoline_onion = peeled_trampoline.next.map(|p| p.into_bytes());
+                let (Some(remaining_trampoline_onion), Some(mut prev_tlc)) =
+                    (peeled_trampoline.next.map(|p| p.into_bytes()), previous_tlc)
+                else {
+                    return Err(TlcErr::new_node_fail(
+                        TlcErrorCode::InvalidOnionPayload,
+                        state.get_public_key(),
+                    ));
+                };
 
-                if let Some(mut prev_tlc) = previous_tlc {
-                    if let Some(shared_secret) = trampoline_outer_shared_secret {
-                        prev_tlc.shared_secret = Some(shared_secret);
-                    }
-                    state
-                        .trampoline_forwarding_tlcs
-                        .entry(payment_hash)
-                        .or_default()
-                        .push(prev_tlc);
+                if let Some(shared_secret) = trampoline_outer_shared_secret {
+                    prev_tlc.shared_secret = Some(shared_secret);
                 }
-
+                // currently we only support single previous tlc in trampoline forwarding,
+                // maybe we need to support multiple previous tlcs in the future
+                let previous_tlcs = vec![prev_tlc];
                 let command = SendPaymentCommand {
                     target_pubkey: Some(next_node_id),
                     amount: Some(amount_to_forward),
@@ -2497,7 +2499,10 @@ where
                     max_parts: Some(1),
                     udt_type_script,
                     allow_self_payment: true,
-                    final_trampoline_onion: remaining_trampoline_onion,
+                    trampoline_context: Some(TrampolineContext {
+                        remaining_trampoline_onion,
+                        previous_tlcs,
+                    }),
                     ..Default::default()
                 };
 
@@ -2811,9 +2816,6 @@ pub struct NetworkActorState<S, C> {
 
     // Inflight payment actors
     inflight_payments: HashMap<Hash256, ActorRef<PaymentActorMessage>>,
-
-    // Trampoline forwarding map: payment_hash -> Vec<PrevTlcInfo>
-    trampoline_forwarding_tlcs: HashMap<Hash256, Vec<PrevTlcInfo>>,
 }
 
 #[derive(Debug, Clone)]
@@ -3968,19 +3970,25 @@ where
             error!("Can't find inflight payment actor");
         }
 
-        if let Some(prev_tlcs) = self.trampoline_forwarding_tlcs.remove(&payment_hash) {
-            let session = self.store.get_payment_session(payment_hash);
-            match session {
-                Some(session) if session.status == PaymentStatus::Success => {
+        // If this payment has associated previous TLCs,
+        // meaning it's a trampoline forwarding payment,
+        // we need to resolve those upstream TLCs based on the payment outcome.
+        let Some(session) = self.store.get_payment_session(payment_hash) else {
+            return;
+        };
+        let trampoline_context = session.request.trampoline_context.as_ref();
+
+        if let Some(context) = trampoline_context {
+            match session.status {
+                PaymentStatus::Success => {
                     let preimage = session
                         .attempts()
                         .find(|a| a.is_success())
                         .and_then(|a| a.preimage);
 
                     if let Some(preimage) = preimage {
-                        // Store preimage globally
                         self.store.insert_preimage(payment_hash, preimage);
-                        for prev_tlc in prev_tlcs {
+                        for prev_tlc in &context.previous_tlcs {
                             let (send, _recv) = oneshot::channel();
                             let rpc_reply = RpcReplyPort::from(send);
                             let command = ChannelCommand::RemoveTlc(
@@ -4003,8 +4011,8 @@ where
                         error!("Payment success but no preimage found for {payment_hash}");
                     }
                 }
-                Some(session) if session.status == PaymentStatus::Failed => {
-                    for prev_tlc in prev_tlcs {
+                PaymentStatus::Failed => {
+                    for prev_tlc in &context.previous_tlcs {
                         let (send, _recv) = oneshot::channel();
                         let rpc_reply = RpcReplyPort::from(send);
                         let shared_secret = prev_tlc.shared_secret.unwrap_or([0u8; 32]);
@@ -4346,7 +4354,6 @@ where
                 funding_timeout_seconds: config.funding_timeout_seconds,
             },
             inflight_payments: Default::default(),
-            trampoline_forwarding_tlcs: Default::default(),
         };
 
         let node_announcement = state.get_or_create_new_node_announcement_message();

crates/fiber-lib/src/fiber/payment.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 
 use super::network::{SendOnionPacketCommand, SendPaymentResponse, ASSUME_NETWORK_MYSELF_ALIVE};
 use super::types::{Hash256, Privkey, Pubkey, TlcErrData};
-use crate::fiber::channel::{ChannelActorStateStore, ProcessingChannelError};
+use crate::fiber::channel::{ChannelActorStateStore, PrevTlcInfo, ProcessingChannelError};
 use crate::fiber::config::{
     DEFAULT_FINAL_TLC_EXPIRY_DELTA, DEFAULT_MAX_PARTS, MAX_PAYMENT_TLC_EXPIRY_LIMIT,
     MIN_TLC_EXPIRY_DELTA, PAYMENT_MAX_PARTS_LIMIT,
@@ -170,7 +170,8 @@ pub struct SendPaymentData {
     /// When set to a non-empty list `[t1, t2, ...]`, routing will only find a path from the
     /// payer to `t1`, and the inner trampoline onion will encode `t1 -> t2 -> ... -> final`.
     pub trampoline_hops: Option<Vec<TrampolineHop>>,
-    pub final_trampoline_onion: Option<Vec<u8>>,
+    #[serde(default)]
+    pub trampoline_context: Option<TrampolineContext>,
     #[serde(skip)]
     pub channel_stats: GraphChannelStat,
 }
@@ -196,7 +197,7 @@ pub struct SendPaymentDataBuilder {
     allow_mpp: bool,
     dry_run: bool,
     trampoline_hops: Option<Vec<TrampolineHop>>,
-    final_trampoline_onion: Option<Vec<u8>>,
+    trampoline_context: Option<TrampolineContext>,
     channel_stats: GraphChannelStat,
 }
 
@@ -308,8 +309,8 @@ impl SendPaymentDataBuilder {
         self
     }
 
-    pub fn final_trampoline_onion(mut self, final_trampoline_onion: Option<Vec<u8>>) -> Self {
-        self.final_trampoline_onion = final_trampoline_onion;
+    pub fn trampoline_context(mut self, trampoline_context: Option<TrampolineContext>) -> Self {
+        self.trampoline_context = trampoline_context;
         self
     }
 
@@ -453,8 +454,8 @@ impl SendPaymentDataBuilder {
             allow_mpp: self.allow_mpp,
             dry_run: self.dry_run,
             trampoline_hops: self.trampoline_hops,
-            final_trampoline_onion: self.final_trampoline_onion,
             channel_stats: self.channel_stats,
+            trampoline_context: self.trampoline_context,
         })
     }
 }
@@ -609,7 +610,7 @@ impl SendPaymentData {
             .allow_mpp(allow_mpp)
             .dry_run(command.dry_run)
             .trampoline_hops(command.trampoline_hops)
-            .final_trampoline_onion(command.final_trampoline_onion)
+            .trampoline_context(command.trampoline_context)
             .channel_stats(Default::default())
             .build()
     }
@@ -1129,12 +1130,21 @@ pub struct SendPaymentCommand {
     /// When set to a non-empty list `[t1, t2, ...]`, routing will only find a path from the
     /// payer to `t1`, and the inner trampoline onion will encode `t1 -> t2 -> ... -> final`.
     pub trampoline_hops: Option<Vec<TrampolineHop>>,
+
+    /// Context for trampoline routing.
+    pub trampoline_context: Option<TrampolineContext>,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, Default)]
+pub struct TrampolineContext {
     /// Optional final trampoline onion packet.
     ///
-    /// When provided, this onion packet will be attached to the payload of the final hop
+    /// When provided, this onion packet will be attached to the payload of the next hop
     /// (which in this context is the next trampoline node).
-    #[serde(skip)]
-    pub final_trampoline_onion: Option<Vec<u8>>,
+    pub remaining_trampoline_onion: Vec<u8>,
+    /// Previous TLCs information for the payment session.
+    /// This is used to associate the outgoing payment with the incoming payment.
+    pub previous_tlcs: Vec<PrevTlcInfo>,
 }
 
 impl SendPaymentCommand {
@@ -1713,9 +1723,10 @@ where
                 Ok(mut hops) => {
                     assert_ne!(hops[0].funding_tx_hash, Hash256::default());
 
-                    if let Some(trampoline_onion) = &session.request.final_trampoline_onion {
+                    if let Some(trampoline) = &session.request.trampoline_context {
                         if let Some(last_hop) = hops.last_mut() {
-                            last_hop.trampoline_onion = Some(trampoline_onion.clone());
+                            last_hop.trampoline_onion =
+                                Some(trampoline.remaining_trampoline_onion.clone());
                         }
                     }
 

crates/fiber-lib/src/fiber/tests/network.rs

@@ -131,6 +131,7 @@ fn test_send_payment_data_trampoline_hops_validation_errors() {
         target_pubkey: Some(target),
         allow_self_payment: false,
         payment_hash: Some(payment_hash),
+        trampoline_context: None,
         final_tlc_expiry_delta: None,
         tlc_expiry_limit: None,
         timeout: None,
@@ -142,7 +143,6 @@ fn test_send_payment_data_trampoline_hops_validation_errors() {
         dry_run: false,
         hop_hints: None,
         custom_records: None,
-        final_trampoline_onion: None,
     };
 
     // Provided but empty.

crates/fiber-lib/src/fiber/tests/payment.rs

@@ -5247,6 +5247,7 @@ async fn test_send_payment_no_preimage_invoice_will_make_payment_failed() {
                 target_pubkey: None,
                 allow_self_payment: true,
                 payment_hash: None,
+                trampoline_context: None,
                 final_tlc_expiry_delta: None,
                 tlc_expiry_limit: None,
                 timeout: None,
@@ -5258,7 +5259,6 @@ async fn test_send_payment_no_preimage_invoice_will_make_payment_failed() {
                 dry_run: false,
                 hop_hints: None,
                 custom_records: None,
-                final_trampoline_onion: None,
             })
             .await
             .unwrap();