From 7c94e61a05458bd505a9af16dec523583cf45647 Mon Sep 17 00:00:00 2001
From: Lucas Santos <lucasgabrielgtx@gmail.com>
Date: Wed, 6 Nov 2024 18:57:26 -0300
Subject: [PATCH] fix: solve glitching balance bug

---
 .../Background/controllers/MainController.ts   | 18 +++++++++++++++++-
 .../Background/controllers/balances/index.ts   |  8 +++++---
 .../Background/controllers/balances/types.ts   |  5 ++++-
 source/scripts/Background/index.ts             |  2 ++
 source/state/store.ts                          |  6 +++++-
 5 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/source/scripts/Background/controllers/MainController.ts b/source/scripts/Background/controllers/MainController.ts
index 8a04dfc4b..43e2c8956 100644
--- a/source/scripts/Background/controllers/MainController.ts
+++ b/source/scripts/Background/controllers/MainController.ts
@@ -125,6 +125,7 @@ class MainController extends KeyringManager {
 
     this.bindMethods();
   }
+
   public setAutolockTimer(minutes: number) {
     store.dispatch(setTimer(minutes));
   }
@@ -1206,6 +1207,7 @@ class MainController extends KeyringManager {
     isBitcoinBased,
     activeNetwork,
     activeAccount,
+    isPolling,
   }: {
     activeAccount: {
       id: number;
@@ -1213,11 +1215,22 @@ class MainController extends KeyringManager {
     };
     activeNetwork: INetwork;
     isBitcoinBased: boolean;
+    isPolling?: boolean;
   }) {
     const { accounts } = store.getState().vault;
 
     const currentAccount = accounts[activeAccount.type][activeAccount.id];
 
+    let internalProvider: CustomJsonRpcProvider | undefined;
+
+    if (isPolling) {
+      const abortController = new AbortController();
+      internalProvider = new CustomJsonRpcProvider(
+        abortController.signal,
+        activeNetwork.url
+      );
+    }
+
     const { currentPromise: balancePromise, cancel } =
       this.cancellablePromises.createCancellablePromise<void>(
         async (resolve, reject) => {
@@ -1226,7 +1239,8 @@ class MainController extends KeyringManager {
               await this.balancesManager.utils.getBalanceUpdatedForAccount(
                 currentAccount,
                 isBitcoinBased,
-                activeNetwork.url
+                activeNetwork.url,
+                internalProvider
               );
 
             const actualUserBalance = isBitcoinBased
@@ -1372,6 +1386,7 @@ class MainController extends KeyringManager {
     store.dispatch(setIsNetworkChanging(false));
     store.dispatch(setIsLoadingBalances(false));
   };
+
   private async configureNetwork(
     network: INetwork,
     chain: string
@@ -1437,6 +1452,7 @@ class MainController extends KeyringManager {
       wallet,
     });
   }
+
   private bindMethods() {
     const proto = Object.getPrototypeOf(this);
     for (const key of Object.getOwnPropertyNames(proto)) {
diff --git a/source/scripts/Background/controllers/balances/index.ts b/source/scripts/Background/controllers/balances/index.ts
index 96df215fb..f9373d9a6 100644
--- a/source/scripts/Background/controllers/balances/index.ts
+++ b/source/scripts/Background/controllers/balances/index.ts
@@ -13,7 +13,8 @@ const BalancesManager = (
   const getBalanceUpdatedForAccount = async (
     currentAccount: IPaliAccount,
     isBitcoinBased: boolean,
-    networkUrl: string
+    networkUrl: string,
+    provider?: CustomJsonRpcProvider
   ) => {
     switch (isBitcoinBased) {
       case true:
@@ -30,8 +31,9 @@ const BalancesManager = (
         }
       case false:
         try {
-          const getEvmBalance =
-            await evmBalanceController.getEvmBalanceForAccount(currentAccount);
+          const getEvmBalance = await EvmBalanceController(
+            provider || web3Provider
+          ).getEvmBalanceForAccount(currentAccount);
 
           return getEvmBalance;
         } catch (evmBalanceError) {
diff --git a/source/scripts/Background/controllers/balances/types.ts b/source/scripts/Background/controllers/balances/types.ts
index 5f78ce85b..34cae1cb5 100644
--- a/source/scripts/Background/controllers/balances/types.ts
+++ b/source/scripts/Background/controllers/balances/types.ts
@@ -1,3 +1,5 @@
+import { CustomJsonRpcProvider } from '@pollum-io/sysweb3-keyring';
+
 import { IPaliAccount } from 'state/vault/types';
 
 export interface IEvmBalanceController {
@@ -15,7 +17,8 @@ export interface IBalancesManagerUtils {
   getBalanceUpdatedForAccount: (
     currentAccount: IPaliAccount,
     isBitcoinBased: boolean,
-    networkUrl: string
+    networkUrl: string,
+    provider?: CustomJsonRpcProvider
   ) => Promise<string>;
 }
 
diff --git a/source/scripts/Background/index.ts b/source/scripts/Background/index.ts
index a12472818..25e4d51ca 100755
--- a/source/scripts/Background/index.ts
+++ b/source/scripts/Background/index.ts
@@ -73,6 +73,7 @@ async function createOffscreen() {
     })
     .catch(() => {});
 }
+
 chrome.runtime.onStartup.addListener(createOffscreen);
 self.onmessage = () => null; // keepAlive
 createOffscreen();
@@ -314,6 +315,7 @@ async function checkForUpdates() {
     walletMethods?.updateUserTransactionsState
   ) {
     walletMethods.updateUserNativeBalance({
+      isPolling: true,
       isBitcoinBased,
       activeNetwork,
       activeAccount,
diff --git a/source/state/store.ts b/source/state/store.ts
index fb67c0ecc..ab279b994 100755
--- a/source/state/store.ts
+++ b/source/state/store.ts
@@ -5,6 +5,7 @@ import {
   Store,
 } from '@reduxjs/toolkit';
 import isEqual from 'lodash/isEqual';
+import logger from 'redux-logger';
 import { thunk } from 'redux-thunk';
 
 import dapp from './dapp';
@@ -27,9 +28,12 @@ const middleware: any = [
 ];
 
 middleware.push(thunk);
-
 const nodeEnv = process.env.NODE_ENV;
 
+if (nodeEnv !== 'production' && nodeEnv !== 'test') {
+  middleware.push(logger as never);
+}
+
 const store: Store<{
   dapp: IDAppState;
   price: IPriceState;