feat(useWallet): add correctNetwork flag (#74)

Dhaiwat10 · web-flow · commit 5e89859585be · 2021-12-07T09:53:25.000+05:30
* feat(useWallet): add correctNetwork flag

* Remove unused import
diff --git a/packages/components/src/components/Address/Address.stories.tsx b/packages/components/src/components/Address/Address.stories.tsx
@@ -1,7 +1,7 @@
 import React from 'react';
 import { Address } from '.';
 
-import { useWallet, Provider } from '@web3-ui/hooks';
+import { useWallet, Provider, NETWORKS } from '@web3-ui/hooks';
 import { Button } from '@chakra-ui/react';
 
 export default {
@@ -25,7 +25,7 @@ const AddressUsingProvider = () => {
 };
 
 export const WithWallet = () => (
-  <Provider network='rinkeby'>
+  <Provider network={NETWORKS.rinkeby}>
     <AddressUsingProvider />
   </Provider>
 );
diff --git a/packages/components/src/components/AddressInput/AddressInput.stories.tsx b/packages/components/src/components/AddressInput/AddressInput.stories.tsx
@@ -1,7 +1,7 @@
 import React from 'react';
 import { AddressInput } from '.';
 import { ethers } from 'ethers';
-import { Provider, useWallet } from '@web3-ui/hooks';
+import { NETWORKS, Provider, useWallet } from '@web3-ui/hooks';
 import { Text } from '@chakra-ui/layout';
 
 export default {
@@ -41,7 +41,7 @@ const Component = ({ ...props }) => {
 export const Default = () => <Component />;
 export const UsingWeb3Hooks = () => {
   return (
-    <Provider network='mainnet'>
+    <Provider network={NETWORKS.mainnet}>
       <WithUseWallet />
     </Provider>
   );
diff --git a/packages/hooks/src/Provider.test.tsx b/packages/hooks/src/Provider.test.tsx
@@ -2,10 +2,11 @@ import React from 'react';
 import { render } from '@testing-library/react';
 
 import { Provider } from './Provider';
+import { NETWORKS } from '.';
 
 describe('Provider', () => {
   it('renders without throwing', () => {
-    const { container } = render(<Provider network='rinkeby' />);
+    const { container } = render(<Provider network={NETWORKS.rinkeby} />);
     expect(container).toBeInTheDocument();
   });
 });
diff --git a/packages/hooks/src/Provider.tsx b/packages/hooks/src/Provider.tsx
@@ -9,21 +9,21 @@ export interface Web3ContextType {
   signer?: JsonRpcSigner | null;
   userAddress?: string | null;
   disconnectWallet?: () => void;
-  network: string;
-  chainId?: number;
+  chainId?: number | null;
   connected: boolean;
   provider?: ethers.providers.Web3Provider | null;
+  correctNetwork: boolean;
 }
 
 export const Web3Context = React.createContext<Web3ContextType | undefined>(undefined);
 
 export interface ProviderProps {
   /**
    * @dev The network you want to connect to.
-   * @example 'mainnet'
+   * @example NETWORKS.mainnet
    * @type string
    */
-  network: string;
+  network: number;
   /**
    * @dev Your Infura project ID. This is required if you want to support WalletConnect.
    * @type string
@@ -39,15 +39,16 @@ export interface ProviderProps {
  */
 export const Provider: React.FC<ProviderProps> = ({ children, network, infuraId }) => {
   const [signer, setSigner] = React.useState<null | JsonRpcSigner>();
-  const [provider, setProvider] = React.useState<ethers.providers.Web3Provider>();
+  const [provider, setProvider] = React.useState<ethers.providers.Web3Provider | null>();
   const [userAddress, setUserAddress] = React.useState<null | string>();
   const [web3Modal, setWeb3Modal] = React.useState<Web3Modal>();
-  const [chainId, setChainId] = React.useState<number>();
+  const [chainId, setChainId] = React.useState<number | null>();
   const [connected, setConnected] = React.useState<boolean>(false);
+  const [correctNetwork, setCorrectNetwork] = React.useState<boolean>(true);
+  const [connection, setConnection] = React.useState<any>();
 
   const connectWallet = React.useCallback(async () => {
     const web3Modal = new Web3Modal({
-      network,
       providerOptions: {
         walletconnect: {
           package: WalletConnectProvider,
@@ -57,16 +58,61 @@ export const Provider: React.FC<ProviderProps> = ({ children, network, infuraId
         },
       },
     });
+    setWeb3Modal(web3Modal);
     const connection = await web3Modal.connect();
+    setConnection(connection);
     const provider = new ethers.providers.Web3Provider(connection);
+    setProvider(provider);
+    const chainId = await provider.getNetwork().then((network) => network.chainId);
+    setChainId(chainId);
+    setCorrectNetwork(chainId === network);
     const signer = provider.getSigner();
-    setWeb3Modal(web3Modal);
     setSigner(signer);
+    setUserAddress(await signer.getAddress());
+    setConnected(true);
+
+    connection.on('chainChanged', onNetworkChange);
+    connection.on('accountsChanged', onAccountsChanged);
+    connection.on('disconnect', onDisconnect);
+  }, [
+    Web3Modal,
+    web3Modal,
+    WalletConnectProvider,
+    network,
+    infuraId,
+    ethers,
+    correctNetwork,
+    connection,
+  ]);
+
+  const onNetworkChange = async (newChainId: string) => {
+    const formattedChainId = +newChainId.split('0x')[1];
+    setChainId(formattedChainId);
+    setCorrectNetwork(formattedChainId === network);
+    const provider = new ethers.providers.Web3Provider(connection);
     setProvider(provider);
+    const signer = provider.getSigner();
+    setSigner(signer);
     setUserAddress(await signer.getAddress());
-    setChainId(await signer.getChainId());
     setConnected(true);
-  }, [Web3Modal, web3Modal, WalletConnectProvider, network, infuraId, ethers]);
+  };
+
+  const onAccountsChanged = async () => {
+    const provider = new ethers.providers.Web3Provider(connection);
+    setProvider(provider);
+    const chainId = await provider.getNetwork().then((network) => network.chainId);
+    setChainId(chainId);
+    setCorrectNetwork(chainId === network);
+    const signer = provider.getSigner();
+    setSigner(signer);
+    setUserAddress(await signer.getAddress());
+    setConnected(true);
+  };
+
+  const onDisconnect = async () => {
+    web3Modal?.clearCachedProvider();
+    disconnectWallet();
+  };
 
   const disconnectWallet = React.useCallback(() => {
     web3Modal?.clearCachedProvider();
@@ -85,8 +131,19 @@ export const Provider: React.FC<ProviderProps> = ({ children, network, infuraId
       provider,
       network,
       chainId,
+      correctNetwork,
     }),
-    [connectWallet, signer, userAddress, web3Modal, connected, provider, network, chainId]
+    [
+      connectWallet,
+      signer,
+      userAddress,
+      web3Modal,
+      connected,
+      provider,
+      network,
+      chainId,
+      correctNetwork,
+    ]
   );
 
   return <Web3Context.Provider value={{ ...value }}>{children}</Web3Context.Provider>;
diff --git a/packages/hooks/src/constants.ts b/packages/hooks/src/constants.ts
@@ -0,0 +1,47 @@
+export const NETWORKS = {
+  mainnet: 1,
+  ropsten: 3,
+  rinkeby: 4,
+  goerli: 5,
+  kovan: 42,
+  BSC: 56,
+  BSCTestnet: 97,
+  xdai: 100,
+  polygon: 137,
+  theta: 361,
+  thetaTestnet: 365,
+  moonriver: 1285,
+  mumbai: 80001,
+  harmony: 1666600000,
+  palm: 11297108109,
+  localhost: 1337,
+  hardhat: 31337,
+  fantom: 250,
+  avalanche: 43114,
+  songbird: 19,
+  moonbaseAlpha: 1287,
+};
+
+export const CHAIN_ID_TO_NETWORK = {
+  1: 'mainnet',
+  3: 'ropsten',
+  4: 'rinkeby',
+  5: 'goerli',
+  42: 'kovan',
+  56: 'BSC',
+  97: 'BSCTestnet',
+  100: 'xdai',
+  137: 'polygon',
+  361: 'theta',
+  365: 'thetaTestnet',
+  1285: 'moonriver',
+  80001: 'mumbai',
+  1666600000: 'harmony',
+  11297108109: 'palm',
+  1337: 'localhost',
+  31337: 'hardhat',
+  250: 'fantom',
+  43114: 'avalanche',
+  19: 'songbird',
+  1287: 'moonbaseAlpha',
+};
diff --git a/packages/hooks/src/hooks/useWallet.ts b/packages/hooks/src/hooks/useWallet.ts
@@ -1,4 +1,5 @@
 import React from 'react';
+import { CHAIN_ID_TO_NETWORK, NETWORKS } from '..';
 import { Web3Context } from '../Provider';
 
 /**
@@ -8,27 +9,40 @@ export function useWallet() {
   const context = React.useContext(Web3Context);
   const [ens, setEns] = React.useState<string>();
 
+  if (!context) {
+    throw new Error('No Web3Context found');
+  }
+
+  const {
+    connectWallet,
+    disconnectWallet,
+    userAddress,
+    chainId,
+    signer,
+    connected,
+    provider,
+    correctNetwork,
+  } = context;
+
   React.useEffect(() => {
-    if (context) {
-      const { userAddress, provider } = context;
-      if (userAddress && provider) {
-        provider.lookupAddress(userAddress).then((address) => {
-          setEns(address as string);
-        });
-      }
+    if (userAddress && provider && chainId === NETWORKS.mainnet) {
+      provider.lookupAddress(userAddress).then((address) => {
+        setEns(address as string);
+      });
     }
-  }, [context?.userAddress, context?.provider]);
+  }, [userAddress, provider]);
 
   return {
-    connectWallet: context?.connectWallet,
-    disconnectWallet: context?.disconnectWallet,
+    connectWallet,
+    disconnectWallet,
     connection: {
-      userAddress: context?.userAddress,
-      network: context?.network,
-      signer: context?.signer,
+      userAddress,
+      network: CHAIN_ID_TO_NETWORK[chainId as number],
+      signer,
       ens,
     },
-    connected: context?.connected,
-    provider: context?.provider,
+    connected,
+    provider,
+    correctNetwork,
   };
 }
diff --git a/packages/hooks/src/index.ts b/packages/hooks/src/index.ts
@@ -2,3 +2,4 @@ export { useWallet } from './hooks';
 export { useTokenBalance } from './hooks/useTokenBalance';
 export { Provider, Web3Context } from './Provider';
 export type { ProviderProps, Web3ContextType } from './Provider';
+export { NETWORKS, CHAIN_ID_TO_NETWORK } from './constants';
diff --git a/packages/hooks/src/stories/ConnectWallet.stories.tsx b/packages/hooks/src/stories/ConnectWallet.stories.tsx
@@ -1,19 +1,26 @@
-import React from 'react';
+import React, { useEffect } from 'react';
 import { Button } from '@chakra-ui/react';
-import { Provider, useWallet } from '..';
+import { NETWORKS, Provider, useWallet } from '..';
 
 export default {
   title: 'Hooks/useWallet',
 };
 
 const DefaultUsingProvider = () => {
-  const { connection, connectWallet, disconnectWallet, connected } = useWallet();
+  const { connection, connectWallet, disconnectWallet, connected, correctNetwork } = useWallet();
+
+  useEffect(() => {
+    if (!correctNetwork) {
+      alert('Please connect to the correct network');
+    }
+  }, [correctNetwork]);
 
   if (connected) {
     return (
       <div>
         <Button onClick={disconnectWallet}>Disconnect wallet</Button>
         <p>{connection.ens || connection.userAddress}</p>
+        <p>Connected to the correct network: {correctNetwork ? 'Yes' : 'no'}</p>
       </div>
     );
   }
@@ -22,7 +29,7 @@ const DefaultUsingProvider = () => {
 };
 
 export const Default = () => (
-  <Provider network='rinkeby'>
+  <Provider network={NETWORKS.rinkeby}>
     <DefaultUsingProvider />
   </Provider>
 );
