From 833bbb1d35482cb64a2eb7f3f651f64c59ade82a Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 5 Oct 2020 12:32:04 +0300
Subject: [PATCH] core: calculate gas for `System.Storage.Put` correctly

If `Put` creates new key, its length should contribute
to a GAS cost of the syscall.
---
 pkg/core/interop_system.go      | 10 ++++-----
 pkg/core/interop_system_test.go | 36 +++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/pkg/core/interop_system.go b/pkg/core/interop_system.go
index efffc3fdb..c1f43d023 100644
--- a/pkg/core/interop_system.go
+++ b/pkg/core/interop_system.go
@@ -389,14 +389,14 @@ func putWithContextAndFlags(ic *interop.Context, stc *StorageContext, key []byte
 		return errors.New("StorageContext is read only")
 	}
 	si := ic.DAO.GetStorageItem(stc.ID, key)
-	if si == nil {
-		si = &state.StorageItem{}
-	}
-	if si.IsConst {
+	if si != nil && si.IsConst {
 		return errors.New("storage item exists and is read-only")
 	}
 	sizeInc := 1
-	if len(value) > len(si.Value) {
+	if si == nil {
+		si = &state.StorageItem{}
+		sizeInc = len(key) + len(value)
+	} else if len(value) > len(si.Value) {
 		sizeInc = len(value) - len(si.Value)
 	}
 	if !ic.VM.AddGas(int64(sizeInc) * StoragePrice) {
diff --git a/pkg/core/interop_system_test.go b/pkg/core/interop_system_test.go
index eac5a680c..365e650b4 100644
--- a/pkg/core/interop_system_test.go
+++ b/pkg/core/interop_system_test.go
@@ -1,6 +1,7 @@
 package core
 
 import (
+	"errors"
 	"math/big"
 	"testing"
 
@@ -327,6 +328,41 @@ func TestBlockchainGetContractState(t *testing.T) {
 	})
 }
 
+func TestStoragePut(t *testing.T) {
+	_, cs, ic, bc := createVMAndContractState(t)
+	defer bc.Close()
+
+	require.NoError(t, ic.DAO.PutContractState(cs))
+
+	initVM := func(t *testing.T, key, value []byte, gas int64) {
+		v := ic.SpawnVM()
+		v.LoadScript(cs.Script)
+		v.GasLimit = gas
+		v.Estack().PushVal(value)
+		v.Estack().PushVal(key)
+		require.NoError(t, storageGetContext(ic))
+	}
+
+	t.Run("create, not enough gas", func(t *testing.T) {
+		initVM(t, []byte{1}, []byte{2, 3}, 2*StoragePrice)
+		err := storagePut(ic)
+		require.True(t, errors.Is(err, errGasLimitExceeded), "got: %v", err)
+	})
+
+	initVM(t, []byte{4}, []byte{5, 6}, 3*StoragePrice)
+	require.NoError(t, storagePut(ic))
+
+	t.Run("update", func(t *testing.T) {
+		t.Run("not enough gas", func(t *testing.T) {
+			initVM(t, []byte{4}, []byte{5, 6, 7, 8}, StoragePrice)
+			err := storagePut(ic)
+			require.True(t, errors.Is(err, errGasLimitExceeded), "got: %v", err)
+		})
+		initVM(t, []byte{4}, []byte{5, 6, 7, 8}, 2*StoragePrice)
+		require.NoError(t, storagePut(ic))
+	})
+}
+
 // getTestContractState returns 2 contracts second of which is allowed to call the first.
 func getTestContractState() (*state.Contract, *state.Contract) {
 	script := []byte{