consensus: check block parts don't exceed maximum block bytes (#5436)

types/part_set.go

@@ -155,6 +155,9 @@ type PartSet struct {
 	parts         []*Part
 	partsBitArray *bits.BitArray
 	count         uint32
+	// a count of the total size (in bytes). Used to ensure that the
+	// part set doesn't exceed the maximum block bytes
+	byteSize int64
 }
 
 // Returns an immutable, full PartSet from the data bytes.
@@ -186,6 +189,7 @@ func NewPartSetFromData(data []byte, partSize uint32) *PartSet {
 		parts:         parts,
 		partsBitArray: partsBitArray,
 		count:         total,
+		byteSize:      int64(len(data)),
 	}
 }
 
@@ -197,6 +201,7 @@ func NewPartSetFromHeader(header PartSetHeader) *PartSet {
 		parts:         make([]*Part, header.Total),
 		partsBitArray: bits.NewBitArray(int(header.Total)),
 		count:         0,
+		byteSize:      0,
 	}
 }
 
@@ -244,6 +249,13 @@ func (ps *PartSet) Count() uint32 {
 	return ps.count
 }
 
+func (ps *PartSet) ByteSize() int64 {
+	if ps == nil {
+		return 0
+	}
+	return ps.byteSize
+}
+
 func (ps *PartSet) Total() uint32 {
 	if ps == nil {
 		return 0
@@ -277,6 +289,7 @@ func (ps *PartSet) AddPart(part *Part) (bool, error) {
 	ps.parts[part.Index] = part
 	ps.partsBitArray.SetIndex(int(part.Index), true)
 	ps.count++
+	ps.byteSize += int64(len(part.Bytes))
 	return true, nil
 }
 

types/part_set_test.go

@@ -17,15 +17,17 @@ const (
 
 func TestBasicPartSet(t *testing.T) {
 	// Construct random data of size partSize * 100
-	data := tmrand.Bytes(testPartSize * 100)
+	nParts := 100
+	data := tmrand.Bytes(testPartSize * nParts)
 	partSet := NewPartSetFromData(data, testPartSize)
 
 	assert.NotEmpty(t, partSet.Hash())
-	assert.EqualValues(t, 100, partSet.Total())
-	assert.Equal(t, 100, partSet.BitArray().Size())
+	assert.EqualValues(t, nParts, partSet.Total())
+	assert.Equal(t, nParts, partSet.BitArray().Size())
 	assert.True(t, partSet.HashesTo(partSet.Hash()))
 	assert.True(t, partSet.IsComplete())
-	assert.EqualValues(t, 100, partSet.Count())
+	assert.EqualValues(t, nParts, partSet.Count())
+	assert.EqualValues(t, testPartSize*nParts, partSet.ByteSize())
 
 	// Test adding parts to a new partSet.
 	partSet2 := NewPartSetFromHeader(partSet.Header())
@@ -49,7 +51,8 @@ func TestBasicPartSet(t *testing.T) {
 	assert.Nil(t, err)
 
 	assert.Equal(t, partSet.Hash(), partSet2.Hash())
-	assert.EqualValues(t, 100, partSet2.Total())
+	assert.EqualValues(t, nParts, partSet2.Total())
+	assert.EqualValues(t, nParts*testPartSize, partSet.ByteSize())
 	assert.True(t, partSet2.IsComplete())
 
 	// Reconstruct data, assert that they are equal.