From 34b1b52784b808eff826d1258ef8fd919666de26 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 24 Mar 2022 17:15:43 +0300
Subject: [PATCH] network: check compressed payload size in `decompress`

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>
---
 pkg/network/compress.go     |  3 +++
 pkg/network/fuzz_test.go    | 26 ++++++++++++++++++++++++++
 pkg/network/message_test.go |  7 +++++++
 3 files changed, 36 insertions(+)
 create mode 100644 pkg/network/fuzz_test.go

diff --git a/pkg/network/compress.go b/pkg/network/compress.go
index a8e73915d..644cad135 100644
--- a/pkg/network/compress.go
+++ b/pkg/network/compress.go
@@ -21,6 +21,9 @@ func compress(source []byte) ([]byte, error) {
 
 // decompress decompresses bytes using lz4.
 func decompress(source []byte) ([]byte, error) {
+	if len(source) < 4 {
+		return nil, errors.New("invalid compressed payload")
+	}
 	length := binary.LittleEndian.Uint32(source[:4])
 	if length > payload.MaxSize {
 		return nil, errors.New("invalid uncompressed payload length")
diff --git a/pkg/network/fuzz_test.go b/pkg/network/fuzz_test.go
new file mode 100644
index 000000000..c6caa171b
--- /dev/null
+++ b/pkg/network/fuzz_test.go
@@ -0,0 +1,26 @@
+//go:build go1.18
+// +build go1.18
+
+package network
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/stretchr/testify/require"
+)
+
+func FuzzMessageDecode(f *testing.F) {
+	for i := 0; i < 100; i++ {
+		seed := make([]byte, rand.Uint32()%1000)
+		rand.Read(seed)
+		f.Add(seed)
+	}
+
+	f.Fuzz(func(t *testing.T, value []byte) {
+		m := new(Message)
+		r := io.NewBinReaderFromBuf(value)
+		require.NotPanics(t, func() { _ = m.Decode(r) })
+	})
+}
diff --git a/pkg/network/message_test.go b/pkg/network/message_test.go
index 38b620189..57cd82d75 100644
--- a/pkg/network/message_test.go
+++ b/pkg/network/message_test.go
@@ -18,6 +18,13 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestMessageDecodeFuzzCases(t *testing.T) {
+	raw := []byte("10\x0200")
+	m := new(Message)
+	r := io.NewBinReaderFromBuf(raw)
+	require.NotPanics(t, func() { _ = m.Decode(r) })
+}
+
 func TestEncodeDecodeVersion(t *testing.T) {
 	// message with tiny payload, shouldn't be compressed
 	expected := NewMessage(CMDVersion, &payload.Version{