[keys] Take ScriptContext into account when converting keys

src/descriptor/dsl.rs

@@ -36,10 +36,11 @@ macro_rules! impl_top_level_sh {
 #[doc(hidden)]
 #[macro_export]
 macro_rules! impl_top_level_pk {
-    ( $descriptor_variant:ident, $key:expr ) => {{
+    ( $descriptor_variant:ident, $ctx:ty, $key:expr ) => {{
-        use $crate::keys::ToDescriptorKey;
+        use $crate::keys::{DescriptorKey, ToDescriptorKey};
+
         $key.to_descriptor_key()
-            .and_then(|key| key.into_key_and_secret())
+            .and_then(|key: DescriptorKey<$ctx>| key.into_key_and_secret())
             .map(|(pk, key_map)| {
                 (
                     $crate::miniscript::Descriptor::<
@@ -198,6 +199,18 @@ macro_rules! impl_node_opcode_three {
 /// }?;
 /// # Ok::<(), Box<dyn std::error::Error>>(())
 /// ```
+///
+/// ------
+///
+/// Native-Segwit single-sig, equivalent to: `wpkh(...)`
+///
+/// ```
+/// # use std::str::FromStr;
+/// let my_key = bitcoin::PrivateKey::from_wif("cVt4o7BGAig1UXywgGSmARhxMdzP5qvQsxKkSsc1XEkw3tDTQFpy")?;
+///
+/// let (descriptor, key_map) = bdk::descriptor!(wpkh ( my_key ) )?;
+/// # Ok::<(), Box<dyn std::error::Error>>(())
+/// ```
 #[macro_export]
 macro_rules! descriptor {
     ( bare ( $( $minisc:tt )* ) ) => ({
@@ -210,19 +223,19 @@ macro_rules! descriptor {
         $crate::impl_top_level_sh!(ShWsh, $( $minisc )*)
     });
     ( pk $key:expr ) => ({
-        $crate::impl_top_level_pk!(Pk, $key)
+        $crate::impl_top_level_pk!(Pk, $crate::miniscript::Legacy, $key)
     });
     ( pkh $key:expr ) => ({
-        $crate::impl_top_level_pk!(Pkh, $key)
+        $crate::impl_top_level_pk!(Pkh,$crate::miniscript::Legacy, $key)
     });
     ( wpkh $key:expr ) => ({
-        $crate::impl_top_level_pk!(Wpkh, $key)
+        $crate::impl_top_level_pk!(Wpkh, $crate::miniscript::Segwitv0, $key)
     });
     ( sh ( wpkh ( $key:expr ) ) ) => ({
         $crate::descriptor!(shwpkh ($( $minisc )*))
     });
     ( shwpkh ( $key:expr ) ) => ({
-        $crate::impl_top_level_pk!(ShWpkh, $key)
+        $crate::impl_top_level_pk!(ShWpkh, $crate::miniscript::Segwitv0, $key)
     });
     ( sh ( $( $minisc:tt )* ) ) => ({
         $crate::impl_top_level_sh!(Sh, $( $minisc )*)
@@ -279,9 +292,7 @@ macro_rules! fragment {
     });
     ( pk_k $key:expr ) => ({
         use $crate::keys::ToDescriptorKey;
-        $key.to_descriptor_key()
+        $key.into_miniscript_and_secret()
-            .and_then(|key| key.into_key_and_secret())
-            .and_then(|(pk, key_map)| Ok(($crate::impl_leaf_opcode_value!(PkK, pk)?.0, key_map)))
     });
     ( pk $key:expr ) => ({
         $crate::fragment!(+c pk_k $key)
@@ -351,21 +362,7 @@ macro_rules! fragment {
             .and_then(|items| $crate::fragment!(thresh_vec $thresh, items))
     });
     ( multi_vec $thresh:expr, $keys:expr ) => ({
-        use $crate::miniscript::descriptor::KeyMap;
+        $crate::keys::make_multi($thresh, $keys)
-        use $crate::keys::{ToDescriptorKey, DescriptorKey};
-
-        $keys.into_iter()
-            .map(|key| key.to_descriptor_key().and_then(DescriptorKey::into_key_and_secret))
-            .collect::<Result<Vec<_>, _>>()
-            .map(|items| items.into_iter().unzip())
-            .and_then(|(keys, key_maps): (Vec<_>, Vec<_>)| {
-                let key_maps = key_maps.into_iter().fold(KeyMap::default(), |mut acc, map| {
-                    acc.extend(map.into_iter());
-                    acc
-                });
-
-                Ok(($crate::impl_leaf_opcode_value_two!(Multi, $thresh, keys)?.0, key_maps))
-            })
     });
     ( multi $thresh:expr $(, $key:expr )+ ) => ({
         use $crate::keys::ToDescriptorKey;
@@ -376,7 +373,7 @@ macro_rules! fragment {
         )*
 
         keys.into_iter().collect::<Result<Vec<_>, _>>()
-            .and_then(|keys| $crate::fragment!(multi_vec $thresh, keys))
+            .and_then(|keys| $crate::keys::make_multi($thresh, keys))
     });
 
 }

src/keys/mod.rs

@@ -24,10 +24,15 @@
 
 //! Key formats
 
+use std::any::TypeId;
+use std::marker::PhantomData;
+
 use bitcoin::util::bip32;
 use bitcoin::{PrivateKey, PublicKey};
 
 use miniscript::descriptor::{DescriptorPublicKey, DescriptorSecretKey, DescriptorXKey, KeyMap};
+pub use miniscript::ScriptContext;
+use miniscript::{Miniscript, Terminal};
 
 use crate::Error;
 
@@ -36,17 +41,20 @@ use crate::Error;
 pub mod bip39;
 
 /// Container for public or secret keys
-pub enum DescriptorKey {
+pub enum DescriptorKey<Ctx: ScriptContext> {
-    Public(DescriptorPublicKey),
+    Public(DescriptorPublicKey, PhantomData<Ctx>),
-    Secret(DescriptorSecretKey),
+    Secret(DescriptorSecretKey, PhantomData<Ctx>),
 }
 
-impl DescriptorKey {
+impl<Ctx: ScriptContext> DescriptorKey<Ctx> {
+    // This method is used internally by `bdk::fragment!` and `bdk::descriptor!`. It has to be
+    // public because it is effectively called by external crates, once the macros are expanded,
+    // but since it is not meant to be part of the public api we hide it from the docs.
     #[doc(hidden)]
     pub fn into_key_and_secret(self) -> Result<(DescriptorPublicKey, KeyMap), Error> {
         match self {
-            DescriptorKey::Public(public) => Ok((public, KeyMap::default())),
+            DescriptorKey::Public(public, _) => Ok((public, KeyMap::default())),
-            DescriptorKey::Secret(secret) => {
+            DescriptorKey::Secret(secret, _) => {
                 let mut key_map = KeyMap::with_capacity(1);
 
                 let public = secret
@@ -60,64 +68,240 @@ impl DescriptorKey {
     }
 }
 
-/// Trait for objects that can be turned into a public or secret [`DescriptorKey`]
+#[derive(Debug, Eq, PartialEq, Copy, Clone)]
-pub trait ToDescriptorKey {
+pub enum ScriptContextEnum {
-    fn to_descriptor_key(self) -> Result<DescriptorKey, Error>;
+    Legacy,
+    Segwitv0,
 }
 
-/// Identity conversion. This is used internally by [`bdk::fragment`]
+impl ScriptContextEnum {
-impl ToDescriptorKey for DescriptorKey {
+    pub fn is_legacy(&self) -> bool {
-    fn to_descriptor_key(self) -> Result<DescriptorKey, Error> {
+        self == &ScriptContextEnum::Legacy
+    }
+
+    pub fn is_segwit_v0(&self) -> bool {
+        self == &ScriptContextEnum::Segwitv0
+    }
+}
+
+pub trait ExtScriptContext: ScriptContext {
+    fn as_enum() -> ScriptContextEnum;
+
+    fn is_legacy() -> bool {
+        Self::as_enum().is_legacy()
+    }
+
+    fn is_segwit_v0() -> bool {
+        Self::as_enum().is_segwit_v0()
+    }
+}
+
+impl<Ctx: ScriptContext + 'static> ExtScriptContext for Ctx {
+    fn as_enum() -> ScriptContextEnum {
+        match TypeId::of::<Ctx>() {
+            t if t == TypeId::of::<miniscript::Legacy>() => ScriptContextEnum::Legacy,
+            t if t == TypeId::of::<miniscript::Segwitv0>() => ScriptContextEnum::Segwitv0,
+            _ => unimplemented!("Unknown ScriptContext type"),
+        }
+    }
+}
+
+/// Trait for objects that can be turned into a public or secret [`DescriptorKey`]
+///
+/// The generic type `Ctx` is used to define the context in which the key is valid: some key
+/// formats, like the mnemonics used by Electrum wallets, encode internally whether the wallet is
+/// legacy or segwit. Thus, trying to turn a valid legacy mnemonic into a `DescriptorKey`
+/// that would become part of a segwit descriptor should fail.
+///
+/// For key types that do care about this, the [`ExtScriptContext`] trait provides some useful
+/// methods that can be used to check at runtime which `Ctx` is being used.
+///
+/// For key types that that do not need to check this at runtime (because they can only work within a
+/// single `Ctx`), the "specialized" trait can be implemented to make the compiler handle the type
+/// checking.
+///
+/// ## Examples
+///
+/// Key type valid in any context:
+///
+/// ```
+/// use bdk::bitcoin::PublicKey;
+///
+/// use bdk::keys::{ScriptContext, ToDescriptorKey, DescriptorKey};
+/// use bdk::Error;
+///
+/// pub struct MyKeyType {
+///     pubkey: PublicKey,
+/// }
+///
+/// impl<Ctx: ScriptContext> ToDescriptorKey<Ctx> for MyKeyType {
+///     fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
+///         self.pubkey.to_descriptor_key()
+///     }
+/// }
+/// ```
+///
+/// Key type that internally encodes in which context it's valid. The context is checked at runtime:
+///
+/// ```
+/// use bdk::bitcoin::PublicKey;
+///
+/// use bdk::keys::{ExtScriptContext, ScriptContext, ToDescriptorKey, DescriptorKey};
+/// use bdk::Error;
+///
+/// pub struct MyKeyType {
+///     is_legacy: bool,
+///     pubkey: PublicKey,
+/// }
+///
+/// impl<Ctx: ScriptContext + 'static> ToDescriptorKey<Ctx> for MyKeyType {
+///     fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
+///         if Ctx::is_legacy() == self.is_legacy {
+///             self.pubkey.to_descriptor_key()
+///         } else {
+///             Err(Error::Generic("Invalid key context".into()))
+///         }
+///     }
+/// }
+/// ```
+///
+/// Key type that can only work within [`miniscript::Segwitv0`] context. Only the specialized version
+/// of the trait is implemented.
+///
+/// This example deliberately fails to compile, to demonstrate how the compiler can catch when keys
+/// are misused. In this case, the "segwit-only" key is used to build a `pkh()` descriptor, which
+/// makes the compiler (correctly) fail.
+///
+/// ```compile_fail
+/// use std::str::FromStr;
+/// use bdk::bitcoin::PublicKey;
+///
+/// use bdk::keys::{ToDescriptorKey, DescriptorKey};
+/// use bdk::Error;
+///
+/// pub struct MySegwitOnlyKeyType {
+///     pubkey: PublicKey,
+/// }
+///
+/// impl ToDescriptorKey<bdk::miniscript::Segwitv0> for MySegwitOnlyKeyType {
+///     fn to_descriptor_key(self) -> Result<DescriptorKey<bdk::miniscript::Segwitv0>, Error> {
+///         self.pubkey.to_descriptor_key()
+///     }
+/// }
+///
+/// let key = MySegwitOnlyKeyType {
+///     pubkey: PublicKey::from_str("...")?,
+/// };
+/// let (descriptor, _) = bdk::descriptor!(pkh ( key ) )?;
+/// //                                    ^^^^^ changing this to `wpkh` would make it compile
+///
+/// # Ok::<_, Box<dyn std::error::Error>>(())
+/// ```
+pub trait ToDescriptorKey<Ctx: ScriptContext>: Sized {
+    /// Turn the key into a [`DescriptorKey`] within the requested [`ScriptContext`]
+    fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error>;
+
+    // Used internally by `bdk::fragment!` to build `pk_k()` fragments
+    #[doc(hidden)]
+    fn into_miniscript_and_secret(
+        self,
+    ) -> Result<(Miniscript<DescriptorPublicKey, Ctx>, KeyMap), Error> {
+        let descriptor_key = self.to_descriptor_key()?;
+        let (key, key_map) = descriptor_key.into_key_and_secret()?;
+
+        Ok((Miniscript::from_ast(Terminal::PkK(key))?, key_map))
+    }
+}
+
+// Used internally by `bdk::fragment!` to build `multi()` fragments
+#[doc(hidden)]
+pub fn make_multi<Pk: ToDescriptorKey<Ctx>, Ctx: ScriptContext>(
+    thresh: usize,
+    pks: Vec<Pk>,
+) -> Result<(Miniscript<DescriptorPublicKey, Ctx>, KeyMap), Error> {
+    let (pks, key_maps): (Vec<_>, Vec<_>) = pks
+        .into_iter()
+        .map(|key| {
+            key.to_descriptor_key()
+                .and_then(DescriptorKey::into_key_and_secret)
+        })
+        .collect::<Result<Vec<_>, _>>()?
+        .into_iter()
+        .unzip();
+
+    let key_map = key_maps
+        .into_iter()
+        .fold(KeyMap::default(), |mut acc, map| {
+            acc.extend(map.into_iter());
+            acc
+        });
+
+    Ok((Miniscript::from_ast(Terminal::Multi(thresh, pks))?, key_map))
+}
+
+/// The "identity" conversion is used internally by some `bdk::fragment`s
+impl<Ctx: ScriptContext> ToDescriptorKey<Ctx> for DescriptorKey<Ctx> {
+    fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
         Ok(self)
     }
 }
 
-impl ToDescriptorKey for DescriptorPublicKey {
+impl<Ctx: ScriptContext> ToDescriptorKey<Ctx> for DescriptorPublicKey {
-    fn to_descriptor_key(self) -> Result<DescriptorKey, Error> {
+    fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
-        Ok(DescriptorKey::Public(self))
+        Ok(DescriptorKey::Public(self, PhantomData))
     }
 }
 
-impl ToDescriptorKey for PublicKey {
+impl<Ctx: ScriptContext> ToDescriptorKey<Ctx> for PublicKey {
-    fn to_descriptor_key(self) -> Result<DescriptorKey, Error> {
+    fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
-        Ok(DescriptorKey::Public(DescriptorPublicKey::PubKey(self)))
+        Ok(DescriptorKey::Public(
+            DescriptorPublicKey::PubKey(self),
+            PhantomData,
+        ))
     }
 }
 
-impl ToDescriptorKey for (bip32::ExtendedPubKey, bip32::DerivationPath) {
+/// This assumes that "is_wildcard" is true, since this is generally the way extended keys are used
-    fn to_descriptor_key(self) -> Result<DescriptorKey, Error> {
+impl<Ctx: ScriptContext> ToDescriptorKey<Ctx> for (bip32::ExtendedPubKey, bip32::DerivationPath) {
-        Ok(DescriptorKey::Public(DescriptorPublicKey::XPub(
+    fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
-            DescriptorXKey {
+        Ok(DescriptorKey::Public(
+            DescriptorPublicKey::XPub(DescriptorXKey {
                 source: None,
                 xkey: self.0,
                 derivation_path: self.1,
                 is_wildcard: true,
-            },
+            }),
-        )))
+            PhantomData,
+        ))
     }
 }
 
-impl ToDescriptorKey for DescriptorSecretKey {
+impl<Ctx: ScriptContext> ToDescriptorKey<Ctx> for DescriptorSecretKey {
-    fn to_descriptor_key(self) -> Result<DescriptorKey, Error> {
+    fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
-        Ok(DescriptorKey::Secret(self))
+        Ok(DescriptorKey::Secret(self, PhantomData))
     }
 }
 
-impl ToDescriptorKey for PrivateKey {
+impl<Ctx: ScriptContext> ToDescriptorKey<Ctx> for PrivateKey {
-    fn to_descriptor_key(self) -> Result<DescriptorKey, Error> {
+    fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
-        Ok(DescriptorKey::Secret(DescriptorSecretKey::PrivKey(self)))
+        Ok(DescriptorKey::Secret(
+            DescriptorSecretKey::PrivKey(self),
+            PhantomData,
+        ))
     }
 }
 
-impl ToDescriptorKey for (bip32::ExtendedPrivKey, bip32::DerivationPath) {
+/// This assumes that "is_wildcard" is true, since this is generally the way extended keys are used
-    fn to_descriptor_key(self) -> Result<DescriptorKey, Error> {
+impl<Ctx: ScriptContext> ToDescriptorKey<Ctx> for (bip32::ExtendedPrivKey, bip32::DerivationPath) {
-        Ok(DescriptorKey::Secret(DescriptorSecretKey::XPrv(
+    fn to_descriptor_key(self) -> Result<DescriptorKey<Ctx>, Error> {
-            DescriptorXKey {
+        Ok(DescriptorKey::Secret(
+            DescriptorSecretKey::XPrv(DescriptorXKey {
                 source: None,
                 xkey: self.0,
                 derivation_path: self.1,
                 is_wildcard: true,
-            },
+            }),
-        )))
+            PhantomData,
+        ))
     }
 }