sender will set custom records for amp

Author: chenyukang

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

@@ -252,6 +252,58 @@ mod tests {
         assert_eq!(children[1], expected_child2);
     }
 
+    #[test]
+    fn test_reconstruct_single_child() {
+        let root = Share::random();
+        let share = root;
+        let desc = ChildDesc::new(share, 1);
+
+        let children = reconstruct_children(&[desc]);
+
+        assert_eq!(children.len(), 1);
+        assert_eq!(children[0].desc, desc);
+
+        // Verify that the child is correctly derived from the reconstructed root
+        let expected_child = derive_child(root, desc);
+        assert_eq!(children[0], expected_child);
+    }
+
+    #[test]
+    fn test_reconstruct_n_children() {
+        let root = Share::random();
+        let n = 100;
+        let mut shares: Vec<Share> = (0..n - 1).map(|_| Share::random()).collect();
+        let final_share = {
+            let mut final_share = root;
+            for share in &shares {
+                final_share.xor_assign(share);
+            }
+            final_share
+        };
+        shares.push(final_share);
+
+        // onion packet will add ChildDesc
+        let descs: Vec<ChildDesc> = shares
+            .iter()
+            .enumerate()
+            .map(|(i, &share)| ChildDesc::new(share, i as u32))
+            .collect();
+
+        // last hop will reconstruct children and derive them
+        let children = reconstruct_children(&descs);
+
+        assert_eq!(children.len(), descs.len());
+        for (i, child) in children.iter().enumerate() {
+            assert_eq!(child.desc, descs[i]);
+        }
+
+        // Verify that each child is correctly derived from the reconstructed root
+        for (i, desc) in descs.iter().enumerate() {
+            let expected_child = derive_child(root, *desc);
+            assert_eq!(children[i], expected_child);
+        }
+    }
+
     #[test]
     fn test_reconstruct_empty_children() {
         let children = reconstruct_children(&[]);

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

@@ -70,6 +70,7 @@ use super::{
 use crate::ckb::config::UdtCfgInfos;
 use crate::ckb::contracts::{check_udt_script, get_udt_whitelist, is_udt_type_auto_accept};
 use crate::ckb::{CkbChainMessage, FundingError, FundingRequest, FundingTx};
+use crate::fiber::amp::{reconstruct_children, ChildDesc, Share};
 use crate::fiber::channel::{
     AddTlcCommand, AddTlcResponse, ChannelEphemeralConfig, ChannelInitializationOperation,
     ShutdownCommand, TxCollaborationCommand, TxUpdateCommand, DEFAULT_COMMITMENT_DELAY_EPOCHS,
@@ -88,7 +89,7 @@ use crate::fiber::payment::SessionRoute;
 use crate::fiber::payment::{Attempt, AttemptStatus, PaymentSession, PaymentStatus};
 use crate::fiber::serde_utils::EntityHex;
 use crate::fiber::types::{
-    FiberChannelMessage, PeeledPaymentOnionPacket, TlcErrPacket, TxSignatures,
+    AMPPaymentData, FiberChannelMessage, PeeledPaymentOnionPacket, TlcErrPacket, TxSignatures,
 };
 use crate::fiber::KeyPair;
 use crate::invoice::{CkbInvoice, CkbInvoiceStatus, InvoiceStore, PreimageStore};
@@ -659,7 +660,7 @@ impl SendPaymentData {
                 .is_some_and(|max_parts| max_parts <= 1 || max_parts > PAYMENT_MAX_PARTS_LIMIT)
         {
             return Err(format!(
-                "invalid max_parts, value should be in range [1, {}]",
+                "invalid max_parts, value should be in range [2, {}]",
                 PAYMENT_MAX_PARTS_LIMIT
             ));
         }
@@ -2384,6 +2385,7 @@ where
         session: &mut PaymentSession,
     ) -> Result<Vec<(Attempt, Vec<PaymentHopData>)>, Error> {
         let graph = self.network_graph.read().await;
+        let amp_mpp = session.request.allow_atomic_mpp;
         let source = graph.get_source_pubkey();
         let active_parts = session.attempts().filter(|a| a.is_active()).count();
         let mut remain_amount = session.remain_amount();
@@ -2464,13 +2466,61 @@ where
             return Err(Error::SendPaymentError(error));
         }
 
+        // generate custom records for amp_mpp
+        if amp_mpp {
+            self.set_custom_records_for_amp_mpp(&mut result, session.payment_hash());
+        }
+
         for (attempt, _) in &result {
             session.append_attempt(attempt.clone());
         }
 
         return Ok(result);
     }
 
+    fn set_custom_records_for_amp_mpp(
+        &self,
+        attempts_routers: &mut [(Attempt, Vec<PaymentHopData>)],
+        payment_hash: Hash256,
+    ) {
+        let root = Share::random();
+        let mut final_secret = root;
+        let mut secrets = vec![];
+        for _i in 0..(attempts_routers.iter().len() - 1) {
+            let secret = Share::random();
+            final_secret = final_secret.xor(&secret);
+            secrets.push(secret);
+        }
+        secrets.push(final_secret);
+
+        let descs: Vec<ChildDesc> = secrets
+            .iter()
+            .enumerate()
+            .map(|(i, &share)| ChildDesc::new(share, i as u32))
+            .collect();
+
+        let children = reconstruct_children(&descs);
+
+        for (i, desc) in descs.iter().enumerate() {
+            let (mut attempt, mut route) = attempts_routers[i].clone();
+
+            let route_len = route.len();
+            let mpp_payment_data =
+                AMPPaymentData::new(payment_hash, i as u16, children.len() as u16, desc.share);
+            let mut custom_records = route
+                .last()
+                .expect("expect last hop data")
+                .custom_records
+                .clone()
+                .unwrap_or_default();
+
+            mpp_payment_data.write(&mut custom_records);
+            route[route_len - 1].custom_records = Some(custom_records);
+            attempt.hash = children[i].hash;
+            attempts_routers[i] = (attempt, route);
+        }
+    }
+
     async fn send_payment_onion_packet(
         &self,
         state: &mut NetworkActorState<S>,

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

@@ -1,6 +1,7 @@
 use crate::{
     ckb::config::{UdtArgInfo, UdtCellDep, UdtCfgInfos, UdtDep, UdtScript},
     fiber::{
+        amp::Share,
         config::AnnouncedNodeName,
         features::FeatureVector,
         gen::{fiber as molecule_fiber, gossip},
@@ -734,7 +735,7 @@ fn test_basic_mpp_custom_records() {
 fn test_amp_custom_records() {
     let mut payment_custom_records = PaymentCustomRecords::default();
     let parent_payment_hash = gen_rand_sha256_hash();
-    let amp_record = AMPPaymentData::new(parent_payment_hash, 0, 3);
+    let amp_record = AMPPaymentData::new(parent_payment_hash, 0, 3, Share::random());
     amp_record.write(&mut payment_custom_records);
 
     let new_amp_record = AMPPaymentData::read(&payment_custom_records).unwrap();

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

@@ -12,6 +12,7 @@ use super::r#gen::fiber::PubNonceOpt;
 use super::serde_utils::{EntityHex, PubNonceAsBytes, SliceBase58, SliceHex};
 use crate::ckb::config::{UdtArgInfo, UdtCellDep, UdtCfgInfos, UdtDep, UdtScript};
 use crate::ckb::contracts::get_udt_whitelist;
+use crate::fiber::amp::Share;
 use crate::fiber::network::USER_CUSTOM_RECORDS_MAX_INDEX;
 use ckb_jsonrpc_types::CellOutput;
 use ckb_types::H256;
@@ -3784,21 +3785,28 @@ impl BasicMppPaymentData {
     }
 }
 
-#[derive(Hash, Eq, PartialEq, Debug)]
+#[derive(Eq, PartialEq, Debug)]
 pub struct AMPPaymentData {
     pub parent_payment_hash: Hash256,
     pub total_amp_count: u16,
     pub index: u16,
+    pub secret: Share,
 }
 
 impl AMPPaymentData {
     pub const CUSTOM_RECORD_KEY: u32 = USER_CUSTOM_RECORDS_MAX_INDEX + 2;
 
-    pub fn new(parent_payment_hash: Hash256, index: u16, total_amp_count: u16) -> Self {
+    pub fn new(
+        parent_payment_hash: Hash256,
+        index: u16,
+        total_amp_count: u16,
+        secret: Share,
+    ) -> Self {
         Self {
             parent_payment_hash,
             total_amp_count,
             index,
+            secret,
         }
     }
 
@@ -3807,6 +3815,7 @@ impl AMPPaymentData {
         vec.extend_from_slice(self.parent_payment_hash.as_ref());
         vec.extend_from_slice(&self.total_amp_count.to_le_bytes());
         vec.extend_from_slice(&self.index.to_le_bytes());
+        vec.extend_from_slice(self.secret.as_bytes());
         vec
     }
 
@@ -3821,16 +3830,18 @@ impl AMPPaymentData {
             .data
             .get(&Self::CUSTOM_RECORD_KEY)
             .and_then(|data| {
-                if data.len() != 32 + 4 {
+                if data.len() != 32 + 4 + 32 {
                     return None;
                 }
                 let parent_hash: [u8; 32] = data[..32].try_into().unwrap();
                 let total_amp_count = u16::from_le_bytes(data[32..34].try_into().unwrap());
-                let index = u16::from_le_bytes(data[34..].try_into().unwrap());
+                let index = u16::from_le_bytes(data[34..36].try_into().unwrap());
+                let secret = Share::new(data[36..68].try_into().unwrap());
                 Some(Self::new(
                     Hash256::from(parent_hash),
                     index,
                     total_amp_count,
+                    secret,
                 ))
             })
     }