From 6cff35c927a912b56e00453834da77f45aa52fc4 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Fri, 2 Oct 2020 19:30:29 +0300
Subject: [PATCH] network: restrict flags size in MerkleBlockPayload

---
 pkg/network/payload/merkleblock.go      |  2 +-
 pkg/network/payload/merkleblock_test.go | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/pkg/network/payload/merkleblock.go b/pkg/network/payload/merkleblock.go
index 0fef31742..9867e4ab0 100644
--- a/pkg/network/payload/merkleblock.go
+++ b/pkg/network/payload/merkleblock.go
@@ -26,7 +26,7 @@ func (m *MerkleBlock) DecodeBinary(br *io.BinReader) {
 	}
 	m.TxCount = txCount
 	br.ReadArray(&m.Hashes, m.TxCount)
-	m.Flags = br.ReadVarBytes()
+	m.Flags = br.ReadVarBytes((txCount + 7) / 8)
 }
 
 // EncodeBinary implements Serializable interface.
diff --git a/pkg/network/payload/merkleblock_test.go b/pkg/network/payload/merkleblock_test.go
index 807367744..30362bb40 100644
--- a/pkg/network/payload/merkleblock_test.go
+++ b/pkg/network/payload/merkleblock_test.go
@@ -54,4 +54,18 @@ func TestMerkleBlock_EncodeDecodeBinary(t *testing.T) {
 		require.NoError(t, err)
 		require.True(t, errors.Is(block.ErrMaxContentsPerBlock, testserdes.DecodeBinary(data, new(MerkleBlock))))
 	})
+
+	t.Run("bad flags size", func(t *testing.T) {
+		b := newDumbBlock()
+		_ = b.Hash()
+		expected := &MerkleBlock{
+			Base:    b,
+			TxCount: 0,
+			Hashes:  []util.Uint256{},
+			Flags:   []byte{1, 2, 3, 4, 5},
+		}
+		data, err := testserdes.EncodeBinary(expected)
+		require.NoError(t, err)
+		require.Error(t, testserdes.DecodeBinary(data, new(MerkleBlock)))
+	})
 }