From 543fd58e9397a3185955d06117f62549402dbd34 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Tue, 29 Sep 2020 10:56:57 +0300
Subject: [PATCH] vm: restrict map key size

---
 docs/compiler.md              |  1 +
 pkg/vm/stackitem/item.go      | 20 ++++++++++++++------
 pkg/vm/stackitem/json.go      |  4 ++--
 pkg/vm/stackitem/json_test.go |  2 +-
 pkg/vm/vm.go                  |  4 ++--
 5 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/docs/compiler.md b/docs/compiler.md
index ebee25c9d..d3166c73b 100644
--- a/docs/compiler.md
+++ b/docs/compiler.md
@@ -24,6 +24,7 @@ a dialect of Go rather than a complete port of the language:
     overhead for all contracts. This can easily be mitigated by first storing values
     in variables and returning the result.
  * lambdas are supported, but closures are not.
+ * maps are supported, but valid map keys are booleans, integers and strings with length <= 64
 
 ## VM API (interop layer)
 Compiler translates interop function calls into NEO VM syscalls or (for custom
diff --git a/pkg/vm/stackitem/item.go b/pkg/vm/stackitem/item.go
index 33292e3f6..ca33f111e 100644
--- a/pkg/vm/stackitem/item.go
+++ b/pkg/vm/stackitem/item.go
@@ -27,6 +27,8 @@ const (
 	// MaxByteArrayComparableSize is the maximum allowed length of ByteArray for Equals method.
 	// It is set to be the maximum uint16 value.
 	MaxByteArrayComparableSize = math.MaxUint16
+	// MaxKeySize is the maximum size of map key.
+	MaxKeySize = 64
 )
 
 // Item represents the "real" value that is pushed on the stack.
@@ -773,8 +775,8 @@ func (i *Map) Convert(typ Type) (Item, error) {
 
 // Add adds key-value pair to the map.
 func (i *Map) Add(key, value Item) {
-	if !IsValidMapKey(key) {
-		panic("wrong key type")
+	if err := IsValidMapKey(key); err != nil {
+		panic(err)
 	}
 	index := i.Index(key)
 	if index >= 0 {
@@ -792,12 +794,18 @@ func (i *Map) Drop(index int) {
 
 // IsValidMapKey checks whether it's possible to use given Item as a Map
 // key.
-func IsValidMapKey(key Item) bool {
+func IsValidMapKey(key Item) error {
 	switch key.(type) {
-	case *Bool, *BigInteger, *ByteArray:
-		return true
+	case *Bool, *BigInteger:
+		return nil
+	case *ByteArray:
+		size := len(key.Value().([]byte))
+		if size > MaxKeySize {
+			return fmt.Errorf("invalid map key size: %d", size)
+		}
+		return nil
 	default:
-		return false
+		return fmt.Errorf("invalid map key of type %s", key.Type())
 	}
 }
 
diff --git a/pkg/vm/stackitem/json.go b/pkg/vm/stackitem/json.go
index 4a4b359a7..a1886ef30 100644
--- a/pkg/vm/stackitem/json.go
+++ b/pkg/vm/stackitem/json.go
@@ -360,8 +360,8 @@ func FromJSONWithTypes(data []byte) (Item, error) {
 			key, err := FromJSONWithTypes(arr[i].Key)
 			if err != nil {
 				return nil, err
-			} else if !IsValidMapKey(key) {
-				return nil, fmt.Errorf("invalid map key of type %s", key.Type())
+			} else if err = IsValidMapKey(key); err != nil {
+				return nil, err
 			}
 			value, err := FromJSONWithTypes(arr[i].Value)
 			if err != nil {
diff --git a/pkg/vm/stackitem/json_test.go b/pkg/vm/stackitem/json_test.go
index bb9e17acd..c8e15dbed 100644
--- a/pkg/vm/stackitem/json_test.go
+++ b/pkg/vm/stackitem/json_test.go
@@ -88,7 +88,7 @@ func TestFromToJSON(t *testing.T) {
 		t.Run("BigNestedArray", getTestDecodeFunc(`[[[[[[[[[[[]]]]]]]]]]]`, nil))
 		t.Run("EncodeRecursive", func(t *testing.T) {
 			// add this item to speed up test a bit
-			item := NewByteArray(make([]byte, MaxSize/100))
+			item := NewByteArray(make([]byte, MaxKeySize))
 			t.Run("Array", func(t *testing.T) {
 				arr := NewArray([]Item{item})
 				arr.Append(arr)
diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 17d6521e6..0aa99207b 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -1668,8 +1668,8 @@ func validateMapKey(key *Element) {
 	if key == nil {
 		panic("no key found")
 	}
-	if !stackitem.IsValidMapKey(key.Item()) {
-		panic("key can't be a collection")
+	if err := stackitem.IsValidMapKey(key.Item()); err != nil {
+		panic(err)
 	}
 }