Add support for proxy contracts

.changeset/modern-goats-talk.md

@@ -0,0 +1,5 @@
+---
+'@graphprotocol/graph-cli': minor
+---
+
+Add EIP-1967/OpenZeppelin proxy contracts support. When proxy contract is detected, user is given an option to use ABI of the implementation contract.

packages/cli/src/command-helpers/contracts.ts

@@ -206,6 +206,51 @@ export class ContractService {
     return null;
   }
 
+  async getProxyImplementation(networkId: string, address: string) {
+    const urls = this.getRpcUrls(networkId);
+    if (!urls.length) {
+      throw new Error(`No JSON-RPC available for ${networkId} in the registry`);
+    }
+
+    const EIP_1967_SLOT = '0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc';
+    const OPEN_ZEPPELIN_SLOT = '0x7050c9e0f4ca769c69bd3a8ef740bc37934f8e2c036e5a723fd8ee048ed3f8c3';
+    const getStorageAt = async (url: string, slot: string) => {
+      const response = await fetch(url, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({
+          jsonrpc: '2.0',
+          method: 'eth_getStorageAt',
+          params: [address, slot, 'latest'],
+          id: 1,
+        }),
+      });
+      const json = await response.json();
+      if (json?.result) {
+        const impl = '0x' + json.result.slice(-40);
+        if (impl !== '0x0000000000000000000000000000000000000000') {
+          return impl;
+        }
+      }
+      return null;
+    };
+
+    for (const url of urls) {
+      for (const slot of [EIP_1967_SLOT, OPEN_ZEPPELIN_SLOT]) {
+        try {
+          const impl = await getStorageAt(url, slot);
+          if (impl) {
+            return impl;
+          }
+        } catch (error) {
+          logger(`Failed to fetch proxy implementation from ${url}: ${error}`);
+        }
+      }
+    }
+
+    throw new Error(`No implementation address found`);
+  }
+
   private async fetchTransactionByHash(networkId: string, txHash: string) {
     const urls = this.getRpcUrls(networkId);
     if (!urls.length) {
@@ -234,6 +279,6 @@ export class ContractService {
       }
     }
 
-    throw new Error(`JSON-RPC is unreachable`);
+    throw new Error(`Failed to fetch tx ${txHash}`);
   }
 }

packages/cli/src/command-helpers/proxy.test.ts

@@ -0,0 +1,91 @@
+import { prompt } from 'gluegun';
+import { describe, expect, it, vi } from 'vitest';
+import EthereumABI from '../protocols/ethereum/abi.js';
+import { ContractService } from './contracts.js';
+import { checkForProxy } from './proxy.js';
+import { loadRegistry } from './registry.js';
+
+// Mock gluegun's prompt
+vi.mock('gluegun', async () => {
+  const actual = await vi.importActual('gluegun');
+  return {
+    ...actual,
+    prompt: {
+      confirm: vi.fn().mockResolvedValue(true),
+    },
+  };
+});
+
+describe('Proxy detection', async () => {
+  const NETWORK = 'mainnet';
+  const registry = await loadRegistry();
+  const contractService = new ContractService(registry);
+
+  interface ProxyTestCase {
+    name: string;
+    type: string;
+    address: string;
+    implementationAddress: string | null;
+    expectedFunctions: string[];
+  }
+
+  const testCases: ProxyTestCase[] = [
+    {
+      name: 'USDC',
+      type: 'EIP-1967 Upgradeable',
+      address: '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48',
+      implementationAddress: '0x43506849d7c04f9138d1a2050bbf3a0c054402dd',
+      expectedFunctions: ['mint(address,uint256)', 'configureMinter(address,uint256)'],
+    },
+    {
+      name: 'BUSD',
+      type: 'OpenZeppelin Unstructured Storage',
+      address: '0x4Fabb145d64652a948d72533023f6E7A623C7C53',
+      implementationAddress: '0x2A3F1A37C04F82aA274f5353834B2d002Db91015',
+      expectedFunctions: ['reclaimBUSD()', 'claimOwnership()'],
+    },
+    {
+      name: 'Gelato',
+      type: 'EIP-2535 Diamond Pattern (not supported)',
+      address: '0x3caca7b48d0573d793d3b0279b5f0029180e83b6',
+      implementationAddress: null,
+      expectedFunctions: [],
+    },
+  ];
+
+  for (const testCase of testCases) {
+    it(`should handle ${testCase.name} ${testCase.type} Proxy`, async () => {
+      const abi = await contractService.getABI(EthereumABI, NETWORK, testCase.address);
+      expect(abi).toBeDefined();
+
+      const { implementationAddress, implementationAbi } = await checkForProxy(
+        contractService,
+        NETWORK,
+        testCase.address,
+        abi!,
+      );
+
+      expect(implementationAddress === testCase.implementationAddress);
+
+      const implFunctions = implementationAbi?.callFunctionSignatures();
+      for (const expectedFunction of testCase.expectedFunctions) {
+        expect(implFunctions).toContain(expectedFunction);
+      }
+    });
+  }
+
+  it('should handle when user declines to use implementation', async () => {
+    vi.mocked(prompt.confirm).mockResolvedValueOnce(false);
+    const abi = await contractService.getABI(EthereumABI, NETWORK, testCases[0].address);
+    expect(abi).toBeDefined();
+
+    const { implementationAddress, implementationAbi } = await checkForProxy(
+      contractService,
+      NETWORK,
+      testCases[0].address,
+      abi!,
+    );
+    expect(implementationAddress).toBeNull();
+    expect(implementationAbi).toBeNull();
+  });
+});

