From 8efffe91a13bb4ffa3a134cc00a2f3d06c634fc5 Mon Sep 17 00:00:00 2001
From: Dr Maxim Orlovsky <orlovsky@lnp-bp.org>
Date: Fri, 6 Oct 2023 10:03:31 +0200
Subject: [PATCH] primitives: implement consensus decoding for ControlBlock

---
 primitives/src/coding.rs  | 46 ++++++++++++++++++++++++++++++++++++---
 primitives/src/taproot.rs | 23 +++++++++++++++++++-
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/primitives/src/coding.rs b/primitives/src/coding.rs
index eb2d0ff4..44ffc84e 100644
--- a/primitives/src/coding.rs
+++ b/primitives/src/coding.rs
@@ -27,9 +27,9 @@ use amplify::hex::{self, FromHex, ToHex};
 use amplify::{confinement, ByteArray, Bytes32, IoError, Wrapper};
 
 use crate::{
-    ControlBlock, InternalPk, LockTime, Outpoint, RedeemScript, Sats, ScriptBytes, ScriptPubkey,
-    SeqNo, SigScript, TapBranchHash, TapScript, Tx, TxIn, TxOut, TxVer, Txid, Vout, Witness,
-    WitnessScript, LIB_NAME_BITCOIN,
+    ControlBlock, InternalPk, InvalidLeafVer, LeafVer, LockTime, Outpoint, Parity, RedeemScript,
+    Sats, ScriptBytes, ScriptPubkey, SeqNo, SigScript, TapBranchHash, TapMerklePath, TapScript, Tx,
+    TxIn, TxOut, TxVer, Txid, Vout, Witness, WitnessScript, LIB_NAME_BITCOIN,
 };
 
 pub type VarIntArray<T> = Confined<Vec<T>, 0, U32>;
@@ -133,7 +133,9 @@ pub enum ConsensusDecodeError {
     #[from(io::Error)]
     Io(IoError),
 
+    #[display(inner)]
     #[from]
+    #[from(InvalidLeafVer)]
     #[from(confinement::Error)]
     Data(ConsensusDataError),
 }
@@ -150,6 +152,17 @@ pub enum ConsensusDataError {
     /// invalid BIP340 (x-only) pubkey data.
     InvalidXonlyPubkey(Bytes32),
 
+    /// taproot Merkle path length exceeds BIP-341 consensus limit of 128
+    /// elements.
+    LongTapMerklePath,
+
+    /// Merkle path in the `PSBT_IN_TAP_TREE` is not encoded correctly.
+    InvalidTapMerklePath,
+
+    #[from]
+    #[display(inner)]
+    InvalidLeafVer(InvalidLeafVer),
+
     #[from]
     #[display(inner)]
     Confined(confinement::Error),
@@ -492,6 +505,33 @@ impl ConsensusEncode for ControlBlock {
     }
 }
 
+impl ConsensusDecode for ControlBlock {
+    fn consensus_decode(reader: &mut impl Read) -> Result<Self, ConsensusDecodeError> {
+        let first_byte = u8::consensus_decode(reader)?;
+        let leaf_version = LeafVer::from_consensus_u8(first_byte & 0xFE)?;
+        let output_key_parity = Parity::from_consensus_u8(first_byte & 0x01).expect("binary value");
+
+        let internal_key = InternalPk::consensus_decode(reader)?;
+
+        let mut buf = vec![];
+        reader.read_to_end(&mut buf)?;
+        let mut iter = buf.chunks_exact(32);
+        let merkle_branch = iter.by_ref().map(TapBranchHash::from_slice_unsafe);
+        let merkle_branch = TapMerklePath::try_from_iter(merkle_branch)
+            .map_err(|_| ConsensusDataError::LongTapMerklePath)?;
+        if !iter.remainder().is_empty() {
+            return Err(ConsensusDataError::InvalidTapMerklePath.into());
+        }
+
+        Ok(ControlBlock {
+            leaf_version,
+            output_key_parity,
+            internal_key,
+            merkle_branch,
+        })
+    }
+}
+
 impl ConsensusEncode for Sats {
     fn consensus_encode(&self, writer: &mut impl Write) -> Result<usize, IoError> {
         self.0.consensus_encode(writer)
diff --git a/primitives/src/taproot.rs b/primitives/src/taproot.rs
index 156cb540..2fc4b5a0 100644
--- a/primitives/src/taproot.rs
+++ b/primitives/src/taproot.rs
@@ -27,7 +27,7 @@ use std::ops::BitXor;
 use std::{cmp, io, slice, vec};
 
 use amplify::confinement::{Confined, U32};
-use amplify::{Bytes32, Wrapper};
+use amplify::{confinement, Bytes32, Wrapper};
 use commit_verify::{DigestExt, Sha256};
 use secp256k1::{Scalar, XOnlyPublicKey};
 use strict_encoding::{
@@ -238,6 +238,27 @@ impl<'a> IntoIterator for &'a TapMerklePath {
     fn into_iter(self) -> Self::IntoIter { self.0.iter() }
 }
 
+impl TapMerklePath {
+    /// Tries to construct a confinement over a collection. Fails if the number
+    /// of items in the collection exceeds one of the confinement bounds.
+    // We can't use `impl TryFrom` due to the conflict with core library blanked
+    // implementation
+    #[inline]
+    pub fn try_from(path: Vec<TapBranchHash>) -> Result<Self, confinement::Error> {
+        Confined::try_from(path).map(Self::from_inner)
+    }
+
+    /// Tries to construct a confinement with a collection of elements taken
+    /// from an iterator. Fails if the number of items in the collection
+    /// exceeds one of the confinement bounds.
+    #[inline]
+    pub fn try_from_iter<I: IntoIterator<Item = TapBranchHash>>(
+        iter: I,
+    ) -> Result<Self, confinement::Error> {
+        Confined::try_from_iter(iter).map(Self::from_inner)
+    }
+}
+
 /// Taproot annex prefix.
 pub const TAPROOT_ANNEX_PREFIX: u8 = 0x50;