[Wallet] Implement auto-locking of masternode collaterals

src/evo/providertx.cpp

@@ -246,3 +246,20 @@ void ProRegPL::ToJson(UniValue& obj) const
     obj.pushKV("operatorReward", (double)nOperatorReward / 100);
     obj.pushKV("inputsHash", inputsHash.ToString());
 }
+
+bool GetProRegCollateral(const CTransactionRef& tx, COutPoint& outRet)
+{
+    if (tx == nullptr) {
+        return false;
+    }
+    if (!tx->IsSpecialTx() || tx->nType != CTransaction::TxType::PROREG) {
+        return false;
+    }
+    ProRegPL pl;
+    if (!GetTxPayload(*tx, pl)) {
+        return false;
+    }
+    outRet = pl.collateralOutpoint.hash.IsNull() ? COutPoint(tx->GetHash(), pl.collateralOutpoint.n)
+                                                 : pl.collateralOutpoint;
+    return true;
+}

src/evo/providertx.h

@@ -70,4 +70,8 @@ class ProRegPL
 
 bool CheckProRegTx(const CTransaction& tx, const CBlockIndex* pindexPrev, CValidationState& state);
 
+// If tx is a ProRegTx, return the collateral outpoint in outRet.
+bool GetProRegCollateral(const CTransactionRef& tx, COutPoint& outRet);
+
+
 #endif  //PIVX_PROVIDERTX_H

src/init.cpp

@@ -22,6 +22,7 @@
 #include "checkpoints.h"
 #include "compat/sanity.h"
 #include "consensus/upgrades.h"
+#include "evo/deterministicmns.h"
 #include "evo/evonotificationinterface.h"
 #include "fs.h"
 #include "httpserver.h"
@@ -1860,6 +1861,7 @@ bool AppInitMain()
     strBudgetMode = gArgs.GetArg("-budgetvotemode", "auto");
 
 #ifdef ENABLE_WALLET
+    // !TODO: remove after complete transition to DMN
     // use only the first wallet here. This section can be removed after transition to DMN
     if (gArgs.GetBoolArg("-mnconflock", DEFAULT_MNCONFLOCK) && !vpwallets.empty() && vpwallets[0]) {
         LOCK(vpwallets[0]->cs_wallet);
@@ -1873,6 +1875,14 @@ bool AppInitMain()
                       mne.getAlias(), mne.getTxHash(), mne.getOutputIndex());
         }
     }
+
+    // automatic lock for DMN
+    if (gArgs.GetBoolArg("-mnconflock", DEFAULT_MNCONFLOCK)) {
+        const auto& mnList = deterministicMNManager->GetListAtChainTip();
+        for (CWallet* pwallet : vpwallets) {
+            pwallet->ScanMasternodeCollateralsAndLock(mnList);
+        }
+    }
 #endif
 
     // lite mode disables all Masternode related functionality

src/wallet/wallet.cpp

@@ -9,6 +9,7 @@
 
 #include "budget/budgetmanager.h"
 #include "coincontrol.h"
+#include "evo/deterministicmns.h"
 #include "init.h"
 #include "guiinterfaceutil.h"
 #include "masternode.h"
@@ -963,7 +964,7 @@ bool CWallet::AddToWallet(const CWalletTx& wtxIn, bool fFlushOnClose)
 {
     LOCK(cs_wallet);
     CWalletDB walletdb(*dbw, "r+", fFlushOnClose);
-    uint256 hash = wtxIn.GetHash();
+    const uint256& hash = wtxIn.GetHash();
 
     // Inserts only if not already there, returns tx inserted or tx found
     std::pair<std::map<uint256, CWalletTx>::iterator, bool> ret = mapWallet.emplace(hash, wtxIn);
@@ -1164,6 +1165,9 @@ bool CWallet::AddToWalletIfInvolvingMe(const CTransactionRef& ptx, const CWallet
             }
         }
 
