Options to change deploy function and proxy contract factory

packages/plugin-hardhat/contracts/AccessManagedProxy.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+// Example of a custom proxy, for testing only.
+
+import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import {AccessManager} from "@openzeppelin/contracts/access/manager/AccessManager.sol";
+import {IAccessManager} from "@openzeppelin/contracts/access/manager/IAccessManager.sol";
+import {IAccessManaged} from "@openzeppelin/contracts/access/manager/IAccessManaged.sol";
+
+contract AccessManagedProxy is ERC1967Proxy {
+    IAccessManager public immutable ACCESS_MANAGER;
+
+    constructor(address implementation, bytes memory _data, IAccessManager manager) payable ERC1967Proxy(implementation, _data) {
+        ACCESS_MANAGER = manager;
+    }
+
+    /**
+     * @dev Checks with the ACCESS_MANAGER if msg.sender is authorized to call the current call's function,
+     * and if so, delegates the current call to `implementation`.
+     *
+     * This function does not return to its internal call site, it will return directly to the external caller.
+     */
+    function _delegate(address implementation) internal virtual override {
+        (bool immediate, ) = ACCESS_MANAGER.canCall(msg.sender, address(this), bytes4(msg.data[0:4]));
+        if (!immediate) revert IAccessManaged.AccessManagedUnauthorized(msg.sender);
+        super._delegate(implementation);
+    }
+}

packages/plugin-hardhat/src/defender/deploy.ts

@@ -10,8 +10,6 @@ import {
   DeployRequestLibraries,
 } from '@openzeppelin/defender-sdk-deploy-client';
 import {
-  Deployment,
-  RemoteDeploymentId,
   getContractNameAndRunValidation,
   UpgradesError,
 } from '@openzeppelin/upgrades-core';
@@ -24,7 +22,12 @@ import UpgradeableBeacon from '@openzeppelin/upgrades-core/artifacts/@openzeppel
 import TransparentUpgradeableProxy from '@openzeppelin/upgrades-core/artifacts/@openzeppelin/contracts-v5/proxy/transparent/TransparentUpgradeableProxy.sol/TransparentUpgradeableProxy.json';
 
 import { getNetwork, parseTxOverrides } from './utils';
-import { DeployTransaction, DefenderDeployOptions, UpgradeOptions, EthersDeployOptions } from '../utils';
+import {
+  DefenderDeployOptions,
+  UpgradeOptions,
+  EthersDeployOptions,
+  DefenderDeployment,
+} from '../utils';
 import debug from '../utils/debug';
 import { getDeployData } from '../utils/deploy-impl';
 import { ContractSourceNotFoundError } from '@openzeppelin/upgrades-core';
@@ -60,7 +63,7 @@ export async function defenderDeploy(
   factory: ContractFactory,
   opts: UpgradeOptions & EthersDeployOptions & DefenderDeployOptions,
   ...args: unknown[]
-): Promise<Required<Deployment & RemoteDeploymentId> & DeployTransaction> {
+): Promise<DefenderDeployment> {
   const client = getDeployClient(hre);
 
   // Override constructor arguments in options with the ones passed as arguments to this function.

packages/plugin-hardhat/src/deploy-beacon-proxy.ts

@@ -80,10 +80,10 @@ export function makeDeployBeaconProxy(
       ]);
     }
 
-    const BeaconProxyFactory = await getBeaconProxyFactory(hre, getSigner(attachTo.runner));
+    const BeaconProxyFactory = await (opts.proxyFactory || getBeaconProxyFactory)(hre, getSigner(attachTo.runner));
     const proxyDeployment: Required<ProxyDeployment> & DeployTransaction & RemoteDeploymentId = Object.assign(
       { kind: opts.kind },
-      await deploy(hre, opts, BeaconProxyFactory, beaconAddress, data),
+      await (opts.deployFunction || deploy)(hre, opts, BeaconProxyFactory, beaconAddress, data),
     );
 
     await manifest.addProxy(proxyDeployment);

packages/plugin-hardhat/src/deploy-contract.ts

@@ -1,12 +1,10 @@
 import { HardhatRuntimeEnvironment } from 'hardhat/types';
 import type { ContractFactory, Contract } from 'ethers';
 