packages/cli/src/command-helpers/proxy.ts

@@ -0,0 +1,55 @@
+import { prompt } from 'gluegun';
+import EthereumABI from '../protocols/ethereum/abi.js';
+import { ContractService } from './contracts.js';
+import { retryWithPrompt } from './retry.js';
+import { withSpinner } from './spinner.js';
+
+export interface CheckForProxyResult {
+  implementationAbi: EthereumABI | null;
+  implementationAddress: string | null;
+}
+
+export async function checkForProxy(
+  contractService: ContractService,
+  network: string,
+  address: string,
+  abi: EthereumABI,
+): Promise<CheckForProxyResult> {
+  let implementationAddress = null;
+  let implementationAbi = null;
+
+  const maybeProxy = abi.callFunctionSignatures()?.includes('upgradeTo(address)');
+  if (maybeProxy) {
+    const impl = await retryWithPrompt(() =>
+      withSpinner(
+        'Fetching proxy implementation address...',
+        'Failed to fetch proxy implementation address',
+        'Warning fetching proxy implementation address',
+        () => contractService.getProxyImplementation(network, address),
+      ),
+    );
+
+    if (impl) {
+      const useImplementation = await prompt
+        .confirm(`Proxy contract detected. Use implementation contract ABI at ${impl}?`, true)
+        .catch(() => false);
+
+      if (useImplementation) {
+        implementationAddress = impl;
+        implementationAbi = await retryWithPrompt(() =>
+          withSpinner(
+            'Fetching implementation contract ABI...',
+            'Failed to fetch implementation ABI',
+            'Warning fetching implementation ABI',
+            () => contractService.getABI(EthereumABI, network, implementationAddress!),
+          ),
+        );
+      }
+    }
+  }
+
+  return {
+    implementationAbi,
+    implementationAddress,
+  };
+}

packages/cli/src/commands/add.ts

@@ -4,6 +4,7 @@ import { Args, Command, Errors, Flags } from '@oclif/core';
 import { ContractService } from '../command-helpers/contracts.js';
 import * as DataSourcesExtractor from '../command-helpers/data-sources.js';
 import { updateNetworksFile } from '../command-helpers/network.js';
