From e06f2a9d8b699ea055b14d14e3192ab6eac249fe Mon Sep 17 00:00:00 2001
From: thunderbiscuit <thunderbiscuit@protonmail.com>
Date: Tue, 22 Oct 2024 12:47:05 -0400
Subject: [PATCH] feat: expose TxBuilder::add_data method

---
 bdk-ffi/src/bdk.udl       |  3 +++
 bdk-ffi/src/error.rs      | 10 ++++++++++
 bdk-ffi/src/tx_builder.rs | 15 +++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/bdk-ffi/src/bdk.udl b/bdk-ffi/src/bdk.udl
index 1c716777..3c22f7b7 100644
--- a/bdk-ffi/src/bdk.udl
+++ b/bdk-ffi/src/bdk.udl
@@ -75,6 +75,7 @@ interface CreateTxError {
   UnknownUtxo(string outpoint);
   MissingNonWitnessUtxo(string outpoint);
   MiniscriptPsbt(string error_message);
+  PushBytesError();
 };
 
 [Error]
@@ -474,6 +475,8 @@ interface TxBuilder {
 
   TxBuilder set_exact_sequence(u32 nsequence);
 
+  TxBuilder add_data(sequence<u8> data);
+
   [Throws=CreateTxError]
   Psbt finish([ByRef] Wallet wallet);
 };
diff --git a/bdk-ffi/src/error.rs b/bdk-ffi/src/error.rs
index 668d7314..2476712e 100644
--- a/bdk-ffi/src/error.rs
+++ b/bdk-ffi/src/error.rs
@@ -1,5 +1,6 @@
 use bitcoin_ffi::OutPoint;
 
+use bdk_core::bitcoin::script::PushBytesError;
 use bdk_electrum::electrum_client::Error as BdkElectrumError;
 use bdk_esplora::esplora_client::{Error as BdkEsploraError, Error};
 use bdk_wallet::bitcoin::address::ParseError as BdkParseError;
@@ -194,6 +195,9 @@ pub enum CreateTxError {
 
     #[error("miniscript psbt error: {error_message}")]
     MiniscriptPsbt { error_message: String },
+
+    #[error("attempt to prepare too many bytes to be pushed into script")]
+    PushBytesError,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -826,6 +830,12 @@ impl From<BdkCreateTxError> for CreateTxError {
     }
 }
 
+impl From<PushBytesError> for CreateTxError {
+    fn from(_: PushBytesError) -> Self {
+        CreateTxError::PushBytesError
+    }
+}
+
 impl From<BdkCreateWithPersistError<chain::rusqlite::Error>> for CreateWithPersistError {
     fn from(error: BdkCreateWithPersistError<chain::rusqlite::Error>) -> Self {
         match error {
diff --git a/bdk-ffi/src/tx_builder.rs b/bdk-ffi/src/tx_builder.rs
index 7ade7c30..fddd9a6f 100644
--- a/bdk-ffi/src/tx_builder.rs
+++ b/bdk-ffi/src/tx_builder.rs
@@ -11,7 +11,9 @@ use bdk_wallet::bitcoin::ScriptBuf as BdkScriptBuf;
 use bdk_wallet::bitcoin::{OutPoint, Sequence, Txid};
 use bdk_wallet::ChangeSpendPolicy;
 
+use bdk_wallet::bitcoin::script::PushBytesBuf;
 use std::collections::HashSet;
+use std::convert::TryFrom;
 use std::str::FromStr;
 use std::sync::Arc;
 
@@ -28,6 +30,7 @@ pub struct TxBuilder {
     pub(crate) drain_wallet: bool,
     pub(crate) drain_to: Option<BdkScriptBuf>,
     pub(crate) sequence: Option<u32>,
+    pub(crate) data: Vec<u8>,
 }
 
 impl TxBuilder {
@@ -44,6 +47,7 @@ impl TxBuilder {
             drain_wallet: false,
             drain_to: None,
             sequence: None,
+            data: Vec::new(),
         }
     }
 
@@ -167,6 +171,13 @@ impl TxBuilder {
         })
     }
 
+    pub(crate) fn add_data(&self, data: Vec<u8>) -> Arc<Self> {
+        Arc::new(TxBuilder {
+            data,
+            ..self.clone()
+        })
+    }
+
     pub(crate) fn finish(&self, wallet: &Arc<Wallet>) -> Result<Arc<Psbt>, CreateTxError> {
         // TODO: I had to change the wallet here to be mutable. Why is that now required with the 1.0 API?
         let mut wallet = wallet.get_wallet();
@@ -205,6 +216,10 @@ impl TxBuilder {
         if let Some(sequence) = self.sequence {
             tx_builder.set_exact_sequence(Sequence(sequence));
         }
+        if !&self.data.is_empty() {
+            let push_bytes = PushBytesBuf::try_from(self.data.clone())?;
+            tx_builder.add_data(&push_bytes);
+        }
 
         let psbt = tx_builder.finish().map_err(CreateTxError::from)?;