Use Ownable in VestingWallet instead of an immutable beneficiary

.changeset/healthy-gorillas-applaud.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': major
+---
+
+`VestingWallet`: Use Ownable2Step instead of an immutable `beneficiary`. The initial owner is set to the benefactor (`msg.sender`) but transferred to the beneficiary address so that unclaimed tokens can be recovered by the benefactor.

contracts/finance/VestingWallet.sol

@@ -6,12 +6,17 @@ import {IERC20} from "../token/ERC20/IERC20.sol";
 import {SafeERC20} from "../token/ERC20/utils/SafeERC20.sol";
 import {Address} from "../utils/Address.sol";
 import {Context} from "../utils/Context.sol";
+import {Ownable2Step, Ownable} from "../access/Ownable2Step.sol";
 
 /**
  * @title VestingWallet
- * @dev This contract handles the vesting of Eth and ERC20 tokens for a given beneficiary. Custody of multiple tokens
- * can be given to this contract, which will release the token to the beneficiary following a given vesting schedule.
- * The vesting schedule is customizable through the {vestedAmount} function.
+ * @dev Handles the vesting of native currency and ERC20 tokens for a given beneficiary who gets the ownership of the
+ * contract by calling {Ownable2Step-acceptOwnership} to accept a 2-step ownership transfer setup at deployment with
+ * {Ownable2Step-transferOwnership}. The initial owner of this contract is the benefactor (`msg.sender`), enabling them
+ * to recover any unclaimed tokens.
+ *
+ * Custody of multiple tokens can be given to this contract, which will release the tokens to the beneficiary following
+ * a given vesting schedule that is customizable through the {vestedAmount} function.
  *
  * Any token transferred to this contract will follow the vesting schedule as if they were locked from the beginning.
  * Consequently, if the vesting has already started, any amount of tokens sent to this contract will (at least partly)
@@ -20,7 +25,7 @@ import {Context} from "../utils/Context.sol";
  * By setting the duration to 0, one can configure this contract to behave like an asset timelock that hold tokens for
  * a beneficiary until a specified time.
  */