+        // If this is a ProRegTx and the wallet owns the collateral, lock the corresponding coin
+        LockIfMyCollateral(ptx);
+
         bool isFromMe = IsFromMe(ptx);
         if (fExisted || IsMine(ptx) || isFromMe || (saplingNoteData && !saplingNoteData->empty())) {
 
@@ -4153,6 +4157,47 @@ void CWallet::AutoCombineDust(CConnman* connman)
     }
 }
 
+void CWallet::LockOutpointIfMine(const CTransactionRef& ptx, const COutPoint& c)
+{
+    AssertLockHeld(cs_wallet);
+    CTxOut txout;
+    if (ptx && c.hash == ptx->GetHash() && c.n < ptx->vout.size()) {
+        // the collateral is an output of this tx
+        txout = ptx->vout[c.n];
+    } else {
+        // the collateral is a reference to an utxo inside this wallet
+        const auto& it = mapWallet.find(c.hash);
+        if (it != mapWallet.end()) {
+            txout = it->second.tx->vout[c.n];
+        }
+    }
+    if (!txout.IsNull() && IsMine(txout) != ISMINE_NO && !IsSpent(c)) {
+        LockCoin(c);
+    }
+}
+
+// Called during Init
+void CWallet::ScanMasternodeCollateralsAndLock(const CDeterministicMNList& mnList)
+{
+    LOCK(cs_wallet);
+
+    LogPrintf("Locking masternode collaterals...\n");
+    mnList.ForEachMN(false, [&](const CDeterministicMNCPtr& dmn) {
+        LockOutpointIfMine(nullptr, dmn->collateralOutpoint);
+    });
+}
+
+// Called from AddToWalletIfInvolvingMe
+void CWallet::LockIfMyCollateral(const CTransactionRef& ptx)
+{
+    AssertLockHeld(cs_wallet);
+
+    COutPoint o;
+    if (GetProRegCollateral(ptx, o)) {
+        LockOutpointIfMine(ptx, o);
+    }
+}
+
 std::string CWallet::GetWalletHelpString(bool showDebug)
 {
     std::string strUsage = HelpMessageGroup(_("Wallet options:"));

src/wallet/wallet.h

@@ -98,6 +98,7 @@ class CScheduler;
 class ScriptPubKeyMan;
 class SaplingScriptPubKeyMan;
 class SaplingNoteData;
+class CDeterministicMNList;
 
 /** (client) version numbers for particular wallet features */
 enum WalletFeature {
@@ -847,6 +848,29 @@ class CWallet : public CCryptoKeyStore, public CValidationInterface
     bool SetStakeSplitThreshold(const CAmount sst);
     CAmount GetStakeSplitThreshold() const { LOCK(cs_wallet); return nStakeSplitThreshold; }
 
+    /*
+     *  Requires cs_wallet lock.
+     *  Lock for spending the coin c, if it's owned by the wallet, it's unspent, and:
+     *  -- If ptx is not null, c is one of the outputs of *ptx
+     *  -- If ptx is null, c is the output of a transaction in mapWallet
+     */
+    void LockOutpointIfMine(const CTransactionRef& ptx, const COutPoint& c);
+
+    /*
+     *  Locks cs_wallet
+     *  Called during Init. If a DMN collateral is found in the wallet,
+     *  lock the corresponding coin, to prevent accidental spending.
+     */
+    void ScanMasternodeCollateralsAndLock(const CDeterministicMNList& mnList);
+
+    /*
+     *  Requires cs_wallet lock.
+     *  Called from AddToWalletIfInvolvingMe. If ptx is a ProRegTx, and the
+     *  collateral (either referenced or created) is owned by this wallet,
+     *  lock the corresponding coin, to prevent accidental spending.
+     */
+    void LockIfMyCollateral(const CTransactionRef& ptx);
+
     //  keystore implementation
     PairResult getNewAddress(CTxDestination& ret, const std::string addressLabel, const std::string purpose,
                                            const CChainParams::Base58Type addrType = CChainParams::PUBKEY_ADDRESS);