audit fixes

README.md

@@ -45,7 +45,7 @@ The MigrationLocker contract is responsible for allowing users to lock their PUS
 
 **Key Features:**
 - Token locking mechanism with unique identifier generation
-- Safety toggles to prevent/allow locking - Owner Controlled
+- Openzeppelin's Pausable library to prevent/allow locking - Owner Controlled
 - Proper access control with Ownable2Step pattern
 - Token burning capability for migrated tokens
 - Emergency fund recovery functionality

script/utils/fetchAndStoreEvents.js

@@ -7,10 +7,16 @@ async function main() {
   const CONTRACT_ADDRESS = LOCKER_CONFIG.CONTRACT_ADDRESS;
   const LOCKER_ABI = LOCKER_CONFIG.ABI;
   const FILTER_EPOCHS = LOCKER_CONFIG.FILTER_EPOCHS;
+  const PUSH_TOKEN_ADDRESS = "0xf418588522d5dd018b425E472991E52EBBeEEEEE";
 
   const provider = ethers.provider;
   const locker = new ethers.Contract(CONTRACT_ADDRESS, LOCKER_ABI, provider);
 
+  // Add PUSH token interface for balance validation
+  const pushToken = new ethers.Contract(PUSH_TOKEN_ADDRESS, [
+    "function balanceOf(address account) view returns (uint256)"
+  ], provider);
+
   // Get current epoch from contract
   const currentEpoch = await locker.epoch();
   console.log(`🔢 Current epoch: ${currentEpoch}`);
@@ -48,6 +54,7 @@ async function main() {
   // Group events by address and combine amounts
   const addressAmounts = {};
   let totalEvents = 0;
+  const epochTotals = {}; // Track total amount per epoch from events
 
   // Process events and filter by epoch on the client side
   for (const event of allEvents) {
@@ -75,6 +82,13 @@ async function main() {
         epoch: eventEpoch.toString()
       };
     }
+
+    // Track total per epoch
+    if (!epochTotals[eventEpoch]) {
+      epochTotals[eventEpoch] = BigInt(0);
+    }
+    epochTotals[eventEpoch] += amount;
+
     totalEvents++;
   }
 
@@ -93,6 +107,48 @@ async function main() {
     console.log(`🔄 Combined ${duplicateCount} duplicate addresses`);
   }
 
+  // Validate that sum of all leaves equals on-chain locked amounts
+  console.log(`\n🔍 Validating totals against on-chain state...`);
+
+  for (const epoch of epochsToProcess) {
+    const offChainTotal = epochTotals[epoch] || BigInt(0);
+
+    // Get on-chain total (incremental balance for this epoch)
+    let onChainTotal;
+    if (epoch === currentEpoch) {
+      // For current epoch: current balance minus balance at start of current epoch
+      const currentBalance = await pushToken.balanceOf(CONTRACT_ADDRESS);
+      const currentEpochStart = await locker.epochStartBlock(currentEpoch);
+      const balanceAtStart = await pushToken.balanceOf(CONTRACT_ADDRESS, {
+        blockTag: Number(currentEpochStart) - 1
+      });
+      onChainTotal = currentBalance - balanceAtStart;
+    } else {
+      // For past epochs: balance at end of epoch minus balance at start of epoch
+      const nextEpochStart = await locker.epochStartBlock(epoch + 1);
+      const epochStart = await locker.epochStartBlock(epoch);
+
+      const endBlockOfEpoch = Number(nextEpochStart) - 1;
+      const startBlockOfEpoch = Number(epochStart) - 1;
+
+      const endBalance = await pushToken.balanceOf(CONTRACT_ADDRESS, { blockTag: endBlockOfEpoch });
+      const startBalance = await pushToken.balanceOf(CONTRACT_ADDRESS, { blockTag: startBlockOfEpoch });
+      onChainTotal = endBalance - startBalance;
+    }
+
+    if (offChainTotal !== onChainTotal) {
+      console.error(`❌ Validation failed for epoch ${epoch}:`);
+      console.error(`   Off-chain total: ${offChainTotal.toString()}`);
+      console.error(`   On-chain total: ${onChainTotal.toString()}`);
+      console.error(`   Difference: ${(offChainTotal > onChainTotal ? offChainTotal - onChainTotal : onChainTotal - offChainTotal).toString()}`);
+      throw new Error(`Funds may be missing for epoch ${epoch}`);
+    }
+
+    console.log(`✅ Epoch ${epoch}: ${offChainTotal.toString()} wei`);
+  }
+
+  console.log(`\n✅ All validations passed!`);
+
   const outputPath = path.join(__dirname, OUTPUT_CONFIG.CLAIMS_PATH);
   fs.mkdirSync(path.dirname(outputPath), { recursive: true });
   fs.writeFileSync(outputPath, JSON.stringify(claims, null, 2));

src/MigrationLocker.sol

@@ -4,12 +4,16 @@ pragma solidity 0.8.29;
 import { PausableUpgradeable } from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
 import { Ownable2StepUpgradeable } from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import { Initializable } from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { IPUSH } from "./interfaces/IPush.sol";
 
 /// @title MigrationLocker
 /// @author Push Chain
 /// @notice Allows users to lock their Push tokens for migration
 contract MigrationLocker is Initializable, Ownable2StepUpgradeable, PausableUpgradeable {
+    using SafeERC20 for IERC20;
+
     /// @notice Indicates the current epoch
     /// @dev    Each specific epoch represents a particular block of time under which all Locked events will be
     ///         recorded to create the merkle tree all user deposits done in that specific epoch.
@@ -33,6 +37,9 @@ contract MigrationLocker is Initializable, Ownable2StepUpgradeable, PausableUpgr
     /// @param epoch The epoch number
     event Locked(address caller, address recipient, uint256 amount, uint256 epoch);
 
+    /// @notice Emitted when a admin initiates a new epoch
+    event NewEpoch(uint256 epoch, uint256 startBlock);
+
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
@@ -49,9 +56,13 @@ contract MigrationLocker is Initializable, Ownable2StepUpgradeable, PausableUpgr
         initiateNewEpoch();
     }
 
+    /// @notice Allows the owner to initiate a new epoch
+    /// @dev The function increments the epoch number and sets the start block for the new epoch
+    /// @dev Emits a NewEpoch event with the epoch number and start block
     function initiateNewEpoch() public onlyOwner {
         epoch++;
         epochStartBlock[epoch] = block.number;
+        emit NewEpoch(epoch, block.number);
     }
 
     /// Pauseable Features
@@ -68,7 +79,7 @@ contract MigrationLocker is Initializable, Ownable2StepUpgradeable, PausableUpgr
     /// @param _recipient The address of the recipient
     /// @dev The recipient address cannot be zero
     /// @dev The function transfers the specified amount of tokens from the user to the contract
-    /// @dev Emits a Locked event with the recipient address, amount, and a unique identifier
+    /// @dev Emits a Locked event with the recipient address, amount, and epoch
     function lock(uint256 _amount, address _recipient) external whenNotPaused {
         uint256 codeLength;
         assembly {
@@ -78,7 +89,44 @@ contract MigrationLocker is Initializable, Ownable2StepUpgradeable, PausableUpgr
             revert("Invalid recipient");
         }
 
-        IPUSH(PUSH_TOKEN).transferFrom(msg.sender, address(this), _amount);
+        IERC20(PUSH_TOKEN).safeTransferFrom(msg.sender, address(this), _amount);
+        emit Locked(msg.sender, _recipient, _amount, epoch);
+    }
+
+    /// @notice Allows users to lock their tokens for migration using permit signature
+    /// @param _amount The amount of tokens to lock
+    /// @param _recipient The address of the recipient
+    /// @param _deadline The deadline for the permit signature
+    /// @param _v The v component of the permit signature
+    /// @param _r The r component of the permit signature
+    /// @param _s The s component of the permit signature
+    /// @dev The recipient address cannot be zero
+    /// @dev The function uses permit to approve tokens and then transfers them to the contract
+    /// @dev Emits a Locked event with the recipient address, amount, and epoch
+    function lockWithPermit(
+        uint256 _amount,
+        address _recipient,
+        uint256 _deadline,
+        uint8 _v,
+        bytes32 _r,
+        bytes32 _s
+    )
+        external
+        whenNotPaused
+    {
+        uint256 codeLength;
+        assembly {
+            codeLength := extcodesize(_recipient)
+        }
+        if (_recipient == address(0) || codeLength > 0) {
+            revert("Invalid recipient");
+        }
+
+        // Use permit to approve tokens for this contract
+        IPUSH(PUSH_TOKEN).permit(msg.sender, address(this), _amount, _deadline, _v, _r, _s);
+
+        // Transfer the approved tokens to this contract
+        IERC20(PUSH_TOKEN).safeTransferFrom(msg.sender, address(this), _amount);
         emit Locked(msg.sender, _recipient, _amount, epoch);
     }
 
@@ -93,7 +141,7 @@ contract MigrationLocker is Initializable, Ownable2StepUpgradeable, PausableUpgr
     function recoverFunds(address _token, address _to, uint256 _amount) external onlyOwner whenNotPaused {
         require(_to != address(0), "Invalid recipient");
 
-        require(_amount > 0 && _amount <= IPUSH(_token).balanceOf(address(this)), "Invalid amount");
-        IPUSH(_token).transfer(_to, _amount);
+        require(_amount > 0 && _amount <= IERC20(_token).balanceOf(address(this)), "Invalid amount");
+        IERC20(_token).safeTransfer(_to, _amount);
     }
 }

