From 314430be1d31f26029305b86d4cbd9aff28056c4 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 1 Jun 2020 10:44:41 +0300
Subject: [PATCH] mpt: implement NEO storage key conversion

---
 cli/server/dump.go           | 31 ++-----------------------------
 pkg/core/mpt/helpers.go      | 36 ++++++++++++++++++++++++++++++++++++
 pkg/core/mpt/helpers_test.go | 30 ++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 29 deletions(-)
 create mode 100644 pkg/core/mpt/helpers_test.go

diff --git a/cli/server/dump.go b/cli/server/dump.go
index ead0b3b33..51ab2399d 100644
--- a/cli/server/dump.go
+++ b/cli/server/dump.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/nspcc-dev/neo-go/pkg/core/mpt"
 	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 )
@@ -33,35 +34,7 @@ func toNeoStorageKey(key []byte) []byte {
 	if len(key) < util.Uint160Size {
 		panic("invalid key in storage")
 	}
-
-	var nkey []byte
-	for i := util.Uint160Size - 1; i >= 0; i-- {
-		nkey = append(nkey, key[i])
-	}
-
-	key = key[util.Uint160Size:]
-
-	index := 0
-	remain := len(key)
-	for remain >= 16 {
-		nkey = append(nkey, key[index:index+16]...)
-		nkey = append(nkey, 0)
-		index += 16
-		remain -= 16
-	}
-
-	if remain > 0 {
-		nkey = append(nkey, key[index:]...)
-	}
-
-	padding := 16 - remain
-	for i := 0; i < padding; i++ {
-		nkey = append(nkey, 0)
-	}
-
-	nkey = append(nkey, byte(padding))
-
-	return nkey
+	return mpt.ToNeoStorageKey(key)
 }
 
 // batchToMap converts batch to a map so that JSON is compatible
diff --git a/pkg/core/mpt/helpers.go b/pkg/core/mpt/helpers.go
index 1c67c6c59..fe59b2917 100644
--- a/pkg/core/mpt/helpers.go
+++ b/pkg/core/mpt/helpers.go
@@ -1,5 +1,7 @@
 package mpt
 
+import "github.com/nspcc-dev/neo-go/pkg/util"
+
 // lcp returns longest common prefix of a and b.
 // Note: it does no allocations.
 func lcp(a, b []byte) []byte {
@@ -33,3 +35,37 @@ func toNibbles(path []byte) []byte {
 	}
 	return result
 }
+
+// ToNeoStorageKey converts storage key to C# neo node's format.
+// Key is expected to be at least 20 bytes in length.
+// our format: script hash in BE + key
+// neo format: script hash in LE + key with 0 between every 16 bytes, padded to len 16.
+func ToNeoStorageKey(key []byte) []byte {
+	const groupSize = 16
+
+	var nkey []byte
+	for i := util.Uint160Size - 1; i >= 0; i-- {
+		nkey = append(nkey, key[i])
+	}
+
+	key = key[util.Uint160Size:]
+
+	index := 0
+	remain := len(key)
+	for remain >= groupSize {
+		nkey = append(nkey, key[index:index+groupSize]...)
+		nkey = append(nkey, 0)
+		index += groupSize
+		remain -= groupSize
+	}
+
+	if remain > 0 {
+		nkey = append(nkey, key[index:]...)
+	}
+
+	padding := groupSize - remain
+	for i := 0; i < padding; i++ {
+		nkey = append(nkey, 0)
+	}
+	return append(nkey, byte(padding))
+}
diff --git a/pkg/core/mpt/helpers_test.go b/pkg/core/mpt/helpers_test.go
new file mode 100644
index 000000000..a27542e8f
--- /dev/null
+++ b/pkg/core/mpt/helpers_test.go
@@ -0,0 +1,30 @@
+package mpt
+
+import (
+	"encoding/hex"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestToNeoStorageKey(t *testing.T) {
+	testCases := []struct{ key, res string }{
+		{
+			"0102030405060708091011121314151617181920",
+			"20191817161514131211100908070605040302010000000000000000000000000000000010",
+		},
+		{
+			"010203040506070809101112131415161718192021222324",
+			"2019181716151413121110090807060504030201212223240000000000000000000000000c",
+		},
+		{
+			"0102030405060708091011121314151617181920212223242526272829303132333435363738",
+			"20191817161514131211100908070605040302012122232425262728293031323334353600373800000000000000000000000000000e",
+		},
+	}
+	for _, tc := range testCases {
+		key, _ := hex.DecodeString(tc.key)
+		res, _ := hex.DecodeString(tc.res)
+		require.Equal(t, res, ToNeoStorageKey(key))
+	}
+}