From 36306a0554ab3e778d83837238545c8727107fa8 Mon Sep 17 00:00:00 2001
From: yukang <moorekang@gmail.com>
Date: Wed, 23 Oct 2024 15:05:46 +0800
Subject: [PATCH] add allow_self_payment for send_payment

---
 src/fiber/graph.rs       |  45 +++++++---
 src/fiber/network.rs     |  26 +++---
 src/fiber/tests/graph.rs | 178 ++++++++++++++++++++++++++++++++++++---
 src/rpc/README.md        |   1 +
 src/rpc/channel.rs       |   4 +
 5 files changed, 216 insertions(+), 38 deletions(-)

diff --git a/src/fiber/graph.rs b/src/fiber/graph.rs
index dd51233c..16f9f441 100644
--- a/src/fiber/graph.rs
+++ b/src/fiber/graph.rs
@@ -510,15 +510,16 @@ where
     /// Returns a list of `PaymentHopData` for all nodes in the route, including the origin and the target node.
     pub fn build_route(
         &self,
-        payment_request: SendPaymentData,
+        payment_data: &SendPaymentData,
     ) -> Result<Vec<PaymentHopData>, GraphError> {
+        let payment_data = payment_data.clone();
         let source = self.get_source_pubkey();
-        let target = payment_request.target_pubkey;
-        let amount = payment_request.amount;
-        let preimage = payment_request.preimage;
-        let payment_hash = payment_request.payment_hash;
-        let udt_type_script = payment_request.udt_type_script;
-        let invoice = payment_request
+        let target = payment_data.target_pubkey;
+        let amount = payment_data.amount;
+        let preimage = payment_data.preimage;
+        let payment_hash = payment_data.payment_hash;
+        let udt_type_script = payment_data.udt_type_script;
+        let invoice = payment_data
             .invoice
             .map(|x| x.parse::<CkbInvoice>().unwrap());
         let hash_algorithm = invoice
@@ -531,12 +532,20 @@ where
             source, target, amount, payment_hash
         );
 
+        let allow_self_payment = payment_data.allow_self_payment;
+        if source == target && !allow_self_payment {
+            return Err(GraphError::PathFind(
+                "source and target are the same and allow_self_payment is not enable".to_string(),
+            ));
+        }
+
         let route = self.find_route(
             source,
             target,
             amount,
-            payment_request.max_fee_amount,
+            payment_data.max_fee_amount,
             udt_type_script,
+            allow_self_payment,
         )?;
         assert!(!route.is_empty());
 
@@ -609,6 +618,7 @@ where
         amount: u128,
         max_fee_amount: Option<u128>,
         udt_type_script: Option<Script>,
+        allow_self: bool,
     ) -> Result<Vec<PathEdge>, GraphError> {
         let started_time = std::time::Instant::now();
         let nodes_len = self.nodes.len();
@@ -624,11 +634,12 @@ where
             ));
         }
 