-contract VestingWallet is Context {
+contract VestingWallet is Context, Ownable2Step {
     event EtherReleased(uint256 amount);
     event ERC20Released(address indexed token, uint256 amount);
 
@@ -31,18 +36,19 @@ contract VestingWallet is Context {
 
     uint256 private _released;
     mapping(address => uint256) private _erc20Released;
-    address private immutable _beneficiary;
     uint64 private immutable _start;
     uint64 private immutable _duration;
 
     /**
-     * @dev Set the beneficiary, start timestamp and vesting duration of the vesting wallet.
+     * @dev Sets the sender as the initial owner, the beneficiary as the pending owner, the start timestamp and the
+     * vesting duration of the vesting wallet.
      */
-    constructor(address beneficiaryAddress, uint64 startTimestamp, uint64 durationSeconds) payable {
+    constructor(address beneficiaryAddress, uint64 startTimestamp, uint64 durationSeconds) payable Ownable(msg.sender) {
         if (beneficiaryAddress == address(0)) {
             revert VestingWalletInvalidBeneficiary(address(0));
         }
-        _beneficiary = beneficiaryAddress;
+        transferOwnership(beneficiaryAddress);
+
         _start = startTimestamp;
         _duration = durationSeconds;
     }
@@ -52,13 +58,6 @@ contract VestingWallet is Context {
      */
     receive() external payable virtual {}
 
-    /**
-     * @dev Getter for the beneficiary address.
-     */
-    function beneficiary() public view virtual returns (address) {
-        return _beneficiary;
-    }
-
     /**
      * @dev Getter for the start timestamp.
      */
@@ -114,23 +113,23 @@ contract VestingWallet is Context {
      *
      * Emits a {EtherReleased} event.
      */
-    function release() public virtual {
+    function release() public virtual onlyOwner {
         uint256 amount = releasable();
         _released += amount;
         emit EtherReleased(amount);
-        Address.sendValue(payable(beneficiary()), amount);
+        Address.sendValue(payable(owner()), amount);
     }
 
     /**
      * @dev Release the tokens that have already vested.
      *
      * Emits a {ERC20Released} event.
      */
-    function release(address token) public virtual {
+    function release(address token) public virtual onlyOwner {
         uint256 amount = releasable(token);
         _erc20Released[token] += amount;
         emit ERC20Released(token, amount);
-        SafeERC20.safeTransfer(IERC20(token), beneficiary(), amount);
+        SafeERC20.safeTransfer(IERC20(token), owner(), amount);
     }
 
     /**

test/finance/VestingWallet.behavior.js

@@ -1,12 +1,14 @@
-const { time } = require('@nomicfoundation/hardhat-network-helpers');
+const { time, setNextBlockBaseFeePerGas } = require('@nomicfoundation/hardhat-network-helpers');
 const { expectEvent } = require('@openzeppelin/test-helpers');
 const { expect } = require('chai');
+const { web3 } = require('hardhat');
+const { expectRevertCustomError } = require('../helpers/customError');
 
 function releasedEvent(token, amount) {
   return token ? ['ERC20Released', { token: token.address, amount }] : ['EtherReleased', { amount }];
 }
 
-function shouldBehaveLikeVesting(beneficiary) {
+function shouldBehaveLikeVesting(beneficiary, accounts) {
   it('check vesting schedule', async function () {
     const [vestedAmount, releasable, ...args] = this.token
       ? ['vestedAmount(address,uint64)', 'releasable(address)', this.token.address]
@@ -22,35 +24,51 @@ function shouldBehaveLikeVesting(beneficiary) {
     }
   });
 
-  it('execute vesting schedule', async function () {
-    const [release, ...args] = this.token ? ['release(address)', this.token.address] : ['release()'];
+  describe('execute vesting schedule', async function () {
+    beforeEach(async function () {
+      [this.release, ...this.args] = this.token ? ['release(address)', this.token.address] : ['release()'];
+    });
 
-    let released = web3.utils.toBN(0);
-    const before = await this.getBalance(beneficiary);
+    it('releases a linearly vested schedule', async function () {
+      let released = web3.utils.toBN(0);
+      const before = await this.getBalance(beneficiary);
+      const releaser = await this.mock.owner();
 
-    {
-      const receipt = await this.mock.methods[release](...args);
+      {
+        // Allows gas price to be 0 so no ETH is spent in the transaction.
+        await setNextBlockBaseFeePerGas(0);
 
-      await expectEvent.inTransaction(receipt.tx, this.mock, ...releasedEvent(this.token, '0'));
+        const receipt = await this.mock.methods[this.release](...this.args, { from: releaser, gasPrice: 0 });
+        await expectEvent.inTransaction(receipt.tx, this.mock, ...releasedEvent(this.token, '0'));
+        await this.checkRelease(receipt, beneficiary, '0');
 
-      await this.checkRelease(receipt, beneficiary, '0');
+        expect(await this.getBalance(beneficiary)).to.be.bignumber.equal(before);
+      }
 
-      expect(await this.getBalance(beneficiary)).to.be.bignumber.equal(before);
-    }
+      for (const timestamp of this.schedule) {
+        await time.setNextBlockTimestamp(timestamp);
+        const vested = this.vestingFn(timestamp);
 
-    for (const timestamp of this.schedule) {
-      await time.setNextBlockTimestamp(timestamp);
-      const vested = this.vestingFn(timestamp);
+        // Allows gas price to be 0 so no ETH is spent in the transaction.
+        await setNextBlockBaseFeePerGas(0);
 
-      const receipt = await this.mock.methods[release](...args);
-      await expectEvent.inTransaction(receipt.tx, this.mock, ...releasedEvent(this.token, vested.sub(released)));
+        const receipt = await this.mock.methods[this.release](...this.args, { from: releaser, gasPrice: 0 });
+        await expectEvent.inTransaction(receipt.tx, this.mock, ...releasedEvent(this.token, vested.sub(released)));
+        await this.checkRelease(receipt, beneficiary, vested.sub(released));
 
-      await this.checkRelease(receipt, beneficiary, vested.sub(released));
+        expect(await this.getBalance(beneficiary)).to.be.bignumber.equal(before.add(vested));
 
-      expect(await this.getBalance(beneficiary)).to.be.bignumber.equal(before.add(vested));
+        released = vested;
+      }
+    });
 
-      released = vested;
-    }
+    it('cannot be released by a non releaser', async function () {
+      await expectRevertCustomError(
+        this.mock.methods[this.release](...this.args, { from: accounts[0] }),
+        'OwnableUnauthorizedAccount',
+        [accounts[0]],
+      );
+    });
   });
 }
 

test/finance/VestingWallet.test.js

@@ -18,6 +18,7 @@ contract('VestingWallet', function (accounts) {
   beforeEach(async function () {
     this.start = (await time.latest()).addn(3600); // in 1 hour
     this.mock = await VestingWallet.new(beneficiary, this.start, duration);
+    await this.mock.acceptOwnership({ from: beneficiary });
   });
 
   it('rejects zero address for beneficiary', async function () {
@@ -28,8 +29,14 @@ contract('VestingWallet', function (accounts) {
     );
   });
 
+  it('sets the initial owner as the sender (benefactor)', async function () {
+    this.mock = await VestingWallet.new(beneficiary, this.start, duration, { from: sender });
+    expect(await this.mock.owner()).to.be.equal(sender);
+    expect(await this.mock.pendingOwner()).to.be.equal(beneficiary);
+  });
+
   it('check vesting contract', async function () {
-    expect(await this.mock.beneficiary()).to.be.equal(beneficiary);
+    expect(await this.mock.owner()).to.be.equal(beneficiary);
     expect(await this.mock.start()).to.be.bignumber.equal(this.start);
     expect(await this.mock.duration()).to.be.bignumber.equal(duration);
     expect(await this.mock.end()).to.be.bignumber.equal(this.start.add(duration));
@@ -50,7 +57,7 @@ contract('VestingWallet', function (accounts) {
         this.checkRelease = () => {};
       });
 
-      shouldBehaveLikeVesting(beneficiary);
+      shouldBehaveLikeVesting(beneficiary, accounts);
     });
 
     describe('ERC20 vesting', function () {
@@ -63,7 +70,7 @@ contract('VestingWallet', function (accounts) {
         await this.token.$_mint(this.mock.address, amount);
       });
 
-      shouldBehaveLikeVesting(beneficiary);
+      shouldBehaveLikeVesting(beneficiary, accounts);
     });
   });
 });