P2TR address from untweaked public key

nlanson · nlanson · commit 803b5fed8a5f · 2021-11-11T14:36:56.000+11:00
Ambiguous TweakedPublicKey and UntweakedPublicKey type aliases and methods to convert

Use structs for Untweaked and Tweaked key type

swap dangerous api to work on tweaked keys

remove unecessary allocations and rename methods

Use type alias for UntweakedPublicKey

TweakedPublicKey::new(...) method added

minor naming and doc changes
diff --git a/src/util/address.rs b/src/util/address.rs
@@ -40,7 +40,7 @@ use core::num::ParseIntError;
 use core::str::FromStr;
 #[cfg(feature = "std")] use std::error;
 
-use secp256k1::schnorrsig;
+use secp256k1::{Secp256k1, Verification};
 use bech32;
 use hashes::Hash;
 use hash_types::{PubkeyHash, WPubkeyHash, ScriptHash, WScriptHash};
@@ -49,7 +49,9 @@ use blockdata::constants::{PUBKEY_ADDRESS_PREFIX_MAIN, SCRIPT_ADDRESS_PREFIX_MAI
 use network::constants::Network;
 use util::base58;
 use util::ecdsa;
+use util::taproot::TapBranchHash;
 use blockdata::script::Instruction;
+use util::schnorr::{TapTweak, UntweakedPublicKey, TweakedPublicKey};
 
 /// Address error.
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -511,13 +513,34 @@ impl Address {
         }
     }
 
-    /// Create a pay to taproot address
-    pub fn p2tr(taptweaked_key: schnorrsig::PublicKey, network: Network) -> Address {
+    /// Create a pay to taproot address from untweaked key
+    pub fn p2tr<C: Verification>(
+        secp: Secp256k1<C>, 
+        internal_key: UntweakedPublicKey, 
+        merkle_root: Option<TapBranchHash>, 
+        network: Network
+    ) -> Address {
         Address {
             network: network,
             payload: Payload::WitnessProgram {
                 version: WitnessVersion::V1,
-                program: taptweaked_key.serialize().to_vec()
+                program: internal_key.tap_tweak(secp, merkle_root).into_inner().serialize().to_vec()
+            }
+        }
+    }
+
+    /// Create a pay to taproot address from a pre-tweaked output key.
+    ///
+    /// This method is not recommended for use and [Address::p2tr()] should be used where possible.
+    pub fn p2tr_tweaked<C: Verification>(
+        output_key: TweakedPublicKey, 
+        network: Network
+    ) -> Address {
+        Address {
+            network: network,
+            payload: Payload::WitnessProgram {
+                version: WitnessVersion::V1,
+                program: output_key.into_inner().serialize().to_vec()
             }
         }
     }
@@ -786,6 +809,7 @@ mod tests {
     use blockdata::script::Script;
     use network::constants::Network::{Bitcoin, Testnet};
     use util::ecdsa::PublicKey;
+    use secp256k1::schnorrsig;
 
     use super::*;
 
@@ -1148,4 +1172,12 @@ mod tests {
         test_addr_type(legacy_payload, LEGACY_EQUIVALENCE_CLASSES);
         test_addr_type(&segwit_payload, SEGWIT_EQUIVALENCE_CLASSES);
     }
+
+    #[test]
+    fn p2tr_from_untweaked(){
+        //Test case from BIP-086
+        let internal_key = schnorrsig::PublicKey::from_str("cc8a4bc64d897bddc5fbc2f670f7a8ba0b386779106cf1223c6fc5d7cd6fc115").unwrap();
+        let address = Address::p2tr(Secp256k1::new(), internal_key,None, Network::Bitcoin);
+        assert_eq!(address.to_string(), "bc1p5cyxnuxmeuwuvkwfem96lqzszd02n6xdcjrs20cac6yqjjwudpxqkedrcr");
+    }
 }
diff --git a/src/util/schnorr.rs b/src/util/schnorr.rs
@@ -17,3 +17,67 @@
 //!
 
 pub use secp256k1::schnorrsig::{PublicKey, KeyPair};
+use secp256k1::{Secp256k1, Verification};
+use hashes::{Hash, HashEngine};
+use util::taproot::{TapBranchHash, TapTweakHash};
+
+/// Untweaked Schnorr public key
+pub type UntweakedPublicKey = PublicKey;
+
+/// Tweaked Schnorr public key
+pub struct TweakedPublicKey(PublicKey);
+
+/// A trait for tweaking Schnorr public keys
+pub trait TapTweak {
+    /// Tweaks an untweaked public key given an untweaked key and optional script tree merkle root.
+    /// 
+    /// This is done by using the equation Q = P + H(P|c)G, where 
+    ///  * Q is the tweaked key
+    ///  * P is the internal key
+    ///  * H is the hash function
+    ///  * c is the commitment data
+    ///  * G is the generator point
+    fn tap_tweak<C: Verification>(&self, secp: Secp256k1<C>, merkle_root: Option<TapBranchHash>) -> TweakedPublicKey;
+
+    /// Directly convert an UntweakedPublicKey to a TweakedPublicKey
+    /// 
+    /// This method is dangerous and can lead to loss of funds if used incorrectly.
+    /// Specifically, in multi-party protocols a peer can provide a value that allows them to steal.
+    fn dangerous_assume_tweaked(self) -> TweakedPublicKey;
+}
+
+impl TapTweak for UntweakedPublicKey {
+    fn tap_tweak<C: Verification>(&self, secp: Secp256k1<C>, merkle_root: Option<TapBranchHash>) -> TweakedPublicKey {
+        // Compute the tweak
+        let mut engine = TapTweakHash::engine();
+        engine.input(&self.serialize());
+        merkle_root.map(|hash| engine.input(&hash));
+        let tweak_value: [u8; 32] = TapTweakHash::from_engine(engine).into_inner();
+            
+
+        //Tweak the internal key by the tweak value
+        let mut output_key = self.clone();
+        let parity = output_key.tweak_add_assign(&secp, &tweak_value).expect("Tap tweak failed");
+        if self.tweak_add_check(&secp, &output_key, parity, tweak_value) {
+            return TweakedPublicKey(output_key);
+        } else { unreachable!("Tap tweak failed") }
+    }
+
+    
+    fn dangerous_assume_tweaked(self) -> TweakedPublicKey {
+        TweakedPublicKey(self)
+    }
+}
+
+
+impl TweakedPublicKey {
+    /// Create a new [TweakedPublicKey] from a [PublicKey]. No tweak is applied.
+    pub fn new(key: PublicKey) -> TweakedPublicKey {
+        TweakedPublicKey(key)
+    }
+    
+    /// Returns the underlying public key 
+    pub fn into_inner(self) -> PublicKey {
+        self.0
+    }
+}
