From b5a5f91b4c8e4ea063afdcf8d80730f92e4658c9 Mon Sep 17 00:00:00 2001
From: Dr Maxim Orlovsky <orlovsky@pandoraprime.ch>
Date: Sat, 2 Jul 2022 20:04:53 +0200
Subject: [PATCH] Finalization of PSBT DBC-related API

---
 Cargo.lock                |  2 +-
 psbt/Cargo.toml           |  2 +-
 psbt/src/commit/lnpbp4.rs | 19 +++++++++++++++++++
 psbt/src/commit/mod.rs    |  2 ++
 psbt/src/commit/tapret.rs | 31 ++++++++++++++++++++++++-------
 5 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c229b42..e4bff60 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -671,7 +671,7 @@ dependencies = [
 
 [[package]]
 name = "psbt"
-version = "0.8.1"
+version = "0.8.3"
 dependencies = [
  "amplify",
  "bitcoin",
diff --git a/psbt/Cargo.toml b/psbt/Cargo.toml
index 1e51700..47e5cdc 100644
--- a/psbt/Cargo.toml
+++ b/psbt/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "psbt"
-version = "0.8.1"
+version = "0.8.3"
 license = "Apache-2.0"
 authors = ["Dr. Maxim Orlovsky <orlovsky@lnp-bp.org>"]
 description = "Partially signed bitcoin transaction v0-2 library (bip174, bip370, bip371)"
diff --git a/psbt/src/commit/lnpbp4.rs b/psbt/src/commit/lnpbp4.rs
index 7224a76..32931d2 100644
--- a/psbt/src/commit/lnpbp4.rs
+++ b/psbt/src/commit/lnpbp4.rs
@@ -14,6 +14,8 @@
 // If not, see <https://opensource.org/licenses/Apache-2.0>.
 
 use amplify::Slice32;
+use bitcoin::hashes::Hash;
+use commit_verify::lnpbp4;
 use commit_verify::lnpbp4::{Message, ProtocolId};
 use strict_encoding::{StrictDecode, StrictEncode};
 
@@ -160,6 +162,23 @@ impl Psbt {
 /// Extension trait for [`Output`] for working with proprietary LNPBP4
 /// keys.
 impl Output {
+    /// Returns [`lnpbp4::MessageMap`] constructed from the proprietary key
+    /// data.
+    pub fn lnpbp4_message_map(&self) -> Result<lnpbp4::MessageMap, Lnpbp4KeyError> {
+        self.proprietary
+            .iter()
+            .filter(|(key, _)| {
+                key.prefix == PSBT_LNPBP4_PREFIX && key.subtype == PSBT_OUT_LNPBP4_MESSAGE
+            })
+            .map(|(key, val)| {
+                Ok((
+                    ProtocolId::from_slice(&key.key).ok_or(Lnpbp4KeyError::InvalidKeyValue)?,
+                    Message::from_slice(val).map_err(|_| Lnpbp4KeyError::InvalidKeyValue)?,
+                ))
+            })
+            .collect()
+    }
+
     /// Returns a valid LNPBP-4 [`Message`] associated with the given
     /// [`ProtocolId`], if any.
     ///
diff --git a/psbt/src/commit/mod.rs b/psbt/src/commit/mod.rs
index fa3c0dd..c4c982f 100644
--- a/psbt/src/commit/mod.rs
+++ b/psbt/src/commit/mod.rs
@@ -12,6 +12,8 @@
 // along with this software.
 // If not, see <https://opensource.org/licenses/Apache-2.0>.
 
+// TODO: Relocate to BP DBC library
+
 //! Managing commitment-related proprietary keys inside PSBT.
 //!
 //! Supports Tapret, Opret, P2C and S2C commitments and LNPBP4 structures used
diff --git a/psbt/src/commit/tapret.rs b/psbt/src/commit/tapret.rs
index c5e72e0..c7d91f3 100644
--- a/psbt/src/commit/tapret.rs
+++ b/psbt/src/commit/tapret.rs
@@ -25,7 +25,6 @@
 //! commitments.
 
 use amplify::Slice32;
-use bitcoin::util::taproot::TaprootMerkleBranch;
 use bitcoin_scripts::taproot::DfsPath;
 use strict_encoding::{StrictDecode, StrictEncode};
 
@@ -93,7 +92,7 @@ impl ProprietaryKeyTapret for ProprietaryKey {}
 
 /// Errors processing tapret-related proprietary PSBT keys and their values.
 #[derive(
-    Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Display, Error
+    Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Display, Error, From
 )]
 #[display(doc_comments)]
 pub enum TapretKeyError {
@@ -104,6 +103,10 @@ pub enum TapretKeyError {
     /// the output is not marked to host tapret commitments. Please first set
     /// PSBT_OUT_TAPRET_HOST flag.
     TapretProhibited,
+
+    /// The key contains invalid value
+    #[from(strict_encoding::Error)]
+    InvalidKeyValue,
 }
 
 /// Error decoding [`DfsPath`] inside PSBT data
@@ -189,8 +192,8 @@ impl Output {
     }
 
     /// Assigns value of the tapreturn commitment to this PSBT output, by
-    /// adding [`PSBT_OUT_TAPRET_COMMITMENT`] proprietary key containing the
-    /// 32-byte commitment as its value.
+    /// adding [`PSBT_OUT_TAPRET_COMMITMENT`] and [`PSBT_OUT_TAPRET_PROOF`]
+    /// proprietary keys containing the 32-byte commitment as its proof.
     ///
     /// # Errors
     ///
@@ -201,6 +204,7 @@ impl Output {
     pub fn set_tapret_commitment(
         &mut self,
         commitment: impl Into<[u8; 32]>,
+        proof: &impl StrictEncode,
     ) -> Result<(), TapretKeyError> {
         if !self.is_tapret_host() {
             return Err(TapretKeyError::TapretProhibited);
@@ -215,6 +219,9 @@ impl Output {
             commitment.into().to_vec(),
         );
 
+        self.proprietary
+            .insert(ProprietaryKey::tapret_proof(), proof.strict_serialize()?);
+
         Ok(())
     }
 
@@ -238,8 +245,18 @@ impl Output {
     /// commitments (having non-32 bytes) will be filtered at the moment of PSBT
     /// deserialization and this function will return `None` only in situations
     /// when the commitment is absent.
-    pub fn tapret_proof(&self) -> Option<TaprootMerkleBranch> {
-        let proof = self.proprietary.get(&ProprietaryKey::tapret_proof())?;
-        TaprootMerkleBranch::from_slice(proof).ok()
+    ///
+    /// Function returns generic type since the real type will create dependency
+    /// on `bp-dpc` crate, which will result in circular dependency with the
+    /// current crate.
+    pub fn tapret_proof<T>(&self) -> Result<Option<T>, TapretKeyError>
+    where
+        T: StrictDecode,
+    {
+        self.proprietary
+            .get(&ProprietaryKey::tapret_proof())
+            .map(T::strict_deserialize)
+            .transpose()
+            .map_err(TapretKeyError::from)
     }
 }