From c8120a139d9260843e75b169b99bef4c85923b2b Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Wed, 6 Oct 2021 16:37:23 +0300
Subject: [PATCH] core: add MaxKeyLength restrictions to MPT's operations

---
 pkg/core/mpt/compat_test.go | 8 ++++++--
 pkg/core/mpt/proof.go       | 4 ++++
 pkg/core/mpt/trie.go        | 6 ++++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/pkg/core/mpt/compat_test.go b/pkg/core/mpt/compat_test.go
index f83273e80..707b7987e 100644
--- a/pkg/core/mpt/compat_test.go
+++ b/pkg/core/mpt/compat_test.go
@@ -40,12 +40,13 @@ func prepareMPTCompat() *Trie {
 // There are some differences, though:
 // 1. In our implementation delete is silent, i.e. we do not return an error is the key is missing or empty.
 //    However, we do return error when contents of hash node are missing from the store
+//    (corresponds to exception in C# implementation). However, if the key is too big, an error is returned
 //    (corresponds to exception in C# implementation).
 // 2. In our implementation put returns error if something goes wrong, while C# implementation throws
 //    an exception and returns nothing.
 // 3. In our implementation get does not immediately return error in case of an empty key. An error is returned
 //    only if value is missing from the storage. C# implementation checks that key is not empty and throws an error
-//    otherwice.
+//    otherwice. However, if the key is too big, an error is returned (corresponds to exception in C# implementation).
 func TestCompatibility(t *testing.T) {
 	mainTrie := prepareMPTCompat()
 
@@ -59,6 +60,7 @@ func TestCompatibility(t *testing.T) {
 		tr.testHas(t, []byte{0xac, 0x01, 0x00}, nil)
 		tr.testHas(t, []byte{0xac, 0x99, 0x10}, nil)
 		tr.testHas(t, []byte{0xac, 0xf1}, nil)
+		tr.testHas(t, make([]byte, MaxKeyLength), nil)
 	})
 
 	t.Run("TryGetResolve", func(t *testing.T) {
@@ -95,7 +97,8 @@ func TestCompatibility(t *testing.T) {
 		require.NoError(t, tr.Delete(nil))
 		require.NoError(t, tr.Delete([]byte{0xac, 0x20}))
 
-		require.Error(t, tr.Delete([]byte{0xac, 0xf1})) // error for can't resolve
+		require.Error(t, tr.Delete([]byte{0xac, 0xf1}))           // error for can't resolve
+		require.Error(t, tr.Delete(make([]byte, MaxKeyLength+1))) // error for too big key
 
 		// In our implementation missing keys are ignored.
 		require.NoError(t, tr.Delete([]byte{0xac}))
@@ -202,6 +205,7 @@ func TestCompatibility(t *testing.T) {
 		testGetProof(t, tr, nil, 0)
 		testGetProof(t, tr, []byte{0xac, 0x01, 0x00}, 0)
 		testGetProof(t, tr, []byte{0xac, 0xf1}, 0)
+		testGetProof(t, tr, make([]byte, MaxKeyLength), 0)
 	})
 
 	t.Run("VerifyProof", func(t *testing.T) {
diff --git a/pkg/core/mpt/proof.go b/pkg/core/mpt/proof.go
index 8308a6f63..d3cfb9134 100644
--- a/pkg/core/mpt/proof.go
+++ b/pkg/core/mpt/proof.go
@@ -2,6 +2,7 @@ package mpt
 
 import (
 	"bytes"
+	"errors"
 
 	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
@@ -13,6 +14,9 @@ import (
 // Proof consist of serialized nodes occurring on path from the root to the leaf of key.
 func (t *Trie) GetProof(key []byte) ([][]byte, error) {
 	var proof [][]byte
+	if len(key) > MaxKeyLength {
+		return nil, errors.New("key is too big")
+	}
 	path := toNibbles(key)
 	r, err := t.getProof(t.root, path, &proof)
 	if err != nil {
diff --git a/pkg/core/mpt/trie.go b/pkg/core/mpt/trie.go
index 59970a585..678dab26a 100644
--- a/pkg/core/mpt/trie.go
+++ b/pkg/core/mpt/trie.go
@@ -49,6 +49,9 @@ func NewTrie(root Node, enableRefCount bool, store *storage.MemCachedStore) *Tri
 
 // Get returns value for the provided key in t.
 func (t *Trie) Get(key []byte) ([]byte, error) {
+	if len(key) > MaxKeyLength {
+		return nil, errors.New("key is too big")
+	}
 	path := toNibbles(key)
 	r, bs, err := t.getWithPath(t.root, path)
 	if err != nil {
@@ -235,6 +238,9 @@ func (t *Trie) putIntoNode(curr Node, path []byte, val Node) (Node, error) {
 // Delete removes key from trie.
 // It returns no error on missing key.
 func (t *Trie) Delete(key []byte) error {
+	if len(key) > MaxKeyLength {
+		return errors.New("key is too big")
+	}
 	path := toNibbles(key)
 	r, err := t.deleteFromNode(t.root, path)
 	if err != nil {