From 6bf130951fa473e36746c99affaff5dc7cd52fe3 Mon Sep 17 00:00:00 2001
From: dustinface <35775977+xdustinface@users.noreply.github.com>
Date: Fri, 27 Nov 2020 15:56:27 +0100
Subject: [PATCH] wallet: Add methods to store governance objects (#3810)

* wallet: Add methods to store governance objects in the walletdb

* wallet: Cache all governance objects in m_gobjects + load them on start

* wallet: Assert cs_wallet is locked in LoadGovernanceObject

Co-authored-by: UdjinM6 <UdjinM6@users.noreply.github.com>

* wallet: Extend a comment

Co-authored-by: PastaPastaPasta <6443210+PastaPastaPasta@users.noreply.github.com>

Co-authored-by: UdjinM6 <UdjinM6@users.noreply.github.com>
Co-authored-by: PastaPastaPasta <6443210+PastaPastaPasta@users.noreply.github.com>
---
 src/wallet/wallet.cpp   | 24 ++++++++++++++++++++++++
 src/wallet/wallet.h     | 11 +++++++++++
 src/wallet/walletdb.cpp | 21 +++++++++++++++++++++
 src/wallet/walletdb.h   |  4 ++++
 4 files changed, 60 insertions(+)

diff --git a/src/wallet/wallet.cpp b/src/wallet/wallet.cpp
index 9f4e60a60c..7f5fc1d6ec 100644
--- a/src/wallet/wallet.cpp
+++ b/src/wallet/wallet.cpp
@@ -5492,6 +5492,30 @@ void CWallet::NotifyChainLock(const CBlockIndex* pindexChainLock, const llmq::CC
     NotifyChainLockReceived(pindexChainLock->nHeight);
 }
 
+bool CWallet::LoadGovernanceObject(const CGovernanceObject& obj)
+{
+    AssertLockHeld(cs_wallet);
+    return m_gobjects.emplace(obj.GetHash(), obj).second;
+}
+
+bool CWallet::WriteGovernanceObject(const CGovernanceObject& obj)
+{
+    AssertLockHeld(cs_wallet);
+    WalletBatch batch(*database);
+    return batch.WriteGovernanceObject(obj) && LoadGovernanceObject(obj);
+}
+
+std::vector<const CGovernanceObject*> CWallet::GetGovernanceObjects()
+{
+    AssertLockHeld(cs_wallet);
+    std::vector<const CGovernanceObject*> vecObjects;
+    vecObjects.reserve(m_gobjects.size());
+    for (auto& obj : m_gobjects) {
+        vecObjects.push_back(&obj.second);
+    }
+    return vecObjects;
+}
+
 CKeyPool::CKeyPool()
 {
     nTime = GetTime();
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
index 26687f19cc..16dea0d75d 100644
--- a/src/wallet/wallet.h
+++ b/src/wallet/wallet.h
@@ -23,6 +23,7 @@
 #include <wallet/walletdb.h>
 #include <wallet/rpcwallet.h>
 
+#include <governance/governance-object.h>
 #include <privatesend/privatesend.h>
 
 #include <algorithm>
@@ -850,6 +851,9 @@ public:
     // Map from Script ID to key metadata (for watch-only keys).
     std::map<CScriptID, CKeyMetadata> m_script_metadata;
 
+    // Map from governance object hash to governance object, they are added by gobject_prepare.
+    std::map<uint256, CGovernanceObject> m_gobjects;
+
     typedef std::map<unsigned int, CMasterKey> MasterKeyMap;
     MasterKeyMap mapMasterKeys;
     unsigned int nMasterKeyMaxID = 0;
@@ -1234,6 +1238,13 @@ public:
     void NotifyTransactionLock(const CTransaction &tx, const llmq::CInstantSendLock& islock) override;
     void NotifyChainLock(const CBlockIndex* pindexChainLock, const llmq::CChainLockSig& clsig) override;
 
+    /** Load a CGovernanceObject into m_gobjects. */
+    bool LoadGovernanceObject(const CGovernanceObject& obj);
+    /** Store a CGovernanceObject in the wallet database. This should only be used by governance objects that are created by this wallet via `gobject prepare`. */
+    bool WriteGovernanceObject(const CGovernanceObject& obj);
+    /** Returns a vector containing pointers to the governance objects in m_gobjects */
+    std::vector<const CGovernanceObject*> GetGovernanceObjects();
+
     /**
      * Blocks until the wallet state is up-to-date to /at least/ the current
      * chain at the time this function is entered
diff --git a/src/wallet/walletdb.cpp b/src/wallet/walletdb.cpp
index d4cf18c65e..93484e8ee1 100644
--- a/src/wallet/walletdb.cpp
+++ b/src/wallet/walletdb.cpp
@@ -10,6 +10,7 @@
 #include <consensus/tx_verify.h>
 #include <consensus/validation.h>
 #include <fs.h>
+#include <governance/governance-object.h>
 #include <protocol.h>
 #include <serialize.h>
 #include <sync.h>
@@ -178,6 +179,11 @@ bool WalletBatch::WritePrivateSendSalt(const uint256& salt)
     return WriteIC(std::string("ps_salt"), salt);
 }
 
+bool WalletBatch::WriteGovernanceObject(const CGovernanceObject& obj)
+{
+    return WriteIC(std::make_pair(std::string("gobject"), obj.GetHash()), obj, false);
+}
+
 CAmount WalletBatch::GetAccountCreditDebit(const std::string& strAccount)
 {
     std::list<CAccountingEntry> entries;
@@ -547,6 +553,21 @@ ReadKeyValue(CWallet* pwallet, CDataStream& ssKey, CDataStream& ssValue,
                 strErr = "Error reading wallet database: LoadHDPubKey failed";
                 return false;
             }
+        } else if (strType == "gobject") {
+            uint256 nObjectHash;
+            CGovernanceObject obj;
+            ssKey >> nObjectHash;
+            ssValue >> obj;
+
+            if (obj.GetHash() != nObjectHash) {
+                strErr = "Invalid governance object: Hash mismatch";
+                return false;
+            }
+
+            if (!pwallet->LoadGovernanceObject(obj)) {
+                strErr = "Invalid governance object: LoadGovernanceObject";
+                return false;
+            }
         }
         else if (strType != "bestblock" && strType != "bestblock_nomerkle"){
             wss.m_unknown_records++;
diff --git a/src/wallet/walletdb.h b/src/wallet/walletdb.h
index 3477c1eb1b..6dcbddfbca 100644
--- a/src/wallet/walletdb.h
+++ b/src/wallet/walletdb.h
@@ -34,6 +34,7 @@ static const bool DEFAULT_FLUSHWALLET = true;
 class CAccount;
 class CAccountingEntry;
 struct CBlockLocator;
+class CGovernanceObject;
 class CKeyPool;
 class CMasterKey;
 class CScript;
@@ -163,6 +164,9 @@ public:
     bool ReadPrivateSendSalt(uint256& salt);
     bool WritePrivateSendSalt(const uint256& salt);
 
+    /** Write a CGovernanceObject to the database */
+    bool WriteGovernanceObject(const CGovernanceObject& obj);
+
     /// Write destination data key,value tuple to database
     bool WriteDestData(const std::string &address, const std::string &key, const std::string &value);
     /// Erase destination data tuple from wallet database