Merge pull request #1876 from nspcc-dev/nep17-fix

core: get native nep17 balance from contract storage

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 {

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

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) {