Add a variant of the Vesting wallet for updating the beneficiary

contracts/vesting/VestingWalletRecoveryLight.sol

@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: UNLICENSED
+// See Forta Network License: https://github.com/forta-network/forta-contracts/blob/master/LICENSE.md
+
+pragma solidity ^0.8.9;
+
+import "@openzeppelin/contracts/utils/Address.sol";
+import "@openzeppelin/contracts/utils/StorageSlot.sol";
+
+/**
+ * This contract is designed for recovery in case the beneficiary was lost.
+ *
+ * This contract replicates a storage layout that is compatible with VestingWallet, VestingWalletV1 and
+ * VestingWalletV2.
+ */
+contract VestingWalletRecoveryLight {
+    /// Storage
+    // Initializable
+    uint8 private _initialized;
+    bool private _initializing;
+    // ContextUpgradeable
+    uint256[50] private __gap_1;
+    // OwnableUpgradeable
+    address private _owner;
+    uint256[49] private __gap_2;
+    // UUPSUpgradeable
+    uint256[50] private __gap_3;
+    // ERC1967UpgradeUpgradeable
+    uint256[50] private __gap_4;
+    // VestingWallet / VestingWalletV1 / VestingWalletV2
+    mapping (address => uint256) private _released;
+    address private _beneficiary;
+    uint256 private _start;
+    uint256 private _cliff;
+    uint256 private _duration;
+    // VestingWalletV1 / VestingWalletV2
+    uint256[45] private __gap_5;
+    // VestingWalletV2
+    uint256 private historicalBalanceMin;
+    uint256[45] private __gap_6;
+
+    /// Constants and Events
+    // ERC1967UpgradeUpgradeable
+    bytes32 internal constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
+    event Upgraded(address indexed implementation);
+    event BeneficiaryUpdate(address newBeneficiary);
+
+    /// @dev The `previousImplementation` address MUST correspond to the real previous implementation.
+    /// Otherwise, upgrading to this implementation contract and calling this function will leave the new implementation uninitialized.
+    function changeBeneficiaryAndRollbackTo(address newBeneficiary, address previousImplementation) external {
+        // restrict to owner, in case the upgrade misses the "andCall" part.
+        require(msg.sender == _owner);
+
+        // change beneficiary
+        _beneficiary = newBeneficiary;
+        emit BeneficiaryUpdate(newBeneficiary);
+
+        // ERC1967Upgrade._setImplementation
+        require(Address.isContract(previousImplementation), "ERC1967: new implementation is not a contract");
+        StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value = previousImplementation;
+        emit Upgraded(previousImplementation);
+    }
+
+    function proxiableUUID() external pure returns (bytes32) {
+        return _IMPLEMENTATION_SLOT;
+    }
+
+    function upgradeTo(address newImplementation) external {
+        require(msg.sender == _owner);
+        StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value = newImplementation;
+    }
+
+    function upgradeToAndCall(address newImplementation, bytes memory data) external {
+        require(msg.sender == _owner);
+        StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value = newImplementation;
+        if (data.length > 0) {
+            Address.functionDelegateCall(newImplementation, data);
+        }
+    }
+}

contracts/vesting/VestingWalletV2.sol

@@ -66,7 +66,7 @@ contract VestingWalletV2 is VestingWalletV1 {
     /**
      * @notice Bridge token to L2, with custom escrow manager on L2.
      * @dev Using a custom escrow manager is needed if the beneficiary isn't valid on the child chain, for example if it
-     * is a smart wallet that doesn't exist at the same address on the child chain. 
+     * is a smart wallet that doesn't exist at the same address on the child chain.
      * If the beneficiary of the contract is a smart wallet valid on both chains, it must be explicitly mentioned as the manager.
      * @dev amount of tokens to be bridged.
      * @dev l2Manager the address that will be controller generated StakinEscrow in L2.
@@ -118,7 +118,7 @@ contract VestingWalletV2 is VestingWalletV1 {
 
     /**
      * @dev Sets historical balance minimum. Only use if there is an imbalance between L1 VestingWallet and L2 StakingEscrow
-     */ 
+     */
     function setHistoricalBalanceMin(uint256 value)
         public
         onlyOwner()
@@ -129,7 +129,7 @@ contract VestingWalletV2 is VestingWalletV1 {
 
     /**
      * @dev Sets  historical balance. Only use if there is an imbalance between L1 VestingWallet and L2 StakingEscrow
-     */ 
+     */
     function updateHistoricalBalanceMin(int256 update)
         public
         onlyOwner()

test/vesting/VestingWallet.recovery.test.js

@@ -0,0 +1,70 @@
+const hre = require('hardhat');
+const { ethers } = hre;
+const { expect } = require('chai');
+const { prepare, deployUpgradeable, performUpgrade } = require('../fixture');
+const utils = require('../../scripts/utils');
+
+const allocation = {
+    start: utils.dateToTimestamp('2021-09-01T00:00:00Z'),
+    cliff: utils.durationToSeconds('1 year'),
+    duration: utils.durationToSeconds('4 years'),
+};
+
+describe('Vesting recovery', function () {
+    prepare();
+
+    for (const { contractName, constructorArgs } of [
+        { contractName: 'VestingWallet' },
+        { contractName: 'VestingWalletV1' },
+        { contractName: 'VestingWalletV2', constructorArgs: Array(4).fill(ethers.constants.AddressZero) },
+    ])
+    describe(contractName, function () {
+        beforeEach(async function () {
+            allocation.beneficiary = this.accounts.user1.address;
+            allocation.newBeneficiary = this.accounts.user2.address;
+            allocation.owner = this.accounts.admin.address;
+
+            this.vesting = await deployUpgradeable(
+                hre,
+                contractName,
+                'uups',
+                [allocation.beneficiary, allocation.owner, allocation.start, allocation.cliff, allocation.duration],
+                { unsafeAllow: 'delegatecall', constructorArgs }
+            );
+        });
+
+        describe('recover beneficiary', function () {
+            beforeEach(async function () {
+                await Promise.all([this.vesting.start(), this.vesting.cliff(), this.vesting.duration(), this.vesting.beneficiary(), this.vesting.owner()]).then(
+                    ([start, cliff, duration, beneficiary, owner]) => {
+                        expect(start).to.be.equal(allocation.start);
+                        expect(cliff).to.be.equal(allocation.cliff);
+                        expect(duration).to.be.equal(allocation.duration);
+                        expect(beneficiary).to.be.equal(allocation.beneficiary);
+                        expect(owner).to.be.equal(allocation.owner);
+                    }
+                );
+            });
+
+            it('perform upgrade', async function () {
+                const implementation = await hre.upgrades.erc1967.getImplementationAddress(this.vesting.address);
+                await performUpgrade(hre, this.vesting, 'VestingWalletRecoveryLight', {
+                    call: { fn: 'changeBeneficiaryAndRollbackTo', args: [allocation.newBeneficiary, implementation] },
+                    unsafeAllow: 'delegatecall'
+                });
+            });
+
+            afterEach(async function () {
+                await Promise.all([this.vesting.start(), this.vesting.cliff(), this.vesting.duration(), this.vesting.beneficiary(), this.vesting.owner()]).then(
+                    ([start, cliff, duration, beneficiary, owner]) => {
+                        expect(start).to.be.equal(allocation.start);
+                        expect(cliff).to.be.equal(allocation.cliff);
+                        expect(duration).to.be.equal(allocation.duration);
+                        expect(beneficiary).to.be.equal(allocation.newBeneficiary);
+                        expect(owner).to.be.equal(allocation.owner);
+                    }
+                );
+            });
+        });
+    });
+});