diff --git a/claim_contracts/src/ClaimableAirdrop.sol b/claim_contracts/src/ClaimableAirdrop.sol index fa3f699fc8..d79f93cf6d 100644 --- a/claim_contracts/src/ClaimableAirdrop.sol +++ b/claim_contracts/src/ClaimableAirdrop.sol @@ -30,9 +30,9 @@ contract ClaimableAirdrop is /// @notice Merkle root of the claimants. bytes32 public claimMerkleRoot; - /// @notice Mapping of the claimants that have claimed the tokens. - /// @dev true if the claimant has claimed the tokens. - mapping(address claimer => bool claimed) public hasClaimed; + /// @notice Mapping tracking claimed leaves. + /// @dev Key is the Merkle leaf itself. + mapping(bytes32 => bool) public hasClaimed; /// @notice Event emitted when a claimant claims the tokens. /// @param to address of the claimant. @@ -78,41 +78,101 @@ contract ClaimableAirdrop is _pause(); } - /// @notice Claim the tokens. - /// @param amount amount of tokens to claim. + /// @notice Claim tokens for a single vesting stage. + /// @param amount amount of tokens to claim for this stage. + /// @param validFrom timestamp from which this stage is claimable. /// @param merkleProof Merkle proof of the claim. function claim( uint256 amount, + uint256 validFrom, bytes32[] calldata merkleProof ) external nonReentrant whenNotPaused { require( - !hasClaimed[msg.sender], - "Account has already claimed the drop" + block.timestamp <= limitTimestampToClaim, + "Drop is no longer claimable" + ); + + _verifyAndMark(amount, validFrom, merkleProof); + + bool success = IERC20(tokenProxy).transferFrom( + tokenDistributor, + msg.sender, + amount ); + require(success, "Failed to transfer funds"); + + emit TokensClaimed(msg.sender, amount); + } + + /// @notice Claim tokens for multiple vesting stages in a single transaction. + /// @param amounts array of token amounts per stage. + /// @param validFroms array of timestamps from which each stage is claimable. + /// @param merkleProofs array of Merkle proofs, one per stage. + function claimBatch( + uint256[] calldata amounts, + uint256[] calldata validFroms, + bytes32[][] calldata merkleProofs + ) external nonReentrant whenNotPaused { require( block.timestamp <= limitTimestampToClaim, "Drop is no longer claimable" ); - - bytes32 leaf = keccak256( - bytes.concat(keccak256(abi.encode(msg.sender, amount))) + uint256 length = amounts.length; + require( + length == validFroms.length && length == merkleProofs.length, + "Array length mismatch" ); - bool verifies = MerkleProof.verify(merkleProof, claimMerkleRoot, leaf); - require(verifies, "Invalid Merkle proof"); + uint256 totalClaimable = 0; - // Done before the transfer call to make sure the reentrancy bug is not possible - hasClaimed[msg.sender] = true; + for (uint256 i = 0; i < length; i++) { + _verifyAndMark(amounts[i], validFroms[i], merkleProofs[i]); + totalClaimable += amounts[i]; + } + + require(totalClaimable > 0, "Nothing to claim"); bool success = IERC20(tokenProxy).transferFrom( tokenDistributor, msg.sender, - amount + totalClaimable ); - require(success, "Failed to transfer funds"); - emit TokensClaimed(msg.sender, amount); + emit TokensClaimed(msg.sender, totalClaimable); + } + + /// @notice Verify a single-stage Merkle proof and mark the stage as claimed. + /// @dev Shared by `claim` and `claimBatch`. Does not perform the token transfer. + /// @param amount amount of tokens for this stage. + /// @param validFrom timestamp from which this stage is claimable. + /// @param merkleProof Merkle proof for this stage. + function _verifyAndMark( + uint256 amount, + uint256 validFrom, + bytes32[] calldata merkleProof + ) internal { + require( + block.timestamp >= validFrom, + "Stage not yet claimable" + ); + + bytes32 leaf = keccak256( + bytes.concat( + keccak256(abi.encode(msg.sender, amount, validFrom)) + ) + ); + + require(!hasClaimed[leaf], "Stage already claimed"); + + bool verifies = MerkleProof.verify( + merkleProof, + claimMerkleRoot, + leaf + ); + require(verifies, "Invalid Merkle proof"); + + hasClaimed[leaf] = true; } /// @notice Update the Merkle root.