From 18a5f63e6812d421c55572882a6dfc208a6f6375 Mon Sep 17 00:00:00 2001
From: Hodor <lost1q84@gmail.com>
Date: Thu, 26 Dec 2024 12:27:41 +0900
Subject: [PATCH] feat: support custom network modify rpc url

---
 src/background/controller/wallet.ts           | 47 +++++++++++++++--
 src/background/service/customTestnet.ts       | 52 +++++++++++++------
 .../AssetList/TestnetChainList.tsx            |  4 ++
 src/ui/views/CustomRPC/index.tsx              |  4 +-
 src/ui/views/SendToken/index.tsx              |  7 ++-
 5 files changed, 92 insertions(+), 22 deletions(-)

diff --git a/src/background/controller/wallet.ts b/src/background/controller/wallet.ts
index d8db06d4111..333da424135 100644
--- a/src/background/controller/wallet.ts
+++ b/src/background/controller/wallet.ts
@@ -1783,12 +1783,43 @@ export class WalletController extends BaseController {
   getGasAccountSig = gasAccountService.getGasAccountSig;
   setGasAccountSig = gasAccountService.setGasAccountSig;
 
-  setCustomRPC = RPCService.setRPC;
-  removeCustomRPC = RPCService.removeCustomRPC;
+  setCustomRPC = (chainEnum: CHAINS_ENUM, url: string) => {
+    RPCService.setRPC(chainEnum, url);
+    const chain = findChain({
+      enum: chainEnum,
+    });
+    if (chain?.isTestnet && RPCService.hasCustomRPC(chainEnum)) {
+      customTestnetService.setCustomRPC({ chainId: chain.id, url: url });
+    }
+  };
+  removeCustomRPC = (chainEnum: CHAINS_ENUM) => {
+    RPCService.removeCustomRPC(chainEnum);
+    const chain = findChain({
+      enum: chainEnum,
+    });
+    if (chain?.isTestnet) {
+      customTestnetService.removeCustomRPC(chain.id);
+    }
+  };
   getAllCustomRPC = RPCService.getAllRPC;
   getCustomRpcByChain = RPCService.getRPCByChain;
   pingCustomRPC = RPCService.ping;
-  setRPCEnable = RPCService.setRPCEnable;
+  setRPCEnable = (chainEnum: CHAINS_ENUM, enable: boolean) => {
+    RPCService.setRPCEnable(chainEnum, enable);
+    const chain = findChain({
+      enum: chainEnum,
+    });
+    if (chain?.isTestnet) {
+      if (enable) {
+        customTestnetService.setCustomRPC({
+          chainId: chain.id,
+          url: RPCService.getRPCByChain(chainEnum).url,
+        });
+      } else {
+        customTestnetService.removeCustomRPC(chain.id);
+      }
+    }
+  };
   validateRPC = async (url: string, chainId: number) => {
     const chain = findChain({
       id: chainId,
@@ -4758,7 +4789,15 @@ export class WalletController extends BaseController {
     return res;
   };
   updateCustomTestnet = customTestnetService.update;
-  removeCustomTestnet = customTestnetService.remove;
+  removeCustomTestnet = (chainId: number) => {
+    const chain = findChain({
+      id: chainId,
+    });
+    if (chain?.enum) {
+      RPCService.removeCustomRPC(chain.enum);
+    }
+    customTestnetService.remove(chainId);
+  };
   getCustomTestnetList = customTestnetService.getList;
 
   getCustomTestnetNonce = async ({
diff --git a/src/background/service/customTestnet.ts b/src/background/service/customTestnet.ts
index 6e091453288..6bee6c8b43b 100644
--- a/src/background/service/customTestnet.ts
+++ b/src/background/service/customTestnet.ts
@@ -6,14 +6,7 @@ import { createPersistStore, withTimeout } from 'background/utils';
 import { BigNumber } from 'bignumber.js';
 import { intToHex } from 'ethereumjs-util';
 import { omit, sortBy } from 'lodash';
-import {
-  Client,
-  createClient,
-  defineChain,
-  erc20Abi,
-  http,
-  isAddress,
-} from 'viem';
+import { createClient, defineChain, erc20Abi, http, isAddress } from 'viem';
 import {
   estimateGas,
   getBalance,
@@ -24,6 +17,8 @@ import {
 } from 'viem/actions';
 import { http as axios } from '../utils/http';
 import { matomoRequestEvent } from '@/utils/matomo-request';
+import RPCService, { RPCServiceStore } from './rpc';
+import { storage } from '../webapi';
 
 const MAX_READ_CONTRACT_TIME = 8000;
 
@@ -89,7 +84,7 @@ class CustomTestnetService {
     customTokenList: [],
   };
 
-  chains: Record<string, Client> = {};
+  chains: Record<string, ReturnType<typeof createClientByChain>> = {};
 
   logos: Record<
     string,
@@ -100,16 +95,22 @@ class CustomTestnetService {
   > = {};
 
   init = async () => {
-    const storage = await createPersistStore<CutsomTestnetServiceStore>({
+    const storageCache = await createPersistStore<CutsomTestnetServiceStore>({
       name: 'customTestnet',
       template: {
         customTestnet: {},
         customTokenList: [],
       },
     });
-    this.store = storage || this.store;
+    this.store = storageCache || this.store;
+    const rpcStorage: RPCServiceStore = await storage.get('rpc');
     Object.values(this.store.customTestnet).forEach((chain) => {
-      const client = createClientByChain(chain);
+      const config =
+        rpcStorage.customRPC[chain.enum] &&
+        rpcStorage.customRPC[chain.enum]?.enable
+          ? { ...chain, rpcUrl: rpcStorage.customRPC[chain.enum].url }
+          : chain;
+      const client = createClientByChain(config);
       this.chains[chain.id] = client;
     });
     this.syncChainList();
@@ -176,12 +177,14 @@ class CustomTestnetService {
         },
       };
     }
-
+    const testnetChain = createTestnetChain(chain);
     this.store.customTestnet = {
       ...this.store.customTestnet,
-      [chain.id]: createTestnetChain(chain),
+      [chain.id]: testnetChain,
     };
-    this.chains[chain.id] = createClientByChain(chain);
+    if (!RPCService.hasCustomRPC(testnetChain.enum)) {
+      this.chains[chain.id] = createClientByChain(chain);
+    }
     this.syncChainList();
 
     if (this.getList().length) {
@@ -677,6 +680,25 @@ class CustomTestnetService {
       return {};
     }
   };
+
+  setCustomRPC = ({ chainId, url }: { chainId: number; url: string }) => {
+    const client = this.getClient(chainId);
+    if (client) {
+      this.chains[chainId] = createClientByChain({
+        ...this.store.customTestnet[chainId],
+        rpcUrl: url,
+      });
+    }
+  };
+
+  removeCustomRPC = (chainId: number) => {
+    const client = this.getClient(chainId);
+    if (client) {
+      this.chains[chainId] = createClientByChain(
+        this.store.customTestnet[chainId]
+      );
+    }
+  };
 }
 
 export const customTestnetService = new CustomTestnetService();
diff --git a/src/ui/views/CommonPopup/AssetList/TestnetChainList.tsx b/src/ui/views/CommonPopup/AssetList/TestnetChainList.tsx
index 2e758dcf536..9507082c4d3 100644
--- a/src/ui/views/CommonPopup/AssetList/TestnetChainList.tsx
+++ b/src/ui/views/CommonPopup/AssetList/TestnetChainList.tsx
@@ -53,6 +53,10 @@ export const TestnetChainList = ({
   //   );
   // }
 
+  if (!chainList?.length) {
+    return null;
+  }
+
   return (
     <div
       className={clsx(
diff --git a/src/ui/views/CustomRPC/index.tsx b/src/ui/views/CustomRPC/index.tsx
index 2060a838011..363d5d58fec 100644
--- a/src/ui/views/CustomRPC/index.tsx
+++ b/src/ui/views/CustomRPC/index.tsx
@@ -1,6 +1,6 @@
 import ThemeIcon from '@/ui/component/ThemeMode/ThemeIcon';
 import { useRabbyDispatch, useRabbySelector } from '@/ui/store';
-import { findChain, findChainByEnum } from '@/utils/chain';
+import { findChain, findChainByEnum, getTestnetChainList } from '@/utils/chain';
 import { matomoRequestEvent } from '@/utils/matomo-request';
 import { CHAINS_ENUM } from '@debank/common';
 import { Button, Switch, message } from 'antd';
@@ -345,7 +345,7 @@ const CustomRPC = () => {
         onChange={handleChainChanged}
         onCancel={handleCancelSelectChain}
         showRPCStatus
-        hideTestnetTab
+        hideTestnetTab={!getTestnetChainList().length}
       />
       <EditRPCModal
         visible={rpcModalVisible}
diff --git a/src/ui/views/SendToken/index.tsx b/src/ui/views/SendToken/index.tsx
index 78f3b49b839..23607d52af0 100644
--- a/src/ui/views/SendToken/index.tsx
+++ b/src/ui/views/SendToken/index.tsx
@@ -1263,7 +1263,12 @@ const SendToken = () => {
       let tokenFromOrder: TokenItem | null = null;
 
       const lastTimeToken = await wallet.getLastTimeSendToken(account.address);
-      if (lastTimeToken) {
+      if (
+        lastTimeToken &&
+        findChain({
+          serverId: lastTimeToken.chain,
+        })
+      ) {
         setCurrentToken(lastTimeToken);
       } else {
         const { firstChain } = await dispatch.chains.getOrderedChainList({