From 5c65d334398cf0ce18879fc6c5f3b5962fad3fe6 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 2 Aug 2021 23:16:12 +0300
Subject: [PATCH] native: move required balance check to token contracts

Which duplicates the check, but deduplicates error path. This check forced
double balance deserialization which is quite costly operation, so we better
do it once.

It's hardly noticeable as of TPS metrics though, maybe some 1-2%%.
---
 pkg/core/native/native_gas.go   |  8 ++++++--
 pkg/core/native/native_neo.go   |  5 +++--
 pkg/core/native/native_nep17.go | 26 ++++----------------------
 3 files changed, 13 insertions(+), 26 deletions(-)

diff --git a/pkg/core/native/native_gas.go b/pkg/core/native/native_gas.go
index 496a22c60..159e7b152 100644
--- a/pkg/core/native/native_gas.go
+++ b/pkg/core/native/native_gas.go
@@ -54,13 +54,17 @@ func newGAS(init int64) *GAS {
 	return g
 }
 
-func (g *GAS) increaseBalance(_ *interop.Context, _ util.Uint160, si *state.StorageItem, amount *big.Int) error {
+func (g *GAS) increaseBalance(_ *interop.Context, _ util.Uint160, si *state.StorageItem, amount *big.Int, checkBal *big.Int) error {
 	acc, err := state.NEP17BalanceFromBytes(*si)
 	if err != nil {
 		return err
 	}
 	if sign := amount.Sign(); sign == 0 {
-		return nil
+		// Requested self-transfer amount can be higher than actual balance.
+		if checkBal != nil && acc.Balance.Cmp(checkBal) < 0 {
+			err = errors.New("insufficient funds")
+		}
+		return err
 	} else if sign == -1 && acc.Balance.Cmp(new(big.Int).Neg(amount)) == -1 {
 		return errors.New("insufficient funds")
 	}
diff --git a/pkg/core/native/native_neo.go b/pkg/core/native/native_neo.go
index e81b658c1..644188aac 100644
--- a/pkg/core/native/native_neo.go
+++ b/pkg/core/native/native_neo.go
@@ -391,12 +391,13 @@ func (n *NEO) getGASPerVote(d dao.DAO, key []byte, index ...uint32) []big.Int {
 	return reward
 }
 
-func (n *NEO) increaseBalance(ic *interop.Context, h util.Uint160, si *state.StorageItem, amount *big.Int) error {
+func (n *NEO) increaseBalance(ic *interop.Context, h util.Uint160, si *state.StorageItem, amount *big.Int, checkBal *big.Int) error {
 	acc, err := state.NEOBalanceFromBytes(*si)
 	if err != nil {
 		return err
 	}
-	if amount.Sign() == -1 && acc.Balance.Cmp(new(big.Int).Neg(amount)) == -1 {
+	if (amount.Sign() == -1 && acc.Balance.Cmp(new(big.Int).Neg(amount)) == -1) ||
+		(amount.Sign() == 0 && checkBal != nil && acc.Balance.Cmp(checkBal) == -1) {
 		return errors.New("insufficient funds")
 	}
 	if err := n.distributeGas(ic, h, acc); err != nil {
diff --git a/pkg/core/native/native_nep17.go b/pkg/core/native/native_nep17.go
index d32d3b925..f08dd0aa1 100644
--- a/pkg/core/native/native_nep17.go
+++ b/pkg/core/native/native_nep17.go
@@ -33,7 +33,7 @@ type nep17TokenNative struct {
 	symbol       string
 	decimals     int64
 	factor       int64
-	incBalance   func(*interop.Context, util.Uint160, *state.StorageItem, *big.Int) error
+	incBalance   func(*interop.Context, util.Uint160, *state.StorageItem, *big.Int, *big.Int) error
 	balFromBytes func(item *state.StorageItem) (*big.Int, error)
 }
 
@@ -173,29 +173,11 @@ func (c *nep17TokenNative) updateAccBalance(ic *interop.Context, acc util.Uint16
 			return errors.New("insufficient funds")
 		}
 		si = state.StorageItem{}
-	} else if amount.Sign() == 0 && requiredBalance != nil {
-		// If amount == 0 then it's either a round trip or an empty transfer. In
-		// case of a round trip account's balance may still be less than actual
-		// transfer's amount, so we need to check it. Other cases are handled by
-		// `incBalance` method.
-		balance, err := c.balFromBytes(&si)
-		if err != nil {
-			return fmt.Errorf("failed to deserialise balance: %w", err)
-		}
-		if balance.Cmp(requiredBalance) < 0 {
-			// Firstly, need to put it back to storage as it affects dumps.
-			err = ic.DAO.PutStorageItem(c.ID, key, si)
-			if err != nil {
-				return err
-			}
-			// Finally, return an error.
-			return errors.New("insufficient funds")
-		}
 	}
 
-	err := c.incBalance(ic, acc, &si, amount)
+	err := c.incBalance(ic, acc, &si, amount, requiredBalance)
 	if err != nil {
-		if si != nil && amount.Sign() < 0 {
+		if si != nil && amount.Sign() <= 0 {
 			_ = ic.DAO.PutStorageItem(c.ID, key, si)
 		}
 		return err
@@ -288,7 +270,7 @@ func (c *nep17TokenNative) addTokens(ic *interop.Context, h util.Uint160, amount
 	if si == nil {
 		si = state.StorageItem{}
 	}
-	if err := c.incBalance(ic, h, &si, amount); err != nil {
+	if err := c.incBalance(ic, h, &si, amount, nil); err != nil {
 		panic(err)
 	}
 	var err error