+import { checkForProxy } from '../command-helpers/proxy.js';
 import { loadRegistry } from '../command-helpers/registry.js';
 import { retryWithPrompt } from '../command-helpers/retry.js';
 import {
@@ -88,7 +89,9 @@ export default class AddCommand extends Command {
 
     let startBlock = startBlockFlag ? parseInt(startBlockFlag).toString() : startBlockFlag;
     let contractName = contractNameFlag || DEFAULT_CONTRACT_NAME;
-    let ethabi = null;
+
+    let ethabi: EthereumABI | null = null;
+    let implAddress = null;
 
     if (sourcifyContractInfo) {
       startBlock ??= sourcifyContractInfo.startBlock;
@@ -112,6 +115,18 @@ export default class AddCommand extends Command {
           ),
         );
         if (!ethabi) throw Error;
+
+        const { implementationAbi, implementationAddress } = await checkForProxy(
+          contractService,
+          network,
+          address,
+          ethabi,
+        );
+        if (implementationAddress) {
+          implAddress = implementationAddress;
+          ethabi = implementationAbi!;
+        }
+        if (!ethabi) throw Error;
       } catch (error) {
         // we cannot ask user to do prompt in test environment
         if (process.env.NODE_ENV !== 'test') {
@@ -136,10 +151,15 @@ export default class AddCommand extends Command {
         }
       }
     }
+    if (!ethabi) {
+      this.error('Failed to load ABI', { exit: 1 });
+    }
 
     try {
       if (isLocalHost) throw Error; // Triggers user prompting without waiting for Etherscan lookup to fail
-      startBlock ||= Number(await contractService.getStartBlock(network, address)).toString();
+      startBlock ||= Number(
+        await contractService.getStartBlock(network, implAddress ?? address),
+      ).toString();
     } catch (error) {
       // we cannot ask user to do prompt in test environment
       if (process.env.NODE_ENV !== 'test') {
@@ -166,7 +186,8 @@ export default class AddCommand extends Command {
       if (isLocalHost) throw Error; // Triggers user prompting without waiting for Etherscan lookup to fail
       if (contractName === DEFAULT_CONTRACT_NAME) {
         contractName =
-          (await contractService.getContractName(network, address)) ?? DEFAULT_CONTRACT_NAME;
+          (await contractService.getContractName(network, implAddress ?? address)) ??
+          DEFAULT_CONTRACT_NAME;
       }
     } catch (error) {
       // not asking user to do prompt in test environment
@@ -266,8 +287,6 @@ export default class AddCommand extends Command {
       'Warning during codegen',
       async () => await system.run(yarn ? 'yarn codegen' : 'npm run codegen'),
     );
-
-    this.exit(0);
   }
 }
 

packages/cli/src/commands/init.ts

@@ -11,6 +11,7 @@ import { DEFAULT_IPFS_URL } from '../command-helpers/ipfs.js';
 import { initNetworksConfig } from '../command-helpers/network.js';
 import { chooseNodeUrl } from '../command-helpers/node.js';
 import { PromptManager } from '../command-helpers/prompt-manager.js';
+import { checkForProxy } from '../command-helpers/proxy.js';
 import { loadRegistry } from '../command-helpers/registry.js';
 import { retryWithPrompt } from '../command-helpers/retry.js';
 import { generateScaffold, writeScaffold } from '../command-helpers/scaffold.js';
@@ -635,6 +636,7 @@ async function processInitForm(
         }
 
         // If ABI is not provided, try to fetch it from Etherscan API
+        let implAddress: string | undefined = undefined;
         if (protocolInstance.hasABIs() && !initAbi) {
           abiFromApi = await retryWithPrompt(() =>
             withSpinner(
@@ -644,18 +646,37 @@ async function processInitForm(
               () => contractService.getABI(protocolInstance.getABI(), network.id, address),
             ),
           );
-          initDebugger.extend('processInitForm')("abiFromEtherscan len: '%s'", abiFromApi?.name);
+          initDebugger.extend('processInitForm')("ABI: '%s'", abiFromApi?.name);
         } else {
           abiFromApi = initAbi;
         }
+
+        if (abiFromApi) {
+          const { implementationAbi, implementationAddress } = await checkForProxy(
+            contractService,
+            network.id,
+            address,
+            abiFromApi,
+          );
+          if (implementationAddress) {
+            implAddress = implementationAddress;
+            abiFromApi = implementationAbi!;
+            initDebugger.extend('processInitForm')(
+              "Impl ABI: '%s', Impl Address: '%s'",
+              abiFromApi?.name,
+              implAddress,
+            );
+          }
+        }
+
         // If startBlock is not provided, try to fetch it from Etherscan API
         if (!initStartBlock) {
           startBlock = await retryWithPrompt(() =>
             withSpinner(
               'Fetching start block from contract API...',
               'Failed to fetch start block',
               'Warning fetching start block',
-              () => contractService.getStartBlock(network.id, address),
+              () => contractService.getStartBlock(network.id, implAddress ?? address),
             ),
           );
           initDebugger.extend('processInitForm')("startBlockFromEtherscan: '%s'", startBlock);
@@ -668,7 +689,7 @@ async function processInitForm(
               'Fetching contract name from contract API...',
               'Failed to fetch contract name',
               'Warning fetching contract name',
-              () => contractService.getContractName(network.id, address),
+              () => contractService.getContractName(network.id, implAddress ?? address),
             ),
           );
           initDebugger.extend('processInitForm')("contractNameFromEtherscan: '%s'", contractName);
@@ -1302,7 +1323,7 @@ async function addAnotherContract(
       name: 'contract',
       initial: ProtocolContract.identifierName(),
       required: true,
-      message: () => `\nContract ${ProtocolContract.identifierName()}`,
+      message: () => `Contract ${ProtocolContract.identifierName()}`,
       validate: value => {
         const { valid, error } = validateContract(value, ProtocolContract);
         return valid ? true : error;