[wallet] Add AddressValidators

Cargo.toml

@@ -66,10 +66,9 @@ rand = "0.7"
 name = "repl"
 required-features = ["cli-utils"]
 [[example]]
-name = "psbt"
-required-features = ["cli-utils"]
-[[example]]
 name = "parse_descriptor"
+[[example]]
+name = "address_validator"
 
 [[example]]
 name = "miniscriptc"

examples/address_validator.rs

@@ -0,0 +1,48 @@
+use std::sync::Arc;
+
+use magical_bitcoin_wallet::bitcoin;
+use magical_bitcoin_wallet::database::MemoryDatabase;
+use magical_bitcoin_wallet::descriptor::HDKeyPaths;
+use magical_bitcoin_wallet::types::ScriptType;
+use magical_bitcoin_wallet::wallet::address_validator::{AddressValidator, AddressValidatorError};
+use magical_bitcoin_wallet::{OfflineWallet, Wallet};
+
+use bitcoin::hashes::hex::FromHex;
+use bitcoin::util::bip32::Fingerprint;
+use bitcoin::{Network, Script};
+
+struct DummyValidator;
+impl AddressValidator for DummyValidator {
+    fn validate(
+        &self,
+        script_type: ScriptType,
+        hd_keypaths: &HDKeyPaths,
+        script: &Script,
+    ) -> Result<(), AddressValidatorError> {
+        let (_, path) = hd_keypaths
+            .values()
+            .find(|(fing, _)| fing == &Fingerprint::from_hex("bc123c3e").unwrap())
+            .ok_or(AddressValidatorError::InvalidScript)?;
+
+        println!(
+            "Validating `{:?}` {} address, script: {}",
+            script_type, path, script
+        );
+
+        Ok(())
+    }
+}
+
+fn main() -> Result<(), magical_bitcoin_wallet::error::Error> {
+    let descriptor = "sh(and_v(v:pk(tpubDDpWvmUrPZrhSPmUzCMBHffvC3HyMAPnWDSAQNBTnj1iZeJa7BZQEttFiP4DS4GCcXQHezdXhn86Hj6LHX5EDstXPWrMaSneRWM8yUf6NFd/*),after(630000)))";
+    let mut wallet: OfflineWallet<_> =
+        Wallet::new_offline(descriptor, None, Network::Regtest, MemoryDatabase::new())?;
+
+    wallet.add_address_validator(Arc::new(Box::new(DummyValidator)));
+
+    wallet.get_new_address()?;
+    wallet.get_new_address()?;
+    wallet.get_new_address()?;
+
+    Ok(())
+}

examples/psbt.rs

@@ -1,50 +0,0 @@
-extern crate base64;
-extern crate magical_bitcoin_wallet;
-
-use std::str::FromStr;
-
-use magical_bitcoin_wallet::bitcoin;
-use magical_bitcoin_wallet::descriptor::*;
-use magical_bitcoin_wallet::psbt::*;
-use magical_bitcoin_wallet::signer::Signer;
-
-use bitcoin::consensus::encode::{deserialize, serialize};
-use bitcoin::util::psbt::PartiallySignedTransaction;
-use bitcoin::SigHashType;
-
-fn main() {
-    let desc = "pkh(tprv8ZgxMBicQKsPd7Uf69XL1XwhmjHopUGep8GuEiJDZmbQz6o58LninorQAfcKZWARbtRtfnLcJ5MQ2AtHcQJCCRUcMRvmDUjyEmNUWwx8UbK/*)";
-
-    let extended_desc = ExtendedDescriptor::from_str(desc).unwrap();
-
-    let psbt_str = "cHNidP8BAFMCAAAAAd9SiQfxXZ+CKjgjRNonWXsnlA84aLvjxtwCmMfRc0ZbAQAAAAD+////ASjS9QUAAAAAF6kUYJR3oB0lS1M0W1RRMMiENSX45IuHAAAAAAABAPUCAAAAA9I7/OqeFeOFdr5VTLnj3UI/CNRw2eWmMPf7qDv6uIF6AAAAABcWABTG+kgr0g44V0sK9/9FN9oG/CxMK/7///+d0ffphPcV6FE9J/3ZPKWu17YxBnWWTJQyRJs3HUo1gwEAAAAA/v///835mYd9DmnjVnUKd2421MDoZmIxvB4XyJluN3SPUV9hAAAAABcWABRfvwFGp+x/yWdXeNgFs9v0duyeS/7///8CFbH+AAAAAAAXqRSEnTOAjJN/X6ZgR9ftKmwisNSZx4cA4fUFAAAAABl2qRTs6pS4x17MSQ4yNs/1GPsfdlv2NIisAAAAACIGApVE9PPtkcqp8Da43yrXGv4nLOotZdyxwJoTWQxuLxIuCAxfmh4JAAAAAAA=";
-    let psbt_buf = base64::decode(psbt_str).unwrap();
-    let mut psbt: PartiallySignedTransaction = deserialize(&psbt_buf).unwrap();
-
-    let signer = PSBTSigner::from_descriptor(&psbt.global.unsigned_tx, &extended_desc).unwrap();
-
-    for (index, input) in psbt.inputs.iter_mut().enumerate() {
-        for (pubkey, (fing, path)) in &input.hd_keypaths {
-            let sighash = input.sighash_type.unwrap_or(SigHashType::All);
-
-            // Ignore the "witness_utxo" case because we know this psbt is a legacy tx
-            if let Some(non_wit_utxo) = &input.non_witness_utxo {
-                let prev_script = &non_wit_utxo.output
-                    [psbt.global.unsigned_tx.input[index].previous_output.vout as usize]
-                    .script_pubkey;
-                let (signature, sighash) = signer
-                    .sig_legacy_from_fingerprint(index, sighash, fing, path, prev_script)
-                    .unwrap()
-                    .unwrap();
-
-                let mut concat_sig = Vec::new();
-                concat_sig.extend_from_slice(&signature.serialize_der());
-                concat_sig.extend_from_slice(&[sighash as u8]);
-
-                input.partial_sigs.insert(*pubkey, concat_sig);
-            }
-        }
-    }
-
-    println!("signed: {}", base64::encode(&serialize(&psbt)));
-}

