From 2952983c94695c87e90b3f3d735a41908bab81a3 Mon Sep 17 00:00:00 2001
From: ChronosX88 <chronosx88@gmail.com>
Date: Tue, 15 Jun 2021 00:45:35 +0300
Subject: [PATCH] Fix some minor mistakes in various places

---
 consensus/consensus.go           | 28 +++++++++-------------------
 consensus/consensus_validator.go |  2 +-
 consensus/miner.go               |  2 +-
 3 files changed, 11 insertions(+), 21 deletions(-)

diff --git a/consensus/consensus.go b/consensus/consensus.go
index 09da50d..b115bc5 100644
--- a/consensus/consensus.go
+++ b/consensus/consensus.go
@@ -5,6 +5,8 @@ import (
 	"math/big"
 	"sync"
 
+	"github.com/Secured-Finance/dione/cache"
+
 	"github.com/asaskevich/EventBus"
 
 	"github.com/Secured-Finance/dione/blockchain"
@@ -92,6 +94,7 @@ func (pcm *PBFTConsensusManager) propose(blk *types3.Block) error {
 		return err
 	}
 	pcm.psb.BroadcastToServiceTopic(prePrepareMsg)
+	pcm.blockPool.AddBlock(blk)
 	pcm.state.status = StateStatusPrePrepared
 	return nil
 }
@@ -152,18 +155,11 @@ func (pcm *PBFTConsensusManager) handlePrepare(message *pubsub.GenericMessage) {
 		Type:      types.ConsensusMessageTypePrepare,
 		From:      message.From,
 		Blockhash: prepare.Blockhash,
-		Signature: prepare.Signature, // TODO check the signature
+		Signature: prepare.Signature,
 	}
 
-	pk, _ := message.From.ExtractPublicKey()
-	ok, err := pk.Verify(cmsg.Blockhash, cmsg.Signature)
-	if err != nil {
-		logrus.Warnf("Failed to verify PREPARE message signature: %s", err.Error())
-		return
-	}
-
-	if !ok {
-		logrus.Errorf("Signature of PREPARE message of peer %s isn't valid!", cmsg.From)
+	if _, err := pcm.blockPool.GetBlock(cmsg.Blockhash); errors.Is(err, cache.ErrNotFound) {
+		logrus.Debugf("received unknown block")
 		return
 	}
 
@@ -171,6 +167,7 @@ func (pcm *PBFTConsensusManager) handlePrepare(message *pubsub.GenericMessage) {
 		logrus.Debugf("received existing prepare msg, dropping...")
 		return
 	}
+
 	if !pcm.validator.Valid(cmsg, nil) {
 		logrus.Warn("received invalid prepare msg, dropping...")
 		return
@@ -204,15 +201,8 @@ func (pcm *PBFTConsensusManager) handleCommit(message *pubsub.GenericMessage) {
 		Signature: commit.Signature, // TODO check the signature
 	}
 
-	pk, _ := message.From.ExtractPublicKey()
-	ok, err := pk.Verify(cmsg.Blockhash, cmsg.Signature)
-	if err != nil {
-		logrus.Warnf("Failed to verify COMMIT message signature: %s", err.Error())
-		return
-	}
-
-	if !ok {
-		logrus.Errorf("Signature of COMMIT message of peer %s isn't valid!", cmsg.From)
+	if _, err := pcm.blockPool.GetBlock(cmsg.Blockhash); errors.Is(err, cache.ErrNotFound) {
+		logrus.Debugf("received unknown block")
 		return
 	}
 
diff --git a/consensus/consensus_validator.go b/consensus/consensus_validator.go
index 1c43a24..571cc36 100644
--- a/consensus/consensus_validator.go
+++ b/consensus/consensus_validator.go
@@ -137,6 +137,7 @@ func NewConsensusValidator(miner *Miner, bc *blockchain.BlockChain) *ConsensusVa
 			for _, v := range msg.Block.Data {
 				wg.Add(1)
 				go func(v *types3.Transaction, c chan error) {
+					defer wg.Done()
 					if err := utils.VerifyTx(msg.Block.Header, v); err != nil {
 						c <- fmt.Errorf("failed to verify tx: %w", err)
 						return
@@ -157,7 +158,6 @@ func NewConsensusValidator(miner *Miner, bc *blockchain.BlockChain) *ConsensusVa
 					} else {
 						logrus.Debugf("Origin chain [%v]/request type[%v] doesn't have any payload validation!", task.OriginChain, task.RequestType)
 					}
-					wg.Done()
 				}(v, result)
 			}
 			go func() {
diff --git a/consensus/miner.go b/consensus/miner.go
index 42a7317..a4527a3 100644
--- a/consensus/miner.go
+++ b/consensus/miner.go
@@ -108,7 +108,7 @@ func (m *Miner) MineBlock(randomness []byte, lastBlockHeader *types2.BlockHeader
 		return nil, nil
 	}
 
-	txs := m.mempool.GetAllTransactions()
+	txs := m.mempool.GetTransactionsForNewBlock()
 	if txs == nil {
 		return nil, fmt.Errorf("there is no txes for processing") // skip new consensus round because there is no transaction for processing
 	}