From 532b15b45e0c5ed27bdfab2f2b595659704f2a6a Mon Sep 17 00:00:00 2001
From: Duke <duke@hush.is>
Date: Sun, 18 May 2025 17:38:03 -0400
Subject: [PATCH] Lock and unlock sapling notes during consolidations

---
 ...asyncrpcoperation_saplingconsolidation.cpp | 32 +++++++++++++++++++
 .../asyncrpcoperation_saplingconsolidation.h  |  5 +++
 2 files changed, 37 insertions(+)

diff --git a/src/wallet/asyncrpcoperation_saplingconsolidation.cpp b/src/wallet/asyncrpcoperation_saplingconsolidation.cpp
index 41455d360..755564492 100644
--- a/src/wallet/asyncrpcoperation_saplingconsolidation.cpp
+++ b/src/wallet/asyncrpcoperation_saplingconsolidation.cpp
@@ -74,6 +74,8 @@ void AsyncRPCOperation_saplingconsolidation::main() {
     }
 
     LogPrintf("%s", s);
+    unlock_notes(); // clean up
+    LogPrint("zrpc", "%s: consolidation input notes unlocked\n", getId());
 }
 
 bool AsyncRPCOperation_saplingconsolidation::main_impl() {
@@ -102,6 +104,11 @@ bool AsyncRPCOperation_saplingconsolidation::main_impl() {
             return true;
         }
 
+        // store sapling inputs so we can correctly lock and unlock them
+        for (auto entry : saplingEntries) {
+            z_sapling_inputs_.push_back(entry);
+        }
+
         if (fConsolidationMapUsed) {
             const vector<string>& v = mapMultiArgs["-consolidatesaplingaddress"];
             for(int i = 0; i < v.size(); i++) {
@@ -233,6 +240,10 @@ bool AsyncRPCOperation_saplingconsolidation::main_impl() {
                 break;
             }
 
+            // Lock shielded input notes
+            lock_notes();
+            LogPrint("zrpc", "%s: consolidation input notes locked\n", getId());
+
             if(pwalletMain->CommitAutomatedTx(tx)) {
                 LogPrintf("%s: Committed consolidation transaction with txid=%s\n",opid, tx.GetHash().ToString());
                 amountConsolidated += actualAmountToSend;
@@ -240,10 +251,15 @@ bool AsyncRPCOperation_saplingconsolidation::main_impl() {
                 numTxCreated++;
             } else {
                 LogPrintf("%s: Consolidation transaction FAILED in CommitTransaction, txid=%s\n",opid , tx.GetHash().ToString());
+
+                unlock_notes();
+                LogPrint("zrpc", "%s: consolidatoin input notes unlocked\n", getId());
+
                 setConsolidationResult(numTxCreated, amountConsolidated, consolidationTxIds);
                 status = false;
                 break;
             }
+
         }
     }
 
@@ -275,3 +291,19 @@ UniValue AsyncRPCOperation_saplingconsolidation::getStatus() const {
     obj.push_back(Pair("target_height", targetHeight_));
     return obj;
 }
+
+// Lock input notes
+void AsyncRPCOperation_saplingconsolidation::lock_notes() {
+    LOCK2(cs_main, pwalletMain->cs_wallet);
+    for (auto note : z_sapling_inputs_) {
+        pwalletMain->LockNote(note.op);
+    }
+}
+
+// Unlock input notes
+void AsyncRPCOperation_saplingconsolidation::unlock_notes() {
+    LOCK2(cs_main, pwalletMain->cs_wallet);
+    for (auto note : z_sapling_inputs_) {
+        pwalletMain->UnlockNote(note.op);
+    }
+}
diff --git a/src/wallet/asyncrpcoperation_saplingconsolidation.h b/src/wallet/asyncrpcoperation_saplingconsolidation.h
index 4cc370431..675a56ad7 100644
--- a/src/wallet/asyncrpcoperation_saplingconsolidation.h
+++ b/src/wallet/asyncrpcoperation_saplingconsolidation.h
@@ -8,6 +8,7 @@
 #include "univalue.h"
 #include "zcash/Address.hpp"
 #include "zcash/zip32.h"
+#include "wallet.h" // for SaplingNoteEntry
 
 //Default fee used for consolidation transactions, in puposhis
 static const CAmount DEFAULT_CONSOLIDATION_FEE = 10000;
@@ -32,10 +33,14 @@ public:
 
     virtual UniValue getStatus() const;
 
+    void lock_notes();
+    void unlock_notes();
+
 private:
     int targetHeight_;
 
     bool main_impl();
+    std::vector<SaplingNoteEntry> z_sapling_inputs_;
 
     void setConsolidationResult(int numTxCreated, const CAmount& amountConsolidated, const std::vector<std::string>& consolidationTxIds);