src/descriptor/mod.rs

@@ -17,15 +17,13 @@ pub mod checksum;
 pub mod error;
 pub mod policy;
 
-// use crate::wallet::utils::AddressType;
-use crate::wallet::signer::SignersContainer;
-
 pub use self::checksum::get_checksum;
 use self::error::Error;
 pub use self::policy::Policy;
+use crate::wallet::signer::SignersContainer;
 
 pub type ExtendedDescriptor = Descriptor<DescriptorPublicKey>;
-type HDKeyPaths = BTreeMap<PublicKey, (Fingerprint, DerivationPath)>;
+pub type HDKeyPaths = BTreeMap<PublicKey, (Fingerprint, DerivationPath)>;
 
 pub trait ExtractPolicy {
     fn extract_policy(
@@ -76,7 +74,6 @@ pub trait DescriptorMeta: Sized {
     fn derive_from_hd_keypaths(&self, hd_keypaths: &HDKeyPaths) -> Option<Self>;
     fn derive_from_psbt_input(&self, psbt_input: &psbt::Input, utxo: Option<TxOut>)
         -> Option<Self>;
-    // fn address_type(&self) -> Option<AddressType>;
 }
 
 pub trait DescriptorScripts {
@@ -258,18 +255,6 @@ impl DescriptorMeta for Descriptor<DescriptorPublicKey> {
             _ => None,
         }
     }
-
-    // fn address_type(&self) -> Option<AddressType> {
-    //     match self {
-    //         Descriptor::Pkh(_) => Some(AddressType::Pkh),
-    //         Descriptor::Wpkh(_) => Some(AddressType::Wpkh),
-    //         Descriptor::ShWpkh(_) => Some(AddressType::ShWpkh),
-    //         Descriptor::Sh(_) => Some(AddressType::Sh),
-    //         Descriptor::Wsh(_) => Some(AddressType::Wsh),
-    //         Descriptor::ShWsh(_) => Some(AddressType::ShWsh),
-    //         _ => None,
-    //     }
-    // }
 }
 
 #[derive(Debug, Clone, Hash, PartialEq, PartialOrd, Eq, Ord, Default)]

src/error.rs

@@ -36,6 +36,7 @@ pub enum Error {
     InvalidOutpoint(OutPoint),
 
     Descriptor(crate::descriptor::error::Error),
+    AddressValidator(crate::wallet::address_validator::AddressValidatorError),
 
     Encode(bitcoin::consensus::encode::Error),
     Miniscript(miniscript::Error),
@@ -66,6 +67,10 @@ macro_rules! impl_error {
 }
 
 impl_error!(crate::descriptor::error::Error, Descriptor);
+impl_error!(
+    crate::wallet::address_validator::AddressValidatorError,
+    AddressValidator
+);
 impl_error!(
     crate::descriptor::policy::PolicyError,
     InvalidPolicyPathError

src/wallet/address_validator.rs

@@ -0,0 +1,63 @@
+use bitcoin::Script;
+
+use crate::descriptor::HDKeyPaths;
+use crate::types::ScriptType;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum AddressValidatorError {
+    UserRejected,
+    ConnectionError,
+    TimeoutError,
+    InvalidScript,
+}
+
+pub trait AddressValidator {
+    fn validate(
+        &self,
+        script_type: ScriptType,
+        hd_keypaths: &HDKeyPaths,
+        script: &Script,
+    ) -> Result<(), AddressValidatorError>;
+}
+
+#[cfg(test)]
+mod test {
+    use std::sync::Arc;
+
+    use super::*;
+    use crate::wallet::test::{get_funded_wallet, get_test_wpkh};
+    use crate::wallet::TxBuilder;
+
+    struct TestValidator;
+    impl AddressValidator for TestValidator {
+        fn validate(
+            &self,
+            _script_type: ScriptType,
+            _hd_keypaths: &HDKeyPaths,
+            _script: &bitcoin::Script,
+        ) -> Result<(), AddressValidatorError> {
+            Err(AddressValidatorError::InvalidScript)
+        }
+    }
+
+    #[test]
+    #[should_panic(expected = "InvalidScript")]
+    fn test_address_validator_external() {
+        let (mut wallet, _, _) = get_funded_wallet(get_test_wpkh());
+        wallet.add_address_validator(Arc::new(Box::new(TestValidator)));
+
+        wallet.get_new_address().unwrap();
+    }
+
+    #[test]
+    #[should_panic(expected = "InvalidScript")]
+    fn test_address_validator_internal() {
+        let (mut wallet, descriptors, _) = get_funded_wallet(get_test_wpkh());
+        wallet.add_address_validator(Arc::new(Box::new(TestValidator)));
+
+        let addr = testutils!(@external descriptors, 10);
+        wallet
+            .create_tx(TxBuilder::from_addressees(vec![(addr, 25_000)]))
+            .unwrap();
+    }
+}

src/wallet/mod.rs

@@ -14,6 +14,7 @@ use miniscript::descriptor::DescriptorPublicKey;
 #[allow(unused_imports)]
 use log::{debug, error, info, trace};
 
+pub mod address_validator;
 pub mod coin_selection;
 pub mod export;
 mod rbf;
@@ -22,7 +23,8 @@ pub mod time;
 pub mod tx_builder;
 pub mod utils;
 
-use signer::{Signer, SignersContainer};
+use address_validator::AddressValidator;
+use signer::{Signer, SignerId, SignersContainer};
 use tx_builder::TxBuilder;
 use utils::{After, FeeRate, IsDust, Older};
 
@@ -33,7 +35,6 @@ use crate::descriptor::{
 };
 use crate::error::Error;
 use crate::psbt::PSBTUtils;
-// use crate::psbt::{utils::PSBTUtils, PSBTSatisfier, PSBTSigner};
 use crate::types::*;
 
 const CACHE_ADDR_BATCH_SIZE: u32 = 100;
@@ -47,6 +48,8 @@ pub struct Wallet<B: Blockchain, D: BatchDatabase> {
     signers: Arc<SignersContainer<DescriptorPublicKey>>,
     change_signers: Arc<SignersContainer<DescriptorPublicKey>>,
 
+    address_validators: Vec<Arc<Box<dyn AddressValidator>>>,
+
     network: Network,
 
     current_height: Option<u32>,
@@ -96,6 +99,7 @@ where
             change_descriptor,
             signers,
             change_signers,
+            address_validators: Vec::new(),
 
             network,
 
@@ -134,6 +138,24 @@ where
             .fold(0, |sum, i| sum + i.txout.value))
     }
 
+    pub fn add_signer(
+        &mut self,
+        script_type: ScriptType,
+        id: SignerId<DescriptorPublicKey>,
+        signer: Arc<Box<dyn Signer>>,
+    ) {
+        let signers = match script_type {
+            ScriptType::External => Arc::make_mut(&mut self.signers),
+            ScriptType::Internal => Arc::make_mut(&mut self.change_signers),
+        };
+
+        signers.add_external(id, signer);
+    }
+
+    pub fn add_address_validator(&mut self, validator: Arc<Box<dyn AddressValidator>>) {
+        self.address_validators.push(validator);
+    }
+
     pub fn create_tx<Cs: coin_selection::CoinSelectionAlgorithm>(
         &self,
         builder: TxBuilder<Cs>,
@@ -724,6 +746,14 @@ where
             self.cache_addresses(script_type, index, CACHE_ADDR_BATCH_SIZE)?;
         }
 
+        let hd_keypaths = descriptor.get_hd_keypaths(index)?;
+        let script = descriptor
+            .derive(&[ChildNumber::from_normal_idx(index).unwrap()])
+            .script_pubkey();
+        for validator in &self.address_validators {
+            validator.validate(script_type, &hd_keypaths, &script)?;
+        }
+
         Ok(index)
     }
 
@@ -1103,21 +1133,21 @@ mod test {
             .is_some());
     }
 
