From 8fe92fab885e0618283336fcd02d7c3b08cfba69 Mon Sep 17 00:00:00 2001
From: Bui Quang Minh <minh.bui@skymavis.com>
Date: Wed, 9 Aug 2023 17:38:49 +0700
Subject: [PATCH] consortium/v2: reject the vote with wrong target number

Currently, a vote with wrong target number is not assembled into block as there
is a check in assembleFinalityVote. In this commit, we move that check to
VerifyVote to reject the vote earlier in vote pool, avoid the vote pool to be
DOSed.
---
 consensus/consortium/v2/consortium.go | 12 +++--
 core/vote/vote_pool_test.go           | 66 ++++++++++++++++++++++++++-
 2 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/consensus/consortium/v2/consortium.go b/consensus/consortium/v2/consortium.go
index ec2b81775f..58e8f80fa0 100644
--- a/consensus/consortium/v2/consortium.go
+++ b/consensus/consortium/v2/consortium.go
@@ -209,6 +209,15 @@ func (c *Consortium) GetRecents(chain consensus.ChainHeaderReader, number uint64
 // VerifyVote check if the finality voter is in the validator set, it assumes the signature is
 // already verified
 func (c *Consortium) VerifyVote(chain consensus.ChainHeaderReader, vote *types.VoteEnvelope) error {
+	header := chain.GetHeaderByHash(vote.Data.TargetHash)
+	if header == nil {
+		return errors.New("header not found")
+	}
+
+	if header.Number.Uint64() != vote.Data.TargetNumber {
+		return errors.New("wrong target number in vote")
+	}
+
 	// Look at the comment assembleFinalityVote in function for the
 	// detailed explanation on the snapshot we need to get to verify the
 	// finality vote.
@@ -1154,9 +1163,6 @@ func (c *Consortium) assembleFinalityVote(header *types.Header, snap *Snapshot)
 						log.Warn("Malformed public key from vote pool", "err", err)
 						continue
 					}
-					if vote.Data.TargetNumber != header.Number.Uint64()-1 {
-						continue
-					}
 					authorized := false
 					for valPosition, validator := range snap.ValidatorsWithBlsPub {
 						if publicKey.Equals(validator.BlsPublicKey) {
diff --git a/core/vote/vote_pool_test.go b/core/vote/vote_pool_test.go
index ada0f3d2e4..738a117518 100644
--- a/core/vote/vote_pool_test.go
+++ b/core/vote/vote_pool_test.go
@@ -367,7 +367,7 @@ func generateVote(
 	secretKey blsCommon.SecretKey,
 ) *types.VoteEnvelope {
 	voteData := types.VoteData{
-		TargetNumber: 1,
+		TargetNumber: uint64(blockNumber),
 		TargetHash:   blockHash,
 	}
 	digest := voteData.Hash()
@@ -469,3 +469,67 @@ func TestVotePoolDosProtection(t *testing.T) {
 		t.Fatalf("Number of future vote per peer, expect %d have %d", 0, votePool.numFutureVotePerPeer["AAAA"])
 	}
 }
+
+type mockPOSAv2 struct {
+	consensus.FastFinalityPoSA
+}
+
+func (p *mockPOSAv2) GetJustifiedNumberAndHash(chain consensus.ChainHeaderReader, header *types.Header) (uint64, common.Hash, error) {
+	parentHeader := chain.GetHeaderByHash(header.ParentHash)
+	if parentHeader == nil {
+		return 0, common.Hash{}, fmt.Errorf("unexpected error")
+	}
+	return parentHeader.Number.Uint64(), parentHeader.Hash(), nil
+}
+
+func (m *mockPOSAv2) VerifyVote(chain consensus.ChainHeaderReader, vote *types.VoteEnvelope) error {
+	header := chain.GetHeaderByHash(vote.Data.TargetHash)
+	if header == nil {
+		return errors.New("header not found")
+	}
+
+	if header.Number.Uint64() != vote.Data.TargetNumber {
+		return errors.New("wrong target number in vote")
+	}
+
+	return nil
+}
+
+func (m *mockPOSAv2) IsActiveValidatorAt(chain consensus.ChainHeaderReader, header *types.Header) bool {
+	return true
+}
+
+func TestVotePoolWrongTargetNumber(t *testing.T) {
+	secretKey, err := bls.RandKey()
+	if err != nil {
+		t.Fatalf("Failed to create secret key, err %s", err)
+	}
+
+	// Create a database pre-initialize with a genesis block
+	db := rawdb.NewMemoryDatabase()
+	genesis := (&core.Genesis{
+		Config:  params.TestChainConfig,
+		Alloc:   core.GenesisAlloc{testAddr: {Balance: big.NewInt(1000000)}},
+		BaseFee: big.NewInt(params.InitialBaseFee),
+	}).MustCommit(db)
+	chain, _ := core.NewBlockChain(db, nil, params.TestChainConfig, ethash.NewFullFaker(), vm.Config{}, nil, nil)
+
+	bs, _ := core.GenerateChain(params.TestChainConfig, genesis, ethash.NewFaker(), db, 1, nil, true)
+	if _, err := chain.InsertChain(bs[:1]); err != nil {
+		panic(err)
+	}
+	mockEngine := &mockPOSAv2{}
+
+	// Create vote pool
+	votePool := NewVotePool(chain, mockEngine, 22)
+
+	// bs[0] is the block 1 so the target block number must be 1.
+	// Here we provide wrong target number 0
+	vote := generateVote(0, bs[0].Hash(), secretKey)
+	votePool.PutVote("AAAA", vote)
+	time.Sleep(100 * time.Millisecond)
+
+	if len(votePool.curVotes) != 0 {
+		t.Fatalf("Current vote length, expect %d have %d", 0, len(votePool.curVotes))
+	}
+}