fix(wallet/descriptor+wallet): expect Threshold type, and handle it

oleonardolima · oleonardolima · commit 97dcf4f13d2d · 2024-06-06T16:05:03.000-03:00
internally
diff --git a/crates/wallet/src/descriptor/dsl.rs b/crates/wallet/src/descriptor/dsl.rs
@@ -703,10 +703,10 @@ macro_rules! fragment {
         $crate::keys::make_pkh($key, &secp)
     });
     ( after ( $value:expr ) ) => ({
-        $crate::impl_leaf_opcode_value!(After, $crate::miniscript::AbsLockTime::from_consensus($value))
+        $crate::impl_leaf_opcode_value!(After, $crate::miniscript::AbsLockTime::from_consensus($value).expect("valid `AbsLockTime`"))
     });
     ( older ( $value:expr ) ) => ({
-        $crate::impl_leaf_opcode_value!(Older, $crate::bitcoin::Sequence($value)) // TODO!!
+        $crate::impl_leaf_opcode_value!(Older, $crate::miniscript::RelLockTime::from_consensus($value).expect("valid `RelLockTime`")) // TODO!!
     });
     ( sha256 ( $hash:expr ) ) => ({
         $crate::impl_leaf_opcode_value!(Sha256, $hash)
@@ -757,7 +757,8 @@ macro_rules! fragment {
             (keys_acc, net_acc)
         });
 
-        $crate::impl_leaf_opcode_value_two!(Thresh, $thresh, items)
+        let thresh = $crate::miniscript::Threshold::new($thresh, items).expect("valid threshold and pks collection");
+        $crate::impl_leaf_opcode_value!(Thresh, thresh)
             .map(|(minisc, _, _)| (minisc, key_maps, valid_networks))
     });
     ( thresh ( $thresh:expr, $( $inner:tt )* ) ) => ({
@@ -769,7 +770,12 @@ macro_rules! fragment {
     ( multi_vec ( $thresh:expr, $keys:expr ) ) => ({
         let secp = $crate::bitcoin::secp256k1::Secp256k1::new();
 
-        $crate::keys::make_multi($thresh, $crate::miniscript::Terminal::Multi, $keys, &secp)
+        let fun = |k, pks| {
+            let thresh = $crate::miniscript::Threshold::new(k, pks).expect("valid threshold and pks collection");
+            $crate::miniscript::Terminal::Multi(thresh)
+        };
+
+        $crate::keys::make_multi($thresh, fun, $keys, &secp)
     });
     ( multi ( $thresh:expr $(, $key:expr )+ ) ) => ({
         $crate::group_multi_keys!( $( $key ),* )
@@ -778,7 +784,12 @@ macro_rules! fragment {
     ( multi_a_vec ( $thresh:expr, $keys:expr ) ) => ({
         let secp = $crate::bitcoin::secp256k1::Secp256k1::new();
 
-        $crate::keys::make_multi($thresh, $crate::miniscript::Terminal::MultiA, $keys, &secp)
+        let fun = |k, pks| {
+            let thresh = $crate::miniscript::Threshold::new(k, pks).expect("valid threshold and pks collection");
+            $crate::miniscript::Terminal::MultiA(thresh)
+        };
+
+        $crate::keys::make_multi($thresh, fun, $keys, &secp)
     });
     ( multi_a ( $thresh:expr $(, $key:expr )+ ) ) => ({
         $crate::group_multi_keys!( $( $key ),* )
diff --git a/crates/wallet/src/descriptor/policy.rs b/crates/wallet/src/descriptor/policy.rs
@@ -40,6 +40,7 @@ use crate::collections::{BTreeMap, HashSet, VecDeque};
 use alloc::string::String;
 use alloc::vec::Vec;
 use core::cmp::max;
+use miniscript::miniscript::limits::{MAX_PUBKEYS_IN_CHECKSIGADD, MAX_PUBKEYS_PER_MULTISIG};
 
 use core::fmt;
 
@@ -53,7 +54,7 @@ use bitcoin::{absolute, key::XOnlyPublicKey, PublicKey, Sequence};
 use miniscript::descriptor::{
     DescriptorPublicKey, ShInner, SinglePub, SinglePubKey, SortedMultiVec, WshInner,
 };
-use miniscript::hash256;
+use miniscript::{hash256, Threshold};
 use miniscript::{
     Descriptor, Miniscript, Satisfier, ScriptContext, SigType, Terminal, ToPublicKey,
 };
@@ -586,30 +587,25 @@ impl Policy {
         Ok(Some(policy))
     }
 
-    fn make_multisig<Ctx: ScriptContext + 'static>(
-        keys: &[DescriptorPublicKey],
+    fn make_multi<Ctx: ScriptContext + 'static>(
+        threshold: &Threshold<DescriptorPublicKey, MAX_PUBKEYS_PER_MULTISIG>,
         signers: &SignersContainer,
         build_sat: BuildSatisfaction,
-        threshold: usize,
         sorted: bool,
         secp: &SecpCtx,
     ) -> Result<Option<Policy>, PolicyError> {
-        if threshold == 0 {
-            return Ok(None);
-        }
-
-        let parsed_keys = keys.iter().map(|k| PkOrF::from_key(k, secp)).collect();
+        let parsed_keys = threshold.iter().map(|k| PkOrF::from_key(k, secp)).collect();
 
         let mut contribution = Satisfaction::Partial {
-            n: keys.len(),
-            m: threshold,
+            n: threshold.n(),
+            m: threshold.k(),
             items: vec![],
             conditions: Default::default(),
             sorted: Some(sorted),
         };
         let mut satisfaction = contribution.clone();
 
-        for (index, key) in keys.iter().enumerate() {
+        for (index, key) in threshold.iter().enumerate() {
             if signers.find(signer_id(key, secp)).is_some() {
                 contribution.add(
                     &Satisfaction::Complete {
@@ -635,7 +631,7 @@ impl Policy {
 
         let mut policy: Policy = SatisfiableItem::Multisig {
             keys: parsed_keys,
-            threshold,
+            threshold: threshold.k(),
         }
         .into();
         policy.contribution = contribution;
@@ -644,6 +640,57 @@ impl Policy {
         Ok(Some(policy))
     }
 
+    fn make_multi_a<Ctx: ScriptContext + 'static>(
+        threshold: &Threshold<DescriptorPublicKey, MAX_PUBKEYS_IN_CHECKSIGADD>,
+        signers: &SignersContainer,
+        build_sat: BuildSatisfaction,
+        sorted: bool,
+        secp: &SecpCtx,
+    ) -> Result<Option<Policy>, PolicyError> {
+        let parsed_keys = threshold.iter().map(|k| PkOrF::from_key(k, secp)).collect();
+
+        let mut contribution = Satisfaction::Partial {
+            n: threshold.n(),
+            m: threshold.k(),
+            items: vec![],
+            conditions: Default::default(),
+            sorted: Some(sorted),
+        };
+        let mut satisfaction = contribution.clone();
+
+        for (index, key) in threshold.iter().enumerate() {
+            if signers.find(signer_id(key, secp)).is_some() {
+                contribution.add(
+                    &Satisfaction::Complete {
+                        condition: Default::default(),
+                    },
+                    index,
+                )?;
+            }
+            if let Some(psbt) = build_sat.psbt() {
+                if Ctx::find_signature(psbt, key, secp) {
+                    satisfaction.add(
+                        &Satisfaction::Complete {
+                            condition: Default::default(),
+                        },
+                        index,
+                    )?;
+                }
+            }
+        }
+        satisfaction.finalize();
+        contribution.finalize();
+
+        let mut policy: Policy = SatisfiableItem::Multisig {
+            keys: parsed_keys,
+            threshold: threshold.k(),
+        }
+        .into();
+        policy.contribution = contribution;
+        policy.satisfaction = satisfaction;
+        Ok(Some(policy))
+    }
+
     /// Return whether or not a specific path in the policy tree is required to unambiguously
     /// create a transaction
     ///
@@ -952,7 +999,10 @@ impl<Ctx: ScriptContext + 'static> ExtractPolicy for Miniscript<DescriptorPublic
                 Some(policy)
             }
             Terminal::Older(value) => {
-                let mut policy: Policy = SatisfiableItem::RelativeTimelock { value: (*value).into() }.into();
+                let mut policy: Policy = SatisfiableItem::RelativeTimelock {
+                    value: (*value).into(),
+                }
+                .into();
                 policy.contribution = Satisfaction::Complete {
                     condition: Condition {
                         timelock: None,
@@ -966,9 +1016,11 @@ impl<Ctx: ScriptContext + 'static> ExtractPolicy for Miniscript<DescriptorPublic
                 } = build_sat
                 {
                     let older = Older::new(Some(current_height), Some(input_max_height), false);
-                    let older_sat = Satisfier::<bitcoin::PublicKey>::check_older(&older, (*value).into());
-                    let inputs_sat = psbt_inputs_sat(psbt)
-                        .all(|sat| Satisfier::<bitcoin::PublicKey>::check_older(&sat, (*value).into()));
+                    let older_sat =
+                        Satisfier::<bitcoin::PublicKey>::check_older(&older, (*value).into());
+                    let inputs_sat = psbt_inputs_sat(psbt).all(|sat| {
+                        Satisfier::<bitcoin::PublicKey>::check_older(&sat, (*value).into())
+                    });
                     if older_sat && inputs_sat {
                         policy.satisfaction = policy.contribution.clone();
                     }
@@ -986,8 +1038,11 @@ impl<Ctx: ScriptContext + 'static> ExtractPolicy for Miniscript<DescriptorPublic
             Terminal::Hash160(hash) => {
                 Some(SatisfiableItem::Hash160Preimage { hash: *hash }.into())
             }
-            Terminal::Multi(k, pks) | Terminal::MultiA(k, pks) => {
-                Policy::make_multisig::<Ctx>(pks, signers, build_sat, *k, false, secp)?
+            Terminal::Multi(threshold) => {
+                Policy::make_multi::<Ctx>(threshold, signers, build_sat, false, secp)?
+            }
+            Terminal::MultiA(threshold) => {
+                Policy::make_multi_a::<Ctx>(threshold, signers, build_sat, false, secp)?
             }
             // Identities
             Terminal::Alt(inner)
@@ -1016,8 +1071,9 @@ impl<Ctx: ScriptContext + 'static> ExtractPolicy for Miniscript<DescriptorPublic
                 a.extract_policy(signers, build_sat, secp)?,
                 b.extract_policy(signers, build_sat, secp)?,
             )?,
-            Terminal::Thresh(k, nodes) => {
-                let mut threshold = *k;
+            Terminal::Thresh(threshold) => {
+                let mut k = threshold.k();
+                let nodes = threshold.data();
                 let mapped: Vec<_> = nodes
                     .iter()
                     .map(|n| n.extract_policy(signers, build_sat, secp))
@@ -1027,13 +1083,13 @@ impl<Ctx: ScriptContext + 'static> ExtractPolicy for Miniscript<DescriptorPublic
                     .collect();
 
                 if mapped.len() < nodes.len() {
-                    threshold = match threshold.checked_sub(nodes.len() - mapped.len()) {
+                    k = match k.checked_sub(nodes.len() - mapped.len()) {
                         None => return Ok(None),
                         Some(x) => x,
                     };
                 }
 
-                Policy::make_thresh(mapped, threshold)?
+                Policy::make_thresh(mapped, k)?
             }
 
             // Unsupported
@@ -1087,13 +1143,10 @@ impl ExtractPolicy for Descriptor<DescriptorPublicKey> {
             build_sat: BuildSatisfaction,
             secp: &SecpCtx,
         ) -> Result<Option<Policy>, Error> {
-            Ok(Policy::make_multisig::<Ctx>(
-                keys.pks.as_ref(),
-                signers,
-                build_sat,
-                keys.k,
-                true,
-                secp,
+            let threshold = Threshold::new(keys.k(), keys.pks().to_vec())
+                .expect("valid threshold and pks collection");
+            Ok(Policy::make_multi::<Ctx>(
+                &threshold, signers, build_sat, true, secp,
             )?)
         }
 
diff --git a/crates/wallet/src/wallet/export.rs b/crates/wallet/src/wallet/export.rs
@@ -166,7 +166,7 @@ impl FullyNodedExport {
         fn check_ms<Ctx: ScriptContext>(
             terminal: &Terminal<String, Ctx>,
         ) -> Result<(), &'static str> {
-            if let Terminal::Multi(_, _) = terminal {
+            if let Terminal::Multi(_) = terminal {
                 Ok(())
             } else {
                 Err("The descriptor contains operators not supported by Bitcoin Core")
