feat(proxy): implements custom upgrade function

[zmalatrax]

This PR addresses issue #478, and potentially #146 . It refactors the proxy upgrading workflow and allows the user to define the upgrade function to be used when upgrading, via an admin contract or not. It mainly solves the issue of upgrading without post-upgrade delegatecall for proxies that only implement upgradeToAndCall or upgradeAndCall (e.g. OpenZeppelin proxies since ~v4.6).

Design Choices

1. The user can specify the upgrade function from its script
2. Arbitrary functions (name, arguments) can be used
3. Specifying the upgrade function must be optional
4. It must be backward-compatible (fallback on current behavior)

   Usage

If one wants to specify the upgrade method to be used, add the upgradeFunction field to your proxy deployment options, fill the methodName and upgradeArgs fields, respectively with the function name and the arguments template (working on the same principle as proxyArgs).

• '{proxy}': Template returning the proxy address (mainly for transparent proxies with admin)
• '{implementation}': The new implementation address once deployed
• '{data}': The calldata, delegating the post-upgrade function if defined
• Other params are pushed as is.

Examples:

• With the new OZ's UUPS proxy:

  const upgradeableDeployment = await deploy("UUPSUpgradeableMock", {
  from: deployer,
  proxy: {
  proxyContract: "UUPSProxy", 
  proxyArgs: ["{implementation}", "{data}"],
  checkProxyAdmin: false,
  checkABIConflict: false,
  upgradeFunction: {
    methodName: "upgradeToAndCall",
    upgradeArgs: ['{implementation}', '{data}']
  },
  execute: {
    init: {
      methodName: "initialize",
      args: [deployer],
    },
  },
  },
  log: true,
  args: [],
  });

• With a custom proxy admin implementing upgradeAndCall(address proxy, address newImplementation, bytes memory data, uint256 bonusArg):

  const upgradeableDeployment = await deploy("MockTransparentImplementation", {
  from: deployer,
  proxy: {
  proxyContract: "ERC1967Proxy",
  viaAdminContract: "CustomProxyAdmin",
  upgradeFunction: {
    methodName: "upgradeAndCall",
    upgradeArgs: ["{proxy}", "{implementation}", "{data}", 27]
  },
  execute: {
    init: {
      methodName: "initialize", // Can only be called once
      args: [deployer],
    },
    onUpdate: {
      methodName: "counter",
      args: []
    }
  },
  },
  log: true,
  args: [],
  });

If upgradeFunction is not specified, hardhat-deploy works as before.

Implementation

To allow the user to specify its upgrade function, the field must be accessible from the DeployOptions. Therefore, the field upgradeFunction has been added to ProxyOptionsBase:

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

As the upgrade function's arguments are not known in advance, the user defines the template arguments. It is inspired by proxyArgs.

It relies on the replaceTemplateArgs function from which a new optional field, proxyAddress, has been added; to allow proxy with admin contracts to also specify a custom upgrade function.

replaceTemplateArgs should return any[] and not string[] as the args could include other types (e.g. numbers).

function replaceTemplateArgs(
proxyArgsTemplate: string[],
{
  implementationAddress,
  proxyAdmin,
  data,
  proxyAddress,
}: {
  // [...]
  proxyAddress?: string;
}): any[] {
const proxyArgs: any[] = [];
  for (let i = 0; i < proxyArgsTemplate.length; i++) {
    const argValue = proxyArgsTemplate[i];
    if (argValue === '{proxy}') {
      if (!proxyAddress) {
        throw new Error(`Expected proxy address but none was specified.`);
      }
      proxyArgs.push(proxyAddress);
    }
    // [...]
  }
}

The execute call to perform the upgrade has been refactored:

--- _deployViaProxy ---

// [...]
let executeReceipt;
if (proxyAdminName) {
  if (oldProxy) {
    throw new Error(`Old Proxy do not support Proxy Admin contracts`);
  }
  if (!currentProxyAdminOwner) {
    throw new Error(`no currentProxyAdminOwner found in ProxyAdmin`);
  }
  executeReceipt = await execute(
    proxyAdminName,
    {...options, from: currentProxyAdminOwner},
    upgradeMethod,
    ...upgradeArgs
  );
} else {
  executeReceipt = await execute(
    name,
    {...options, from},
    upgradeMethod,
    ...upgradeArgs
  );
}
if (!executeReceipt) {
  throw new Error(`could not execute ${upgradeMethod}`);
}

upgradeMethod & upgradeArgsTemplate are computed in _getProxyInfo: upgradeArgsTemplate will then allow computing updateArgs, passed to execute.

• If upgradeFunction has been specified, the given values are used.

  
  --- _getProxyInfo ---

if (options.proxy.upgradeFunction) { upgradeMethod = options.proxy.upgradeFunction.methodName; upgradeArgsTemplate = options.proxy.upgradeFunction.upgradeArgs; }

- Otherwise, it fallbacks on the current behavior, refactored as so:
```typescript
--- _getProxyInfo ---

if (!upgradeMethod) {
  if (viaAdminContract) {
    if (updateMethod) {
      upgradeMethod = 'upgradeAndCall';
      upgradeArgsTemplate = ['{proxy}', '{implementation}', '{data}'];
    } else {
      upgradeMethod = 'upgrade';
      upgradeArgsTemplate = ['{proxy}', '{implementation}'];
    }
  } else if (updateMethod) {
    upgradeMethod = 'upgradeToAndCall';
    upgradeArgsTemplate = ['{implementation}', '{data}'];
  } else {
    upgradeMethod = 'upgradeTo';
    upgradeArgsTemplate = ['{implementation}'];
  }
}

The oldProxy case, to check if changeImplementation is implemented has been left in _deployViaProxy, before computing upgradeArgs and performing the upgrade:

--- _deployViaProxy ---

let proxy = await getDeploymentOrNUll(proxyName);
if (!proxy) { // [...]
} else {
  // [...]
  const oldProxy = proxy.abi.find(
    (frag: {name: string}) => frag.name === 'changeImplementation'
  );
  if (oldProxy) {
    upgradeMethod = 'changeImplementation';
    upgradeArgsTemplate = ['{implementation}', '{data}'];
  }
  let proxyAddress = proxy.address;
  let upgradeArgs = replaceTemplateArgs(upgradeArgsTemplate, {
    implementationAddress: implementation.address,
    proxyAdmin,
    data,
    proxyAddress,
  });
  if (!upgradeMethod) {
    throw new Error(`No upgrade method found, cannot make upgrades`);
  }
  let executeReceipt;
  // [...]
}

[Alivers]

@wighawag Hi, could you please review this feature ❤️ . Openzeppelin has removed upgradeTo function in 5.0.0, so we really need this custom upgrade function feature. Thank you!

[wighawag]

Thanks, available in v0.11.38