From 58ab3b5ad114cfde421349adb9382bbf26c0e38e Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Thu, 1 Apr 2021 18:16:43 +0300
Subject: [PATCH] core: get native nep17 balance from contract storage

Not from DAO's items with storage.STNEP17Balances prefix, because
changes are reflected there only after notifications processing. And
this happens only after the transaction script is executed, but there
might be cases when we need to get the balance that was updated
earlier during the same transaction processing.

Affects storage dumps.
---
 pkg/core/native/native_gas.go   |  9 +++++++++
 pkg/core/native/native_neo.go   |  9 +++++++++
 pkg/core/native/native_nep17.go | 24 +++++++++++++++---------
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/pkg/core/native/native_gas.go b/pkg/core/native/native_gas.go
index 61e3b79bb..84916686e 100644
--- a/pkg/core/native/native_gas.go
+++ b/pkg/core/native/native_gas.go
@@ -34,6 +34,7 @@ func newGAS() *GAS {
 	nep17.decimals = 8
 	nep17.factor = GASFactor
 	nep17.incBalance = g.increaseBalance
+	nep17.balFromBytes = g.balanceFromBytes
 
 	g.nep17TokenNative = *nep17
 
@@ -59,6 +60,14 @@ func (g *GAS) increaseBalance(_ *interop.Context, _ util.Uint160, si *state.Stor
 	return nil
 }
 
+func (g *GAS) balanceFromBytes(si *state.StorageItem) (*big.Int, error) {
+	acc, err := state.NEP17BalanceStateFromBytes(*si)
+	if err != nil {
+		return nil, err
+	}
+	return &acc.Balance, err
+}
+
 // Initialize initializes GAS contract.
 func (g *GAS) Initialize(ic *interop.Context) error {
 	if err := g.nep17TokenNative.Initialize(ic); err != nil {
diff --git a/pkg/core/native/native_neo.go b/pkg/core/native/native_neo.go
index 75b2a7a95..0580680fa 100644
--- a/pkg/core/native/native_neo.go
+++ b/pkg/core/native/native_neo.go
@@ -107,6 +107,7 @@ func newNEO() *NEO {
 	nep17.decimals = 0
 	nep17.factor = 1
 	nep17.incBalance = n.increaseBalance
+	nep17.balFromBytes = n.balanceFromBytes
 
 	n.nep17TokenNative = *nep17
 	n.votesChanged.Store(true)
@@ -404,6 +405,14 @@ func (n *NEO) increaseBalance(ic *interop.Context, h util.Uint160, si *state.Sto
 	return nil
 }
 
+func (n *NEO) balanceFromBytes(si *state.StorageItem) (*big.Int, error) {
+	acc, err := state.NEOBalanceStateFromBytes(*si)
+	if err != nil {
+		return nil, err
+	}
+	return &acc.Balance, err
+}
+
 func (n *NEO) distributeGas(ic *interop.Context, h util.Uint160, acc *state.NEOBalanceState) error {
 	if ic.Block == nil || ic.Block.Index == 0 {
 		return nil
diff --git a/pkg/core/native/native_nep17.go b/pkg/core/native/native_nep17.go
index 2d3cde642..bbc55bec7 100644
--- a/pkg/core/native/native_nep17.go
+++ b/pkg/core/native/native_nep17.go
@@ -2,6 +2,7 @@ package native
 
 import (
 	"errors"
+	"fmt"
 	"math"
 	"math/big"
 
@@ -29,10 +30,11 @@ func makeAccountKey(h util.Uint160) []byte {
 // nep17TokenNative represents NEP-17 token contract.
 type nep17TokenNative struct {
 	interop.ContractMD
-	symbol     string
-	decimals   int64
-	factor     int64
-	incBalance func(*interop.Context, util.Uint160, *state.StorageItem, *big.Int) error
+	symbol       string
+	decimals     int64
+	factor       int64
+	incBalance   func(*interop.Context, util.Uint160, *state.StorageItem, *big.Int) error
+	balFromBytes func(item *state.StorageItem) (*big.Int, error)
 }
 
 // totalSupplyKey is the key used to store totalSupply value.
@@ -221,12 +223,16 @@ func (c *nep17TokenNative) TransferInternal(ic *interop.Context, from, to util.U
 
 func (c *nep17TokenNative) balanceOf(ic *interop.Context, args []stackitem.Item) stackitem.Item {
 	h := toUint160(args[0])
-	bs, err := ic.DAO.GetNEP17Balances(h)
-	if err != nil {
-		panic(err)
+	key := makeAccountKey(h)
+	si := ic.DAO.GetStorageItem(c.ID, key)
+	if si == nil {
+		return stackitem.NewBigInteger(big.NewInt(0))
 	}
-	balance := bs.Trackers[c.ID].Balance
-	return stackitem.NewBigInteger(&balance)
+	balance, err := c.balFromBytes(&si)
+	if err != nil {
+		panic(fmt.Errorf("can not deserialize balance state: %w", err))
+	}
+	return stackitem.NewBigInteger(balance)
 }
 
 func (c *nep17TokenNative) mint(ic *interop.Context, h util.Uint160, amount *big.Int, callOnPayment bool) {