diff --git a/beacon-chain/p2p/encoder/ssz.go b/beacon-chain/p2p/encoder/ssz.go index 5a0ac1cdc..001ff674f 100644 --- a/beacon-chain/p2p/encoder/ssz.go +++ b/beacon-chain/p2p/encoder/ssz.go @@ -99,14 +99,14 @@ func (e SszNetworkEncoder) doDecode(b []byte, to interface{}) error { // DecodeGossip decodes the bytes to the protobuf gossip message provided. func (e SszNetworkEncoder) DecodeGossip(b []byte, to interface{}) error { - var err error + size, err := snappy.DecodedLen(b) + if uint64(size) > MaxGossipSize { + return errors.Errorf("gossip message exceeds max gossip size: %d bytes > %d bytes", size, MaxGossipSize) + } b, err = snappy.Decode(nil /*dst*/, b) if err != nil { return err } - if uint64(len(b)) > MaxGossipSize { - return errors.Errorf("gossip message exceeds max gossip size: %d bytes > %d bytes", len(b), MaxGossipSize) - } return e.doDecode(b, to) } diff --git a/beacon-chain/p2p/encoder/ssz_test.go b/beacon-chain/p2p/encoder/ssz_test.go index a0e74bead..8f44d3697 100644 --- a/beacon-chain/p2p/encoder/ssz_test.go +++ b/beacon-chain/p2p/encoder/ssz_test.go @@ -2,6 +2,7 @@ package encoder_test import ( "bytes" + "encoding/binary" "fmt" "testing" @@ -19,10 +20,13 @@ func TestSszNetworkEncoder_RoundTrip(t *testing.T) { testRoundTripWithGossip(t, e) } -func TestSszNetworkEncoder_RoundTrip_Snappy(t *testing.T) { +func TestSszNetworkEncoder_FailsSnappyLength(t *testing.T) { e := &encoder.SszNetworkEncoder{} - testRoundTripWithLength(t, e) - testRoundTripWithGossip(t, e) + att := &testpb.TestSimpleMessage{} + data := make([]byte, 32) + binary.PutUvarint(data, encoder.MaxGossipSize+32) + err := e.DecodeGossip(data, att) + require.ErrorContains(t, "gossip message exceeds max gossip size", err) } func testRoundTripWithLength(t *testing.T, e *encoder.SszNetworkEncoder) {