From 62ae44a131825e8ebdd83b053e94208ec15b7509 Mon Sep 17 00:00:00 2001
From: Greg Pfeil <greg@electriccoin.co>
Date: Mon, 20 Mar 2023 09:40:42 -0600
Subject: [PATCH] Ensure that a WalletTxBuilder tx balances
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This also removes the code that stops adding spends if they ever go
`>= targetAmount`. The included note limiting and change calculation should
ensure that it’s always `==` at the end, and we don’t want paper over a mistake
in those earlier calculations.

There are existing tests that fail if either
- the newly-added Orchard increment is missing or
- the assertion is applied when there’s Sprout change.
---
 src/wallet/wallet_tx_builder.cpp | 28 +++++++++++++---------------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/src/wallet/wallet_tx_builder.cpp b/src/wallet/wallet_tx_builder.cpp
index 455feff14..9c78c1921 100644
--- a/src/wallet/wallet_tx_builder.cpp
+++ b/src/wallet/wallet_tx_builder.cpp
@@ -543,8 +543,7 @@ TransactionBuilderResult TransactionEffects::ApproveAndBuild(
 
     // Track the total of notes that we've added to the builder. This
     // shouldn't strictly be necessary, given `spendable.LimitToAmount`
-    CAmount sum = 0;
-    CAmount targetAmount = payments.Total() + fee;
+    CAmount totalSpend = 0;
 
     // Create Sapling outpoints
     std::vector<SaplingOutPoint> saplingOutPoints;
@@ -559,10 +558,7 @@ TransactionBuilderResult TransactionEffects::ApproveAndBuild(
         assert(wallet.GetSaplingExtendedSpendingKey(t.address, saplingKey));
         saplingKeys.push_back(saplingKey);
 
-        sum += t.note.value();
-        if (sum >= targetAmount) {
-            break;
-        }
+        totalSpend += t.note.value();
     }
 
     // Fetch Sapling anchor and witnesses, and Orchard Merkle paths.
@@ -591,6 +587,8 @@ TransactionBuilderResult TransactionEffects::ApproveAndBuild(
             return TransactionBuilderResult(
                 strprintf("Failed to add Orchard note to transaction (check %s for details)", GetDebugLogPath())
             );
+        } else {
+            totalSpend += spendInfo.second.Value();
         }
     }
 
@@ -633,10 +631,7 @@ TransactionBuilderResult TransactionEffects::ApproveAndBuild(
         const CTxOut& txOut = out.tx->vout[out.i];
         builder.AddTransparentInput(COutPoint(out.tx->GetHash(), out.i), txOut.scriptPubKey, txOut.nValue);
 
-        sum += txOut.nValue;
-        if (sum >= targetAmount) {
-            break;
-        }
+        totalSpend += txOut.nValue;
     }
 
     // Find Sprout witnesses
@@ -673,18 +668,21 @@ TransactionBuilderResult TransactionEffects::ApproveAndBuild(
 
         builder.AddSproutInput(sk, t.note, vSproutWitnesses[i].value());
 
-        sum += t.note.value();
-        if (sum >= targetAmount) {
-            break;
-        }
+        totalSpend += t.note.value();
     }
 
+    // TODO: We currently can’t store Sprout change in `Payments`, so we only validate the
+    //       spend/output balance in the case that `TransactionBuilder` doesn’t need to
+    //       (re)calculate the change. In future, we shouldn’t rely on `TransactionBuilder` ever
+    //       calculating change.
     if (changeAddr.has_value()) {
         std::visit(match {
             [&](const SproutPaymentAddress& addr) {
                 builder.SendChangeToSprout(addr);
             },
-            [](const RecipientAddress&) { }
+            [&](const RecipientAddress&) {
+                assert(totalSpend == payments.Total() + fee);
+            }
         }, changeAddr.value());
     }