Add support for custom upgrade functions

[zmalatrax]

Is your feature request related to a problem? Please describe.

Current versions of OpenZeppelin's (OZ) Proxies (ERC1967Utils.sol & UUPSUpgradeable.sol) only implements the upgradeToAndCall(address newImplementation, bytes memory data) function. The ProxyAdmin contract only implements upgradeAndCall(address proxy, address implementation, bytes memory data).

When trying to upgrade the implementation contract of a Transparent or UUPS proxy one without performing a post-upgrade call through hardhat-deploy (the equivalent of the old upgradeTo(address newImplementation) included in hardhat-deploy), the plugin defaults to calling upgrade or upgradeTo (depending if a proxy admin has been defined or not) which are not implemented, resulting in the error method not implemented....

Here is the deployment script used when trying to initialize then upgrade a UUPS proxy, with the implementation only having upgradeToAndCall(address newImplementation, bytes memory data) defined:

const upgradeableDeployment = await deploy("UUPSUpgradeableMock", {
  from: deployer,
  proxy: {
  proxyContract: "UUPSProxy", //Custom UUPS Proxy adapted from OpenZeppelin's one
  /* proxyArgs, checkProxyAdmin & checkABIConflict must be specified
  because of using a "custom" proxy, which is not in the included contracts such as "UUPS" */
  proxyArgs: ["{implementation}", "{data}"],
  checkProxyAdmin: false,
  checkABIConflict: false,
  execute: {
      init: {
        methodName: "initialize",
        args: [deployer],
      },
    },
  },
  log: true,
  args: [],
});

Fig.1: hardhat-deploy deployment script of a custom UUPS proxy implementing upgradeToAndCall only for upgrading

It nicely deploys and initializes the proxy but, when modifying the implementation and running the script again (which should perform the upgrade), it errors with Error: No method named "upgradeTo" on contract deployed as "UUPSUpgradeableMock"

Fig.2: hardhat-deploy logs when trying to update the UUPSUpgradeableMock with no updateMethod defined

Note: In helper.ts, the variables updateMethod & updateArgs are referring to the post-upgrade function that is delegatecalled to the new implementation once the upgrade is done. However, their name is misleading (to my understanding, update ~= upgrade) and should be renamed for better semantics, such as postUpgradeMethod & postUpgradeArgs

Describe the solution you'd like

• Allow the user to define which function to call when performing an upgrade:

  // ProxyOptions
  type ProxyOptionsBase = {
  // [...]
  upgradeFunction?: {
  methodName: string;
  upgradeArgs: any[];
  };
  }

• Default to the current behavior if none was given.

This way, it would allow the use of newer versions of OZ's proxies but also custom proxies/admin proxies to be supported by hardhat-deploy while having backward compatibility with the current workflow. It should solve #146.

For the issue I faced, it would be resolved with this script:

const upgradeableDeployment = await deploy("UUPSUpgradeableMock", {
  from: deployer,
  proxy: {
    proxyContract: "UUPSProxy", //Custom UUPS Proxy adapted from OpenZeppelin's one
    /* proxyArgs, checkProxyAdmin & checkABIConflict must be specified
    because of using a "custom" proxy, which is not in the included contracts such as "UUPS" */
    proxyArgs: ["{implementation}", "{data}"],
    checkProxyAdmin: false,
    checkABIConflict: false,
    upgradeFunction: {
      methodName: "upgradeToAndCall",
      upgradeArgs: ['{implementation}', '{data}']
    }
    execute: {
      init: {
        methodName: "initialize",
        args: [deployer],
      },
    },
  },
  log: true,
  args: [],
});

As no updateMethod is defined, data defaults to '0x' thus not performing a delegatecall after the successful upgrade.

Describe alternatives you've considered

• A current workaround to successfully perform upgrades with new OZ contracts is by having a post-upgrade function be called. Calling a view or pure function allows not interfering with the system state when upgrading (e.g. a public variable getter function) but it is sketchy and inadequate (gas bad).

const upgradeableDeployment = await deploy("UUPSUpgradeableMock", {
  from: deployer,
  proxy: {
  proxyContract: "UUPSProxy", //Custom UUPS Proxy adapted from OpenZeppelin's one
  /* proxyArgs, checkProxyAdmin & checkABIConflict must be specified
  because of using a "custom" proxy, which is not in the included contracts such as "UUPS" */
  proxyArgs: ["{implementation}", "{data}"],
  checkProxyAdmin: false,
  checkABIConflict: false,
  execute: {
      init: {
        methodName: "initialize",
        args: [deployer],
      },
      onUpdate: {
        methodName: "counter",
        args: []
      }
    },
  },
  log: true,
  args: [],
});

• If updateMethod is undefined, check the abi to find upgradeTo, or upgrade if it is via an admin. If not found (i.e. not defined), make a call to upgradeToAndCall with data = '0x'. However, doing so doesn't give the user total freedom when designing its custom proxy. Plus, it adds complexity when refactoring the upgrading workflow could simplify it.

Additional Context

• UUPSUpgradeableMock.sol

[jaybuidl]

Can be marked as closed :)