diff --git a/contracts/airdrop/MerkleDistributor.sol b/contracts/airdrop/MerkleDistributor.sol index 47815b3b..baff5d9d 100644 --- a/contracts/airdrop/MerkleDistributor.sol +++ b/contracts/airdrop/MerkleDistributor.sol @@ -6,202 +6,230 @@ import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; import "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol"; import "@openzeppelin/contracts/access/Ownable.sol"; + /** * @title MerkleDistributor - * @notice Handles token airdrops from an unlimited amount of token rewards - * @dev Based on https://github.com/Uniswap/merkle-distributor but modified to handle multiple airdrops concurrently + * @notice Handles token airdrops with version-controlled distributions + * @dev Allows for multiple versiond of a token ditribution and the ability to pause and withdraw unclaimed tokens */ contract MerkleDistributor is Ownable { + using SafeERC20 for IERC20; struct Distribution { + uint256 version; address token; bool isPaused; - uint256 expiryTimestamp; bytes32 merkleRoot; uint256 totalAmount; - uint256 version; - mapping(address => uint256) claimed; - mapping(address => uint256) lastClaimedVersion; + bool isWithdrawn; + } + + struct Claim { + uint256 claimedAmount; + uint256 versionClaimed; } + address[] public tokens; - mapping(address => Distribution) public distributions; - event Claimed(address indexed token, uint256 index, address indexed account, uint256 amount); - event DistributionAdded(uint256 indexed tokenIndex, address indexed token, uint256 totalAmount, uint256 expiryTimestamp); - event DistributionUpdated(address indexed token, uint256 additionalAmount, uint256 expiryTimestamp); - event SetExpiryTimestamp(address indexed token, uint256 expiryTimestamp); + mapping(address => mapping(uint256 => Distribution)) public distributions; + mapping(address => mapping(address => mapping(uint256 => Claim))) public claims; + mapping(address => uint256) public latestVersion; + + event Claimed(address indexed token, uint256 version, address indexed account, uint256 amount); + event DistributionAdded(address indexed token, uint256 version, uint256 totalAmount); + event DistributionUpdated(address indexed token, uint256 version, uint256 additionalAmount); + event DistributionWithdrawn(address indexed token, uint256 version, uint256 amount); + /** + * @notice modifier to check if a distribution exists + * @param _token token address + **/ modifier distributionExists(address _token) { - require(distributions[_token].token != address(0), "MerkleDistributor: Distribution does not exist."); + // check if distribution exists with latest version of it + require(latestVersion[_token] > 0, "MerkleDistributor: No distributions exist for this token."); + require(distributions[_token][latestVersion[_token]].token != address(0), "MerkleDistributor: No distributions exist for this token."); _; } /** - * @notice returns the total amount that an account has claimed from a distribution + * @notice returns the total claimed amount for a token + * @param _token token address + * @param _account account to check + **/ + function getTotalClaimed(address _token, address _account) public view returns (uint256) { + uint256 _latestVersion = latestVersion[_token]; + uint256 totalClaimed; + for (uint256 i = 1; i <= _latestVersion; i++) { + totalClaimed += claims[_token][_account][i].claimedAmount; + } + return totalClaimed; + } + + /** + * @notice returns the distribution details for a specific version of a token distribution * @param _token token address - * @param _account address of the account to return claimed amount for + * @param _version version of the distribution **/ - function getClaimed(address _token, address _account) external view distributionExists(_token) returns (uint256) { - return distributions[_token].claimed[_account]; + function getDistribution(address _token, uint256 _version) public view returns (Distribution memory) { + return distributions[_token][_version]; } /** - * @notice add multiple token distributions - * @param _tokens the list of token addresses to add - * @param _merkleRoots list of merkle roots for each distribution - * @param _totalAmounts list of total distribution amounts for each token - * @param _expiryTimestamps list of expiry timestamps for each distribution + * @notice adds multiple distributions at once + * @param _tokens token addresses + * @param _merkleRoots merkle roots of the distributions + * @param _totalAmounts total amounts of tokens to be distributed **/ function addDistributions( - address[] calldata _tokens, - bytes32[] calldata _merkleRoots, - uint256[] calldata _totalAmounts, - uint256[] calldata _expiryTimestamps + address[] memory _tokens, + bytes32[] memory _merkleRoots, + uint256[] memory _totalAmounts ) external onlyOwner { - require( - _tokens.length == _merkleRoots.length && _tokens.length == _totalAmounts.length, - "MerkleDistributor: Array lengths need to match." - ); - + require(_tokens.length == _merkleRoots.length, "MerkleDistributor: Invalid input length."); + require(_tokens.length == _totalAmounts.length, "MerkleDistributor: Invalid input length."); for (uint256 i = 0; i < _tokens.length; i++) { - addDistribution(_tokens[i], _merkleRoots[i], _totalAmounts[i], _expiryTimestamps[i]); + addDistribution(_tokens[i], _merkleRoots[i], _totalAmounts[i]); } } - + /** - * @notice add a token distribution - * @param _token token address - * @param _merkleRoot merkle root for token distribution - * @param _totalAmount total distribution amount - * @param _expiryTimestamp timestamp when unclaimed tokens can be withdrawn + * @notice updates multiple distributions at once + * @param _tokens token addresses + * @param _merkleRoots merkle roots of the distributions + * @param _additionalAmounts additional amounts of tokens to be distributed **/ - function addDistribution( - address _token, - bytes32 _merkleRoot, - uint256 _totalAmount, - uint256 _expiryTimestamp - ) public onlyOwner { - require(distributions[_token].token == address(0), "MerkleDistributor: Distribution is already added."); - require(IERC20(_token).balanceOf(address(this)) >= _totalAmount, "MerkleDistributor: Insufficient balance."); - - tokens.push(_token); - distributions[_token].token = _token; - distributions[_token].merkleRoot = _merkleRoot; - distributions[_token].totalAmount = _totalAmount; - distributions[_token].expiryTimestamp = _expiryTimestamp; - distributions[_token].version++; - - emit DistributionAdded(tokens.length - 1, _token, _totalAmount, _expiryTimestamp); + function updateDistributions( + address[] memory _tokens, + bytes32[] memory _merkleRoots, + uint256[] memory _additionalAmounts + ) external onlyOwner { + require(_tokens.length == _merkleRoots.length, "MerkleDistributor: Invalid input length."); + require(_tokens.length == _additionalAmounts.length, "MerkleDistributor: Invalid input length."); + for (uint256 i = 0; i < _tokens.length; i++) { + updateDistribution(_tokens[i], _merkleRoots[i], _additionalAmounts[i]); + } } /** - * @notice update multiple token distributions - * @param _tokens the list of token addresses to update - * @param _merkleRoots list of updated merkle roots for the distributions - * @param _additionalAmounts list of total additional distribution amounts for each token - * @param _expiryTimestamps list of updated expiry timestamps for each distribution + * @notice adds multiple distribution versions at once + * @param _tokens token addresses + * @param _merkleRoots merkle roots of the distributions + * @param _totalAmounts total amounts of tokens to be distributed **/ - function updateDistributions( - address[] calldata _tokens, - bytes32[] calldata _merkleRoots, - uint256[] calldata _additionalAmounts, - uint256[] calldata _expiryTimestamps + function addDistributionVersions( + address[] memory _tokens, + bytes32[] memory _merkleRoots, + uint256[] memory _totalAmounts ) external onlyOwner { - require( - _tokens.length == _merkleRoots.length && _tokens.length == _additionalAmounts.length, - "MerkleDistributor: Array lengths need to match." - ); - + require(_tokens.length == _merkleRoots.length, "MerkleDistributor: Invalid input length."); + require(_tokens.length == _totalAmounts.length, "MerkleDistributor: Invalid input length."); for (uint256 i = 0; i < _tokens.length; i++) { - updateDistribution(_tokens[i], _merkleRoots[i], _additionalAmounts[i], _expiryTimestamps[i]); + addDistributionVersion(_tokens[i], _merkleRoots[i], _totalAmounts[i]); } } - + /** - * @notice update a token distribution - * @dev merkle root should be updated to reflect additional amount - the amount for each - * account should be incremented by any additional allocation and any new accounts should be added - * to the tree + * @notice adds a new distribution * @param _token token address - * @param _merkleRoot updated merkle root for token distribution - * @param _additionalAmount total additional distribution amount - * @param _expiryTimestamp timestamp when unclaimed tokens can be withdrawn + * @param _merkleRoot merkle root of the distribution + * @param _totalAmount total amount of tokens to be distributed **/ - function updateDistribution( + function addDistribution( address _token, bytes32 _merkleRoot, - uint256 _additionalAmount, - uint256 _expiryTimestamp - ) public onlyOwner distributionExists(_token) { - require(!distributions[_token].isPaused, "MerkleDistributor: Distribution is paused."); - require(IERC20(_token).balanceOf(address(this)) >= _additionalAmount, "MerkleDistributor: Insufficient balance."); - require(_expiryTimestamp >= distributions[_token].expiryTimestamp, "MerkleDistributor: Invalid expiry timestamp."); - - distributions[_token].merkleRoot = _merkleRoot; - distributions[_token].totalAmount += _additionalAmount; - distributions[_token].expiryTimestamp = _expiryTimestamp; - distributions[_token].version++; + uint256 _totalAmount + ) public onlyOwner { + require(distributions[_token][1].token == address(0), "MerkleDistributor: Distribution already exists."); + require(IERC20(_token).balanceOf(address(this)) >= _totalAmount, "MerkleDistributor: Insufficient balance."); - emit DistributionUpdated(_token, _additionalAmount, _expiryTimestamp); - } + uint256 _version = 1; - /** - * @notice claim multiple token distributions - * @param _tokens list of token address - * @param _indexes list of indexes of the claims within the distributions - * @param _account address of the account to claim for - * @param _amounts list of lifetime amounts of the tokens allocated to account - * @param _merkleProofs list of merkle proofs for the token claims - **/ - function claimDistributions( - address[] calldata _tokens, - uint256[] calldata _indexes, - address _account, - uint256[] calldata _amounts, - bytes32[][] calldata _merkleProofs - ) external { - require( - _tokens.length == _indexes.length && _tokens.length == _amounts.length && _tokens.length == _merkleProofs.length, - "MerkleDistributor: Array lengths need to match." - ); + Distribution storage distribution = distributions[_token][_version]; + distribution.token = _token; + distribution.version = _version; + distribution.merkleRoot = _merkleRoot; + distribution.totalAmount = _totalAmount; + latestVersion[_token] = _version; - for (uint256 i = 0; i < _tokens.length; i++) { - claimDistribution(_tokens[i], _indexes[i], _account, _amounts[i], _merkleProofs[i]); - } + emit DistributionAdded(_token, _version, _totalAmount); } /** - * @notice claim a token distribution - * @param _token token address - * @param _index index of the claim within the distribution - * @param _account address of the account to claim for - * @param _amount lifetime amount of the token allocated to account - * @param _merkleProof the merkle proof for the token claim - **/ + * @notice Claims tokens from the latest distribution version for a given token + * @param _token token address + * @param _index index of the claim + * @param _account account to claim tokens for + * @param _amount amount of tokens to claim + * @param _merkleProof merkle proof + **/ function claimDistribution( address _token, uint256 _index, address _account, uint256 _amount, bytes32[] calldata _merkleProof - ) public distributionExists(_token) { - require(!distributions[_token].isPaused, "MerkleDistributor: Distribution is paused."); - require(_amount > 0, "MerkleDistributor: No claimable tokens."); - Distribution storage distribution = distributions[_token]; - + ) external { + uint256 _latestVersion = latestVersion[_token]; + require(_latestVersion > 0, "MerkleDistributor: No distributions exist for this token."); + Distribution storage distribution = distributions[_token][_latestVersion]; + require(!distribution.isPaused, "MerkleDistributor: Distribution is paused."); + Claim storage claim = claims[_token][_account][_latestVersion]; bytes32 node = keccak256(abi.encodePacked(_index, _account, _amount)); require(MerkleProof.verify(_merkleProof, distribution.merkleRoot, node), "MerkleDistributor: Invalid proof."); - require( - distribution.lastClaimedVersion[_account] < distribution.version, - "MerkleDistributor: Already claimed for current version." - ); + require(claim.claimedAmount < _amount, "MerkleDistributor: Tokens claimed for the latest version."); + uint256 claimableAmount = _amount - claim.claimedAmount; + claim.claimedAmount = _amount; + claim.versionClaimed = _latestVersion; + IERC20(_token).safeTransfer(_account, claimableAmount); + emit Claimed(_token, _latestVersion, _account, claimableAmount); + } + + + /* + * @notice updates an existing distribution on its current version + * @param _token token address + * @param _merkleRoot merkle root of the distribution + * @param _additionalAmount additional amount of tokens to be distributed + **/ + function updateDistribution( + address _token, + bytes32 _merkleRoot, + uint256 _additionalAmount + ) public onlyOwner distributionExists(_token) { + uint256 _latestVersion = latestVersion[_token]; + require(IERC20(_token).balanceOf(address(this)) >= _additionalAmount, "MerkleDistributor: Insufficient balance."); + Distribution storage distribution = distributions[_token][_latestVersion]; + distribution.totalAmount += _additionalAmount; + distribution.merkleRoot = _merkleRoot; + + emit DistributionUpdated(_token, _latestVersion, _additionalAmount); + } + + /** + * @notice Adds a new version of a token distribution automatically incrementing the version + * @param _token token address + * @param _merkleRoot merkle root of the distribution + * @param _totalAmount total amount of tokens to be distributed + **/ + function addDistributionVersion( + address _token, + bytes32 _merkleRoot, + uint256 _totalAmount + ) public onlyOwner distributionExists(_token) { + require(distributions[_token][latestVersion[_token]].isWithdrawn, "MerkleDistributor: Latest version is not withdrawn."); + uint256 _version = latestVersion[_token] + 1; + require(IERC20(_token).balanceOf(address(this)) >= _totalAmount, "MerkleDistributor: Insufficient balance."); - distribution.claimed[_account] += _amount; - distribution.lastClaimedVersion[_account] = distribution.version; - IERC20(_token).safeTransfer(_account, _amount); + Distribution storage distribution = distributions[_token][_version]; + distribution.token = _token; + distribution.version = _version; + distribution.merkleRoot = _merkleRoot; + distribution.totalAmount = _totalAmount; - emit Claimed(_token, _index, _account, _amount); + latestVersion[_token] = _version; + + emit DistributionAdded(_token, _version, _totalAmount); } /** @@ -209,51 +237,52 @@ contract MerkleDistributor is Ownable { * @dev merkle root should be updated to reflect current state of claims - the amount for each * account should be equal to it's claimed amount * @param _token token address - * @param _merkleRoot updated merkle root - * @param _totalAmount updated total amount **/ function withdrawUnclaimedTokens( - address _token, - bytes32 _merkleRoot, - uint256 _totalAmount + address _token ) external onlyOwner distributionExists(_token) { - require(distributions[_token].isPaused, "MerkleDistributor: Distribution is not paused."); + uint256 _version = latestVersion[_token]; + + require(!distributions[_token][_version].isWithdrawn, "MerkleDistributor: Distribution is already withdrawn."); + require(distributions[_token][_version].isPaused, "MerkleDistributor: Distribution is not paused."); IERC20 token = IERC20(_token); uint256 balance = token.balanceOf(address(this)); if (balance > 0) { token.safeTransfer(msg.sender, balance); } + distributions[_token][_version].isWithdrawn = true; - distributions[_token].merkleRoot = _merkleRoot; - distributions[_token].totalAmount = _totalAmount; - distributions[_token].isPaused = false; + emit DistributionWithdrawn(_token, _version, balance); } + /** * @notice pauses a token distribution for withdrawal of unclaimed tokens * @dev must be called before withdrawUnlclaimedTokens to ensure state doesn't change * while the new merkle root is calculated * @param _token token address **/ - function pauseForWithdrawal(address _token) external onlyOwner distributionExists(_token) { - require( - distributions[_token].expiryTimestamp <= block.timestamp, - "MerkleDistributor: Expiry timestamp not reached." - ); - require(!distributions[_token].isPaused, "MerkleDistributor: Already paused."); - distributions[_token].isPaused = true; + function pauseForWithdrawal(address _token) external onlyOwner distributionExists(_token) { + uint256 _version = latestVersion[_token]; + + require(!distributions[_token][_version].isPaused, "MerkleDistributor: Distribution is already paused."); + + distributions[_token][_version].isPaused = true; } /** - * @notice sets the timestamp when unclaimed tokens can be withdrawn for a token distribution + * @notice unpauses a token distribution * @param _token token address - * @param _expiryTimestamp expiry timestamp **/ - function setExpiryTimestamp(address _token, uint256 _expiryTimestamp) external onlyOwner distributionExists(_token) { - require(!distributions[_token].isPaused, "MerkleDistributor: Distribution is paused."); - require(_expiryTimestamp >= distributions[_token].expiryTimestamp, "MerkleDistributor: Invalid expiry timestamp."); - distributions[_token].expiryTimestamp = _expiryTimestamp; - emit SetExpiryTimestamp(_token, _expiryTimestamp); + function unpause(address _token) external onlyOwner distributionExists(_token) { + uint256 _version = latestVersion[_token]; + + require(distributions[_token][_version].isPaused, "MerkleDistributor: Distribution is not paused."); + + distributions[_token][_version].isPaused = false; } + + } + diff --git a/test/foundry/Base.t.sol b/test/foundry/Base.t.sol index 50bd3812..8c6cd036 100644 --- a/test/foundry/Base.t.sol +++ b/test/foundry/Base.t.sol @@ -10,8 +10,10 @@ abstract contract BaseTest is Test, Utils { Users internal users; uint256 internal network; MerkleDistributor public merkleDistributor; + address owner; function init(bool _fork) public { + owner = address(this); if (_fork) { network = vm.createSelectFork(vm.rpcUrl("ethereum")); } else { diff --git a/test/foundry/unit/MerkleDistributor.t.sol b/test/foundry/unit/MerkleDistributor.t.sol index 4a9e15e6..7a4ad3ad 100644 --- a/test/foundry/unit/MerkleDistributor.t.sol +++ b/test/foundry/unit/MerkleDistributor.t.sol @@ -4,19 +4,353 @@ pragma solidity ^0.8.15; import {MerkleDistributor} from "../../../contracts/airdrop/MerkleDistributor.sol"; import {BaseTest} from "../Base.t.sol"; import {ERC677} from "../../../contracts/core/tokens/base/ERC677.sol"; +import "forge-std/console.sol"; + contract MerkleDistributorTest is BaseTest { bool internal _fork = false; + bytes32 _merkleRoot_1 = 0xca99ea02947aea2b4e36e85b6f48ee9bda45ad8a7c13460a30642b131557434a; + uint256 _totalAmount = 14337996000000007761321915; + + address account1 = 0x0000000000002D534FF79e9C69e7Fcc742f0BE83; + uint256 index1 = 0; + uint256 amount1 = 14813531288235778; + bytes32[] validProof1 = [ + bytes32(0xf7b20929116e3c7302692e96bafa244dafa5caceb5577d906648fd399a5ff7bf), + bytes32(0xf41c8fbc357d7ef386de1dfacf7e0021a4c3761bc036ea971f9686a58455fad7), + bytes32(0x1608c0ceba9ec48bf1d1d247d9bec0b23df765759576507814f73fcf74071878), + bytes32(0x5900dd8a9a5b8c1bc6b3fe5427e66ad095bfc8c3f91edf5060921732c297490f), + bytes32(0x2c1e0b76aef43bebef5252e0ebe2d0b748a4dc1a1ca5a1dd222a7636f9cd6d5d), + bytes32(0x4bd1864308465547ac58aae2c6a5a2280440ff685fd646e94038edae57064523), + bytes32(0x2b66dec807f531149f860a0b08643355a7b478140383defed8f4e7017e87aaf2), + bytes32(0xb397cde1180e15d034bd2f7199b204d20db212962fffd5e6a2c7a9ce0c081efc), + bytes32(0xb811da3159e2fd9604a749951af9bf9949341775787f5ede1a933191531a78c7), + bytes32(0x973eca14272bc83b99355e1f7902f15e4dbd97711b788358afd0900b7bde9b3b) + ]; + + address account2 = 0x0000000000007F150Bd6f54c40A34d7C3d5e9f56; + uint256 index2 = 1; + uint256 amount2 = 23035; + bytes32[] validProof2 = [ + bytes32(0x112589f019ac15c6ec94f62fd7d4328ccac3a463da609d32fb582721d817fef8), + bytes32(0x01ffa77fc27fc339e7026a6d2c01227f8e6d066cde2f01d9611f3a589dcbefe7), + bytes32(0x9092c60f0cb2295aa969cbf3a9d33ce8a7ace9f07e3913f133a1c6bb5299cc1c), + bytes32(0x9a11c9d8787464f3fc4787d69205f0229872b2e77f82984c9a33e4e512f1ba58), + bytes32(0xb5e5534c0c35caf0aabaa93d1e5741381b15894aa49174e0e2a80474c1547dce), + bytes32(0x30e1ac74c2856c807e6b05508da347450d44b2a6ecd71b7d1763c41482352f2f), + bytes32(0x3ec7e4400e750832247c5d7d4bd320610aef76c91d9114c5aa0c93183219aef2), + bytes32(0x53a34af34e8826d4d5c24b75a026bbb5bac5a0ac7eea7e1df28e683397f95691), + bytes32(0x123ccedfe2e059b503b08f2cbfc76156cc27bfc0f609b39ec53455bb99cc1629), + bytes32(0x411a1a030a9b4c0462c4b1cb9b10506a7488eea01db61f9d49055aa501f1d3a7), + bytes32(0xcdb1fad7f35ac9b56ea7d04044d9c2e7e47551fa1e3b2bac80472f322c5b89b9), + bytes32(0x4ad3fc2220a0a5bb66a43a5dfa9e04b0df7f357d28944d40f88ddc1e3084bdfa) + ]; + + + + + bytes32 _merkleRoot_2 = 0x52ec042bbcf8a478270012ae0e4a658ff3921314df95f0487b1f05b6a46c3af8; + uint256 _totalAmount_2 = 1065156117468259689325813; + + address account1_2 = 0x00393D62F17B07e64f7cdcDF9BdC2fd925b20Bba; + uint256 index1_2 = 0; + uint256 amount1_2 = 1840233889215604467618; + bytes32[] validProof1_2 = [ + bytes32(0x8f69123df82b69d232b7ad69c5d95b77ab0f8d4525d3544bbae543c629f32408), + bytes32(0x24c37992ee0700c6c1d640cd5233f3949dca5a543f1b5fa6ba776b3e669c629a), + bytes32(0xf78ab8f97474318c199ac417bac751a62c1e1cd7b03613782b5a66ba9225e379), + bytes32(0x539821bea843ee787e20860247831f7824608ac9ba72afeb3fcadad36640c58c), + bytes32(0x96f883d35c39bb71c6a30c68159cc151bb5c7528d9bd8010797736126d7b46ef), + bytes32(0x3db7b50759069a96c99303b5a5cdbdab240a83f1c98372efbe09cd9fdc6e02f6), + bytes32(0xcfb78726e966fe4d350da6ec55a084f45dd8bb3cc9ce257e2cc62a206bc7037d), + bytes32(0x44c8dfb10fac55138806452ecdf93237c0e3ee7213d40780e4d1d8e3c678fde8), + bytes32(0x357c7629bf032666620838427aba400dff312480f88e276b5f0c77cd2aa67e1c) + ]; + + + + uint256 tokenSupply = 1000000000; + uint256 _version1 = 1; + uint256 _version2 = 2; + uint256 _version3 = 3; + + ERC677 _testToken = new ERC677("Token", "TKN", tokenSupply); + ERC677 _testToken2 = new ERC677("Coin", "COI", tokenSupply); + + + function setUp() public { BaseTest.init(_fork); + owner = address(this); + } + + // Setup the MerkleDistributor contract + function _setupTokenDistribution(bytes32 _merkleRoot) internal { + _testToken.transfer(address(merkleDistributor), _totalAmount); + merkleDistributor.addDistribution(address(_testToken), _merkleRoot, _totalAmount); } - function test_claimDistribution_EmptyProof() public { - ERC677 _testToken = new ERC677("Token", "TKN", 1000000); - merkleDistributor.addDistribution(address(_testToken), bytes32(""), 0, 0); - bytes32[] memory _proof = new bytes32[](0); + + function test_Success_addDistribution() public { + _testToken.transfer(address(merkleDistributor), _totalAmount); + merkleDistributor.addDistribution(address(_testToken), _merkleRoot_1, _totalAmount); + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).version, + _version1 + ); + _testToken2.transfer(address(merkleDistributor), _totalAmount); + merkleDistributor.addDistribution(address(_testToken2), _merkleRoot_2, _totalAmount); + assertEq( + merkleDistributor.getDistribution(address(_testToken2), _version1).version, + _version1 + ); + } + + function test_Revert_addDistribution_NotOwner() public { + _testToken.transfer(address(merkleDistributor), _totalAmount); + vm.startPrank(account1); + vm.expectRevert("Ownable: caller is not the owner"); + merkleDistributor.addDistribution(address(_testToken), _merkleRoot_1, _totalAmount); + } + + function test_Revert_addDistribution_InsufficientBalance() public { + vm.expectRevert("MerkleDistributor: Insufficient balance."); + merkleDistributor.addDistribution(address(_testToken), _merkleRoot_1, _totalAmount); + } + + function test_Revert_addDistribution_DistributionExists() public { + _testToken.transfer(address(merkleDistributor), _totalAmount); + merkleDistributor.addDistribution(address(_testToken), _merkleRoot_1, _totalAmount); + vm.expectRevert("MerkleDistributor: Distribution already exists."); + merkleDistributor.addDistribution(address(_testToken), _merkleRoot_1, _totalAmount); + } + + + function test_Success_claimDistribution() public { + _setupTokenDistribution(_merkleRoot_1); + // Attempt to claim with the mock proof + vm.startPrank(account1); + merkleDistributor.claimDistribution(address(_testToken), index1, account1, amount1, validProof1); // Removed version parameter + vm.stopPrank(); + vm.startPrank(account2); + // Verify claim for the latest version (which is 1 in this case) + uint256 claimedAmount = merkleDistributor.getTotalClaimed(address(_testToken), account1); // Explicitly check version 1 + assertEq(claimedAmount, amount1); + // Verify token balance + assertEq(_testToken.balanceOf(account1), amount1); + assertEq(_testToken.balanceOf(address(merkleDistributor)), _totalAmount - amount1); + } + + // function test_Success_claimDistribution_merkleRoot_2() public { + // _setupTokenDistribution(_merkleRoot_2); + // // Attempt to claim with the mock proof + // vm.startPrank(account1_2); + // merkleDistributor.claimDistribution(address(_testToken), index1_2, account1_2, amount1_2, validProof1_2); // Removed version parameter + // vm.stopPrank(); + // // Verify claim for the latest version (which is 1 in this case) + // uint256 claimedAmount = merkleDistributor.getTotalClaimed(address(_testToken), account1_2); // Explicitly check version 1 + // assertEq(claimedAmount, amount1_2); + // // Verify token balance + // assertEq(_testToken.balanceOf(account1_2), amount1_2); + // assertEq(_testToken.balanceOf(address(merkleDistributor)), _totalAmount - amount1_2); + // } + + function test_Revert_claimDistribution_InvalidProof() public { + _setupTokenDistribution(_merkleRoot_1); vm.expectRevert("MerkleDistributor: Invalid proof."); - merkleDistributor.claimDistribution(address(_testToken), 0, users.user1, 10, _proof); + merkleDistributor.claimDistribution(address(_testToken), index1, account1, amount1, validProof2); + } + + function test_Revert_claimDistribution_DistributionPaused() public { + _setupTokenDistribution(_merkleRoot_1); + merkleDistributor.pauseForWithdrawal(address(_testToken)); + vm.expectRevert("MerkleDistributor: Distribution is paused."); + merkleDistributor.claimDistribution(address(_testToken), index1, account1, amount1, validProof1); + } + + function test_Revert_claimDistribution_TokensClaimed() public { + _setupTokenDistribution(_merkleRoot_1); + vm.startPrank(account1); + merkleDistributor.claimDistribution(address(_testToken), index1, account1, amount1, validProof1); + vm.stopPrank(); + vm.expectRevert("MerkleDistributor: Tokens claimed for the latest version."); + merkleDistributor.claimDistribution(address(_testToken), index1, account1, amount1, validProof1); + } + + + // addDistributionVersion + function test_addDistributionVersion() public { + _setupTokenDistribution(_merkleRoot_1); + // Verify the version + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).version, + _version1 + ); + //Pause distribution first + merkleDistributor.pauseForWithdrawal(address(_testToken)); + // Withdraw the tokens first + merkleDistributor.withdrawUnclaimedTokens( + address(_testToken) + ); + // Make sure that the tokens were withdrawn + assertEq( + _testToken.balanceOf(address(merkleDistributor)), + 0 + ); + // Make sure distribution is withdrawn + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).isWithdrawn, + true + ); + // send more tokens to create a new version + _testToken.transfer(address(merkleDistributor), _totalAmount); + // Add a new version + merkleDistributor.addDistributionVersion(address(_testToken), _merkleRoot_1, _totalAmount); + // Verify the version + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version2).version, + _version2 + ); + } + + function test_Revert_addDistributionVersion_InsufficientBalance() public { + _setupTokenDistribution(_merkleRoot_1); + //Pause distribution first + merkleDistributor.pauseForWithdrawal(address(_testToken)); + // Withdraw the tokens first + merkleDistributor.withdrawUnclaimedTokens( + address(_testToken) + ); + // Make sure that the tokens were withdrawn + assertEq( + _testToken.balanceOf(address(merkleDistributor)), + 0 + ); + // Make sure distribution is withdrawn + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).isWithdrawn, + true + ); + // send more tokens to create a new version + vm.expectRevert("MerkleDistributor: Insufficient balance."); + merkleDistributor.addDistributionVersion(address(_testToken), _merkleRoot_1, _totalAmount); + } + + function test_Revert_addDistributionVersion_LatestVersionNotWithdrawn() public { + _setupTokenDistribution(_merkleRoot_1); + + vm.expectRevert("MerkleDistributor: Latest version is not withdrawn."); + merkleDistributor.addDistributionVersion(address(_testToken), _merkleRoot_1, _totalAmount); + } + + function test_Revert_addDistributionversion_DistributionDoesNotExist() public { + vm.expectRevert("MerkleDistributor: No distributions exist for this token."); + merkleDistributor.addDistributionVersion(address(_testToken), _merkleRoot_1, _totalAmount); + } + + + + function test_Success_updateDistribution() public { + _setupTokenDistribution(_merkleRoot_1); + bytes32 _newMerkleRoot = 0x0; + uint256 _additionalAmount = 100000000000000000000000000; + _testToken.transfer(address(merkleDistributor), _additionalAmount); + merkleDistributor.updateDistribution(address(_testToken), _newMerkleRoot, _additionalAmount); + // verify new distribution details + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).merkleRoot, + _newMerkleRoot + ); + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).totalAmount, + _totalAmount + _additionalAmount + ); + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).version, + _version1 + ); + } + + function test_Revert_updateDistribution_revertInsufficientBalance() public { + _setupTokenDistribution(_merkleRoot_1); + bytes32 _newMerkleRoot = 0x0; + uint256 _additionalAmount = 100000000000000000000000000; + vm.expectRevert("MerkleDistributor: Insufficient balance."); + merkleDistributor.updateDistribution(address(_testToken), _newMerkleRoot, _additionalAmount); + } + + function test_Success_pauseForWithdrawl() public { + _setupTokenDistribution(_merkleRoot_1); + merkleDistributor.pauseForWithdrawal(address(_testToken)); + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).isPaused, + true + ); + } + + function test_Revert_pauseForWithdrawl_DistributionDoesNotExist() public { + vm.expectRevert("MerkleDistributor: No distributions exist for this token."); + merkleDistributor.pauseForWithdrawal(address(_testToken)); + } + + function test_Revert_pauseForWithdrawl_AlreadyPaused() public { + _setupTokenDistribution(_merkleRoot_1); + merkleDistributor.pauseForWithdrawal(address(_testToken)); + vm.expectRevert("MerkleDistributor: Distribution is already paused."); + merkleDistributor.pauseForWithdrawal(address(_testToken)); + } + + function test_Success_unpause() public { + _setupTokenDistribution(_merkleRoot_1); + merkleDistributor.pauseForWithdrawal(address(_testToken)); + merkleDistributor.unpause(address(_testToken)); + assertEq( + merkleDistributor.getDistribution(address(_testToken), _version1).isPaused, + false + ); + } + + function test_Revert_unpause_DistributionNotPaused() public { + _setupTokenDistribution(_merkleRoot_1); + vm.expectRevert("MerkleDistributor: Distribution is not paused."); + merkleDistributor.unpause(address(_testToken)); + } + + function test_Sucess_withdrawUnclaimedTokens() public { + uint256 _tokenSupply_3 = 1000000000; + ERC677 _testToken_3 = new ERC677("Test", "TST", _tokenSupply_3); + uint256 _start_owner_token_balance = _testToken_3.balanceOf(owner); + console.log("Owner start balance: ", _testToken_3.balanceOf(owner)); + _testToken_3.transfer(address(merkleDistributor), _totalAmount); + console.log("Owner balance after transfer to merkledistributor: ", _testToken_3.balanceOf(owner)); + console.log("Merkle distributor balance after transfer to it",_testToken_3.balanceOf(address(merkleDistributor))); + merkleDistributor.addDistribution(address(_testToken_3), _merkleRoot_1, _totalAmount); + + assertEq( + _testToken_3.balanceOf(address(merkleDistributor)), + _totalAmount + ); + + merkleDistributor.pauseForWithdrawal(address(_testToken_3)); + merkleDistributor.withdrawUnclaimedTokens(address(_testToken_3)); + + assertEq( + _testToken_3.balanceOf(address(merkleDistributor)), + 0 + ); + assertEq( + _testToken_3.balanceOf(owner), + _start_owner_token_balance + ); + } + + function test_Revert_withdrawUnclaimedTokens_NotOwner() public { + _setupTokenDistribution(_merkleRoot_1); + merkleDistributor.pauseForWithdrawal(address(_testToken)); + vm.startPrank(account1); + vm.expectRevert("Ownable: caller is not the owner"); + merkleDistributor.withdrawUnclaimedTokens(address(_testToken)); } + }