add initializable version of vesting (#95)

arr00 · web-flow · commit ad85c34d2f98 · 2025-07-08T15:03:08.000-06:00
* add initializable version of vesting

* update tests
diff --git a/contracts/finance/VestingWalletCliffConfidential.sol b/contracts/finance/VestingWalletCliffConfidential.sol
@@ -1,11 +1,11 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.8.20;
+pragma solidity ^0.8.24;
 
 import {VestingWalletConfidential} from "./VestingWalletConfidential.sol";
 import {euint64} from "@fhevm/solidity/lib/FHE.sol";
 
-contract VestingWalletCliffConfidential is VestingWalletConfidential {
-    uint64 private _cliff;
+abstract contract VestingWalletCliffConfidential is VestingWalletConfidential {
+    uint64 private immutable _cliff;
 
     /// @dev The specified cliff duration is larger than the vesting duration.
     error InvalidCliffDuration(uint64 cliffSeconds, uint64 durationSeconds);
@@ -14,18 +14,7 @@ contract VestingWalletCliffConfidential is VestingWalletConfidential {
      * @dev Set the duration of the cliff, in seconds. The cliff starts at the vesting
      * start timestamp (see {VestingWalletConfidential}) and ends `cliffSeconds` later.
      */
-    constructor() {
-        _disableInitializers();
-    }
-
-    function initialize(
-        address executor_,
-        address beneficiary,
-        uint64 startTimestamp,
-        uint64 durationSeconds,
-        uint64 cliffSeconds
-    ) public virtual initializer {
-        VestingWalletConfidential.initialize(executor_, beneficiary, startTimestamp, durationSeconds);
+    constructor(uint64 cliffSeconds) {
         if (cliffSeconds > duration()) {
             revert InvalidCliffDuration(cliffSeconds, duration());
         }
diff --git a/contracts/finance/VestingWalletCliffConfidentialInitializable.sol b/contracts/finance/VestingWalletCliffConfidentialInitializable.sol
@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {VestingWalletConfidential} from "./VestingWalletConfidential.sol";
+import {VestingWalletCliffConfidential} from "./VestingWalletCliffConfidential.sol";
+import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+
+contract VestingWalletCliffConfidentialInitializable is Initializable, VestingWalletCliffConfidential {
+    address private _executor;
+    uint64 private _start;
+    uint64 private _duration;
+    uint64 private _cliff;
+
+    constructor() VestingWalletCliffConfidential(0) VestingWalletConfidential(address(0), address(1), 0, 0) {}
+
+    function initialize(
+        address executor_,
+        address beneficiary,
+        uint64 startTimestamp,
+        uint64 durationSeconds,
+        uint64 cliffSeconds
+    ) public virtual initializer {
+        require(beneficiary != address(0), Ownable.OwnableInvalidOwner(address(0)));
+        _transferOwnership(beneficiary);
+
+        _executor = executor_;
+        _start = startTimestamp;
+        _duration = durationSeconds;
+
+        if (cliffSeconds > duration()) {
+            revert InvalidCliffDuration(cliffSeconds, duration());
+        }
+        _cliff = start() + cliffSeconds;
+    }
+
+    /// @inheritdoc VestingWalletConfidential
+    function executor() public view virtual override returns (address) {
+        return _executor;
+    }
+
+    /// @inheritdoc VestingWalletConfidential
+    function start() public view virtual override returns (uint64) {
+        return _start;
+    }
+
+    /// @inheritdoc VestingWalletConfidential
+    function duration() public view virtual override returns (uint64) {
+        return _duration;
+    }
+
+    /// @inheritdoc VestingWalletCliffConfidential
+    function cliff() public view virtual override returns (uint64) {
+        return _cliff;
+    }
+}
diff --git a/contracts/finance/VestingWalletConfidential.sol b/contracts/finance/VestingWalletConfidential.sol
@@ -1,8 +1,8 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.8.20;
+pragma solidity ^0.8.24;
 
 import {Address} from "@openzeppelin/contracts/utils/Address.sol";
-import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 import {FHE, ebool, euint64} from "@fhevm/solidity/lib/FHE.sol";
 
 import {IConfidentialFungibleToken} from "../interfaces/IConfidentialFungibleToken.sol";
@@ -24,29 +24,24 @@ import {TFHESafeMath} from "../utils/TFHESafeMath.sol";
  * NOTE: When using this contract with any token whose balance is adjusted automatically (i.e. a rebase token), make
  * sure to account the supply/balance adjustment in the vesting schedule to ensure the vested amount is as intended.
  */
-contract VestingWalletConfidential is OwnableUpgradeable {
+abstract contract VestingWalletConfidential is Ownable {
     mapping(address token => euint64) private _tokenReleased;
-    uint64 private _start;
-    uint64 private _duration;
-    address private _executor;
+    uint64 private immutable _start;
+    uint64 private immutable _duration;
+    address private immutable _executor;
 
     event VestingWalletConfidentialTokenReleased(address indexed token, euint64 amount);
     event VestingWalletCallExecuted(address indexed target, uint256 value, bytes data);
 
     error VestingWalletConfidentialInvalidDuration();
     error VestingWalletConfidentialOnlyExecutor();
 
-    constructor() {
-        _disableInitializers();
-    }
-
-    function initialize(
+    constructor(
         address executor_,
         address beneficiary,
         uint64 startTimestamp,
         uint64 durationSeconds
-    ) public virtual initializer {
-        __Ownable_init(beneficiary);
+    ) Ownable(beneficiary) {
         _start = startTimestamp;
         _duration = durationSeconds;
         _executor = executor_;
diff --git a/contracts/mocks/Create2Factory.sol b/contracts/mocks/Create2Factory.sol
@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.26;
+
+import {Clones} from "@openzeppelin/contracts/proxy/Clones.sol";
+
+contract Create2Factory {
+    event Deployed(address clone);
+
+    function create2(address impl, bytes memory data) public returns (address) {
+        address deployedTo = Clones.clone(impl);
+        (bool success, ) = deployedTo.call(data);
+        if (!success) {
+            assembly {
+                returndatacopy(0, 0, returndatasize())
+                revert(0, returndatasize())
+            }
+        }
+
+        emit Deployed(deployedTo);
+        return deployedTo;
+    }
+}
diff --git a/contracts/mocks/VestingWalletCliffConfidentialInitializableMock.sol b/contracts/mocks/VestingWalletCliffConfidentialInitializableMock.sol
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {VestingWalletCliffConfidentialInitializable} from "../finance/VestingWalletCliffConfidentialInitializable.sol";
+import {FHE} from "@fhevm/solidity/lib/FHE.sol";
+import {ZamaConfig} from "@fhevm/solidity/config/ZamaConfig.sol";
+
+abstract contract VestingWalletCliffConfidentialInitializableMock is VestingWalletCliffConfidentialInitializable {
+    function initialize(
+        address executor_,
+        address beneficiary,
+        uint64 startTimestamp,
+        uint64 durationSeconds,
+        uint64 cliffSeconds
+    ) public override {
+        super.initialize(executor_, beneficiary, startTimestamp, durationSeconds, cliffSeconds);
+
+        FHE.setCoprocessor(ZamaConfig.getSepoliaConfig());
+        FHE.setDecryptionOracle(ZamaConfig.getSepoliaOracleAddress());
+    }
+}
diff --git a/test/finance/VestingWalletCliffConfidential.test.ts b/test/finance/VestingWalletCliffConfidential.test.ts
@@ -2,62 +2,126 @@ import { shouldBehaveLikeVestingConfidential } from './VestingWalletConfidential
 import { FhevmType } from '@fhevm/hardhat-plugin';
 import { time } from '@nomicfoundation/hardhat-network-helpers';
 import { expect } from 'chai';
+import { EventLog } from 'ethers';
 import { ethers, fhevm } from 'hardhat';
 
 const name = 'ConfidentialFungibleToken';
 const symbol = 'CFT';
 const uri = 'https://example.com/metadata';
 
-describe('VestingWalletCliffConfidential', function () {
-  beforeEach(async function () {
-    const accounts = (await ethers.getSigners()).slice(3);
-    const [holder, recipient, operator] = accounts;
-
-    const token = await ethers.deployContract('$ConfidentialFungibleTokenMock', [name, symbol, uri]);
-
-    const encryptedInput = await fhevm
-      .createEncryptedInput(await token.getAddress(), holder.address)
-      .add64(1000)
-      .encrypt();
-
-    const currentTime = await time.latest();
-    const schedule = [currentTime + 60, currentTime + 60 * 121];
-    const vesting = await ethers.deployContract('$VestingWalletCliffConfidentialMock', [
-      operator,
-      recipient,
-      currentTime + 60,
-      60 * 60 * 2 /* 2 hours */,
-      60 * 60 /* 1 hour */,
-    ]);
-
-    await (token as any)
-      .connect(holder)
-      ['$_mint(address,bytes32,bytes)'](vesting.target, encryptedInput.handles[0], encryptedInput.inputProof);
-
-    Object.assign(this, { accounts, holder, recipient, operator, token, vesting, schedule, vestingAmount: 1000 });
-  });
+for (const useInitializable of [false, true]) {
+  describe(`VestingWalletCliffConfidential${useInitializable ? 'Initializable' : ''}`, function () {
+    beforeEach(async function () {
+      const accounts = (await ethers.getSigners()).slice(3);
+      const [holder, recipient, operator] = accounts;
 
-  it('should release nothing before cliff', async function () {
-    await time.increaseTo(this.schedule[0] + 60);
-    await this.vesting.release(this.token);
+      const token = await ethers.deployContract('$ConfidentialFungibleTokenMock', [name, symbol, uri]);
 
-    const balanceOfHandle = await this.token.balanceOf(this.recipient);
-    await expect(
-      fhevm.userDecryptEuint(FhevmType.euint64, balanceOfHandle, this.token.target, this.recipient),
-    ).to.eventually.equal(0);
-  });
+      const encryptedInput = await fhevm
+        .createEncryptedInput(await token.getAddress(), holder.address)
+        .add64(1000)
+        .encrypt();
 
-  it('should fail construction if cliff is longer than duration', async function () {
-    await expect(
-      ethers.deployContract('$VestingWalletCliffConfidentialMock', [
-        this.operator,
-        this.recipient,
-        (await time.latest()) + 60,
-        60 * 10,
-        60 * 60,
-      ]),
-    ).to.be.revertedWithCustomError(this.vesting, 'InvalidCliffDuration');
-  });
+      const currentTime = await time.latest();
+      const schedule = [currentTime + 60, currentTime + 60 * 121];
+
+      let vesting;
+      let factory;
+      let impl;
+
+      if (!useInitializable) {
+        vesting = await ethers.deployContract('$VestingWalletCliffConfidentialMock', [
+          operator,
+          recipient,
+          currentTime + 60,
+          60 * 60 * 2 /* 2 hours */,
+          60 * 60 /* 1 hour */,
+        ]);
+      } else {
+        impl = await ethers.deployContract('$VestingWalletCliffConfidentialInitializableMock');
+        factory = await ethers.deployContract('Create2Factory');
+
+        const callData = await impl.initialize.populateTransaction(
+          operator,
+          recipient,
+          currentTime + 60,
+          60 * 60 * 2 /* 2 hours */,
+          60 * 60 /* 1 hour */,
+        );
+        const cloneTx = (await (await factory.create2(impl.target, callData.data)).wait())!;
+        const cloneAddress = (cloneTx.logs[2] as EventLog).args[0];
+
+        vesting = await ethers.getContractAt('$VestingWalletCliffConfidentialInitializableMock', cloneAddress);
+      }
 
-  shouldBehaveLikeVestingConfidential();
-});
+      await (token as any)
+        .connect(holder)
+        ['$_mint(address,bytes32,bytes)'](vesting.target, encryptedInput.handles[0], encryptedInput.inputProof);
+
+      Object.assign(this, {
+        accounts,
+        holder,
+        recipient,
+        operator,
+        token,
+        vesting,
+        schedule,
+        vestingAmount: 1000,
+        factory,
+        impl,
+      });
+    });
+
+    it('should release nothing before cliff', async function () {
+      await time.increaseTo(this.schedule[0] + 60);
+      await this.vesting.release(this.token);
+
+      const balanceOfHandle = await this.token.balanceOf(this.recipient);
+      await expect(
+        fhevm.userDecryptEuint(FhevmType.euint64, balanceOfHandle, this.token.target, this.recipient),
+      ).to.eventually.equal(0);
+    });
+
+    it('should fail construction if cliff is longer than duration', async function () {
+      if (!useInitializable) {
+        await expect(
+          ethers.deployContract('$VestingWalletCliffConfidentialMock', [
+            this.operator,
+            this.recipient,
+            (await time.latest()) + 60,
+            60 * 10,
+            60 * 60,
+          ]),
+        ).to.be.revertedWithCustomError(this.vesting, 'InvalidCliffDuration');
+      } else {
+        const callData = await this.impl.initialize.populateTransaction(
+          this.operator,
+          this.recipient,
+          (await time.latest()) + 60,
+          60 * 10,
+          60 * 60,
+        );
+        await expect(this.factory.create2(this.impl.target, callData.data)).to.be.revertedWithCustomError(
+          this.vesting,
+          'InvalidCliffDuration',
+        );
+      }
+    });
+
+    if (useInitializable) {
+      it('cannot reinitialize', async function () {
+        await expect(
+          this.vesting.initialize(
+            this.operator,
+            this.recipient,
+            (await time.latest()) + 60,
+            60 * 60 * 2 /* 2 hours */,
+            60 * 60 /* 1 hour */,
+          ),
+        ).to.be.revertedWithCustomError(this.vesting, 'InvalidInitialization');
+      });
+    }
+
+    shouldBehaveLikeVestingConfidential();
+  });
+}
