From 6a150368674046f796f5c37755896f16d8345fbc Mon Sep 17 00:00:00 2001
From: Daniela Brozzoni <danielabrozzoni@protonmail.com>
Date: Mon, 6 Jun 2022 21:45:13 +0200
Subject: [PATCH] Restrict `drain_to` usage

Before this commit, you could create a transaction with `drain_to` set
without specifying recipients, nor `drain_wallet`, nor `utxos`. What would
happen is that BDK would pick one input from the wallet and send
that one to `drain_to`, which is quite weird.
This PR restricts the usage of `drain_to`: if you want to use it as a
change output, you need to set recipients as well. If you want to send
a specific utxo to the `drain_to` address, you specify it through
`add_utxos`. If you want to drain the whole wallet, you set
`drain_wallet`. In any other case, if `drain_to` is set, we return a
`NoRecipients` error.

Fixes #620
---
 src/wallet/mod.rs        | 43 +++++++++++++++++++++++++++++++++++++++-
 src/wallet/tx_builder.rs |  4 ++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs
index 6d2ff385..50ea6fe7 100644
--- a/src/wallet/mod.rs
+++ b/src/wallet/mod.rs
@@ -791,7 +791,14 @@ where
         let drain_val = (coin_selection.selected_amount() - outgoing).saturating_sub(fee_amount);
 
         if tx.output.is_empty() {
-            if params.drain_to.is_some() {
+            // Uh oh, our transaction has no outputs.
+            // We allow this when:
+            // - We have a drain_to address and the utxos we must spend (this happens,
+            // for example, when we RBF)
+            // - We have a drain_to address and drain_wallet set
+            // Otherwise, we don't know who we should send the funds to, and how much
+            // we should send!
+            if params.drain_to.is_some() && (params.drain_wallet || !params.utxos.is_empty()) {
                 if drain_val.is_dust(&drain_output.script_pubkey) {
                     return Err(Error::InsufficientFunds {
                         needed: drain_output.script_pubkey.dust_value().as_sat(),
@@ -2176,6 +2183,40 @@ pub(crate) mod test {
         assert_eq!(drain_output.value, 30_000 - details.fee.unwrap_or(0));
     }
 
+    #[test]
+    fn test_create_tx_drain_to_and_utxos() {
+        let (wallet, _, _) = get_funded_wallet(get_test_wpkh());
+        let addr = wallet.get_address(New).unwrap();
+        let utxos: Vec<_> = wallet
+            .get_available_utxos()
+            .unwrap()
+            .into_iter()
+            .map(|(u, _)| u.outpoint)
+            .collect();
+        let mut builder = wallet.build_tx();
+        builder
+            .drain_to(addr.script_pubkey())
+            .add_utxos(&utxos)
+            .unwrap();
+        let (psbt, details) = builder.finish().unwrap();
+
+        assert_eq!(psbt.unsigned_tx.output.len(), 1);
+        assert_eq!(
+            psbt.unsigned_tx.output[0].value,
+            50_000 - details.fee.unwrap_or(0)
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "NoRecipients")]
+    fn test_create_tx_drain_to_no_drain_wallet_no_utxos() {
+        let (wallet, _, _) = get_funded_wallet(get_test_wpkh());
+        let drain_addr = wallet.get_address(New).unwrap();
+        let mut builder = wallet.build_tx();
+        builder.drain_to(drain_addr.script_pubkey());
+        builder.finish().unwrap();
+    }
+
     #[test]
     fn test_create_tx_default_fee_rate() {
         let (wallet, _, _) = get_funded_wallet(get_test_wpkh());
diff --git a/src/wallet/tx_builder.rs b/src/wallet/tx_builder.rs
index 59dedcf8..6ffdd1e3 100644
--- a/src/wallet/tx_builder.rs
+++ b/src/wallet/tx_builder.rs
@@ -574,6 +574,9 @@ impl<'a, D: BatchDatabase, Cs: CoinSelectionAlgorithm<D>> TxBuilder<'a, D, Cs, C
     /// difference is that it is valid to use `drain_to` without setting any ordinary recipients
     /// with [`add_recipient`] (but it is perfectly fine to add recipients as well).
     ///
+    /// If you choose not to set any recipients, you should either provide the utxos that the
+    /// transaction should spend via [`add_utxos`], or set [`drain_wallet`] to spend all of them.
+    ///
     /// When bumping the fees of a transaction made with this option, you probably want to
     /// use [`allow_shrinking`] to allow this output to be reduced to pay for the extra fees.
     ///
@@ -604,6 +607,7 @@ impl<'a, D: BatchDatabase, Cs: CoinSelectionAlgorithm<D>> TxBuilder<'a, D, Cs, C
     ///
     /// [`allow_shrinking`]: Self::allow_shrinking
     /// [`add_recipient`]: Self::add_recipient
+    /// [`add_utxos`]: Self::add_utxos
     /// [`drain_wallet`]: Self::drain_wallet
     pub fn drain_to(&mut self, script_pubkey: Script) -> &mut Self {
         self.params.drain_to = Some(script_pubkey);