From e8b0458f16747a5c1cc5eea361da945241adb105 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 21 Dec 2017 14:04:05 -0600
Subject: [PATCH] check for overflow and underflow while choosing proposer

Refs #919
---
 types/validator_set.go      | 45 +++++++++++++++++++++++++++++++++----
 types/validator_set_test.go | 44 ++++++++++++++++++++++++++++++++++--
 2 files changed, 83 insertions(+), 6 deletions(-)

diff --git a/types/validator_set.go b/types/validator_set.go
index 134e4e06e..9aaa6830d 100644
--- a/types/validator_set.go
+++ b/types/validator_set.go
@@ -48,12 +48,17 @@ func NewValidatorSet(vals []*Validator) *ValidatorSet {
 }
 
 // incrementAccum and update the proposer
-// TODO: mind the overflow when times and votingPower shares too large.
 func (valSet *ValidatorSet) IncrementAccum(times int) {
 	// Add VotingPower * times to each validator and order into heap.
 	validatorsHeap := cmn.NewHeap()
 	for _, val := range valSet.Validators {
-		val.Accum += val.VotingPower * int64(times) // TODO: mind overflow
+		res, overflow := signedMulWithOverflowCheck(val.VotingPower, int64(times))
+		// check for overflow both multiplication and sum
+		if !overflow && val.Accum <= mostPositive-res {
+			val.Accum += res
+		} else {
+			val.Accum = mostPositive
+		}
 		validatorsHeap.Push(val, accumComparable{val})
 	}
 
@@ -63,7 +68,13 @@ func (valSet *ValidatorSet) IncrementAccum(times int) {
 		if i == times-1 {
 			valSet.Proposer = mostest
 		}
-		mostest.Accum -= int64(valSet.TotalVotingPower())
+
+		// mind underflow
+		if mostest.Accum >= mostNegative+valSet.TotalVotingPower() {
+			mostest.Accum -= valSet.TotalVotingPower()
+		} else {
+			mostest.Accum = mostNegative
+		}
 		validatorsHeap.Update(mostest, accumComparable{mostest})
 	}
 }
@@ -117,7 +128,13 @@ func (valSet *ValidatorSet) Size() int {
 func (valSet *ValidatorSet) TotalVotingPower() int64 {
 	if valSet.totalVotingPower == 0 {
 		for _, val := range valSet.Validators {
-			valSet.totalVotingPower += val.VotingPower
+			// mind overflow
+			if valSet.totalVotingPower <= mostPositive-val.VotingPower {
+				valSet.totalVotingPower += val.VotingPower
+			} else {
+				valSet.totalVotingPower = mostPositive
+				return valSet.totalVotingPower
+			}
 		}
 	}
 	return valSet.totalVotingPower
@@ -425,3 +442,23 @@ func RandValidatorSet(numValidators int, votingPower int64) (*ValidatorSet, []*P
 	sort.Sort(PrivValidatorsByAddress(privValidators))
 	return valSet, privValidators
 }
+
+const mostNegative int64 = -mostPositive - 1
+const mostPositive int64 = 1<<63 - 1
+
+func signedMulWithOverflowCheck(a, b int64) (int64, bool) {
+	if a == 0 || b == 0 {
+		return 0, false
+	}
+	if a == 1 {
+		return b, false
+	}
+	if b == 1 {
+		return a, false
+	}
+	if a == mostNegative || b == mostNegative {
+		return -1, true
+	}
+	c := a * b
+	return c, c/b != a
+}
diff --git a/types/validator_set_test.go b/types/validator_set_test.go
index 572b7b007..c65f507fa 100644
--- a/types/validator_set_test.go
+++ b/types/validator_set_test.go
@@ -5,8 +5,9 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/tendermint/go-crypto"
-	"github.com/tendermint/go-wire"
+	"github.com/stretchr/testify/assert"
+	crypto "github.com/tendermint/go-crypto"
+	wire "github.com/tendermint/go-wire"
 	cmn "github.com/tendermint/tmlibs/common"
 )
 
@@ -190,6 +191,45 @@ func TestProposerSelection3(t *testing.T) {
 	}
 }
 
+func TestValidatorSetIncrementAccumOverflows(t *testing.T) {
+	// NewValidatorSet calls IncrementAccum(1)
+	vset := NewValidatorSet([]*Validator{
+		// too much voting power
+		0: {Address: []byte("a"), VotingPower: mostPositive, Accum: 0},
+		// too big accum
+		1: {Address: []byte("b"), VotingPower: 10, Accum: mostPositive},
+		// almost too big accum
+		2: {Address: []byte("c"), VotingPower: 10, Accum: mostPositive - 5},
+	})
+
+	assert.Equal(t, int64(0), vset.Validators[0].Accum, "0") // because we decrement val with most voting power
+	assert.Equal(t, mostPositive, vset.Validators[1].Accum, "1")
+	assert.Equal(t, mostPositive, vset.Validators[2].Accum, "2")
+}
+
+func TestValidatorSetIncrementAccumUnderflows(t *testing.T) {
+	// NewValidatorSet calls IncrementAccum(1)
+	vset := NewValidatorSet([]*Validator{
+		0: {Address: []byte("a"), VotingPower: mostPositive, Accum: mostNegative},
+		1: {Address: []byte("b"), VotingPower: 1, Accum: mostNegative},
+	})
+
+	vset.IncrementAccum(5)
+
+	assert.Equal(t, mostNegative, vset.Validators[0].Accum, "0")
+	assert.Equal(t, mostNegative, vset.Validators[1].Accum, "1")
+}
+
+func TestValidatorSetTotalVotingPowerOverflows(t *testing.T) {
+	vset := NewValidatorSet([]*Validator{
+		{Address: []byte("a"), VotingPower: mostPositive, Accum: 0},
+		{Address: []byte("b"), VotingPower: mostPositive, Accum: 0},
+		{Address: []byte("c"), VotingPower: mostPositive, Accum: 0},
+	})
+
+	assert.Equal(t, mostPositive, vset.TotalVotingPower())
+}
+
 func BenchmarkValidatorSetCopy(b *testing.B) {
 	b.StopTimer()
 	vset := NewValidatorSet([]*Validator{})