-        if source == target {
+        if source == target && !allow_self {
             return Err(GraphError::PathFind(
                 "source and target are the same".to_string(),
             ));
         }
+
         let Some(source_node) = self.nodes.get(&source) else {
             return Err(GraphError::PathFind(format!(
                 "source node not found: {:?}",
@@ -641,6 +652,7 @@ where
                 &target
             )));
         };
+
         // initialize the target node
         nodes_heap.push(NodeHeapElement {
             node_id: target,
@@ -652,19 +664,21 @@ where
             next_hop: None,
             incoming_cltv_height: 0,
         });
+        let route_to_self = source == target;
         while let Some(cur_hop) = nodes_heap.pop() {
-            if cur_hop.node_id == source {
-                break;
-            }
             nodes_visited += 1;
 
             for (from, channel_info, channel_update) in self.get_node_inbounds(cur_hop.node_id) {
-                edges_expanded += 1;
+                if from == target && !route_to_self {
+                    continue;
+                }
                 // if charge inbound fees for exit hop
                 if udt_type_script != channel_info.announcement_msg.udt_type_script {
                     continue;
                 }
 
+                edges_expanded += 1;
+
                 let fee_rate = channel_update.fee_rate;
                 let next_hop_received_amount = cur_hop.amount_received;
                 let fee = calculate_tlc_forward_fee(next_hop_received_amount, fee_rate as u128);
@@ -748,7 +762,7 @@ where
         }
 
         let mut current = source_node.node_id;
-        while current != target {
+        loop {
             if let Some(elem) = distances.get(&current) {
                 let next_hop = elem.next_hop.as_ref().expect("next_hop is none");
                 result.push(PathEdge {
@@ -759,6 +773,9 @@ where
             } else {
                 break;
             }
+            if current == target {
+                break;
+            }
         }
 
         info!(
diff --git a/src/fiber/network.rs b/src/fiber/network.rs
index 9a186ccc..6cb2055a 100644
--- a/src/fiber/network.rs
+++ b/src/fiber/network.rs
@@ -276,6 +276,8 @@ pub struct SendPaymentCommand {
     // udt type script
     #[serde_as(as = "Option<EntityHex>")]
     pub udt_type_script: Option<Script>,
+    // allow self payment, default is false
+    pub allow_self_payment: bool,
 }
 
 #[serde_as]
@@ -293,10 +295,11 @@ pub struct SendPaymentData {
     #[serde_as(as = "Option<EntityHex>")]
     pub udt_type_script: Option<Script>,
     pub preimage: Option<Hash256>,
+    pub allow_self_payment: bool,
 }
 
 impl SendPaymentData {
-    pub fn new(command: SendPaymentCommand) -> Result<SendPaymentData, String> {
+    pub fn new(command: SendPaymentCommand, source: Pubkey) -> Result<SendPaymentData, String> {
         let invoice = command
             .invoice
             .as_ref()
@@ -336,6 +339,10 @@ impl SendPaymentData {
             "target_pubkey",
         )?;
 
+        if !command.allow_self_payment && target == source {
+            return Err("allow_self_payment is not enable, can not pay self".to_string());
+        }
+
         let amount = validate_field(
             command.amount,
             invoice.as_ref().and_then(|i| i.amount()),
@@ -388,6 +395,7 @@ impl SendPaymentData {
             keysend,
             udt_type_script,
             preimage,
+            allow_self_payment: command.allow_self_payment,
         })
     }
 }
@@ -2211,12 +2219,7 @@ where
         let mut error = None;
         while payment_session.can_retry() {
             payment_session.retried_times += 1;
-            let hops_infos = match self
-                .network_graph
-                .read()
-                .await
-                .build_route(payment_data.clone())
-            {
+            let hops_infos = match self.network_graph.read().await.build_route(&payment_data) {
                 Err(e) => {
                     error!("Failed to build route: {:?}", e);
                     error = Some(format!("Failed to build route: {:?}", payment_hash));
@@ -2280,10 +2283,11 @@ where
         state: &mut NetworkActorState<S>,
         payment_request: SendPaymentCommand,
     ) -> Result<SendPaymentResponse, Error> {
-        let payment_data = SendPaymentData::new(payment_request.clone()).map_err(|e| {
-            error!("Failed to validate payment request: {:?}", e);
-            Error::InvalidParameter(format!("Failed to validate payment request: {:?}", e))
-        })?;
+        let payment_data = SendPaymentData::new(payment_request.clone(), state.get_public_key())
+            .map_err(|e| {
+                error!("Failed to validate payment request: {:?}", e);
+                Error::InvalidParameter(format!("Failed to validate payment request: {:?}", e))
+            })?;
 
         // initialize the payment session in db and begin the payment process lifecycle
         if let Some(payment_session) = self.store.get_payment_session(payment_data.payment_hash) {
diff --git a/src/fiber/tests/graph.rs b/src/fiber/tests/graph.rs
index c2457e39..e4adf007 100644
--- a/src/fiber/tests/graph.rs
+++ b/src/fiber/tests/graph.rs
@@ -1,7 +1,8 @@
+use crate::fiber::types::Pubkey;
 use crate::{
     fiber::{
         graph::{ChannelInfo, GraphError, NetworkGraph, NodeInfo, PathEdge},
-        network::{get_chain_hash, SendPaymentData},
+        network::{get_chain_hash, SendPaymentCommand, SendPaymentData},
         types::{ChannelAnnouncement, ChannelUpdate, Hash256, NodeAnnouncement},
     },
     store::Store,
@@ -189,7 +190,7 @@ impl MockNetworkGraph {
         let source = self.keys[source].into();
         let target = self.keys[target].into();
         self.graph
-            .find_route(source, target, amount, Some(max_fee), None)
+            .find_route(source, target, amount, Some(max_fee), None, false)
     }
 
     pub fn find_route_udt(
@@ -202,8 +203,26 @@ impl MockNetworkGraph {
     ) -> Result<Vec<PathEdge>, GraphError> {
         let source = self.keys[source].into();
         let target = self.keys[target].into();
-        self.graph
-            .find_route(source, target, amount, Some(max_fee), Some(udt_type_script))
+        self.graph.find_route(
+            source,
+            target,
+            amount,
+            Some(max_fee),
+            Some(udt_type_script),
+            false,
+        )
+    }
+
+    pub fn build_route_with_expect(&self, payment_data: &SendPaymentData, expect: Vec<usize>) {
+        let route = self.graph.build_route(payment_data);
+        assert!(route.is_ok());
+        let route = route.unwrap();
+        let nodes = route.iter().filter_map(|x| x.next_hop).collect::<Vec<_>>();
+        let expecptected_nodes: Vec<Pubkey> = expect
+            .iter()
+            .map(|x| self.keys[*x].into())
+            .collect::<Vec<_>>();
+        assert_eq!(nodes, expecptected_nodes);
     }
 }
 
@@ -509,6 +528,7 @@ fn test_graph_find_path_err() {
         100,
         Some(1000),
         None,
+        false,
     );
     assert!(route.is_err());
 
@@ -518,6 +538,7 @@ fn test_graph_find_path_err() {
         100,
         Some(1000),
         None,
+        false,
     );
     assert!(route.is_err());
 }
@@ -531,7 +552,7 @@ fn test_graph_build_route_three_nodes() {
     let node2 = network.keys[2];
     let node3 = network.keys[3];
     // Test build route from node1 to node3
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node3.into(),
         amount: 100,
         payment_hash: Hash256::default(),
@@ -543,6 +564,7 @@ fn test_graph_build_route_three_nodes() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     eprintln!("return {:?}", route);
     assert!(route.is_ok());
@@ -569,7 +591,7 @@ fn test_graph_build_route_exceed_max_htlc_value() {
     let node3 = network.keys[3];
 
     // Test build route from node1 to node3 with amount exceeding max_htlc_value
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node3.into(),
         amount: 100, // Exceeds max_htlc_value of 50
         payment_hash: Hash256::default(),
@@ -581,6 +603,7 @@ fn test_graph_build_route_exceed_max_htlc_value() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     assert!(route.is_err());
 }
@@ -594,7 +617,7 @@ fn test_graph_build_route_below_min_htlc_value() {
     let node3 = network.keys[3];
 
     // Test build route from node1 to node3 with amount below min_htlc_value
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node3.into(),
         amount: 10, // Below min_htlc_value of 50
         payment_hash: Hash256::default(),
@@ -606,6 +629,7 @@ fn test_graph_build_route_below_min_htlc_value() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     assert!(route.is_err());
 }
@@ -638,7 +662,7 @@ fn test_graph_mark_failed_channel() {
 
     network.mark_channel_failed(2, 3);
     // Test build route from node1 to node3
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node3.into(),
         amount: 100,
         payment_hash: Hash256::default(),
@@ -650,6 +674,7 @@ fn test_graph_mark_failed_channel() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     eprintln!("return {:?}", route);
     assert!(route.is_err());
@@ -658,7 +683,7 @@ fn test_graph_mark_failed_channel() {
     network.add_edge(5, 3, Some(500), Some(2));
 
     // Test build route from node1 to node3
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node3.into(),
         amount: 100,
         payment_hash: Hash256::default(),
@@ -670,6 +695,7 @@ fn test_graph_mark_failed_channel() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     eprintln!("return {:?}", route);
     assert!(route.is_ok());
@@ -686,7 +712,7 @@ fn test_graph_mark_failed_node() {
     let node4 = network.keys[4];
 
     // Test build route from node1 to node3
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node3.into(),
         amount: 100,
         payment_hash: Hash256::default(),
@@ -698,12 +724,13 @@ fn test_graph_mark_failed_node() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     eprintln!("return {:?}", route);
     assert!(route.is_ok());
 
     // Test build route from node1 to node4 should be Ok
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node4.into(),
         amount: 100,
         payment_hash: Hash256::default(),
@@ -715,6 +742,7 @@ fn test_graph_mark_failed_node() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     eprintln!("return {:?}", route);
     assert!(route.is_ok());
@@ -722,7 +750,7 @@ fn test_graph_mark_failed_node() {
     network.mark_node_failed(2);
 
     // Test build route from node1 to node3
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node3.into(),
         amount: 100,
         payment_hash: Hash256::default(),
@@ -734,12 +762,13 @@ fn test_graph_mark_failed_node() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     eprintln!("return {:?}", route);
     assert!(route.is_err());
 
     // Test build route from node1 to node4
-    let route = network.graph.build_route(SendPaymentData {
+    let route = network.graph.build_route(&SendPaymentData {
         target_pubkey: node4.into(),
         amount: 100,
         payment_hash: Hash256::default(),
@@ -751,7 +780,130 @@ fn test_graph_mark_failed_node() {
         keysend: false,
         udt_type_script: None,
         preimage: None,
+        allow_self_payment: false,
     });
     eprintln!("return {:?}", route);
     assert!(route.is_err());
 }
+
+#[test]
+fn test_graph_payment_self_default_is_false() {
+    let mut network = MockNetworkGraph::new(5);
+    network.add_edge(0, 2, Some(500), Some(2));
+    network.add_edge(2, 3, Some(500), Some(2));
+    network.add_edge(2, 4, Some(500), Some(2));
+
+    let node0 = network.keys[0];
+
+    // node0 is the source node
+    let command = SendPaymentCommand {
+        target_pubkey: Some(node0.into()),
+        amount: Some(100),
+        payment_hash: Some(Hash256::default()),
+        final_cltv_delta: Some(100),
+        invoice: None,
+        timeout: Some(10),
+        max_fee_amount: Some(1000),
+        max_parts: None,
+        keysend: Some(false),
+        udt_type_script: None,
+        allow_self_payment: false,
+    };
+    let payment_data = SendPaymentData::new(command, node0.into());
+    let error = payment_data.unwrap_err().to_string();
+    assert!(error.contains("can not pay self"));
+
+    let route = network.graph.build_route(&SendPaymentData {
+        target_pubkey: node0.into(),
+        amount: 100,
+        payment_hash: Hash256::default(),
+        invoice: None,
+        final_cltv_delta: Some(100),
+        timeout: Some(10),
+        max_fee_amount: Some(1000),
+        max_parts: None,
+        keysend: false,
+        udt_type_script: None,
+        preimage: None,
+        allow_self_payment: false,
+    });
+    assert!(route.is_err());
+}
+
+#[test]
+fn test_graph_payment_pay_single_path() {
+    let mut network = MockNetworkGraph::new(9);
+    network.add_edge(0, 2, Some(500), Some(2));
+    network.add_edge(2, 3, Some(500), Some(2));
+    network.add_edge(2, 4, Some(500), Some(2));
+    network.add_edge(4, 5, Some(500), Some(2));
+    network.add_edge(5, 6, Some(500), Some(2));
+
+    let node0 = network.keys[0];
+    // node0 is the source node
+    let command = SendPaymentCommand {
+        target_pubkey: Some(network.keys[6].into()),
+        amount: Some(100),
+        payment_hash: Some(Hash256::default()),
+        final_cltv_delta: Some(100),
+        invoice: None,
+        timeout: Some(10),
+        max_fee_amount: Some(1000),
+        max_parts: None,
+        keysend: Some(false),
+        udt_type_script: None,
+        allow_self_payment: true,
+    };
+    let payment_data = SendPaymentData::new(command, node0.into());
+    assert!(payment_data.is_ok());
+    let payment_data = payment_data.unwrap();
+
+    network.add_edge(4, 0, Some(1000), Some(2));
+    network.build_route_with_expect(&payment_data, vec![2, 4, 5, 6]);
+}
+
+#[test]
+fn test_graph_payment_pay_self_will_ok() {
+    let mut network = MockNetworkGraph::new(9);
+    network.add_edge(0, 2, Some(500), Some(2));
+    network.add_edge(2, 3, Some(500), Some(2));
+    network.add_edge(2, 4, Some(500), Some(2));
+    network.add_edge(4, 5, Some(500), Some(2));
+    network.add_edge(5, 6, Some(500), Some(2));
+
+    let node0 = network.keys[0];
+
+    // node0 is the source node
+    let command = SendPaymentCommand {
+        target_pubkey: Some(network.keys[0].into()),
+        amount: Some(100),
+        payment_hash: Some(Hash256::default()),
+        final_cltv_delta: Some(100),
+        invoice: None,
+        timeout: Some(10),
+        max_fee_amount: Some(1000),
+        max_parts: None,
+        keysend: Some(false),
+        udt_type_script: None,
+        allow_self_payment: true,
+    };
+    let payment_data = SendPaymentData::new(command, node0.into());
+    assert!(payment_data.is_ok());
+    let payment_data = payment_data.unwrap();
+
+    let route = network.graph.build_route(&payment_data);
+    // since we don't have a route to node0, it will fail
+    assert!(route.is_err());
+
+    // add a long path
+    network.add_edge(6, 0, Some(500), Some(2));
+    network.build_route_with_expect(&payment_data, vec![2, 4, 5, 6, 0]);
+
+    // now add another shorter path
+    network.add_edge(4, 0, Some(1000), Some(2));
+    network.build_route_with_expect(&payment_data, vec![2, 4, 0]);
+
+    // now add another shorter path
+    network.add_edge(2, 0, Some(1000), Some(2));
+    network.build_route_with_expect(&payment_data, vec![2, 0]);
+}
diff --git a/src/rpc/README.md b/src/rpc/README.md
index 197021f1..fb792761 100644
--- a/src/rpc/README.md
+++ b/src/rpc/README.md
@@ -182,6 +182,7 @@ Sends a payment to a peer.
 - `timeout` (type: `Option<u64>`): The payment timeout in seconds. If the payment is not completed within this time, it will be cancelled.
 - `max_fee_amount` (type: `Option<u128>`): The maximum fee amounts in shannons that the sender is willing to pay.
 - `max_parts` (type: `Option<u64>`): Max parts for the payment, only used for multi-part payments.
+- `allow_self_payment` (type: `Option<bool>`): Allow self payment, if it's true path finding may construct a payment router that target to the same node, default is false.
 
 Note `target_pubkey`, `amount`, `payment_hash` should be consistent with the invoice. If `invoice` is provided, the `target_pubkey`, `amount`, `payment_hash` can be omitted.
 
diff --git a/src/rpc/channel.rs b/src/rpc/channel.rs
index 300b06e8..73bdc301 100644
--- a/src/rpc/channel.rs
+++ b/src/rpc/channel.rs
@@ -226,6 +226,9 @@ pub(crate) struct SendPaymentCommandParams {
 
     // udt type script for the payment
     udt_type_script: Option<Script>,
+
+    // allow self payment, default is false
+    allow_self_payment: Option<bool>,
 }
 
 #[rpc(server)]
@@ -522,6 +525,7 @@ where
                     max_parts: params.max_parts,
                     keysend: params.keysend,
                     udt_type_script: params.udt_type_script.clone().map(|s| s.into()),
+                    allow_self_payment: params.allow_self_payment.unwrap_or(false),
                 },
                 rpc_reply,
             ))