From 3f1b940a3e15b6dc24c02034dc4259a8ebd87363 Mon Sep 17 00:00:00 2001 From: YaroShkvorets Date: Wed, 8 Jan 2025 21:15:50 -0500 Subject: [PATCH] refactor proxy check code into helper function --- packages/cli/src/command-helpers/proxy.ts | 55 +++++++++++++++++++++++ packages/cli/src/commands/add.ts | 40 +++++------------ packages/cli/src/commands/init.ts | 47 +++++++------------ 3 files changed, 83 insertions(+), 59 deletions(-) create mode 100644 packages/cli/src/command-helpers/proxy.ts diff --git a/packages/cli/src/command-helpers/proxy.ts b/packages/cli/src/command-helpers/proxy.ts new file mode 100644 index 00000000..969cf880 --- /dev/null +++ b/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 { + 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, + }; +} diff --git a/packages/cli/src/commands/add.ts b/packages/cli/src/commands/add.ts index 2f39b7e4..e40b77e8 100644 --- a/packages/cli/src/commands/add.ts +++ b/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 { @@ -101,35 +102,18 @@ export default class AddCommand extends Command { ), ); if (!ethabi) throw Error; - const isProxy = ethabi.callFunctionSignatures()?.includes('upgradeTo(address)'); - if (isProxy) { - 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) { - implAddress = impl; - ethabi = await retryWithPrompt(() => - withSpinner( - 'Fetching implementation contract ABI...', - 'Failed to fetch implementation ABI', - 'Warning fetching implementation ABI', - () => contractService.getABI(EthereumABI, network, implAddress!), - ), - ); - } - 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') { diff --git a/packages/cli/src/commands/init.ts b/packages/cli/src/commands/init.ts index d25a3699..f119ba6e 100644 --- a/packages/cli/src/commands/init.ts +++ b/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'; @@ -622,38 +623,22 @@ async function processInitForm( () => contractService.getABI(protocolInstance.getABI(), network.id, address), ), ); - initDebugger.extend('processInitForm')("abiFromEtherscan len: '%s'", abiFromApi?.name); - const isProxy = abiFromApi?.callFunctionSignatures().includes('upgradeTo(address)'); - if (isProxy) { - const impl = await retryWithPrompt(() => - withSpinner( - 'Fetching proxy implementation address...', - 'Failed to fetch proxy implementation address', - 'Warning fetching proxy implementation address', - () => contractService.getProxyImplementation(network.id, address), - ), + initDebugger.extend('processInitForm')("ABI: '%s'", abiFromApi?.name); + if (abiFromApi) { + const { implementationAbi, implementationAddress } = await checkForProxy( + contractService, + network.id, + address, + abiFromApi, ); - initDebugger.extend('processInitForm')("proxyImplementation: '%s'", impl); - if (impl) { - const useImplementation = await prompt - .confirm( - `Proxy contract detected. Use implementation contract ABI at ${impl}?`, - true, - ) - .catch(() => false); - - if (useImplementation) { - implAddress = impl; - abiFromApi = await retryWithPrompt(() => - withSpinner( - 'Fetching implementation contract ABI...', - 'Failed to fetch implementation ABI', - 'Warning fetching implementation ABI', - () => - contractService.getABI(protocolInstance.getABI(), network.id, implAddress!), - ), - ); - } + if (implementationAddress) { + implAddress = implementationAddress; + abiFromApi = implementationAbi!; + initDebugger.extend('processInitForm')( + "Impl ABI: '%s', Impl Address: '%s'", + abiFromApi?.name, + implAddress, + ); } } }