From 0bcf57603c8a9cdb92adb5c3c0d863176a324df6 Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Thu, 17 Dec 2020 17:35:32 +0100
Subject: [PATCH] Effective amounts add missing checks

---
 db/statedb/txprocessors.go      | 22 +++++++++++++++++++--
 db/statedb/txprocessors_test.go | 34 +++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/db/statedb/txprocessors.go b/db/statedb/txprocessors.go
index d14a5b6..6c2648b 100644
--- a/db/statedb/txprocessors.go
+++ b/db/statedb/txprocessors.go
@@ -1107,6 +1107,20 @@ func (s *StateDB) computeEffectiveAmounts(tx *common.L1Tx) {
 			tx.EffectiveAmount = big.NewInt(0)
 			return
 		}
+
+		// check if tx.TokenID==receiver.TokenID
+		accReceiver, err := s.GetAccount(tx.ToIdx)
+		if err != nil {
+			log.Debugf("EffectiveAmount & EffectiveDepositAmount = 0: can not get account for tx.ToIdx: %d", tx.ToIdx)
+			tx.EffectiveDepositAmount = big.NewInt(0)
+			tx.EffectiveAmount = big.NewInt(0)
+			return
+		}
+		if tx.TokenID != accReceiver.TokenID {
+			log.Debugf("EffectiveAmount = 0: tx TokenID (%d) != receiver account TokenID (%d)", tx.TokenID, accReceiver.TokenID)
+			tx.EffectiveAmount = big.NewInt(0)
+			return
+		}
 		return
 	}
 
@@ -1159,8 +1173,12 @@ func (s *StateDB) computeEffectiveAmounts(tx *common.L1Tx) {
 		return
 	}
 	if accSender.TokenID != accReceiver.TokenID {
-		log.Debugf("EffectiveAmount & EffectiveDepositAmount = 0: sender account TokenID (%d) != receiver account TokenID (%d)", accSender.TokenID, accReceiver.TokenID)
-		tx.EffectiveDepositAmount = big.NewInt(0)
+		log.Debugf("EffectiveAmount = 0: sender account TokenID (%d) != receiver account TokenID (%d)", accSender.TokenID, accReceiver.TokenID)
+		tx.EffectiveAmount = big.NewInt(0)
+		return
+	}
+	if tx.TokenID != accReceiver.TokenID {
+		log.Debugf("EffectiveAmount & EffectiveDepositAmount = 0: tx TokenID (%d) != receiver account TokenID (%d)", tx.TokenID, accReceiver.TokenID)
 		tx.EffectiveAmount = big.NewInt(0)
 		return
 	}
diff --git a/db/statedb/txprocessors_test.go b/db/statedb/txprocessors_test.go
index 87a218e..1d34999 100644
--- a/db/statedb/txprocessors_test.go
+++ b/db/statedb/txprocessors_test.go
@@ -160,6 +160,40 @@ func TestComputeEffectiveAmounts(t *testing.T) {
 	sdb.computeEffectiveAmounts(&tx)
 	assert.Equal(t, big.NewInt(8), tx.EffectiveDepositAmount)
 	assert.Equal(t, big.NewInt(0), tx.EffectiveAmount)
+
+	// CreateAccountDepositTransfer for TokenID=1 when receiver does not
+	// have an account for that TokenID, expect that the
+	// EffectiveDepositAmount=DepositAmount, but EffectiveAmount==0
+	tx = common.L1Tx{
+		FromIdx:       0,
+		ToIdx:         257,
+		Amount:        big.NewInt(8),
+		DepositAmount: big.NewInt(8),
+		FromEthAddr:   tc.Users["A"].Addr,
+		TokenID:       2,
+		UserOrigin:    true,
+		Type:          common.TxTypeCreateAccountDepositTransfer,
+	}
+	sdb.computeEffectiveAmounts(&tx)
+	assert.Equal(t, big.NewInt(8), tx.EffectiveDepositAmount)
+	assert.Equal(t, big.NewInt(0), tx.EffectiveAmount)
+
+	// DepositTransfer for TokenID=1 when receiver does not have an account
+	// for that TokenID, expect that the
+	// EffectiveDepositAmount=DepositAmount, but EffectiveAmount=0
+	tx = common.L1Tx{
+		FromIdx:       258,
+		ToIdx:         256,
+		Amount:        big.NewInt(8),
+		DepositAmount: big.NewInt(8),
+		FromEthAddr:   tc.Users["C"].Addr,
+		TokenID:       1,
+		UserOrigin:    true,
+		Type:          common.TxTypeDepositTransfer,
+	}
+	sdb.computeEffectiveAmounts(&tx)
+	assert.Equal(t, big.NewInt(8), tx.EffectiveDepositAmount)
+	assert.Equal(t, big.NewInt(0), tx.EffectiveAmount)
 }
 
 func TestProcessTxsBalances(t *testing.T) {