[keys] Add a way to restrict the networks in which keys are valid

Thanks to the `ToWalletDescriptor` trait we can also very easily validate the checksum
for descriptors that are loaded from strings, if they contain one. Fixes #20.

src/descriptor/mod.rs

@@ -35,7 +35,7 @@ use bitcoin::hashes::hash160;
 use bitcoin::secp256k1::Secp256k1;
 use bitcoin::util::bip32::{ChildNumber, DerivationPath, Fingerprint};
 use bitcoin::util::psbt;
-use bitcoin::{PublicKey, Script, TxOut};
+use bitcoin::{Network, PublicKey, Script, TxOut};
 
 use miniscript::descriptor::{DescriptorPublicKey, DescriptorXKey, InnerXKey, KeyMap};
 pub use miniscript::{
@@ -50,6 +50,7 @@ pub mod policy;
 pub use self::checksum::get_checksum;
 use self::error::Error;
 pub use self::policy::Policy;
+use crate::keys::{KeyError, ToDescriptorKey, ValidNetworks};
 use crate::wallet::signer::SignersContainer;
 
 /// Alias for a [`Descriptor`] that can contain extended keys using [`DescriptorPublicKey`]
@@ -62,32 +63,124 @@ pub type ExtendedDescriptor = Descriptor<DescriptorPublicKey>;
 /// [`psbt::Output`]: bitcoin::util::psbt::Output
 pub type HDKeyPaths = BTreeMap<PublicKey, (Fingerprint, DerivationPath)>;
 
-/// Trait for types which can be converted into an [`ExtendedDescriptor`] and a [`KeyMap`] usable by a wallet
+/// Trait for types which can be converted into an [`ExtendedDescriptor`] and a [`KeyMap`] usable by a wallet in a specific [`Network`]
 pub trait ToWalletDescriptor {
-    fn to_wallet_descriptor(self) -> Result<(ExtendedDescriptor, KeyMap), Error>;
+    fn to_wallet_descriptor(
+        self,
+        network: Network,
+    ) -> Result<(ExtendedDescriptor, KeyMap), KeyError>;
 }
 
 impl ToWalletDescriptor for &str {
-    fn to_wallet_descriptor(self) -> Result<(ExtendedDescriptor, KeyMap), Error> {
-        Ok(ExtendedDescriptor::parse_secret(self)?)
+    fn to_wallet_descriptor(
+        self,
+        network: Network,
+    ) -> Result<(ExtendedDescriptor, KeyMap), KeyError> {
+        let descriptor = if self.contains("#") {
+            let parts: Vec<&str> = self.splitn(2, "#").collect();
+            if !get_checksum(parts[0])
+                .ok()
+                .map(|computed| computed == parts[1])
+                .unwrap_or(false)
+            {
+                return Err(KeyError::InvalidChecksum);
+            }
+
+            parts[0]
+        } else {
+            self
+        };
+
+        ExtendedDescriptor::parse_secret(descriptor)?.to_wallet_descriptor(network)
     }
 }
 
 impl ToWalletDescriptor for &String {
-    fn to_wallet_descriptor(self) -> Result<(ExtendedDescriptor, KeyMap), Error> {
-        self.as_str().to_wallet_descriptor()
-    }
-}
-
-impl ToWalletDescriptor for (ExtendedDescriptor, KeyMap) {
-    fn to_wallet_descriptor(self) -> Result<(ExtendedDescriptor, KeyMap), Error> {
-        Ok(self)
+    fn to_wallet_descriptor(
+        self,
+        network: Network,
+    ) -> Result<(ExtendedDescriptor, KeyMap), KeyError> {
+        self.as_str().to_wallet_descriptor(network)
     }
 }
 
 impl ToWalletDescriptor for ExtendedDescriptor {
-    fn to_wallet_descriptor(self) -> Result<(ExtendedDescriptor, KeyMap), Error> {
-        (self, KeyMap::default()).to_wallet_descriptor()
+    fn to_wallet_descriptor(
+        self,
+        network: Network,
+    ) -> Result<(ExtendedDescriptor, KeyMap), KeyError> {
+        (self, KeyMap::default()).to_wallet_descriptor(network)
+    }
+}
+
+impl ToWalletDescriptor for (ExtendedDescriptor, KeyMap) {
+    fn to_wallet_descriptor(
+        self,
+        network: Network,
+    ) -> Result<(ExtendedDescriptor, KeyMap), KeyError> {
+        use crate::keys::DescriptorKey;
+
+        // check the network for the keys
+        let translated = self.0.translate_pk(
+            |pk| {
+                let (pk, _, networks) = if self.0.is_witness() {
+                    let desciptor_key: DescriptorKey<miniscript::Segwitv0> =
+                        pk.clone().to_descriptor_key()?;
+                    desciptor_key.extract()?
+                } else {
+                    let desciptor_key: DescriptorKey<miniscript::Legacy> =
+                        pk.clone().to_descriptor_key()?;
+                    desciptor_key.extract()?
+                };
+
+                if networks.contains(&network) {
+                    Ok(pk)
+                } else {
+                    Err(KeyError::InvalidNetwork)
+                }
+            },
+            |pkh| Ok::<_, KeyError>(*pkh),
+        )?;
+
+        Ok((translated, self.1))
+    }
+}
+
+impl ToWalletDescriptor for (ExtendedDescriptor, KeyMap, ValidNetworks) {
+    fn to_wallet_descriptor(
+        self,
+        network: Network,
+    ) -> Result<(ExtendedDescriptor, KeyMap), KeyError> {
+        let valid_networks = &self.2;
+
+        // fixup the network for keys that need it
+        let translated = self.0.translate_pk(
+            |pk| {
+                if valid_networks.contains(&network) {
+                    // workaround for xpubs generated by other key types, like bip39: since when the
+                    // conversion is made one network has to be chosen, what we generally choose
+                    // "mainnet", but then override the set of valid networks to specify that all of
+                    // them are valid. here we reset the network to make sure the wallet struct gets a
+                    // descriptor with the right network everywhere.
+                    let pk = match pk {
+                        DescriptorPublicKey::XPub(ref xpub) => {
+                            let mut xpub = xpub.clone();
+                            xpub.xkey.network = network;
+
+                            DescriptorPublicKey::XPub(xpub)
+                        }
+                        other @ _ => other.clone(),
+                    };
+
+                    Ok(pk)
+                } else {
+                    Err(KeyError::InvalidNetwork)
+                }
+            },
+            |pkh| Ok::<_, KeyError>(*pkh),
+        )?;
+
+        Ok((translated, self.1))
     }
 }
 
@@ -355,7 +448,7 @@ mod test {
 
     use bitcoin::consensus::encode::deserialize;
     use bitcoin::hashes::hex::FromHex;
-    use bitcoin::util::psbt;
+    use bitcoin::util::{bip32, psbt};
 
     use super::*;
     use crate::psbt::PSBTUtils;
@@ -467,4 +560,25 @@ mod test {
             .derive_from_psbt_input(&psbt.inputs[0], psbt.get_utxo_for(0))
             .is_some());
     }
+
+    #[test]
+    fn test_to_wallet_descriptor_fixup_networks() {
+        use crate::keys::{any_network, ToDescriptorKey};
+
+        let xpub = bip32::ExtendedPubKey::from_str("xpub6ERApfZwUNrhLCkDtcHTcxd75RbzS1ed54G1LkBUHQVHQKqhMkhgbmJbZRkrgZw4koxb5JaHWkY4ALHY2grBGRjaDMzQLcgJvLJuZZvRcEL").unwrap();
+        let path = bip32::DerivationPath::from_str("m/0").unwrap();
+
+        // here `to_descriptor_key` will set the valid networks for the key to only mainnet, since
+        // we are using an "xpub"
+        let key = (xpub, path).to_descriptor_key().unwrap();
+        // override it with any. this happens in some key conversions, like bip39
+        let key = key.override_valid_networks(any_network());
+
+        // make a descriptor out of it
+        let desc = crate::descriptor!(wpkh(key)).unwrap();
+        // this should conver the key that supports "any_network" to the right network (testnet)
+        let (wallet_desc, _) = desc.to_wallet_descriptor(Network::Testnet).unwrap();
+
+        assert_eq!(wallet_desc.to_string(), "wpkh(tpubDEnoLuPdBep9bzw5LoGYpsxUQYheRQ9gcgrJhJEcdKFB9cWQRyYmkCyRoTqeD4tJYiVVgt6A3rN6rWn9RYhR9sBsGxji29LYWHuKKbdb1ev/0/*)");
+    }
 }