Adding threshold error on exceeding 1000 ways

Signed-off-by: Harshil Jani <[email protected]>

Author: Harshil-Jani

src/descriptor/mod.rs

@@ -571,33 +571,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(),
             },
@@ -617,6 +617,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

@@ -163,7 +163,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) => {
                 let a = left.count_assets();

src/lib.rs

@@ -568,6 +568,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

@@ -412,7 +412,7 @@ impl<Ctx: ScriptContext> Miniscript<DescriptorPublicKey, Ctx> {
     }
 
     /// Get the total number of assets possible
-    pub fn count_assets(&self) -> u64 {
+    pub fn count_assets(&self) -> u32 {
         self.node.count_assets()
     }
 }

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()
 }
@@ -60,7 +62,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()
@@ -138,7 +140,7 @@ pub fn combine_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();
 
@@ -149,10 +151,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];
                 }
             }
@@ -164,9 +166,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)
 }