feat(contracts): apply audit results about vesting wallet (#342)

kroma-network · Jul 12, 2024 · 0247393 · 0247393
1 parent 6db98fc
commit 0247393
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 9 deletions.
diff --git a/packages/contracts/.gas-snapshot b/packages/contracts/.gas-snapshot
@@ -164,14 +164,15 @@ KromaPortal_Test:test_receive_succeeds() (gas: 127638)
 KromaPortal_Test:test_simple_isOutputFinalized_succeeds() (gas: 35393)
 KromaPortal_Test:test_unpause_onlyGuardian_reverts() (gas: 46358)
 KromaPortal_Test:test_unpause_succeeds() (gas: 31932)
-KromaVestingWalletTest:test_constructor_succeeds() (gas: 13938)
-KromaVestingWalletTest:test_constructor_zeroValues_reverts() (gas: 72905)
-KromaVestingWalletTest:test_initialize_succeeds() (gas: 21631)
-KromaVestingWalletTest:test_release_afterFullyVested_succeeds() (gas: 79730)
-KromaVestingWalletTest:test_release_notBeneficiary_reverts() (gas: 28574)
-KromaVestingWalletTest:test_release_succeeds() (gas: 116326)
-KromaVestingWalletTest:test_release_tokenAfterFullyVested_succeeds() (gas: 95049)
-KromaVestingWalletTest:test_release_token_succeeds() (gas: 154458)
+KromaVestingWalletTest:test_constructor_succeeds() (gas: 13916)
+KromaVestingWalletTest:test_constructor_zeroValues_reverts() (gas: 72941)
+KromaVestingWalletTest:test_initialize_durationNotMultiple_reverts() (gas: 1609487)
+KromaVestingWalletTest:test_initialize_succeeds() (gas: 21609)
+KromaVestingWalletTest:test_release_afterFullyVested_succeeds() (gas: 79771)
+KromaVestingWalletTest:test_release_notBeneficiary_reverts() (gas: 28552)
+KromaVestingWalletTest:test_release_succeeds() (gas: 115505)
+KromaVestingWalletTest:test_release_tokenAfterFullyVested_succeeds() (gas: 95024)
+KromaVestingWalletTest:test_release_token_succeeds() (gas: 153617)
 L1BlockBedrock_Test:test_updateValues_succeeds() (gas: 65631)
 L1BlockEcotone_Test:test_setL1BlockValuesEcotone_isDepositor_succeeds() (gas: 80519)
 L1BlockEcotone_Test:test_setL1BlockValuesEcotone_notDepositor_fails() (gas: 7621)

diff --git a/packages/contracts/contracts/test/KromaVestingWallet.t.sol b/packages/contracts/contracts/test/KromaVestingWallet.t.sol
@@ -76,6 +76,18 @@ contract KromaVestingWalletTest is CommonTest {
         assertEq(vestingWallet.duration(), durationSec);
     }
 
+    function test_initialize_durationNotMultiple_reverts() external {
+        vestingWallet = KromaVestingWallet(payable(address(new Proxy(multisig))));
+        KromaVestingWallet vestingWalletImpl = new KromaVestingWallet(cliffDivider, vestingCycle);
+
+        vm.prank(multisig);
+        vm.expectRevert("Proxy: delegatecall to new implementation contract failed");
+        toProxy(address(vestingWallet)).upgradeToAndCall(
+            address(vestingWalletImpl),
+            abi.encodeCall(vestingWallet.initialize, (beneficiary, startTime, durationSec + 1))
+        );
+    }
+
     function test_release_token_succeeds() external {
         // Ensure test env is set properly
         assertEq(token.balanceOf(beneficiary), 0);

diff --git a/packages/contracts/contracts/universal/KromaVestingWallet.sol b/packages/contracts/contracts/universal/KromaVestingWallet.sol
@@ -59,6 +59,11 @@ contract KromaVestingWallet is VestingWalletUpgradeable {
         uint64 _startTimestamp,
         uint64 _durationSeconds
     ) public initializer {
+        require(
+            _durationSeconds % VESTING_CYCLE == 0,
+            "KromaVestingWallet: duration should be multiple of vesting cycle"
+        );
+
         __VestingWallet_init(_beneficiary, _startTimestamp, _durationSeconds);
     }
 
@@ -94,7 +99,7 @@ contract KromaVestingWallet is VestingWalletUpgradeable {
     ) internal view override returns (uint256) {
         if (timestamp < start()) {
             return 0;
-        } else if (timestamp > start() + duration()) {
+        } else if (timestamp >= start() + duration()) {
             return totalAllocation;
         } else {
             // At the start date, cliff amount of the assets are immediately vested.