add atomic custom records

Author: chenyukang

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

@@ -6,14 +6,15 @@ use crate::{
         gen::{fiber as molecule_fiber, gossip},
         hash_algorithm::HashAlgorithm,
         types::{
-            pack_hop_data, secp256k1_instance, unpack_hop_data, AddTlc, BroadcastMessageID, Cursor,
-            Hash256, NodeAnnouncement, NodeId, PaymentHopData, PeeledOnionPacket, Privkey, Pubkey,
-            TlcErr, TlcErrPacket, TlcErrorCode, NO_SHARED_SECRET,
+            pack_hop_data, secp256k1_instance, unpack_hop_data, AMPDataRecord, AddTlc,
+            BroadcastMessageID, Cursor, Hash256, NodeAnnouncement, NodeId, PaymentHopData,
+            PeeledOnionPacket, Privkey, Pubkey, TlcErr, TlcErrPacket, TlcErrorCode,
+            NO_SHARED_SECRET,
         },
         PaymentCustomRecords,
     },
     gen_deterministic_fiber_private_key, gen_rand_channel_outpoint, gen_rand_fiber_private_key,
-    gen_rand_fiber_public_key, now_timestamp_as_millis_u64,
+    gen_rand_fiber_public_key, gen_rand_sha256_hash, now_timestamp_as_millis_u64,
 };
 use ckb_hash::blake2b_256;
 use ckb_jsonrpc_types::OutPoint;
@@ -701,3 +702,15 @@ fn test_serde_node_id() {
         "to NodeId"
     );
 }
+
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
+#[cfg_attr(not(target_arch = "wasm32"), test)]
+fn test_amp_custom_records() {
+    let mut payment_custom_records = PaymentCustomRecords::default();
+    let parent_payment_hash = gen_rand_sha256_hash();
+    let amp_record = AMPDataRecord::new(parent_payment_hash, 0, 3);
+    amp_record.write(&mut payment_custom_records);
+
+    let new_amp_record = AMPDataRecord::read(&payment_custom_records).unwrap();
+    assert_eq!(new_amp_record, amp_record);
+}

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

@@ -3786,6 +3786,58 @@ impl PaymentDataRecord {
     }
 }
 
+#[derive(Hash, Eq, PartialEq, Debug)]
+pub struct AMPDataRecord {
+    pub parent_payment_hash: Hash256,
+    pub total_amp_count: u16,
+    pub index: u16,
+}
+
+impl AMPDataRecord {
+    pub const CUSTOM_RECORD_KEY: u32 = 65537;
+
+    pub fn new(parent_payment_hash: Hash256, index: u16, total_amp_count: u16) -> Self {
+        Self {
+            parent_payment_hash,
+            total_amp_count,
+            index,
+        }
+    }
+
+    fn to_vec(&self) -> Vec<u8> {
+        let mut vec = Vec::new();
+        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
+    }
+
+    pub fn write(&self, custom_records: &mut PaymentCustomRecords) {
+        custom_records
+            .data
+            .insert(Self::CUSTOM_RECORD_KEY, self.to_vec());
+    }
+
+    pub fn read(custom_records: &PaymentCustomRecords) -> Option<Self> {
+        custom_records
+            .data
+            .get(&Self::CUSTOM_RECORD_KEY)
+            .and_then(|data| {
+                if data.len() != 32 + 4 {
+                    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());
+                Some(Self::new(
+                    Hash256::from(parent_hash),
+                    index,
+                    total_amp_count,
+                ))
+            })
+    }
+}
+
 #[serde_as]
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct PaymentHopData {