-    fn get_test_wpkh() -> &'static str {
+    pub(crate) fn get_test_wpkh() -> &'static str {
         "wpkh(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW)"
     }
 
-    fn get_test_single_sig_csv() -> &'static str {
+    pub(crate) fn get_test_single_sig_csv() -> &'static str {
         // and(pk(Alice),older(6))
         "wsh(and_v(v:pk(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW),older(6)))"
     }
 
-    fn get_test_single_sig_cltv() -> &'static str {
+    pub(crate) fn get_test_single_sig_cltv() -> &'static str {
         // and(pk(Alice),after(100000))
         "wsh(and_v(v:pk(cVpPVruEDdmutPzisEsYvtST1usBR3ntr8pXSyt6D2YYqXRyPcFW),after(100000)))"
     }
 
-    fn get_funded_wallet(
+    pub(crate) fn get_funded_wallet(
         descriptor: &str,
     ) -> (
         OfflineWallet<MemoryDatabase>,

src/wallet/signer.rs

@@ -1,6 +1,7 @@
 use std::any::Any;
 use std::collections::HashMap;
 use std::fmt;
+use std::sync::Arc;
 
 use bitcoin::blockdata::opcodes;
 use bitcoin::blockdata::script::Builder as ScriptBuilder;
@@ -17,7 +18,7 @@ use crate::descriptor::XKeyUtils;
 
 /// Identifier of a signer in the `SignersContainers`. Used as a key to find the right signer among
 /// many of them
-#[derive(Debug, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum SignerId<Pk: MiniscriptKey> {
     PkHash(<Pk as MiniscriptKey>::Hash),
     Fingerprint(Fingerprint),
@@ -150,8 +151,8 @@ impl Signer for PrivateKey {
 }
 
 /// Container for multiple signers
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
-pub struct SignersContainer<Pk: MiniscriptKey>(HashMap<SignerId<Pk>, Box<dyn Signer>>);
+pub struct SignersContainer<Pk: MiniscriptKey>(HashMap<SignerId<Pk>, Arc<Box<dyn Signer>>>);
 
 impl SignersContainer<DescriptorPublicKey> {
     pub fn as_key_map(&self) -> KeyMap {
@@ -189,11 +190,12 @@ impl From<KeyMap> for SignersContainer<DescriptorPublicKey> {
                             .public_key(&Secp256k1::signing_only())
                             .to_pubkeyhash(),
                     ),
-                    Box::new(private_key),
+                    Arc::new(Box::new(private_key)),
+                ),
+                DescriptorSecretKey::XPrv(xprv) => container.add_external(
+                    SignerId::from(xprv.root_fingerprint()),
+                    Arc::new(Box::new(xprv)),
                 ),
-                DescriptorSecretKey::XPrv(xprv) => {
-                    container.add_external(SignerId::from(xprv.root_fingerprint()), Box::new(xprv))
-                }
             };
         }
 
@@ -212,13 +214,13 @@ impl<Pk: MiniscriptKey> SignersContainer<Pk> {
     pub fn add_external(
         &mut self,
         id: SignerId<Pk>,
-        signer: Box<dyn Signer>,
+        signer: Arc<Box<dyn Signer>>,
-    ) -> Option<Box<dyn Signer>> {
+    ) -> Option<Arc<Box<dyn Signer>>> {
         self.0.insert(id, signer)
     }
 
     /// Removes a signer from the container and returns it
-    pub fn remove(&mut self, id: SignerId<Pk>) -> Option<Box<dyn Signer>> {
+    pub fn remove(&mut self, id: SignerId<Pk>) -> Option<Arc<Box<dyn Signer>>> {
         self.0.remove(&id)
     }
 
@@ -228,7 +230,7 @@ impl<Pk: MiniscriptKey> SignersContainer<Pk> {
     }
 
     /// Finds the signer with a given id in the container
-    pub fn find(&self, id: SignerId<Pk>) -> Option<&Box<dyn Signer>> {
+    pub fn find(&self, id: SignerId<Pk>) -> Option<&Arc<Box<dyn Signer>>> {
         self.0.get(&id)
     }
 }