From 4a664931b44570007010243878a58e577a092e25 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Tue, 9 Nov 2021 10:47:38 -0500
Subject: [PATCH] consensus: add some more checks to vote counting (#7253)
 (#7262)

(cherry picked from commit b3b90f820cd2001e3504142b9b22bb723129ed1b)

Co-authored-by: Callum Waters <cmwaters19@gmail.com>
---
 internal/consensus/peer_state.go | 8 +++++++-
 types/vote_set.go                | 3 +++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/internal/consensus/peer_state.go b/internal/consensus/peer_state.go
index 73e61f21c..6a64e8e10 100644
--- a/internal/consensus/peer_state.go
+++ b/internal/consensus/peer_state.go
@@ -193,7 +193,10 @@ func (ps *PeerState) PickVoteToSend(votes types.VoteSetReader) (*types.Vote, boo
 	}
 
 	if index, ok := votes.BitArray().Sub(psVotes).PickRandom(); ok {
-		return votes.GetByIndex(int32(index)), true
+		vote := votes.GetByIndex(int32(index))
+		if vote != nil {
+			return vote, true
+		}
 	}
 
 	return nil, false
@@ -358,6 +361,9 @@ func (ps *PeerState) BlockPartsSent() int {
 
 // SetHasVote sets the given vote as known by the peer
 func (ps *PeerState) SetHasVote(vote *types.Vote) {
+	if vote == nil {
+		return
+	}
 	ps.mtx.Lock()
 	defer ps.mtx.Unlock()
 
diff --git a/types/vote_set.go b/types/vote_set.go
index b064f2c07..e014ae7bb 100644
--- a/types/vote_set.go
+++ b/types/vote_set.go
@@ -372,6 +372,9 @@ func (voteSet *VoteSet) GetByIndex(valIndex int32) *Vote {
 	}
 	voteSet.mtx.Lock()
 	defer voteSet.mtx.Unlock()
+	if int(valIndex) >= len(voteSet.votes) {
+		return nil
+	}
 	return voteSet.votes[valIndex]
 }