Atomic mpp

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

@@ -0,0 +1,153 @@
+use crate::fiber::{hash_algorithm::HashAlgorithm, types::Hash256};
+use bitcoin::hashes::{sha256::Hash as Sha256, Hash as _};
+use rand::RngCore;
+use serde::{Deserialize, Serialize};
+
+/// AmpSecret represents an n-of-n sharing of a secret 32-byte value.
+/// The secret can be recovered by XORing all n shares together.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub struct AmpSecret([u8; 32]);
+
+impl AmpSecret {
+    /// Create a new AmpSecret from a 32-byte array
+    pub fn new(bytes: [u8; 32]) -> Self {
+        Self(bytes)
+    }
+
+    /// Create a zero AmpSecret
+    pub fn zero() -> Self {
+        Self([0u8; 32])
+    }
+
+    /// Generate a random AmpSecret
+    pub fn random() -> Self {
+        let mut rng = rand::thread_rng();
+        let mut bytes = [0u8; 32];
+        rng.fill_bytes(&mut bytes);
+        Self(bytes)
+    }
+
+    /// XOR this AmpSecret with another AmpSecret, storing the result in self
+    pub fn xor_assign(&mut self, other: &AmpSecret) {
+        for (a, b) in self.0.iter_mut().zip(other.0.iter()) {
+            *a ^= b;
+        }
+    }
+
+    /// XOR two shares and return the result
+    pub fn xor(&self, other: &AmpSecret) -> AmpSecret {
+        let mut result = *self;
+        result.xor_assign(other);
+        result
+    }
+
+    /// Get the underlying bytes
+    pub fn as_bytes(&self) -> &[u8; 32] {
+        &self.0
+    }
+
+    /// Convert to bytes
+    pub fn to_bytes(self) -> [u8; 32] {
+        self.0
+    }
+
+    /// generate a random AmpSecret sequence
+    pub fn gen_random_sequence(root: AmpSecret, n: u16) -> Vec<AmpSecret> {
+        let mut shares: Vec<AmpSecret> = (0..n - 1).map(|_| AmpSecret::random()).collect();
+
+        let mut final_secret = root;
+        for share in &shares {
+            final_secret.xor_assign(share);
+        }
+        shares.push(final_secret);
+        shares
+    }
+}
+
+impl From<[u8; 32]> for AmpSecret {
+    fn from(bytes: [u8; 32]) -> Self {
+        Self(bytes)
+    }
+}
+
+impl AsRef<[u8]> for AmpSecret {
+    fn as_ref(&self) -> &[u8] {
+        &self.0
+    }
+}
+
+/// AmpChildDesc is the meta data for a child payment derived from the root seed.
+/// It contains the index of the child and the share used in the derivation.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct AmpChildDesc {
+    pub index: u16,
+    pub secret: AmpSecret,
+}
+
+impl AmpChildDesc {
+    pub fn new(index: u16, secret: AmpSecret) -> Self {
+        Self { index, secret }
+    }
+}
+
+/// Child is a payment hash and preimage pair derived from the root seed and ChildDesc.
+/// In addition to the derived values, a Child carries all information required in
+/// the derivation apart from the root seed (unless n=1).
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct AmpChild {
+    /// Preimage is the child payment preimage that can be used to settle the HTLC
+    pub preimage: Hash256,
+
+    /// Hash is the child payment hash that to be carried by the HTLC
+    pub hash: Hash256,
+}
+
+impl AmpChild {
+    pub fn new(preimage: Hash256, hash: Hash256) -> Self {
+        Self { preimage, hash }
+    }
+
+    /// DeriveChild computes the child preimage and child hash for a given (root, share, index) tuple.
+    /// The derivation is defined as:
+    ///   child_preimage = SHA256(root || share || be16(index))
+    ///   child_hash = SHA256(child_preimage)
+    pub fn derive_child(
+        root: AmpSecret,
+        desc: AmpChildDesc,
+        hash_algorithm: HashAlgorithm,
+    ) -> AmpChild {
+        let index_bytes = desc.index.to_be_bytes();
+
+        // Compute child_preimage as SHA256(root || share || child_index)
+        let mut preimage_data = Vec::with_capacity(32 + 32 + 2);
+        preimage_data.extend_from_slice(root.as_bytes());
+        preimage_data.extend_from_slice(desc.secret.as_bytes());
+        preimage_data.extend_from_slice(&index_bytes);
+
+        let preimage_hash = Sha256::hash(&preimage_data);
+        let preimage: Hash256 = preimage_hash.to_byte_array().into();
+
+        // this is the payment hash for HTLC
+        let hash: Hash256 = hash_algorithm.hash(preimage.as_ref()).into();
+        AmpChild::new(preimage, hash)
+    }
+
+    /// ReconstructChildren derives the set of children hashes and preimages from the
+    /// provided descriptors.
+    pub fn reconstruct_amp_children(
+        child_descs: &[AmpChildDesc],
+        hash_algorithm: HashAlgorithm,
+    ) -> Vec<AmpChild> {
+        // Recompute the root by XORing the provided shares
+        let mut root = AmpSecret::zero();
+        for desc in child_descs {
+            root.xor_assign(&desc.secret);
+        }
+
+        // With the root computed, derive the child hashes and preimages
+        child_descs
+            .iter()
+            .map(|data| Self::derive_child(root, data.clone(), hash_algorithm))
+            .collect()
+    }
+}

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

