From 64ebb9b50953cf3d3c6a80a2b95854bb57f40c0c Mon Sep 17 00:00:00 2001
From: Christophe Deveaux <contact@christophedeveaux.com>
Date: Thu, 15 Sep 2022 11:42:02 +0200
Subject: [PATCH] fix: update to correct L2 when switching networks (#429)

* refactor: Detect corrects L1 and L2 when switching network

* Restore unused component, fix undefined check

* Properly override chainIdToDefaultL2ChainId in RegisterLocalNetwork
---
 .../src/hooks/useNetworksAndSigners.tsx       | 45 ++++++++++++++-----
 .../arb-token-bridge-ui/src/util/networks.ts  | 24 +++++-----
 2 files changed, 47 insertions(+), 22 deletions(-)

diff --git a/packages/arb-token-bridge-ui/src/hooks/useNetworksAndSigners.tsx b/packages/arb-token-bridge-ui/src/hooks/useNetworksAndSigners.tsx
index cdf7002bc3..1b09c973e8 100644
--- a/packages/arb-token-bridge-ui/src/hooks/useNetworksAndSigners.tsx
+++ b/packages/arb-token-bridge-ui/src/hooks/useNetworksAndSigners.tsx
@@ -24,7 +24,7 @@ import {
 import { L1Network, L2Network, getL1Network, getL2Network } from '@arbitrum/sdk'
 import { useWallet } from '@arbitrum/use-wallet'
 import { useLatest } from 'react-use'
-
+import { useHistory } from 'react-router-dom'
 import { chainIdToDefaultL2ChainId, rpcURLs } from '../util/networks'
 import { trackEvent } from '../util/AnalyticsUtils'
 import { modalProviderOpts } from '../util/modelProviderOpts'
@@ -118,7 +118,7 @@ export function NetworksAndSignersProvider(
 ): JSX.Element {
   const { selectedL2ChainId } = props
   const { provider, account, network, connect } = useWallet()
-
+  const history = useHistory()
   const [result, setResult] = useState<UseNetworksAndSignersResult>({
     status: defaultStatus
   })
@@ -156,18 +156,39 @@ export function NetworksAndSignersProvider(
     async (web3Provider: Web3Provider, address: string) => {
       const providerChainId = (await web3Provider.getNetwork()).chainId
 
+      // If provider is not supported, display warning message
+      if (!(providerChainId in chainIdToDefaultL2ChainId)) {
+        console.error(`Provider chainId not supported: ${providerChainId}`)
+        setResult({ status: UseNetworksAndSignersStatus.NOT_SUPPORTED })
+        return
+      }
+
+      /***
+       * Case 1: selectedChainId is undefined => set it to provider's default L2
+       * Case 2: selectedChainId is defined but not supported by provider => reset query params -> case 1
+       * Case 3: selectedChainId is defined and supported, continue
+       */
       let _selectedL2ChainId = selectedL2ChainId
-      if (_selectedL2ChainId === undefined) {
-        // If l2ChainId is undefined, use a default L2 based on the connected provider chainid
-        _selectedL2ChainId = chainIdToDefaultL2ChainId[providerChainId]
-        if (_selectedL2ChainId === undefined) {
-          console.error(`Unknown provider chainId: ${providerChainId}`)
-          setResult({ status: UseNetworksAndSignersStatus.NOT_SUPPORTED })
-          return
-        }
+      const providerSupportedL2 = chainIdToDefaultL2ChainId[providerChainId]
+
+      // Case 1: use a default L2 based on the connected provider chainid
+      _selectedL2ChainId = _selectedL2ChainId || providerSupportedL2[0]
+      if (typeof _selectedL2ChainId === 'undefined') {
+        console.error(`Unknown provider chainId: ${providerChainId}`)
+        setResult({ status: UseNetworksAndSignersStatus.NOT_SUPPORTED })
+        return
+      }
+
+      // Case 2: L2 is not supported by provider
+      if (!providerSupportedL2.includes(_selectedL2ChainId)) {
+        history.replace({
+          pathname: '/'
+        })
+        return
       }
 
-      getL1Network(web3Provider, _selectedL2ChainId!)
+      // Case 3
+      getL1Network(web3Provider, _selectedL2ChainId)
         .then(async l1Network => {
           // Web3Provider is connected to an L1 network. We instantiate a provider for the L2 network.
           const l2Provider = new JsonRpcProvider(rpcURLs[_selectedL2ChainId!])
@@ -232,7 +253,7 @@ export function NetworksAndSignersProvider(
             })
         })
     },
-    [latestResult, selectedL2ChainId]
+    [latestResult, selectedL2ChainId, history]
   )
 
   useEffect(() => {
diff --git a/packages/arb-token-bridge-ui/src/util/networks.ts b/packages/arb-token-bridge-ui/src/util/networks.ts
index c7adc7fd0d..d1ee21ccd2 100644
--- a/packages/arb-token-bridge-ui/src/util/networks.ts
+++ b/packages/arb-token-bridge-ui/src/util/networks.ts
@@ -57,14 +57,14 @@ export const l2LptGatewayAddresses: { [chainId: number]: string } = {
 }
 
 // Default L2 Chain to use for a certain chainId
-export const chainIdToDefaultL2ChainId: { [chainId: number]: number } = {
-  [ChainId.Mainnet]: ChainId.ArbitrumOne,
-  [ChainId.ArbitrumOne]: ChainId.ArbitrumOne,
-  [ChainId.Rinkeby]: ChainId.ArbitrumRinkeby,
-  [ChainId.ArbitrumRinkeby]: ChainId.ArbitrumRinkeby,
-  [ChainId.Goerli]: ChainId.ArbitrumGoerli,
-  [ChainId.ArbitrumGoerli]: ChainId.ArbitrumGoerli,
-  [ChainId.ArbitrumNova]: ChainId.ArbitrumNova
+export const chainIdToDefaultL2ChainId: { [chainId: number]: ChainId[] } = {
+  [ChainId.ArbitrumGoerli]: [ChainId.ArbitrumGoerli],
+  [ChainId.ArbitrumNova]: [ChainId.ArbitrumNova],
+  [ChainId.ArbitrumOne]: [ChainId.ArbitrumOne],
+  [ChainId.ArbitrumRinkeby]: [ChainId.ArbitrumRinkeby],
+  [ChainId.Goerli]: [ChainId.ArbitrumGoerli],
+  [ChainId.Mainnet]: [ChainId.ArbitrumOne, ChainId.ArbitrumNova],
+  [ChainId.Rinkeby]: [ChainId.ArbitrumRinkeby]
 }
 
 export function registerLocalNetwork() {
@@ -88,8 +88,12 @@ export function registerLocalNetwork() {
 
     rpcURLs[customL1Network.chainID] = customL1Network.rpcURL
     rpcURLs[customL2Network.chainID] = customL2Network.rpcURL
-    chainIdToDefaultL2ChainId[customL1Network.chainID] = customL2Network.chainID
-    chainIdToDefaultL2ChainId[customL2Network.chainID] = customL2Network.chainID
+    chainIdToDefaultL2ChainId[customL1Network.chainID] = [
+      customL2Network.chainID
+    ]
+    chainIdToDefaultL2ChainId[customL2Network.chainID] = [
+      customL2Network.chainID
+    ]
 
     addCustomNetwork({ customL1Network, customL2Network })
   } catch (error: any) {