-import { deploy, DeployContractOptions, DeployTransaction } from './utils';
+import { deploy, DeployContractOptions, EthersOrDefenderDeployment } from './utils';
 import { DeployData, getDeployData } from './utils/deploy-impl';
 import { enableDefender } from './defender/utils';
 import {
-  Deployment,
-  RemoteDeploymentId,
   getContractNameAndRunValidation,
   inferProxyKind,
   UpgradesError,
@@ -30,7 +28,7 @@ async function deployNonUpgradeableContract(
     assertNonUpgradeable(deployData);
   }
 
-  const deployment: Required<Deployment> & DeployTransaction & RemoteDeploymentId = await deploy(
+  const deployment: EthersOrDefenderDeployment = await deploy(
     hre,
     opts,
     Contract,

packages/plugin-hardhat/src/deploy-proxy.ts

@@ -51,6 +51,7 @@ export function makeDeployProxy(hre: HardhatRuntimeEnvironment, defenderModule:
 
     const contractInterface = ImplFactory.interface;
     const data = getInitializerData(contractInterface, args, opts.initializer);
+    const deployFn = opts.deployFunction || deploy;
 
     if (await manifest.getAdmin()) {
       if (kind === 'uups') {
@@ -79,8 +80,8 @@ export function makeDeployProxy(hre: HardhatRuntimeEnvironment, defenderModule:
           throw new InitialOwnerUnsupportedKindError(kind);
         }
 
-        const ProxyFactory = await getProxyFactory(hre, signer);
-        proxyDeployment = Object.assign({ kind }, await deploy(hre, opts, ProxyFactory, impl, data));
+        const ProxyFactory = await (opts.proxyFactory || getProxyFactory)(hre, signer);
+        proxyDeployment = Object.assign({ kind }, await deployFn(hre, opts, ProxyFactory, impl, data));
         break;
       }
 
@@ -95,10 +96,13 @@ export function makeDeployProxy(hre: HardhatRuntimeEnvironment, defenderModule:
           );
         }
 
-        const TransparentUpgradeableProxyFactory = await getTransparentUpgradeableProxyFactory(hre, signer);
+        const TransparentUpgradeableProxyFactory = await (opts.proxyFactory || getTransparentUpgradeableProxyFactory)(
+          hre,
+          signer,
+        );
         proxyDeployment = Object.assign(
           { kind },
-          await deploy(hre, opts, TransparentUpgradeableProxyFactory, impl, initialOwner, data),
+          await deployFn(hre, opts, TransparentUpgradeableProxyFactory, impl, initialOwner, data),
         );
         break;
       }

packages/plugin-hardhat/src/utils/deploy.ts

@@ -8,13 +8,17 @@ export interface DeployTransaction {
   deployTransaction?: ethers.TransactionResponse;
 }
 
+// Defender always includes RemoteDeploymentId, while ethers always includes DeployTransaction
+export type EthersOrDefenderDeployment = Required<Deployment> & DeployTransaction & RemoteDeploymentId;
+export type DefenderDeployment = Required<Deployment & RemoteDeploymentId> & DeployTransaction;
+export type EthersDeployment = Required<Deployment & DeployTransaction> & RemoteDeploymentId;
+
 export async function deploy(
   hre: HardhatRuntimeEnvironment,
   opts: UpgradeOptions & EthersDeployOptions & DefenderDeployOptions,
   factory: ContractFactory,
   ...args: unknown[]
-): Promise<Required<Deployment> & DeployTransaction & RemoteDeploymentId> {
-  // defender always includes RemoteDeploymentId, while ethers always includes DeployTransaction
+): Promise<EthersOrDefenderDeployment> {
   if (opts?.useDefenderDeploy) {
     return await defenderDeploy(hre, factory, opts, ...args);
   } else {
@@ -28,7 +32,7 @@ export async function deploy(
 async function ethersDeploy(
   factory: ContractFactory,
   ...args: ContractMethodArgs<unknown[]>
-): Promise<Required<Deployment & DeployTransaction> & RemoteDeploymentId> {
+): Promise<EthersDeployment> {
   const contractInstance = await factory.deploy(...args);
 
   const deployTransaction = contractInstance.deploymentTransaction();

packages/plugin-hardhat/src/utils/options.ts

@@ -6,7 +6,23 @@ import {
   ValidationOptions,
   withValidationDefaults,
 } from '@openzeppelin/upgrades-core';
-import { Overrides } from 'ethers';
+import { ContractFactory, Overrides } from 'ethers';
+import { EthersOrDefenderDeployment } from './deploy';
+
+/**
+ * Options for customizing the factory or deploy functions
+ */
+export type DeployFactoryOpts = {
+  /**
+   * Allows to customize the proxyFactory used instead of the ones defined in utils/factories.ts
+   */
+  proxyFactory?: () => Promise<ContractFactory>;
+
+  /**
+   * Allows to customize the deploy function used instead of utils/deploy.ts:deploy
+   */
+  deployFunction?: () => Promise<EthersOrDefenderDeployment>;
+};
 
 /**
  * Options for functions that can deploy an implementation contract.
@@ -91,6 +107,7 @@ export type InitialOwner = {
 
 export type DeployBeaconProxyOptions = EthersDeployOptions &
   DeployOpts &
+  DeployFactoryOpts &
   ProxyKindOption &
   Initializer &
   DefenderDeployOptions;
@@ -101,7 +118,11 @@ export type DeployContractOptions = Omit<StandaloneOptions, 'txOverrides'> & //
   DefenderDeployOptions & {
     unsafeAllowDeployContract?: boolean;
   };
-export type DeployProxyOptions = StandaloneOptions & Initializer & InitialOwner & DefenderDeployOptions;
+export type DeployProxyOptions = StandaloneOptions &
+  DeployFactoryOpts &
+  Initializer &
+  InitialOwner &
+  DefenderDeployOptions;
 export type ForceImportOptions = ProxyKindOption;
 export type PrepareUpgradeOptions = UpgradeOptions & GetTxResponse & DefenderDeployOptions;
 export type UpgradeBeaconOptions = UpgradeOptions & DefenderDeploy;

packages/plugin-hardhat/test/uups-accessmanaged-proxy.js

@@ -0,0 +1,57 @@
+const test = require('ava');
+
+const { ethers, upgrades } = require('hardhat');
+const { deploy } = require('../dist/utils/deploy');
+
+test.before(async t => {
+  t.context.Greeter = await ethers.getContractFactory('GreeterProxiable');
+  t.context.GreeterV2 = await ethers.getContractFactory('GreeterV2Proxiable');
+  t.context.GreeterV3 = await ethers.getContractFactory('GreeterV3Proxiable');
+  t.context.AccessManagedProxy = await ethers.getContractFactory('AccessManagedProxy');
+  const AccessManager = await ethers.getContractFactory('AccessManager');
+  const [admin, anon] = await ethers.getSigners();
+  t.context.anon = anon;
+  t.context.admin = admin;
+  t.context.acMgr = await AccessManager.deploy(admin);
+});
+
+async function getAccessManagedProxy(hre, signer) {
+  return hre.ethers.getContractFactory('AccessManagedProxy', signer);
+}
+
+async function deployWithExtraProxyArgs(hre, opts, factory, ...args) {
+  const allArgs = [...args, ...(opts.proxyExtraConstructorArgs || [])];
+  return deploy(hre, opts, factory, ...allArgs);
+}
+
+test('accessmanaged proxy / customization of deploy and factory functions', async t => {
+  const { Greeter, GreeterV2, GreeterV3, acMgr, anon, admin } = t.context;
+
+  const greeter = await upgrades.deployProxy(Greeter, ['Hello, Hardhat!'], {
+    kind: 'uups',
+    proxyExtraConstructorArgs: [await acMgr.getAddress()],
+    deployFunction: deployWithExtraProxyArgs,
+    proxyFactory: getAccessManagedProxy,
+  });
+
+  // By default it calls from admin address, so, it works fine
+  let greet = await greeter.connect(admin).greet();
+  t.is(greet, 'Hello, Hardhat!');
+  // But fails when called from other user
+  let e = await t.throwsAsync(() => greeter.connect(anon).greet());
+  t.true(e.message.includes('AccessManagedUnauthorized'), e.message);
+
+  // Upgrades work well, because the call executed from the default signer that is the admin
+  const greeter2 = await upgrades.upgradeProxy(greeter, GreeterV2);
+  await greeter2.waitForDeployment();
+  await greeter2.resetGreeting();
+
+  // Upgrades don't break the access control
+  e = await t.throwsAsync(() => greeter2.connect(anon).resetGreeting());
+  t.true(e.message.includes('AccessManagedUnauthorized'), e.message);
+
+  const greeter3ImplAddr = await upgrades.prepareUpgrade(await greeter.getAddress(), GreeterV3);
+  const greeter3 = GreeterV3.attach(greeter3ImplAddr);
+  const version3 = await greeter3.version();
+  t.is(version3, 'V3');
+});