From 8934426cbd0ea7b3ddcf090dd4335a29ce242b09 Mon Sep 17 00:00:00 2001
From: Stefan <stefand@status.im>
Date: Mon, 6 Feb 2023 19:05:58 +0200
Subject: [PATCH] fix: request token balance for overridden tokens

Fix balance for tokens overrides are not requested in case of fall back
when etherscan bulk fetch not working.

Updates status-desktop: #9091
---
 services/wallet/token/token.go      | 33 +++++++++++----
 services/wallet/token/token_test.go | 64 +++++++++++++++++++++++++----
 2 files changed, 81 insertions(+), 16 deletions(-)

diff --git a/services/wallet/token/token.go b/services/wallet/token/token.go
index 01aad13f5..b50e16785 100644
--- a/services/wallet/token/token.go
+++ b/services/wallet/token/token.go
@@ -44,6 +44,7 @@ func (t *Token) IsNative() bool {
 	return t.Address == nativeChainAddress
 }
 
+// Manager is used for accessing token store. It changes the token store based on overridden tokens
 type Manager struct {
 	db             *sql.DB
 	RPCClient      *rpc.Client
@@ -56,27 +57,41 @@ func NewTokenManager(
 	networkManager *network.Manager,
 ) *Manager {
 	tokenManager := &Manager{db, RPCClient, networkManager}
-	// Check the networks' custom tokens to see if we must update the tokenStore
-	networks := networkManager.GetConfiguredNetworks()
+
+	overrideTokensInPlace(networkManager.GetConfiguredNetworks(), tokenStore)
+
+	return tokenManager
+}
+
+// overrideTokensInPlace overrides tokens in the store with the ones from the networks
+// BEWARE: overridden tokens will have their original address removed and replaced by the one in networks
+func overrideTokensInPlace(networks []params.Network, store map[uint64]map[common.Address]*Token) {
 	for _, network := range networks {
 		if len(network.TokenOverrides) == 0 {
 			continue
 		}
 
+		// Map from original address to overridden address
+		overriddenMap := make(map[common.Address]common.Address, len(network.TokenOverrides))
+		tokensMap, ok := store[network.ChainID]
+		if !ok {
+			continue
+		}
 		for _, overrideToken := range network.TokenOverrides {
-			tokensMap, ok := tokenStore[network.ChainID]
-			if !ok {
-				continue
-			}
 			for _, token := range tokensMap {
 				if token.Symbol == overrideToken.Symbol {
-					token.Address = overrideToken.Address
+					overriddenMap[token.Address] = overrideToken.Address
 				}
 			}
 		}
-	}
+		for originalAddress, newAddress := range overriddenMap {
+			newToken := *tokensMap[originalAddress]
+			tokensMap[newAddress] = &newToken
+			newToken.Address = newAddress
 
-	return tokenManager
+			delete(tokensMap, originalAddress)
+		}
+	}
 }
 
 func (tm *Manager) inStore(address common.Address, chainID uint64) bool {
diff --git a/services/wallet/token/token_test.go b/services/wallet/token/token_test.go
index 43c83de24..067b42ce1 100644
--- a/services/wallet/token/token_test.go
+++ b/services/wallet/token/token_test.go
@@ -1,8 +1,6 @@
 package token
 
 import (
-	"io/ioutil"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -10,17 +8,14 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 
 	"github.com/status-im/status-go/appdatabase"
-	"github.com/status-im/status-go/sqlite"
+	"github.com/status-im/status-go/params"
 )
 
 func setupTestTokenDB(t *testing.T) (*Manager, func()) {
-	tmpfile, err := ioutil.TempFile("", "wallet-token-tests-")
-	require.NoError(t, err)
-	db, err := appdatabase.InitializeDB(tmpfile.Name(), "wallet-token-tests", sqlite.ReducedKDFIterationsNumber)
+	db, err := appdatabase.InitializeDB(":memory:", "wallet-token-tests", 1)
 	require.NoError(t, err)
 	return &Manager{db, nil, nil}, func() {
 		require.NoError(t, db.Close())
-		require.NoError(t, os.Remove(tmpfile.Name()))
 	}
 }
 
@@ -56,3 +51,58 @@ func TestCustoms(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 0, len(rst))
 }
+
+func TestTokenOverride(t *testing.T) {
+	networks := []params.Network{
+		{
+			ChainID:   1,
+			ChainName: "TestChain1",
+			TokenOverrides: []params.TokenOverride{
+				{
+					Symbol:  "SNT",
+					Address: common.Address{11},
+				},
+			},
+		}, {
+			ChainID:   2,
+			ChainName: "TestChain2",
+			TokenOverrides: []params.TokenOverride{
+				{
+					Symbol:  "STT",
+					Address: common.Address{33},
+				},
+			},
+		},
+	}
+	testTokenStore := map[uint64]map[common.Address]*Token{
+		1: {
+			common.Address{1}: {
+				Address: common.Address{1},
+				Symbol:  "SNT",
+			},
+			common.Address{2}: {
+				Address: common.Address{2},
+				Symbol:  "TNT",
+			},
+		},
+		2: {
+			common.Address{3}: {
+				Address: common.Address{3},
+				Symbol:  "STT",
+			},
+			common.Address{4}: {
+				Address: common.Address{4},
+				Symbol:  "TTT",
+			},
+		},
+	}
+	overrideTokensInPlace(networks, testTokenStore)
+	_, found := testTokenStore[1][common.Address{1}]
+	require.False(t, found)
+	require.Equal(t, common.Address{11}, testTokenStore[1][common.Address{11}].Address)
+	require.Equal(t, common.Address{2}, testTokenStore[1][common.Address{2}].Address)
+	_, found = testTokenStore[2][common.Address{3}]
+	require.False(t, found)
+	require.Equal(t, common.Address{33}, testTokenStore[2][common.Address{33}].Address)
+	require.Equal(t, common.Address{4}, testTokenStore[2][common.Address{4}].Address)
+}