@@ -0,0 +1,124 @@
+use serde::{Deserialize, Serialize};
+
+use crate::fiber::{
+    amp::{AmpChildDesc, AmpSecret},
+    network::USER_CUSTOM_RECORDS_MAX_INDEX,
+    types::Hash256,
+    PaymentCustomRecords,
+};
+
+#[derive(Eq, PartialEq, Debug)]
+/// Bolt04 basic MPP payment data record
+pub struct BasicMppPaymentData {
+    pub payment_secret: Hash256,
+    pub total_amount: u128,
+}
+
+impl BasicMppPaymentData {
+    // record type for payment data record in bolt04
+    // custom records key from 65536 is reserved for internal usage
+    pub const CUSTOM_RECORD_KEY: u32 = USER_CUSTOM_RECORDS_MAX_INDEX + 1;
+
+    pub fn new(payment_secret: Hash256, total_amount: u128) -> Self {
+        Self {
+            payment_secret,
+            total_amount,
+        }
+    }
+
+    fn to_vec(&self) -> Vec<u8> {
+        let mut vec = Vec::new();
+        vec.extend_from_slice(self.payment_secret.as_ref());
+        vec.extend_from_slice(&self.total_amount.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 + 16 {
+                    return None;
+                }
+                let secret: [u8; 32] = data[..32].try_into().unwrap();
+                let payment_secret = Hash256::from(secret);
+                let total_amount = u128::from_le_bytes(data[32..].try_into().unwrap());
+                Some(Self::new(payment_secret, total_amount))
+            })
+    }
+}
+
+#[derive(Clone, Eq, PartialEq, Debug, Serialize, Deserialize)]
+pub struct AmpPaymentData {
+    pub total_amp_count: u16,
+    pub payment_hash: Hash256,
+    pub child_desc: AmpChildDesc,
+    pub total_amount: u128,
+}
+
+impl AmpPaymentData {
+    pub const CUSTOM_RECORD_KEY: u32 = USER_CUSTOM_RECORDS_MAX_INDEX + 2;
+
+    pub fn new(
+        payment_hash: Hash256,
+        total_amp_count: u16,
+        child_desc: AmpChildDesc,
+        total_amount: u128,
+    ) -> Self {
+        Self {
+            payment_hash,
+            total_amp_count,
+            child_desc,
+            total_amount,
+        }
+    }
+
+    fn to_vec(&self) -> Vec<u8> {
+        let mut vec = Vec::new();
+        vec.extend_from_slice(self.payment_hash.as_ref());
+        vec.extend_from_slice(&self.total_amp_count.to_le_bytes());
+        vec.extend_from_slice(&self.child_desc.index.to_le_bytes());
+        vec.extend_from_slice(self.child_desc.secret.as_bytes());
+        vec.extend_from_slice(&self.total_amount.to_le_bytes());
+        vec
+    }
+
+    pub fn index(&self) -> u16 {
+        self.child_desc.index
+    }
+
+    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 + 32 + 16 {
+                    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..36].try_into().unwrap());
+                let secret = AmpSecret::new(data[36..68].try_into().unwrap());
+                let total_amount = u128::from_le_bytes(data[68..].try_into().unwrap());
+                Some(Self::new(
+                    Hash256::from(parent_hash),
+                    total_amp_count,
+                    AmpChildDesc::new(index, secret),
+                    total_amount,
+                ))
+            })
+    }
+}