src/MigrationRelease.sol

@@ -27,7 +27,6 @@ contract MigrationRelease is Initializable, Ownable2StepUpgradeable, PausableUpg
     uint256 public constant VESTING_RATIO = 75;
 
     uint256 public totalReleased;
-    bool public isClaimPaused;
 
     mapping(bytes32 => uint256) public instantClaimTime;
 
@@ -83,7 +82,7 @@ contract MigrationRelease is Initializable, Ownable2StepUpgradeable, PausableUpg
     /// @dev calculates the instant amount based on the INSTANT_RATIO
     /// @dev updates the instantClaimTime mapping and totalReleased variable
     /// @dev transfers the instant amount to the recipient, reverting if the transfer fails
-    /// @dev emits a ReleasedInstant event with the recipient address, amount, and release time
+    /// @dev emits a ReleasedInstant event with the recipient address, amount, and epoch
 
     function releaseInstant(
         address _recipient,
@@ -95,11 +94,8 @@ contract MigrationRelease is Initializable, Ownable2StepUpgradeable, PausableUpg
         whenNotPaused
     {
         bytes32 leaf = keccak256(abi.encodePacked(_recipient, _amount, _epoch));
-        require(
-            verifyAddress(_recipient, _amount, _epoch, _merkleProof) && instantClaimTime[leaf] == 0,
-            "Not Whitelisted or already Claimed"
-        );
-        uint256 instantAmount = (_amount * INSTANT_RATIO) / 10; //Instantly relaese 7.5 times the amount
+        require(verifyAddress(leaf, _merkleProof) && instantClaimTime[leaf] == 0, "Not Whitelisted or already Claimed");
+        uint256 instantAmount = (_amount * INSTANT_RATIO) / 10; //Instantly release 7.5 times the amount
 
         instantClaimTime[leaf] = block.timestamp;
         totalReleased += instantAmount;
@@ -117,11 +113,11 @@ contract MigrationRelease is Initializable, Ownable2StepUpgradeable, PausableUpg
     /// @dev calculates the vested amount based on the VESTING_RATIO
     /// @dev updates the claimedvested mapping and totalReleased variable
     /// @dev transfers the vested amount to the recipient, reverting if the transfer fails
-    /// @dev emits a ReleasedVested event with the recipient address, amount, and release time
+    /// @dev emits a ReleasedVested event with the recipient address, amount, and epoch
 
     function releaseVested(address _recipient, uint256 _amount, uint256 _epoch) external whenNotPaused {
         bytes32 leaf = keccak256(abi.encodePacked(_recipient, _amount, _epoch));
-        if (claimedvested[leaf] == true) {
+        if (claimedvested[leaf]) {
             revert("Already Claimed");
         }
 
@@ -144,26 +140,15 @@ contract MigrationRelease is Initializable, Ownable2StepUpgradeable, PausableUpg
         require(res, "Transfer failed");
     }
 
-    function verifyAddress(
-        address recipient,
-        uint256 amount,
-        uint256 _epoch,
-        bytes32[] calldata _merkleProof
-    )
-        private
-        view
-        returns (bool)
-    {
-        bytes32 leaf = keccak256(abi.encodePacked(recipient, amount, _epoch));
-        return MerkleProof.verify(_merkleProof, merkleRoot, leaf);
+    function verifyAddress(bytes32 _leaf, bytes32[] calldata _merkleProof) private view returns (bool) {
+        return MerkleProof.verify(_merkleProof, merkleRoot, _leaf);
     }
 
     function recoverFunds(address _token, address _to, uint256 _amount) external onlyOwner whenNotPaused {
         require(_to != address(0), "Invalid recipient");
 
         if (_token == address(0)) {
             transferFunds(_to, _amount);
-            return;
         } else {
             require(_amount > 0 && _amount <= IERC20(_token).balanceOf(address(this)), "Invalid amount");
             IERC20(_token).safeTransfer(_to, _amount);

test/MigrationRelease.t.sol

@@ -149,7 +149,6 @@ contract MigrationReleaseTest is Test {
     //////////////////////////////////////////////////////////////*/
 
     function testInitialization() public {
-        assertEq(release.isClaimPaused(), false);
         assertEq(release.owner(), owner);
         assertEq(release.merkleRoot(), merkleRoot);
         assertEq(address(release).balance, 10_000 ether);