From f93bc0f40477bfb61ee329121c8156fc8fa144d8 Mon Sep 17 00:00:00 2001
From: random-zebra <random.zebra@protonmail.com>
Date: Thu, 11 Feb 2021 06:03:49 +0100
Subject: [PATCH] [Budget] Validate proposal votes from deterministic
 masternodes

---
 src/budget/budgetmanager.cpp | 41 ++++++++++++++++++++++++++++++++----
 1 file changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/budget/budgetmanager.cpp b/src/budget/budgetmanager.cpp
index dbffb3b1ffdc0..91926b59d4691 100644
--- a/src/budget/budgetmanager.cpp
+++ b/src/budget/budgetmanager.cpp
@@ -5,6 +5,7 @@
 
 #include "budget/budgetmanager.h"
 
+#include "evo/deterministicmns.h"
 #include "masternode-sync.h"
 #include "masternodeman.h"
 #include "net_processing.h"
@@ -533,7 +534,7 @@ void CBudgetManager::VoteOnFinalizedBudgets()
             LogPrintf("%s: Error submitting vote - %s\n", __func__, strError);
             continue;
         }
-        LogPrint(BCLog::MNBUDGET,"%s: new finalized budget vote signed: %s\n", __func__, vote.GetHash().ToString());
+        LogPrint(BCLog::MNBUDGET, "%s: new finalized budget vote signed: %s\n", __func__, vote.GetHash().ToString());
         AddSeenFinalizedBudgetVote(vote);
         vote.Relay();
     }
@@ -800,10 +801,18 @@ void CBudgetManager::RemoveStaleVotesOnProposal(CBudgetProposal* prop)
 
     auto it = prop->mapVotes.begin();
     while (it != prop->mapVotes.end()) {
-        CMasternode* pmn = mnodeman.Find(it->first);
-        (*it).second.SetValid(pmn && pmn->IsEnabled());
+        auto mnList = deterministicMNManager->GetListAtChainTip();
+        auto dmn = mnList.GetMNByCollateral(it->first);
+        if (dmn) {
+            (*it).second.SetValid(mnList.IsMNValid(dmn));
+        } else {
+            // -- Legacy System (!TODO: remove after enforcement) --
+            CMasternode* pmn = mnodeman.Find(it->first);
+            (*it).second.SetValid(pmn && pmn->IsEnabled());
+        }
         ++it;
     }
+
     LogPrint(BCLog::MNBUDGET, "Cleaned proposal votes for %s. After: YES=%d, NO=%d\n",
             prop->GetName(), prop->GetYeas(), prop->GetNays());
 }
@@ -967,6 +976,30 @@ int CBudgetManager::ProcessProposalVote(CBudgetVote& vote, CNode* pfrom)
     }
 
     const CTxIn& voteVin = vote.GetVin();
+
+    // See if this vote was signed with a deterministic masternode
+    auto mnList = deterministicMNManager->GetListAtChainTip();
+    auto dmn = mnList.GetMNByCollateral(voteVin.prevout);
+    if (dmn) {
+        if (!vote.CheckSignature(dmn->pdmnState->keyIDVoting)) {
+            LogPrint(BCLog::MNBUDGET, "mvote - signature invalid from dmn %s\n", dmn->proTxHash.ToString());
+            return 100;
+        }
+        AddSeenProposalVote(vote);
+        std::string strError = "masternode not valid or PoSe banned";
+        if (mnList.IsMNValid(dmn) && UpdateProposal(vote, pfrom, strError)) {
+            vote.Relay();
+            const uint256& voteID = vote.GetHash();
+            masternodeSync.AddedBudgetItem(voteID);
+            LogPrint(BCLog::MNBUDGET, "mvote - new vote (%s) for proposal %s from dmn %s\n",
+                    voteID.ToString(), vote.GetProposalHash().ToString(), dmn->proTxHash.ToString());
+        } else {
+            LogPrint(BCLog::MNBUDGET, "mvote - error: %s", strError);
+        }
+        return 0;
+    }
+
+    // -- Legacy System (!TODO: remove after enforcement) --
     CMasternode* pmn = mnodeman.Find(voteVin.prevout);
     if (!pmn) {
         LogPrint(BCLog::MNBUDGET, "mvote - unknown masternode - vin: %s\n", voteVin.ToString());
@@ -978,7 +1011,7 @@ int CBudgetManager::ProcessProposalVote(CBudgetVote& vote, CNode* pfrom)
 
     if (!vote.CheckSignature(pmn->pubKeyMasternode.GetID())) {
         if (masternodeSync.IsSynced()) {
-            LogPrintf("mvote - signature invalid\n");
+            LogPrint(BCLog::MNBUDGET, "mvote - signature invalid\n");
             return 20;
         }
         // it could just be a non-synced masternode