From c6831db44dbac49af16246eefcdcab71b927aa93 Mon Sep 17 00:00:00 2001
From: port <108868128+portdeveloper@users.noreply.github.com>
Date: Wed, 20 Mar 2024 21:52:59 +0300
Subject: [PATCH] Use and show implementation address

---
 .../scaffold-eth/Contract/ContractUI.tsx      | 13 +++++++++---
 .../pages/[contractAddress]/[network].tsx     | 20 +++++++++++++++++--
 packages/nextjs/pages/index.tsx               |  6 +++++-
 packages/nextjs/services/store/store.ts       |  4 ++++
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/packages/nextjs/components/scaffold-eth/Contract/ContractUI.tsx b/packages/nextjs/components/scaffold-eth/Contract/ContractUI.tsx
index 592bab93..28f19c69 100644
--- a/packages/nextjs/components/scaffold-eth/Contract/ContractUI.tsx
+++ b/packages/nextjs/components/scaffold-eth/Contract/ContractUI.tsx
@@ -62,8 +62,11 @@ const mainNetworks = getTargetNetworks();
  **/
 export const ContractUI = ({ className = "", initialContractData }: ContractUIProps) => {
   const [refreshDisplayVariables, triggerRefreshDisplayVariables] = useReducer(value => !value, false);
-  const mainChainId = useAbiNinjaState(state => state.mainChainId);
-  const mainNetwork = mainNetworks.find(network => network.id === mainChainId);
+  const { implementationAddress, chainId } = useAbiNinjaState(state => ({
+    chainId: state.mainChainId,
+    implementationAddress: state.implementationAddress,
+  }));
+  const mainNetwork = mainNetworks.find(network => network.id === chainId);
   const networkColor = useNetworkColor(mainNetwork);
   const router = useRouter();
   const { network } = router.query as { network?: string };
@@ -124,7 +127,7 @@ export const ContractUI = ({ className = "", initialContractData }: ContractUIPr
   const { data: contractNameData, isLoading: isContractNameLoading } = useContractRead({
     address: initialContractData.address,
     abi: initialContractData.abi,
-    chainId: mainChainId,
+    chainId: chainId,
     functionName: "name",
   });
 
@@ -219,6 +222,10 @@ export const ContractUI = ({ className = "", initialContractData }: ContractUIPr
                       <span className="font-medium text-base mr-4"> {displayContractName} </span>
                       <Address address={initialContractData.address} />
                     </div>
+                    <div className="flex items-center gap-1">
+                      <span className="font-medium text-base mr-4 text-green-600">Implementation Address</span>
+                      <Address address={implementationAddress} />
+                    </div>
                     <div className="flex items-center gap-1">
                       <span className="text-sm font-bold">Balance:</span>
                       <Balance address={initialContractData.address} className="h-1.5 min-h-[0.375rem] px-0" />
diff --git a/packages/nextjs/pages/[contractAddress]/[network].tsx b/packages/nextjs/pages/[contractAddress]/[network].tsx
index b3b4beb9..d2cb6647 100644
--- a/packages/nextjs/pages/[contractAddress]/[network].tsx
+++ b/packages/nextjs/pages/[contractAddress]/[network].tsx
@@ -2,7 +2,7 @@ import { useEffect, useState } from "react";
 import Link from "next/link";
 import { useRouter } from "next/router";
 import { ParsedUrlQuery } from "querystring";
-import { Abi, isAddress } from "viem";
+import { Abi, Address, isAddress } from "viem";
 import * as chains from "viem/chains";
 import { ExclamationTriangleIcon } from "@heroicons/react/24/outline";
 import { Footer } from "~~/components/Footer";
@@ -11,6 +11,7 @@ import { MiniHeader } from "~~/components/MiniHeader";
 import { ContractUI } from "~~/components/scaffold-eth";
 import { useAbiNinjaState } from "~~/services/store/store";
 import { fetchContractABIFromAnyABI, fetchContractABIFromEtherscan } from "~~/utils/abi";
+import detectProxyTarget from "~~/utils/abi.ninja/proxyContracts";
 
 interface ParsedQueryContractDetailsPage extends ParsedUrlQuery {
   contractAddress: string;
@@ -22,6 +23,15 @@ type ContractData = {
   address: string;
 };
 
+const getImplementationAddress = async (proxyAddress: Address) => {
+  try {
+    const target = await detectProxyTarget(proxyAddress);
+    return target;
+  } catch (e) {
+    console.error(e);
+  }
+};
+
 const ContractDetailPage = () => {
   const router = useRouter();
   const { contractAddress, network } = router.query as ParsedQueryContractDetailsPage;
@@ -33,10 +43,12 @@ const ContractDetailPage = () => {
     contractAbi: storedAbi,
     setMainChainId,
     chainId,
+    setImplementationAddress,
   } = useAbiNinjaState(state => ({
     contractAbi: state.contractAbi,
     setMainChainId: state.setMainChainId,
     chainId: state.mainChainId,
+    setImplementationAddress: state.setImplementationAddress,
   }));
 
   const getNetworkName = (chainId: number) => {
@@ -73,7 +85,11 @@ const ContractDetailPage = () => {
         }
 
         try {
-          const abi = await fetchContractABIFromAnyABI(contractAddress, parsedNetworkId);
+          const implementationAddress = await getImplementationAddress(contractAddress);
+          if (implementationAddress) {
+            setImplementationAddress(implementationAddress);
+          }
+          const abi = await fetchContractABIFromAnyABI(implementationAddress || contractAddress, parsedNetworkId);
           if (!abi) throw new Error("Got empty or undefined ABI from AnyABI");
           setContractData({ abi, address: contractAddress });
           setError(null);
diff --git a/packages/nextjs/pages/index.tsx b/packages/nextjs/pages/index.tsx
index 7f0f7e60..24c0d71e 100644
--- a/packages/nextjs/pages/index.tsx
+++ b/packages/nextjs/pages/index.tsx
@@ -46,9 +46,10 @@ const Home: NextPage = () => {
     chainId: parseInt(network),
   });
 
-  const { setContractAbi, setAbiContractAddress } = useAbiNinjaState(state => ({
+  const { setContractAbi, setAbiContractAddress, setImplementationAddress } = useAbiNinjaState(state => ({
     setContractAbi: state.setContractAbi,
     setAbiContractAddress: state.setAbiContractAddress,
+    setImplementationAddress: state.setImplementationAddress,
   }));
 
   const [isAbiAvailable, setIsAbiAvailable] = useState(false);
@@ -60,6 +61,9 @@ const Home: NextPage = () => {
       setIsFetchingAbi(true);
       try {
         const implementationAddress = await getImplementationAddress(verifiedContractAddress);
+        if (implementationAddress) {
+          setImplementationAddress(implementationAddress);
+        }
         const abi = await fetchContractABIFromAnyABI(
           implementationAddress || verifiedContractAddress,
           parseInt(network),
diff --git a/packages/nextjs/services/store/store.ts b/packages/nextjs/services/store/store.ts
index 55ce9b39..c0482fb2 100644
--- a/packages/nextjs/services/store/store.ts
+++ b/packages/nextjs/services/store/store.ts
@@ -17,6 +17,8 @@ type AbiNinjaState = {
   setContractAbi: (newAbi: Abi) => void;
   abiContractAddress: Address;
   setAbiContractAddress: (newAbiContractAddress: Address) => void;
+  implementationAddress: Address;
+  setImplementationAddress: (newImplementationAddress: Address) => void;
 };
 
 export const useGlobalState = create<GlobalState>(set => ({
@@ -33,4 +35,6 @@ export const useAbiNinjaState = create<AbiNinjaState>(set => ({
   setContractAbi: (newAbi: Abi): void => set({ contractAbi: newAbi }),
   abiContractAddress: "",
   setAbiContractAddress: (newAddress: Address): void => set({ abiContractAddress: newAddress }),
+  implementationAddress: "",
+  setImplementationAddress: (newAddress: Address): void => set({ implementationAddress: newAddress }),
 }));