Setting Maximum threshold for obtaining plan

For the planning module we are considering that total possible ways to
spend should always be less than 1000. This protects the network from
any DDoS Attack if in case you receive a very large descriptor enough
that it has 1000 possible spend paths.

Signed-off-by: Harshil Jani <harshiljani2002@gmail.com>

Author: Harshil-Jani

src/descriptor/mod.rs

@@ -549,33 +549,33 @@ impl Descriptor<DefiniteDescriptorKey> {
 
 impl Descriptor<DescriptorPublicKey> {
     /// Count total possible assets for a given descriptor.
-    pub fn count_assets(&self) -> u64 {
+    pub fn count_assets(&self) -> u32 {
         match self {
             Descriptor::Bare(k) => k.as_inner().count_assets(),
             Descriptor::Pkh(_) => 1,
             Descriptor::Wpkh(_) => 1,
             Descriptor::Sh(k) => match k.as_inner() {
                 ShInner::Wsh(k) => match k.as_inner() {
                     WshInner::SortedMulti(k) => {
-                        let n = k.pks.len() as u64;
-                        let k = k.k as u64;
-                        k_of_n(k, n)
+                        let n = k.pks.len() as u32;
+                        let k = k.k as u32;
+                        k_of_n(k, n).unwrap()
                     }
                     WshInner::Ms(k) => k.count_assets(),
                 },
                 ShInner::Wpkh(_) => 1,
                 ShInner::SortedMulti(k) => {
-                    let n = k.clone().pks.len() as u64;
-                    let k = k.clone().k as u64;
-                    k_of_n(k, n)
+                    let n = k.clone().pks.len() as u32;
+                    let k = k.clone().k as u32;
+                    k_of_n(k, n).unwrap()
                 }
                 ShInner::Ms(k) => k.count_assets(),
             },
             Descriptor::Wsh(k) => match k.as_inner() {
                 WshInner::SortedMulti(k) => {
-                    let n = k.clone().pks.len() as u64;
-                    let k = k.clone().k as u64;
-                    k_of_n(k, n)
+                    let n = k.clone().pks.len() as u32;
+                    let k = k.clone().k as u32;
+                    k_of_n(k, n).unwrap()
                 }
                 WshInner::Ms(k) => k.count_assets(),
             },
@@ -595,6 +595,10 @@ impl Descriptor<DescriptorPublicKey> {
 
     /// Get all possible assets for a given descriptor
     pub fn all_assets(&self) -> Result<Vec<Assets>, Error> {
+        let threshold = self.count_assets();
+        if threshold >= 1000 {
+            return Err(Error::MaxAssetThresholdExceeded);
+        }
         match self {
             Descriptor::Bare(k) => Ok(k.as_inner().all_assets()),
             Descriptor::Pkh(k) => {

src/descriptor/tr.rs

@@ -168,7 +168,7 @@ impl TapTree<DescriptorPublicKey> {
     }
 
     /// Get total possible assets for TapTree
-    pub fn count_assets(&self) -> u64 {
+    pub fn count_assets(&self) -> u32 {
         match self {
             TapTree::Tree { left, right, height: _ } => {
                 let a = left.count_assets();

src/lib.rs

@@ -509,6 +509,9 @@ pub enum Error {
 const MAX_RECURSION_DEPTH: u32 = 402;
 // https://github.com/bitcoin/bips/blob/master/bip-0141.mediawiki
 const MAX_SCRIPT_SIZE: u32 = 10000;
+// For the planning module we are considering that total possible ways to spend
+// should be less than 1000
+const MAX_ASSET_THRESHOLD: u32 = 1000;
 
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

src/miniscript/astelem.rs

@@ -598,7 +598,7 @@ impl<Pk: MiniscriptKey, Ctx: ScriptContext> Terminal<Pk, Ctx> {
 
 impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
     /// Count total possible assets
-    pub fn count_assets(&self) -> u64 {
+    pub fn count_assets(&self) -> u32 {
         match self {
             Terminal::True => 0,
             Terminal::False => 0,
@@ -642,17 +642,17 @@ impl<Ctx: ScriptContext> Terminal<DescriptorPublicKey, Ctx> {
                 for ms in ms_v {
                     count_array.push(ms.count_assets());
                 }
-                let products = get_combinations_product(&count_array, *k as u64);
-                let mut total_count: u64 = 0;
+                let products = get_combinations_product(&count_array, *k as u32);
+                let mut total_count: u32 = 0;
                 for product in products {
                     total_count += product;
                 }
                 total_count
             }
             Terminal::Multi(k, dpk) | Terminal::MultiA(k, dpk) => {
-                let k: u64 = *k as u64;
-                let n: u64 = dpk.len() as u64;
-                k_of_n(k, n)
+                let k: u32 = *k as u32;
+                let n: u32 = dpk.len() as u32;
+                k_of_n(k, n).unwrap()
             }
         }
     }

src/miniscript/mod.rs

@@ -391,7 +391,7 @@ impl<Ctx: ScriptContext> Miniscript<DescriptorPublicKey, Ctx> {
     pub fn all_assets(&self) -> Vec<Assets> { self.node.all_assets() }
 
     /// Get the total number of assets possible
-    pub fn count_assets(&self) -> u64 { self.node.count_assets() }
+    pub fn count_assets(&self) -> u32 { self.node.count_assets() }
 }
 
 impl<Pk: MiniscriptKey, Ctx: ScriptContext> ForEachKey<Pk> for Miniscript<Pk, Ctx> {

src/util.rs

@@ -10,7 +10,9 @@ use crate::miniscript::context;
 use crate::miniscript::satisfy::Placeholder;
 use crate::plan::Assets;
 use crate::prelude::*;
-use crate::{DescriptorPublicKey, MiniscriptKey, ScriptContext, ToPublicKey};
+use crate::{
+    DescriptorPublicKey, Error, MiniscriptKey, ScriptContext, ToPublicKey, MAX_ASSET_THRESHOLD,
+};
 pub(crate) fn varint_len(n: usize) -> usize { bitcoin::VarInt(n as u64).len() }
 
 pub(crate) trait ItemSize {
@@ -56,7 +58,7 @@ pub(crate) fn witness_to_scriptsig(witness: &[Vec<u8>]) -> ScriptBuf {
         } else {
             let push = <&PushBytes>::try_from(wit.as_slice())
                 .expect("All pushes in miniscript are <73 bytes");
-            b = b.push_slice(push)
+            b = b.push_slice(push);
         }
     }
     b.into_script()
@@ -129,12 +131,11 @@ pub fn combine_assets(
     combine_assets(k, dpk_v, index + 1, current_assets.clone(), all_assets);
     let mut new_asset = current_assets;
     new_asset = new_asset.add(dpk_v[index].clone());
-    println!("{:#?}", new_asset);
     combine_assets(k - 1, dpk_v, index + 1, new_asset, all_assets)
 }
 
 // Do product of K combinations
-pub fn get_combinations_product(values: &[u64], k: u64) -> Vec<u64> {
+pub fn get_combinations_product(values: &[u32], k: u32) -> Vec<u32> {
     let mut products = Vec::new();
     let n = values.len();
 
@@ -145,10 +146,10 @@ pub fn get_combinations_product(values: &[u64], k: u64) -> Vec<u64> {
     // Using bitwise operations to generate combinations
     let max_combinations = 1u32 << n;
     for combination_bits in 1..max_combinations {
-        if combination_bits.count_ones() as usize == k as usize {
+        if (combination_bits.count_ones() as usize) == (k as usize) {
             let mut product = 1;
             for i in 0..n {
-                if combination_bits & (1u32 << i) != 0 {
+                if (combination_bits & (1u32 << i)) != 0 {
                     product *= values[i];
                 }
             }
@@ -160,9 +161,20 @@ pub fn get_combinations_product(values: &[u64], k: u64) -> Vec<u64> {
 }
 
 // ways to select k things out of n
-pub fn k_of_n(k: u64, n: u64) -> u64 {
-    if k == 0 || k == n {
-        return 1;
+pub fn k_of_n(k: u32, n: u32) -> Result<u32, Error> {
+    let mut k = k;
+    if k > n - k {
+        k = n - k;
     }
-    k_of_n(k - 1, n - 1) + k_of_n(k, n - 1)
+
+    let mut result = 1;
+    for i in 0..k {
+        result *= n - i;
+        result /= i + 1;
+        if result > MAX_ASSET_THRESHOLD.into() {
+            return Err(Error::MaxAssetThresholdExceeded);
+        }
+    }
+
+    Ok(result)
 }