From 1490ff36ee76c2c0cc48323f30f4c923702a5594 Mon Sep 17 00:00:00 2001 From: sanket1729 Date: Fri, 8 Oct 2021 15:03:49 -0700 Subject: [PATCH] Add support for verifying merkle proofs --- src/util/taproot.rs | 91 +++++++++++++++++++++++++++++++++++++++------ 1 file changed, 80 insertions(+), 11 deletions(-) diff --git a/src/util/taproot.rs b/src/util/taproot.rs index 1ce264d1..474e2e61 100644 --- a/src/util/taproot.rs +++ b/src/util/taproot.rs @@ -113,6 +113,20 @@ impl TapTweakHash { } } +impl TapLeafHash { + /// function to compute leaf hash from components + pub fn from_script(script: &Script, ver: LeafVersion) -> TapLeafHash { + let mut eng = TapLeafHash::engine(); + ver.as_u8() + .consensus_encode(&mut eng) + .expect("engines don't err"); + script + .consensus_encode(&mut eng) + .expect("engines don't err"); + TapLeafHash::from_engine(eng) + } +} + /// Maximum depth of a Taproot Tree Script spend path // https://github.com/bitcoin/bitcoin/blob/e826b22da252e0599c61d21c98ff89f366b3120f/src/script/interpreter.h#L229 pub const TAPROOT_CONTROL_MAX_NODE_COUNT: usize = 128; @@ -133,6 +147,8 @@ pub const TAPROOT_CONTROL_BASE_SIZE: usize = 33; pub const TAPROOT_CONTROL_MAX_SIZE: usize = TAPROOT_CONTROL_BASE_SIZE + TAPROOT_CONTROL_NODE_SIZE * TAPROOT_CONTROL_MAX_NODE_COUNT; +// type alias for versioned tap script corresponding merkle proof +type ScriptMerkleProofMap = BTreeMap<(Script, LeafVersion), BTreeSet>; /// Data structure for representing Taproot spending information. /// Taproot output corresponds to a combination of a /// single public key condition (known the internal key), and zero or more @@ -164,7 +180,7 @@ pub struct TaprootSpendInfo { /// appears in multiple branches of the tree. In all cases, keeping one should /// be enough for spending funds, but we keep all of the paths so that /// a full tree can be constructed again from spending data if required. - script_map: BTreeMap<(Script, LeafVersion), BTreeSet>, + script_map: ScriptMerkleProofMap, } impl TaprootSpendInfo { @@ -268,7 +284,7 @@ impl TaprootSpendInfo { /// Output key(the key used in script pubkey) from Spend data. See also /// [`TaprootSpendInfo::output_key_parity`] - pub fn output_key(&self) -> schnorr::PublicKey { + pub fn output_key(&self) -> schnorr::PublicKey { self.output_key } @@ -302,6 +318,29 @@ impl TaprootSpendInfo { } info } + + /// Access the internal script map + pub fn as_script_map(&self) -> &ScriptMerkleProofMap { + &self.script_map + } + + /// Obtain a [`ControlBlock`] for particular script with the given version. + /// Returns [`None`] if the script is not contained in the [`TaprootSpendInfo`] + /// If there are multiple ControlBlocks possible, this returns the shortest one. + pub fn control_block(&self, script_ver: &(Script, LeafVersion)) -> Option { + let merkle_branch_set = self.script_map.get(script_ver)?; + // Choose the smallest one amongst the multiple script maps + let smallest = merkle_branch_set + .iter() + .min_by(|x, y| x.0.len().cmp(&y.0.len())) + .expect("Invariant: Script map key must contain non-empty set value"); + Some(ControlBlock { + internal_key: self.internal_key, + output_key_parity: self.output_key_parity, + leaf_version: LeafVersion::default(), + merkle_branch: smallest.clone(), + }) + } } /// Builder for building taproot iteratively. Users can specify tap leaf or omitted/hidden @@ -521,17 +560,11 @@ impl LeafInfo { // Compute a leaf hash for the given leaf fn hash(&self) -> sha256::Hash { - let mut eng = TapLeafHash::engine(); - self.ver - .as_u8() - .consensus_encode(&mut eng) - .expect("engines don't err"); - self.script - .consensus_encode(&mut eng) - .expect("engines don't err"); - sha256::Hash::from_engine(eng) + let leaf_hash = TapLeafHash::from_script(&self.script, self.ver); + sha256::Hash::from_inner(leaf_hash.into_inner()) } } + /// The Merkle proof for inclusion of a tree in a taptree hash // The type of hash is sha256::Hash because the vector might contain // both TapBranchHash and TapLeafHash @@ -673,6 +706,42 @@ impl ControlBlock { .expect("writers don't error"); buf } + + /// Verify that a control block is correct proof for a given output key and script + /// This only checks that script is contained inside the taptree described by + /// output key, full verification must also execute the script with witness data + pub fn verify_taproot_commitment( + &self, + secp: &Secp256k1, + output_key: &schnorr::PublicKey, + script: &Script, + ) -> bool { + // compute the script hash + // Initially the curr_hash is the leaf hash + let leaf_hash = TapLeafHash::from_script(&script, self.leaf_version); + let mut curr_hash = TapBranchHash::from_inner(leaf_hash.into_inner()); + // Verify the proof + for elem in self.merkle_branch.as_inner() { + let mut eng = TapBranchHash::engine(); + if curr_hash.as_inner() < elem.as_inner() { + eng.input(&curr_hash); + eng.input(elem); + } else { + eng.input(elem); + eng.input(&curr_hash); + } + // Recalculate the curr hash as parent hash + curr_hash = TapBranchHash::from_engine(eng); + } + // compute the taptweak + let tweak = TapTweakHash::from_key_and_tweak(self.internal_key, Some(curr_hash)); + self.internal_key.tweak_add_check( + secp, + output_key, + self.output_key_parity, + tweak.into_inner(), + ) + } } /// The leaf version for tapleafs