From 14d98811a59843c70f706513994f25618a8a7f64 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Thu, 2 Nov 2023 14:01:36 +0300
Subject: [PATCH] smartcontract: restrict maximum NEF file size on
 deserialisation

Port https://github.com/neo-project/neo/pull/2939.

Signed-off-by: Anna Shaleva <shaleva.ann@nspcc.ru>
---
 pkg/smartcontract/nef/nef.go      |  9 ++++++---
 pkg/smartcontract/nef/nef_test.go | 27 ++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/pkg/smartcontract/nef/nef.go b/pkg/smartcontract/nef/nef.go
index dd427a656..cd7f92d13 100644
--- a/pkg/smartcontract/nef/nef.go
+++ b/pkg/smartcontract/nef/nef.go
@@ -4,10 +4,12 @@ import (
 	"bytes"
 	"encoding/binary"
 	"errors"
+	"fmt"
 
 	"github.com/nspcc-dev/neo-go/pkg/config"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 )
 
 // NEO Executable Format 3 (NEF3)
@@ -31,8 +33,6 @@ import (
 const (
 	// Magic is a magic File header constant.
 	Magic uint32 = 0x3346454E
-	// MaxScriptLength is the maximum allowed contract script length.
-	MaxScriptLength = 512 * 1024
 	// MaxSourceURLLength is the maximum allowed source URL length.
 	MaxSourceURLLength = 256
 	// compilerFieldSize is the length of `Compiler` File header field in bytes.
@@ -139,7 +139,7 @@ func (n *File) DecodeBinary(r *io.BinReader) {
 		r.Err = errInvalidReserved
 		return
 	}
-	n.Script = r.ReadVarBytes(MaxScriptLength)
+	n.Script = r.ReadVarBytes(stackitem.MaxSize)
 	if r.Err == nil && len(n.Script) == 0 {
 		r.Err = errors.New("empty script")
 		return
@@ -165,6 +165,9 @@ func (n File) Bytes() ([]byte, error) {
 // FileFromBytes returns a NEF File deserialized from the given bytes.
 func FileFromBytes(source []byte) (File, error) {
 	result := File{}
+	if len(source) > stackitem.MaxSize {
+		return result, fmt.Errorf("invalid NEF file size: expected %d at max, got %d", stackitem.MaxSize, len(source))
+	}
 	r := io.NewBinReaderFromBuf(source)
 	result.DecodeBinary(r)
 	if r.Err != nil {
diff --git a/pkg/smartcontract/nef/nef_test.go b/pkg/smartcontract/nef/nef_test.go
index 133513276..9d0e70f7e 100644
--- a/pkg/smartcontract/nef/nef_test.go
+++ b/pkg/smartcontract/nef/nef_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/callflag"
 	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 	"github.com/stretchr/testify/require"
 )
 
@@ -49,7 +50,7 @@ func TestEncodeDecodeBinary(t *testing.T) {
 	})
 
 	t.Run("invalid script length", func(t *testing.T) {
-		newScript := make([]byte, MaxScriptLength+1)
+		newScript := make([]byte, stackitem.MaxSize+1)
 		expected.Script = newScript
 		expected.Checksum = expected.CalculateChecksum()
 		checkDecodeError(t, expected)
@@ -126,6 +127,30 @@ func TestBytesFromBytes(t *testing.T) {
 	require.Equal(t, expected, actual)
 }
 
+func TestNewFileFromBytesLimits(t *testing.T) {
+	expected := File{
+		Header: Header{
+			Magic:    Magic,
+			Compiler: "best compiler version 1",
+		},
+		Tokens: []MethodToken{{
+			Hash:       random.Uint160(),
+			Method:     "someMethod",
+			ParamCount: 3,
+			HasReturn:  true,
+			CallFlag:   callflag.WriteStates,
+		}},
+		Script: make([]byte, stackitem.MaxSize-100),
+	}
+	expected.Checksum = expected.CalculateChecksum()
+
+	bytes, err := expected.Bytes()
+	require.NoError(t, err)
+	_, err = FileFromBytes(bytes)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "invalid NEF file size")
+}
+
 func TestMarshalUnmarshalJSON(t *testing.T) {
 	expected := &File{
 		Header: Header{