no-panic

Author: technicallyty

types/validator_set.go

@@ -323,28 +323,39 @@ func (vals *ValidatorSet) Size() int {
 }
 
 // updateTotalVotingPower forces recalculation of the set's total voting power.
-// Panics if total voting power is bigger than MaxTotalVotingPower.
-func (vals *ValidatorSet) updateTotalVotingPower() {
+// Returns an error if total voting power exceeds MaxTotalVotingPower.
+func (vals *ValidatorSet) updateTotalVotingPower() error {
 	sum := int64(0)
 	for _, val := range vals.Validators {
 		// mind overflow
 		sum = safeAddClip(sum, val.VotingPower)
 		if sum > MaxTotalVotingPower {
-			panic(fmt.Sprintf(
-				"Total voting power should be guarded to not exceed %v; got: %v",
-				MaxTotalVotingPower,
-				sum))
+			return fmt.Errorf("total voting power %d exceeds maximum %d", sum, MaxTotalVotingPower)
 		}
 	}
 
 	vals.totalVotingPower = sum
+	return nil
+}
+
+// TotalVotingPowerSafe returns the sum of the voting powers of all validators,
+// or an error if the total exceeds MaxTotalVotingPower.
+func (vals *ValidatorSet) TotalVotingPowerSafe() (int64, error) {
+	if vals.totalVotingPower == 0 {
+		if err := vals.updateTotalVotingPower(); err != nil {
+			return 0, err
+		}
+	}
+	return vals.totalVotingPower, nil
 }
 
 // TotalVotingPower returns the sum of the voting powers of all validators.
 // It recomputes the total voting power if required.
 func (vals *ValidatorSet) TotalVotingPower() int64 {
 	if vals.totalVotingPower == 0 {
-		vals.updateTotalVotingPower()
+		if err := vals.updateTotalVotingPower(); err != nil {
+			panic(err)
+		}
 	}
 	return vals.totalVotingPower
 }
@@ -682,7 +693,9 @@ func (vals *ValidatorSet) updateWithChangeSet(changes []*Validator, allowDeletes
 	// Should go after additions.
 	vals.checkAllKeysHaveSameType()
 
-	vals.updateTotalVotingPower() // will panic if total voting power > MaxTotalVotingPower
+	if err = vals.updateTotalVotingPower(); err != nil {
+		panic(err)
+	}
 
 	// Scale and center.
 	vals.RescalePriorities(PriorityWindowSizeFactor * vals.TotalVotingPower())
@@ -952,7 +965,10 @@ func ValidatorSetFromProto(vp *cmtproto.ValidatorSet) (*ValidatorSet, error) {
 	// power hence we need to recompute it.
 	// FIXME: We should look to remove TotalVotingPower from proto or add it in the validators hash
 	// so we don't have to do this
-	vals.TotalVotingPower()
+	// NOTE: Use TotalVotingPowerSafe to return error instead of panicking on invalid input.
+	if _, err := vals.TotalVotingPowerSafe(); err != nil {
+		return nil, err
+	}
 
 	return vals, vals.ValidateBasic()
 }
@@ -977,7 +993,9 @@ func ValidatorSetFromExistingValidators(valz []*Validator) (*ValidatorSet, error
 	}
 	vals.checkAllKeysHaveSameType()
 	vals.Proposer = vals.findPreviousProposer()
-	vals.updateTotalVotingPower()
+	if err := vals.updateTotalVotingPower(); err != nil {
+		return nil, err
+	}
 	sort.Sort(ValidatorsByVotingPower(vals.Validators))
 	return vals, nil
 }

types/validator_set_test.go

@@ -15,6 +15,7 @@ import (
 	cmtproto "github.com/cometbft/cometbft/api/cometbft/types/v1"
 	"github.com/cometbft/cometbft/crypto"
 	"github.com/cometbft/cometbft/crypto/ed25519"
+	cryptoenc "github.com/cometbft/cometbft/crypto/encoding"
 	"github.com/cometbft/cometbft/crypto/secp256k1"
 	cmtrand "github.com/cometbft/cometbft/internal/rand"
 	cmtmath "github.com/cometbft/cometbft/libs/math"
@@ -471,6 +472,25 @@ func TestValidatorSetTotalVotingPowerPanicsOnOverflow(t *testing.T) {
 	assert.Panics(t, shouldPanic)
 }
 
+func TestValidatorSetFromProtoReturnsErrorOnOverflow(t *testing.T) {
+	// ValidatorSetFromProto should return an error instead of panicking when total voting power exceeds MaxTotalVotingPower.
+	pubKey := ed25519.GenPrivKey().PubKey()
+	pkProto, err := cryptoenc.PubKeyToProto(pubKey)
+	require.NoError(t, err)
+
+	protoVals := &cmtproto.ValidatorSet{
+		Validators: []*cmtproto.Validator{
+			{Address: pubKey.Address(), PubKey: pkProto, VotingPower: math.MaxInt64, ProposerPriority: 0},
+			{Address: pubKey.Address(), PubKey: pkProto, VotingPower: math.MaxInt64, ProposerPriority: 0},
+		},
+		Proposer: &cmtproto.Validator{Address: pubKey.Address(), PubKey: pkProto, VotingPower: math.MaxInt64, ProposerPriority: 0},
+	}
+
+	_, err = ValidatorSetFromProto(protoVals)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "exceeds maximum")
+}
+
 func TestAvgProposerPriority(t *testing.T) {
 	// Create Validator set without calling IncrementProposerPriority:
 	tcs := []struct {
@@ -833,7 +853,8 @@ func verifyValidatorSet(t *testing.T, valSet *ValidatorSet) {
 
 	// verify that the set's total voting power has been updated
 	tvp := valSet.totalVotingPower
-	valSet.updateTotalVotingPower()
+	err := valSet.updateTotalVotingPower()
+	require.NoError(t, err)
 	expectedTvp := valSet.TotalVotingPower()
 	assert.Equal(t, expectedTvp, tvp,
 		"expected TVP %d. Got %d, valSet=%s", expectedTvp, tvp, valSet)