From bdb9748c1b6f05a0551c8abff660dd639dcbe642 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 2 Aug 2021 13:28:12 +0300
Subject: [PATCH] native/std: restrict amount of items in JSON deserialization

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>
---
 pkg/core/native/std.go        |  2 +-
 pkg/vm/stackitem/json.go      | 20 ++++++++++++++++++--
 pkg/vm/stackitem/json_test.go | 21 +++++++++++++++++++--
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/pkg/core/native/std.go b/pkg/core/native/std.go
index 2155883b5..26fc49115 100644
--- a/pkg/core/native/std.go
+++ b/pkg/core/native/std.go
@@ -204,7 +204,7 @@ func (s *Std) jsonDeserialize(_ *interop.Context, args []stackitem.Item) stackit
 		panic(err)
 	}
 
-	item, err := stackitem.FromJSON(data)
+	item, err := stackitem.FromJSON(data, stackitem.MaxDeserialized)
 	if err != nil {
 		panic(err)
 	}
diff --git a/pkg/vm/stackitem/json.go b/pkg/vm/stackitem/json.go
index 2f6088f1e..51a89b43f 100644
--- a/pkg/vm/stackitem/json.go
+++ b/pkg/vm/stackitem/json.go
@@ -15,6 +15,7 @@ import (
 type decoder struct {
 	json.Decoder
 
+	count int
 	depth int
 }
 
@@ -157,8 +158,11 @@ func itemToJSONString(it Item) ([]byte, error) {
 //   null -> Null
 //   array -> Array
 //   map -> Map, keys are UTF-8
-func FromJSON(data []byte) (Item, error) {
-	d := decoder{Decoder: *json.NewDecoder(bytes.NewReader(data))}
+func FromJSON(data []byte, maxCount int) (Item, error) {
+	d := decoder{
+		Decoder: *json.NewDecoder(bytes.NewReader(data)),
+		count:   maxCount,
+	}
 	if item, err := d.decode(); err != nil {
 		return nil, err
 	} else if _, err := d.Token(); err != gio.EOF {
@@ -173,6 +177,12 @@ func (d *decoder) decode() (Item, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	d.count--
+	if d.count < 0 && tok != json.Delim('}') && tok != json.Delim(']') {
+		return nil, errTooBigElements
+	}
+
 	switch t := tok.(type) {
 	case json.Delim:
 		switch t {
@@ -190,6 +200,7 @@ func (d *decoder) decode() (Item, error) {
 			d.depth--
 			return item, err
 		default:
+			d.count++
 			// no error above means corresponding closing token
 			// was encountered for map or array respectively
 			return nil, nil
@@ -234,6 +245,11 @@ func (d *decoder) decodeMap() (*Map, error) {
 		if !ok {
 			return m, nil
 		}
+
+		d.count--
+		if d.count < 0 {
+			return nil, errTooBigElements
+		}
 		val, err := d.decode()
 		if err != nil {
 			return nil, err
diff --git a/pkg/vm/stackitem/json_test.go b/pkg/vm/stackitem/json_test.go
index a6e661115..254ac0222 100644
--- a/pkg/vm/stackitem/json_test.go
+++ b/pkg/vm/stackitem/json_test.go
@@ -11,7 +11,7 @@ import (
 
 func getTestDecodeFunc(js string, expected ...interface{}) func(t *testing.T) {
 	return func(t *testing.T) {
-		actual, err := FromJSON([]byte(js))
+		actual, err := FromJSON([]byte(js), 20)
 		if expected[0] == nil {
 			require.Error(t, err)
 			return
@@ -50,6 +50,14 @@ func TestFromToJSON(t *testing.T) {
 			NewArray([]Item{NewBigInteger(big.NewInt(1)), NewByteArray([]byte("test")), NewBool(true), Null{}})))
 		t.Run("Nested", getTestDecodeFunc(`[[],[{},null]]`,
 			NewArray([]Item{NewArray([]Item{}), NewArray([]Item{NewMap(), Null{}})})))
+		t.Run("ManyElements", func(t *testing.T) {
+			js := `[1, 2, 3]` // 3 elements + array itself
+			_, err := FromJSON([]byte(js), 4)
+			require.NoError(t, err)
+
+			_, err = FromJSON([]byte(js), 3)
+			require.True(t, errors.Is(err, errTooBigElements), err)
+		})
 	})
 	t.Run("Map", func(t *testing.T) {
 		small := NewMap()
@@ -64,6 +72,15 @@ func TestFromToJSON(t *testing.T) {
 		m := NewMap()
 		m.Add(NewByteArray([]byte("\t")), NewBool(true))
 		t.Run("escape keys", getTestDecodeFunc(`{"\t":true}`, m))
+
+		t.Run("ManyElements", func(t *testing.T) {
+			js := `{"a":1,"b":3}` // 4 elements + map itself
+			_, err := FromJSON([]byte(js), 5)
+			require.NoError(t, err)
+
+			_, err = FromJSON([]byte(js), 4)
+			require.True(t, errors.Is(err, errTooBigElements), err)
+		})
 	})
 	t.Run("Invalid", func(t *testing.T) {
 		t.Run("Empty", getTestDecodeFunc(``, nil))
@@ -114,7 +131,7 @@ func testToJSON(t *testing.T, expectedErr error, item Item) {
 	}
 	require.NoError(t, err)
 
-	actual, err := FromJSON(data)
+	actual, err := FromJSON(data, 1024)
 	require.NoError(t, err)
 	require.Equal(t, item, actual)
 }