Merge branch 'policy-refactor'

benma · benma · commit 50b05b6c8424 · 2024-06-12T22:23:28.000+02:00
diff --git a/src/rust/bitbox02-rust/src/hww/api/bitcoin.rs b/src/rust/bitbox02-rust/src/hww/api/bitcoin.rs
@@ -25,6 +25,7 @@ mod payment_request;
 mod policies;
 mod registration;
 mod script;
+mod script_configs;
 pub mod signmsg;
 pub mod signtx;
 
diff --git a/src/rust/bitbox02-rust/src/hww/api/bitcoin/common.rs b/src/rust/bitbox02-rust/src/hww/api/bitcoin/common.rs
@@ -24,6 +24,7 @@ use pb::btc_script_config::SimpleType;
 pub use pb::btc_sign_init_request::FormatUnit;
 pub use pb::{BtcCoin, BtcOutputType};
 
+use super::script_configs::{ValidatedScriptConfig, ValidatedScriptConfigWithKeypath};
 use super::{multisig, params::Params, script};
 
 use sha2::{Digest, Sha256};
@@ -155,8 +156,8 @@ impl Payload {
         keypath: &[u32],
     ) -> Result<Self, Error> {
         let witness_script = policy.witness_script_at_keypath(keypath)?;
-        match policy {
-            super::policies::ParsedPolicy::Wsh { .. } => Ok(Payload {
+        match &policy.descriptor {
+            super::policies::Descriptor::Wsh { .. } => Ok(Payload {
                 data: Sha256::digest(witness_script).to_vec(),
                 output_type: BtcOutputType::P2wsh,
             }),
@@ -169,39 +170,19 @@ impl Payload {
         xpub_cache: &mut Bip32XpubCache,
         params: &Params,
         keypath: &[u32],
-        script_config_account: &pb::BtcScriptConfigWithKeypath,
+        script_config_account: &ValidatedScriptConfigWithKeypath,
     ) -> Result<Self, Error> {
-        match script_config_account {
-            pb::BtcScriptConfigWithKeypath {
-                script_config:
-                    Some(pb::BtcScriptConfig {
-                        config: Some(pb::btc_script_config::Config::SimpleType(simple_type)),
-                    }),
-                ..
-            } => {
-                let simple_type = pb::btc_script_config::SimpleType::try_from(*simple_type)?;
-                Self::from_simple(xpub_cache, params, simple_type, keypath)
+        match &script_config_account.config {
+            ValidatedScriptConfig::SimpleType(simple_type) => {
+                Self::from_simple(xpub_cache, params, *simple_type, keypath)
             }
-            pb::BtcScriptConfigWithKeypath {
-                script_config:
-                    Some(pb::BtcScriptConfig {
-                        config: Some(pb::btc_script_config::Config::Multisig(multisig)),
-                    }),
-                ..
-            } => Self::from_multisig(
+            ValidatedScriptConfig::Multisig(multisig) => Self::from_multisig(
                 params,
                 multisig,
                 keypath[keypath.len() - 2],
                 keypath[keypath.len() - 1],
             ),
-            pb::BtcScriptConfigWithKeypath {
-                script_config:
-                    Some(pb::BtcScriptConfig {
-                        config: Some(pb::btc_script_config::Config::Policy(policy)),
-                    }),
-                ..
-            } => Self::from_policy(&super::policies::parse(policy)?, keypath),
-            _ => Err(Error::InvalidInput),
+            ValidatedScriptConfig::Policy(policy) => Self::from_policy(policy, keypath),
         }
     }
 
diff --git a/src/rust/bitbox02-rust/src/hww/api/bitcoin/policies.rs b/src/rust/bitbox02-rust/src/hww/api/bitcoin/policies.rs
@@ -137,6 +137,47 @@ fn get_change_and_address_index<R: core::convert::AsRef<str>, T: core::iter::Ite
     Err(Error::InvalidInput)
 }
 
+/// Check that it is impossible to create a derivation with duplicate pubkeys, assuming all the keys
+/// in the key vector are distinct.
+///
+/// Even though the rust-miniscript library checks for duplicate keys, it does so on the raw
+/// miniscript, which would not catch e.g. that `wsh(or_b(pk(@0/<0;1>/*),s:pk(@0/<2;1>/*)))` has a
+/// duplicate change derivation if we derive at the receive path.
+///
+/// Also checks that each key is used, e.g. if there are 3 keys in the key vector, @0, @1 and @2
+/// must be present.
+fn check_dups<R: core::convert::AsRef<str>, T: core::iter::Iterator<Item = R>>(
+    pubkeys: T,
+    keys: &[pb::KeyOriginInfo],
+) -> Result<(), Error> {
+    // in "@key_index/<left;right>", keeps track of (key_index,left) and
+    // (key_index,right) to check for duplicates.
+    let mut derivations_seen: Vec<(usize, u32)> = Vec::new();
+
+    let mut keys_seen: Vec<bool> = vec![false; keys.len()];
+
+    for pk in pubkeys {
+        let (key_index, multipath_index_left, multipath_index_right) =
+            parse_wallet_policy_pk(pk.as_ref()).or(Err(Error::InvalidInput))?;
+
+        if derivations_seen.contains(&(key_index, multipath_index_left)) {
+            return Err(Error::InvalidInput);
+        }
+        derivations_seen.push((key_index, multipath_index_left));
+        if derivations_seen.contains(&(key_index, multipath_index_right)) {
+            return Err(Error::InvalidInput);
+        }
+        derivations_seen.push((key_index, multipath_index_right));
+
+        *keys_seen.get_mut(key_index).ok_or(Error::InvalidInput)? = true;
+    }
+
+    if !keys_seen.into_iter().all(|b| b) {
+        return Err(Error::InvalidInput);
+    }
+    Ok(())
+}
+
 struct WalletPolicyPkTranslator<'a> {
     keys: &'a [pb::KeyOriginInfo],
     is_change: bool,
@@ -174,67 +215,30 @@ impl<'a> miniscript::Translator<String, bitcoin::PublicKey, Error>
 
 /// See `ParsedPolicy`.
 #[derive(Debug)]
-pub struct Wsh<'a> {
-    policy: &'a Policy,
+pub struct Wsh {
     miniscript_expr: miniscript::Miniscript<String, miniscript::Segwitv0>,
 }
 
-/// Result of `parse()`.
+/// See `ParsedPolicy`.
 #[derive(Debug)]
-pub enum ParsedPolicy<'a> {
+pub enum Descriptor {
     // `wsh(...)` policies
-    Wsh(Wsh<'a>),
+    Wsh(Wsh),
     // `tr(...)` Taproot etc. in the future.
 }
 
-impl<'a> ParsedPolicy<'a> {
-    fn get_policy(&self) -> &Policy {
-        match self {
-            Self::Wsh(Wsh { policy, .. }) => policy,
-        }
-    }
+/// Result of `parse()`.
+#[derive(Debug)]
+pub struct ParsedPolicy<'a> {
+    policy: &'a Policy,
+    pub descriptor: Descriptor,
+}
 
-    /// Check that it is impossible to create a derivation with duplicate pubkeys, assuming all the
-    /// keys in the key vector are distinct.
-    ///
-    /// Even though the rust-miniscript library checks for duplicate keys, it does so on the raw
-    /// miniscript, which would not catch e.g. that `wsh(or_b(pk(@0/<0;1>/*),s:pk(@0/<2;1>/*)))` has
-    /// a duplicate change derivation if we derive at the receive path.
-    ///
-    /// Also checks that each key is used, e.g. if there are 3 keys in the key vector, @0, @1 and @2
-    /// must be present.
+impl<'a> ParsedPolicy<'a> {
     fn validate_keys(&self) -> Result<(), Error> {
-        match self {
-            Self::Wsh(Wsh {
-                policy,
-                miniscript_expr,
-            }) => {
-                // in "@key_index/<left;right>", keeps track of (key_index,left) and
-                // (key_index,right) to check for duplicates.
-                let mut derivations_seen: Vec<(usize, u32)> = Vec::new();
-
-                let mut keys_seen: Vec<bool> = vec![false; policy.keys.len()];
-
-                for pk in miniscript_expr.iter_pk() {
-                    let (key_index, multipath_index_left, multipath_index_right) =
-                        parse_wallet_policy_pk(&pk).or(Err(Error::InvalidInput))?;
-
-                    if derivations_seen.contains(&(key_index, multipath_index_left)) {
-                        return Err(Error::InvalidInput);
-                    }
-                    derivations_seen.push((key_index, multipath_index_left));
-                    if derivations_seen.contains(&(key_index, multipath_index_right)) {
-                        return Err(Error::InvalidInput);
-                    }
-                    derivations_seen.push((key_index, multipath_index_right));
-
-                    *keys_seen.get_mut(key_index).ok_or(Error::InvalidInput)? = true;
-                }
-
-                if !keys_seen.into_iter().all(|b| b) {
-                    return Err(Error::InvalidInput);
-                }
-                Ok(())
+        match &self.descriptor {
+            Descriptor::Wsh(Wsh { miniscript_expr }) => {
+                check_dups(miniscript_expr.iter_pk(), &self.policy.keys)
             }
         }
     }
@@ -248,7 +252,7 @@ impl<'a> ParsedPolicy<'a> {
     pub fn validate(&self, coin: BtcCoin) -> Result<(), Error> {
         check_enabled(coin)?;
 
-        let policy = self.get_policy();
+        let policy = self.policy;
 
         if policy.keys.len() > MAX_KEYS {
             return Err(Error::InvalidInput);
@@ -297,13 +301,10 @@ impl<'a> ParsedPolicy<'a> {
     /// wsh(and_v(v:pk(@0/0/5),pk(@1/20/5))).
     /// The same derived using `is_change=true` derives: wsh(and_v(v:pk(@0/1/5),pk(@1/21/5)))
     pub fn witness_script(&self, is_change: bool, address_index: u32) -> Result<Vec<u8>, Error> {
-        match self {
-            Self::Wsh(Wsh {
-                policy,
-                miniscript_expr,
-            }) => {
+        match &self.descriptor {
+            Descriptor::Wsh(Wsh { miniscript_expr }) => {
                 let mut translator = WalletPolicyPkTranslator {
-                    keys: policy.keys.as_ref(),
+                    keys: self.policy.keys.as_ref(),
                     is_change,
                     address_index,
                 };
@@ -322,27 +323,27 @@ impl<'a> ParsedPolicy<'a> {
     /// derived using keypath m/48'/1'/0'/3'/11/5 derives:
     /// wsh(and_v(v:pk(@0/11/5),pk(@1/21/5))).
     pub fn witness_script_at_keypath(&self, keypath: &[u32]) -> Result<Vec<u8>, Error> {
-        match self {
-            Self::Wsh(Wsh {
-                policy,
-                miniscript_expr,
-            }) => {
-                let (is_change, address_index) =
-                    get_change_and_address_index(miniscript_expr.iter_pk(), &policy.keys, keypath)?;
+        match &self.descriptor {
+            Descriptor::Wsh(Wsh { miniscript_expr }) => {
+                let (is_change, address_index) = get_change_and_address_index(
+                    miniscript_expr.iter_pk(),
+                    &self.policy.keys,
+                    keypath,
+                )?;
                 self.witness_script(is_change, address_index)
             }
         }
     }
 
     /// Returns true if the address-level keypath points to a change address.
     pub fn is_change_keypath(&self, keypath: &[u32]) -> Result<bool, Error> {
-        match self {
-            Self::Wsh(Wsh {
-                policy,
-                miniscript_expr,
-            }) => {
-                let (is_change, _) =
-                    get_change_and_address_index(miniscript_expr.iter_pk(), &policy.keys, keypath)?;
+        match &self.descriptor {
+            Descriptor::Wsh(Wsh { miniscript_expr }) => {
+                let (is_change, _) = get_change_and_address_index(
+                    miniscript_expr.iter_pk(),
+                    &self.policy.keys,
+                    keypath,
+                )?;
                 Ok(is_change)
             }
         }
@@ -364,10 +365,10 @@ pub fn parse(policy: &Policy) -> Result<ParsedPolicy, Error> {
                 miniscript::Miniscript::from_str(&desc[4..desc.len() - 1])
                     .or(Err(Error::InvalidInput))?;
 
-            Ok(ParsedPolicy::Wsh(Wsh {
+            Ok(ParsedPolicy {
                 policy,
-                miniscript_expr,
-            }))
+                descriptor: Descriptor::Wsh(Wsh { miniscript_expr }),
+            })
         }
         _ => Err(Error::InvalidInput),
     }
@@ -595,10 +596,9 @@ mod tests {
     fn test_parse_wsh_miniscript() {
         // Parse a valid example and check that the keys are collected as is as strings.
         let policy = make_policy("wsh(pk(@0/**))", &[]);
-        match parse(&policy).unwrap() {
-            ParsedPolicy::Wsh(Wsh {
-                ref miniscript_expr,
-                ..
+        match &parse(&policy).unwrap().descriptor {
+            Descriptor::Wsh(Wsh {
+                miniscript_expr, ..
             }) => {
                 assert_eq!(
                     miniscript_expr.iter_pk().collect::<Vec<String>>(),
@@ -609,10 +609,9 @@ mod tests {
 
         // Parse another valid example and check that the keys are collected as is as strings.
         let policy = make_policy("wsh(or_b(pk(@0/**),s:pk(@1/**)))", &[]);
-        match parse(&policy).unwrap() {
-            ParsedPolicy::Wsh(Wsh {
-                ref miniscript_expr,
-                ..
+        match &parse(&policy).unwrap().descriptor {
+            Descriptor::Wsh(Wsh {
+                miniscript_expr, ..
             }) => {
                 assert_eq!(
                     miniscript_expr.iter_pk().collect::<Vec<String>>(),
@@ -770,8 +769,8 @@ mod tests {
         assert!(parse(&pol).unwrap().validate(coin).is_ok());
 
         // Duplicate path, one time in change, one time in receive. While the keys technically are
-        // never duplicate in the final miniscript with the pubkeys inserted, we still prohibit, as
-        // it does not look like there would be a sane use case for this and would likely be an
+        // never duplicate in the final miniscript with the pubkeys inserted, we still prohibit it,
+        // as it does not look like there would be a sane use case for this and would likely be an
         // accident.
         let pol = make_policy(
             "wsh(or_b(pk(@0/<0;1>/*),s:pk(@0/<1;2>/*)))",
diff --git a/src/rust/bitbox02-rust/src/hww/api/bitcoin/script_configs.rs b/src/rust/bitbox02-rust/src/hww/api/bitcoin/script_configs.rs
@@ -0,0 +1,31 @@
+// Copyright 2024 Shift Crypto AG
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::pb;
+use pb::btc_script_config::{Multisig, SimpleType};
+
+use super::policies::ParsedPolicy;
+
+/// Parsed and validated form of `pb::BtcScriptConfig`.
+pub enum ValidatedScriptConfig<'a> {
+    SimpleType(SimpleType),
+    Multisig(&'a Multisig),
+    Policy(ParsedPolicy<'a>),
+}
+
+/// Parsed and validated form of `pb::BtcScriptConfigWithKeypath`.
+pub struct ValidatedScriptConfigWithKeypath<'a> {
+    pub keypath: &'a [u32],
+    pub config: ValidatedScriptConfig<'a>,
+}
diff --git a/src/rust/bitbox02-rust/src/hww/api/bitcoin/signtx.rs b/src/rust/bitbox02-rust/src/hww/api/bitcoin/signtx.rs
