Skip to content

Commit

Permalink
feat(distribution): add tests for claim and sweep
Browse files Browse the repository at this point in the history
  • Loading branch information
@robertu7
robertu7 committed Dec 25, 2023
1 parent 516e941 commit ea3e38d
Show file tree
Hide file tree
Showing 4 changed files with 172 additions and 15 deletions.
21 changes: 15 additions & 6 deletions .gas-snapshot
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,13 +70,22 @@ CurationTest:testCannotCurateNativeTokenZeroAddress() (gas: 16488)
CurationTest:testERC20Curation() (gas: 59908)
CurationTest:testNativeTokenCuration() (gas: 60085)
CurationTest:testNativeTokenCurationToContractAcceptor() (gas: 37466)
DistributionTest:testCannotDropIfInsufficientAllowance(uint256) (runs: 256, μ: 158670, ~: 158683)
DistributionTest:testCannotDropIfInsufficientBalance(uint256) (runs: 256, μ: 160868, ~: 161107)
DistributionTest:testCannotSetAdminByAdmin() (gas: 17364)
DistributionTest:testCannotClaimIfAlreadyClaimed() (gas: 284072)
DistributionTest:testCannotClaimIfInsufficientBalance() (gas: 391953)
DistributionTest:testCannotClaimIfInvalidProof() (gas: 244569)
DistributionTest:testCannotClaimIfInvalidTreeId() (gas: 243424)
DistributionTest:testCannotDropByAttacker() (gas: 11045)
DistributionTest:testCannotDropIfInsufficientAllowance(uint256) (runs: 256, μ: 158715, ~: 158729)
DistributionTest:testCannotDropIfInsufficientBalance(uint256) (runs: 256, μ: 160955, ~: 161196)
DistributionTest:testCannotDropIfZeroAmount() (gas: 150081)
DistributionTest:testCannotSetAdminByAdmin() (gas: 17342)
DistributionTest:testCannotSetAdminByAttacker() (gas: 11111)
DistributionTest:testClaim() (gas: 402817)
DistributionTest:testDrop() (gas: 354444)
DistributionTest:testSetAdmin() (gas: 15417)
DistributionTest:testCannotSweepByAttacker() (gas: 228843)
DistributionTest:testCannotSweepIfZeroBalance() (gas: 230636)
DistributionTest:testClaim() (gas: 410328)
DistributionTest:testDrop() (gas: 354674)
DistributionTest:testSetAdmin() (gas: 20182)
DistributionTest:testSweep() (gas: 250986)
LogbookNFTSVGTest:testTokenURI(uint8,uint8,uint16) (runs: 256, μ: 2019505, ~: 1310779)
LogbookTest:testClaim() (gas: 135608)
LogbookTest:testDonate(uint96) (runs: 256, μ: 155485, ~: 156936)
Expand Down
17 changes: 11 additions & 6 deletions src/Billboard/Distribution.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ contract Distribution is IDistribution, Ownable {
mapping(uint256 => mapping(string => mapping(address => bool))) public hasClaimed;

constructor(address token_, address admin_) {
require(token_ != address(0), "zero address");
require(token_ != address(0), "Zero address");

admin = admin_;
token = token_;
Expand Down Expand Up @@ -55,6 +55,8 @@ contract Distribution is IDistribution, Ownable {

/// @inheritdoc IDistribution
function drop(bytes32 merkleRoot_, uint256 amount_) external payable isFromAdmin returns (uint256 treeId_) {
require(amount_ > 0, "Zero amount");

// Transfer
SafeERC20.safeTransferFrom(IERC20(token), msg.sender, address(this), amount_);

Expand All @@ -77,17 +79,20 @@ contract Distribution is IDistribution, Ownable {
uint256 amount_,
bytes32[] calldata merkleProof_
) external {
require(!hasClaimed[treeId_][cid_][account_], "already claimed.");
require(!hasClaimed[treeId_][cid_][account_], "Already claimed");

bytes32 _root = merkleRoots[treeId_];
require(_root != bytes32(0), "Invalid tree ID");

// Verify the merkle proof
bytes32 _leaf = keccak256(bytes.concat(keccak256(abi.encode(cid_, account_, amount_))));
require(MerkleProof.verify(merkleProof_, merkleRoots[treeId_], _leaf), "invalid proof");
require(MerkleProof.verify(merkleProof_, _root, _leaf), "Invalid proof");

// Mark it as claimed first for to prevent reentrancy
hasClaimed[treeId_][cid_][account_] = true;

// Transfer
require(IERC20(token).transfer(account_, amount_), "failed token transfer");
require(IERC20(token).transfer(account_, amount_), "Failed token transfer");

// Update the balance for the tree
balances[treeId_] -= amount_;
Expand All @@ -99,10 +104,10 @@ contract Distribution is IDistribution, Ownable {
function sweep(uint256 treeId_, address target_) external isFromAdmin {
uint256 _balance = balances[treeId_];

require(_balance > 0, "zero balance");
require(_balance > 0, "Zero balance");

// Transfer
require(IERC20(token).transfer(target_, _balance), "failed token transfer");
require(IERC20(token).transfer(target_, _balance), "Failed token transfer");

// Update the balance for the tree
balances[treeId_] = 0;
Expand Down
2 changes: 1 addition & 1 deletion src/Billboard/IDistribution.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ interface IDistribution {
* @param account_ Address of claim
* @param amount_ Amount of claim
*/
event Claim(string cid_, address indexed account_, uint256 amount_);
event Claim(string indexed cid_, address indexed account_, uint256 amount_);

/**
* @dev Emitted when admin is changed.
Expand Down
147 changes: 145 additions & 2 deletions src/test/Billboard/DistributionTest.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,18 @@ import "@openzeppelin/contracts/token/ERC721/IERC721.sol";
import "./DistributionTestBase.t.sol";

contract DistributionTest is DistributionTestBase {
//////////////////////////////
/// Access control
//////////////////////////////

function testSetAdmin() public {
vm.prank(OWNER);
distribution.setAdmin(ADMIN);
assertEq(distribution.admin(), ADMIN);

vm.expectEmit(true, true, false, false);
emit IDistribution.AdminChanged(ADMIN, USER_ALICE);

distribution.setAdmin(USER_ALICE);
assertEq(distribution.admin(), USER_ALICE);
}

function testCannotSetAdminByAdmin() public {
Expand All @@ -29,6 +37,10 @@ contract DistributionTest is DistributionTestBase {
distribution.setAdmin(USER_ALICE);
}

//////////////////////////////
/// Drop
//////////////////////////////

function testDrop() public {
// drop#1
uint256 _amount = 1510000000000000000;
Expand All @@ -46,6 +58,19 @@ contract DistributionTest is DistributionTestBase {
assertEq(usdt.balanceOf(address(distribution)), _amount * 2);
}

function testCannotDropByAttacker() public {
vm.prank(ATTACKER);
vm.expectRevert("Admin");
distribution.drop(TREE_1_ROOT, 1);
}

function testCannotDropIfZeroAmount() public {
deal(address(usdt), ADMIN, 0);
vm.prank(ADMIN);
vm.expectRevert("Zero amount");
distribution.drop(TREE_1_ROOT, 0);
}

function testCannotDropIfInsufficientAllowance(uint256 amount_) public {
vm.assume(amount_ > 0);
deal(address(usdt), ADMIN, amount_);
Expand All @@ -69,12 +94,17 @@ contract DistributionTest is DistributionTestBase {
distribution.drop(TREE_1_ROOT, amount_ + 1);
}

//////////////////////////////
/// Claim
//////////////////////////////
function testClaim() public {
// drop#1
uint256 _amount = 1510000000000000000;
drop(_amount);

// claim#Alice
vm.expectEmit(true, true, false, false);
emit IDistribution.Claim(TREE_1_CIDS[USER_ALICE], USER_ALICE, TREE_1_AMOUNTS[USER_ALICE]);
uint256 balanceAlce = address(USER_ALICE).balance;
distribution.claim(
1,
Expand All @@ -84,13 +114,18 @@ contract DistributionTest is DistributionTestBase {
TREE_1_PROOFS[USER_ALICE]
);
assertEq(usdt.balanceOf(address(USER_ALICE)), balanceAlce + TREE_1_AMOUNTS[USER_ALICE]);
assertEq(usdt.balanceOf(address(distribution)), _amount - TREE_1_AMOUNTS[USER_ALICE]);

// claim#Bob
vm.expectEmit(true, true, false, false);
emit IDistribution.Claim(TREE_1_CIDS[USER_BOB], USER_BOB, TREE_1_AMOUNTS[USER_BOB]);
uint256 balanceBob = address(USER_BOB).balance;
distribution.claim(1, TREE_1_CIDS[USER_BOB], USER_BOB, TREE_1_AMOUNTS[USER_BOB], TREE_1_PROOFS[USER_BOB]);
assertEq(usdt.balanceOf(address(USER_BOB)), balanceBob + TREE_1_AMOUNTS[USER_BOB]);

// claim#Charlie
vm.expectEmit(true, true, false, false);
emit IDistribution.Claim(TREE_1_CIDS[USER_CHARLIE], USER_CHARLIE, TREE_1_AMOUNTS[USER_CHARLIE]);
uint256 balanceCharlie = address(USER_CHARLIE).balance;
distribution.claim(
1,
Expand All @@ -104,4 +139,112 @@ contract DistributionTest is DistributionTestBase {
// check balance
assertEq(address(distribution).balance, 0);
}

function testCannotClaimIfAlreadyClaimed() public {
// drop#1
uint256 _amount = 1510000000000000000;
drop(_amount);

// claim#Alice
distribution.claim(
1,
TREE_1_CIDS[USER_ALICE],
USER_ALICE,
TREE_1_AMOUNTS[USER_ALICE],
TREE_1_PROOFS[USER_ALICE]
);

// claim#Alice again
vm.expectRevert("Already claimed");
distribution.claim(
1,
TREE_1_CIDS[USER_ALICE],
USER_ALICE,
TREE_1_AMOUNTS[USER_ALICE],
TREE_1_PROOFS[USER_ALICE]
);
}

function testCannotClaimIfInvalidProof() public {
// drop#1
uint256 _amount = 1510000000000000000;
drop(_amount);

// claim#Alice
vm.expectRevert("Invalid proof");
distribution.claim(1, TREE_1_CIDS[USER_ALICE], USER_ALICE, TREE_1_AMOUNTS[USER_ALICE], TREE_1_PROOFS[USER_BOB]);
}

function testCannotClaimIfInvalidTreeId() public {
// drop#1
uint256 _amount = 1510000000000000000;
drop(_amount);

// claim#Alice
vm.expectRevert("Invalid tree ID");
distribution.claim(
2,
TREE_1_CIDS[USER_ALICE],
USER_ALICE,
TREE_1_AMOUNTS[USER_ALICE],
TREE_1_PROOFS[USER_ALICE]
);
}

function testCannotClaimIfInsufficientBalance() public {
// drop#1
uint256 _amount = 1510000000000000000;
drop(_amount);
deal(address(usdt), address(distribution), 0);
assertEq(usdt.balanceOf(address(distribution)), 0);

// claim#Alice
vm.expectRevert("ERC20: transfer amount exceeds balance");
distribution.claim(
1,
TREE_1_CIDS[USER_ALICE],
USER_ALICE,
TREE_1_AMOUNTS[USER_ALICE],
TREE_1_PROOFS[USER_ALICE]
);
}

//////////////////////////////
/// Sweep
//////////////////////////////
function testSweep() public {
// drop
uint256 _amount = 1510000000000000000;
drop(_amount);

// sweep
uint256 prevBalance = usdt.balanceOf(ADMIN);
vm.prank(ADMIN);
distribution.sweep(1, ADMIN);
assertEq(usdt.balanceOf(ADMIN), prevBalance + _amount);
assertEq(usdt.balanceOf(address(distribution)), 0);
assertEq(distribution.balances(1), 0);
}

function testCannotSweepByAttacker() public {
// drop
uint256 _amount = 1510000000000000000;
drop(_amount);

// sweep
vm.prank(ATTACKER);
vm.expectRevert("Admin");
distribution.sweep(1, ADMIN);
}

function testCannotSweepIfZeroBalance() public {
// drop
uint256 _amount = 1510000000000000000;
drop(_amount);

// sweep
vm.prank(ADMIN);
vm.expectRevert("Zero balance");
distribution.sweep(2, ADMIN);
}
}

0 comments on commit ea3e38d

Please sign in to comment.