libs/bits: validate BitArray in FromProto (#5720)

Closes #5705

libs/bits/bit_array.go

@@ -2,7 +2,9 @@ package bits
 
 import (
 	"encoding/binary"
+	"errors"
 	"fmt"
+	"math"
 	"regexp"
 	"strings"
 	"sync"
@@ -27,7 +29,7 @@ func NewBitArray(bits int) *BitArray {
 	}
 	return &BitArray{
 		Bits:  bits,
-		Elems: make([]uint64, (bits+63)/64),
+		Elems: make([]uint64, numElems(bits)),
 	}
 }
 
@@ -100,7 +102,7 @@ func (bA *BitArray) copy() *BitArray {
 }
 
 func (bA *BitArray) copyBits(bits int) *BitArray {
-	c := make([]uint64, (bits+63)/64)
+	c := make([]uint64, numElems(bits))
 	copy(c, bA.Elems)
 	return &BitArray{
 		Bits:  bits,
@@ -418,27 +420,45 @@ func (bA *BitArray) UnmarshalJSON(bz []byte) error {
 	return nil
 }
 
-// ToProto converts BitArray to protobuf
+// ToProto converts BitArray to protobuf. It returns nil if BitArray is
+// nil/empty.
+//
+// XXX: It does not copy the array.
 func (bA *BitArray) ToProto() *tmprotobits.BitArray {
-	if bA == nil || len(bA.Elems) == 0 {
+	if bA == nil ||
+		(len(bA.Elems) == 0 && bA.Bits == 0) { // empty
 		return nil
 	}
 
-	return &tmprotobits.BitArray{
-		Bits:  int64(bA.Bits),
-		Elems: bA.Elems,
-	}
+	return &tmprotobits.BitArray{Bits: int64(bA.Bits), Elems: bA.Elems}
 }
 
-// FromProto sets a protobuf BitArray to the given pointer.
-func (bA *BitArray) FromProto(protoBitArray *tmprotobits.BitArray) {
+// FromProto sets BitArray to the given protoBitArray. It returns an error if
+// protoBitArray is invalid.
+//
+// XXX: It does not copy the array.
+func (bA *BitArray) FromProto(protoBitArray *tmprotobits.BitArray) error {
 	if protoBitArray == nil {
-		bA = nil
-		return
+		return nil
+	}
+
+	// Validate protoBitArray.
+	if protoBitArray.Bits < 0 {
+		return errors.New("negative Bits")
+	}
+	// #[32bit]
+	if protoBitArray.Bits > math.MaxInt32 { // prevent overflow on 32bit systems
+		return errors.New("too many Bits")
+	}
+	if got, exp := len(protoBitArray.Elems), numElems(int(protoBitArray.Bits)); got != exp {
+		return fmt.Errorf("invalid number of Elems: got %d, but exp %d", got, exp)
 	}
 
 	bA.Bits = int(protoBitArray.Bits)
-	if len(protoBitArray.Elems) > 0 {
-		bA.Elems = protoBitArray.Elems
-	}
+	bA.Elems = protoBitArray.Elems
+	return nil
+}
+
+func numElems(bits int) int {
+	return (bits + 63) / 64
 }

libs/bits/bit_array_test.go

@@ -4,12 +4,14 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"math"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	tmrand "github.com/tendermint/tendermint/libs/rand"
+	tmprotobits "github.com/tendermint/tendermint/proto/tendermint/libs/bits"
 )
 
 func randBitArray(bits int) (*BitArray, []byte) {
@@ -266,7 +268,7 @@ func TestJSONMarshalUnmarshal(t *testing.T) {
 	}
 }
 
-func TestBitArrayProtoBuf(t *testing.T) {
+func TestBitArrayToFromProto(t *testing.T) {
 	testCases := []struct {
 		msg     string
 		bA1     *BitArray
@@ -280,11 +282,41 @@ func TestBitArrayProtoBuf(t *testing.T) {
 	for _, tc := range testCases {
 		protoBA := tc.bA1.ToProto()
 		ba := new(BitArray)
-		ba.FromProto(protoBA)
+		err := ba.FromProto(protoBA)
 		if tc.expPass {
+			assert.NoError(t, err)
 			require.Equal(t, tc.bA1, ba, tc.msg)
 		} else {
 			require.NotEqual(t, tc.bA1, ba, tc.msg)
 		}
 	}
 }
+
+func TestBitArrayFromProto(t *testing.T) {
+	testCases := []struct {
+		pbA    *tmprotobits.BitArray
+		resA   *BitArray
+		expErr bool
+	}{
+		0: {nil, &BitArray{}, false},
+		1: {&tmprotobits.BitArray{}, &BitArray{}, false},
+
+		2: {&tmprotobits.BitArray{Bits: 1, Elems: make([]uint64, 1)}, &BitArray{Bits: 1, Elems: make([]uint64, 1)}, false},
+
+		3: {&tmprotobits.BitArray{Bits: -1, Elems: make([]uint64, 1)}, &BitArray{}, true},
+		4: {&tmprotobits.BitArray{Bits: math.MaxInt32 + 1, Elems: make([]uint64, 1)}, &BitArray{}, true},
+		5: {&tmprotobits.BitArray{Bits: 1, Elems: make([]uint64, 2)}, &BitArray{}, true},
+	}
+
+	for i, tc := range testCases {
+		bA := new(BitArray)
+		err := bA.FromProto(tc.pbA)
+		if tc.expErr {
+			assert.Error(t, err, "#%d", i)
+			assert.Equal(t, tc.resA, bA, "#%d", i)
+		} else {
+			assert.NoError(t, err, "#%d", i)
+			assert.Equal(t, tc.resA, bA, "#%d", i)
+		}
+	}
+}