Skip to content

Commit

Permalink
feat(distribution): use share_ instead of amount_ on claim
Browse files Browse the repository at this point in the history
  • Loading branch information
@robertu7
robertu7 committed Jan 24, 2024
1 parent 7fb2c41 commit 8ce7758
Show file tree
Hide file tree
Showing 6 changed files with 91 additions and 77 deletions.
32 changes: 16 additions & 16 deletions .gas-snapshot
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,22 +70,22 @@ CurationTest:testCannotCurateNativeTokenZeroAddress() (gas: 16488)
CurationTest:testERC20Curation() (gas: 59908)
CurationTest:testNativeTokenCuration() (gas: 60085)
CurationTest:testNativeTokenCurationToContractAcceptor() (gas: 37466)
DistributionTest:testCannotClaimIfAlreadyClaimed() (gas: 284069)
DistributionTest:testCannotClaimIfInsufficientBalance() (gas: 393827)
DistributionTest:testCannotClaimIfInvalidProof() (gas: 244573)
DistributionTest:testCannotClaimIfInvalidTreeId() (gas: 243432)
DistributionTest:testCannotDropByAttacker() (gas: 11045)
DistributionTest:testCannotDropIfInsufficientAllowance(uint256) (runs: 256, μ: 213134, ~: 213153)
DistributionTest:testCannotDropIfInsufficientBalance(uint256) (runs: 256, μ: 215359, ~: 215614)
DistributionTest:testCannotDropIfZeroAmount() (gas: 150081)
DistributionTest:testCannotSetAdminByAdmin() (gas: 17377)
DistributionTest:testCannotSetAdminByAttacker() (gas: 11111)
DistributionTest:testCannotSweepByAttacker() (gas: 228850)
DistributionTest:testCannotSweepIfZeroBalance() (gas: 230643)
DistributionTest:testClaim() (gas: 410304)
DistributionTest:testDrop() (gas: 354688)
DistributionTest:testSetAdmin() (gas: 20217)
DistributionTest:testSweep() (gas: 251040)
DistributionTest:testCannotClaimIfAlreadyClaimed() (gas: 302159)
DistributionTest:testCannotClaimIfInsufficientBalance() (gas: 411915)
DistributionTest:testCannotClaimIfInvalidProof() (gas: 263211)
DistributionTest:testCannotClaimIfInvalidTreeId() (gas: 261207)
DistributionTest:testCannotDropByAttacker() (gas: 11092)
DistributionTest:testCannotDropIfInsufficientAllowance(uint256) (runs: 256, μ: 230891, ~: 230907)
DistributionTest:testCannotDropIfInsufficientBalance(uint256) (runs: 256, μ: 233150, ~: 233389)
DistributionTest:testCannotDropIfZeroAmount() (gas: 150119)
DistributionTest:testCannotSetAdminByAdmin() (gas: 17333)
DistributionTest:testCannotSetAdminByAttacker() (gas: 11089)
DistributionTest:testCannotSweepByAttacker() (gas: 246643)
DistributionTest:testCannotSweepIfZeroBalance() (gas: 248436)
DistributionTest:testClaim() (gas: 430983)
DistributionTest:testDrop() (gas: 390240)
DistributionTest:testSetAdmin() (gas: 20195)
DistributionTest:testSweep() (gas: 268852)
LogbookNFTSVGTest:testTokenURI(uint8,uint8,uint16) (runs: 256, μ: 2019505, ~: 1310779)
LogbookTest:testClaim() (gas: 135608)
LogbookTest:testDonate(uint96) (runs: 256, μ: 155485, ~: 156936)
Expand Down
6 changes: 3 additions & 3 deletions scripts/billboard/generate-merkle-tree.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,17 @@ const values = [
[
"Qmf5z5DKcwNWYUP9udvnSCTN2Se4A8kpZJY7JuUVFEqdGU",
"0x0000000000000000000000000000000000000066",
"1000000000000000000", // 1 USDT
"1000", // 10%
],
[
"QmSAwncsWGXeqwrL5USBzQXvjqfH1nFfARLGM91sfd4NZe",
"0x0000000000000000000000000000000000000067",
"500000000000000000", // 0.5 USDT
"2055", // 20.55%
],
[
"QmUQQSeWxcqoNLKroGtz137c7QBWpzbNr9RcqDtVzZxJ3x",
"0x0000000000000000000000000000000000000068",
"10000000000000000", // 0.01 USDT
"6945", // 69.45%
],
];

Expand Down
18 changes: 14 additions & 4 deletions src/Billboard/Distribution.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,9 @@ contract Distribution is IDistribution, Ownable {
// treeId_ => balance_
mapping(uint256 => uint256) public balances;

// treeId_ => totalAmount_
mapping(uint256 => uint256) public totalAmounts;

// treeId_ => cid_ => account_
mapping(uint256 => mapping(string => mapping(address => bool))) public hasClaimed;

Expand Down Expand Up @@ -57,7 +60,7 @@ contract Distribution is IDistribution, Ownable {
//////////////////////////////

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

// Set the merkle root
Expand All @@ -67,8 +70,9 @@ contract Distribution is IDistribution, Ownable {

emit Drop(treeId_, amount_);

// Set the balance for the tree
// Set the balance & total amount for the tree
balances[treeId_] = amount_;
totalAmounts[treeId_] = amount_;

// Transfer
SafeERC20.safeTransferFrom(IERC20(token), msg.sender, address(this), amount_);
Expand All @@ -79,7 +83,7 @@ contract Distribution is IDistribution, Ownable {
uint256 treeId_,
string calldata cid_,
address account_,
uint256 amount_,
uint256 share_,
bytes32[] calldata merkleProof_
) external {
require(!hasClaimed[treeId_][cid_][account_], "Already claimed");
Expand All @@ -88,12 +92,14 @@ contract Distribution is IDistribution, Ownable {
require(_root != bytes32(0), "Invalid tree ID");

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

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

uint256 amount_ = calculateAmount(share_, totalAmounts[treeId_]);

emit Claim(cid_, account_, amount_);

// Update the balance for the tree
Expand All @@ -115,4 +121,8 @@ contract Distribution is IDistribution, Ownable {
// Transfer
require(IERC20(token).transfer(target_, _balance), "Failed token transfer");
}

function calculateAmount(uint256 share_, uint256 totalAmount_) public pure returns (uint256 amount) {
amount = (totalAmount_ * share_) / 10000;
}
}
8 changes: 4 additions & 4 deletions src/Billboard/IDistribution.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,24 +44,24 @@ interface IDistribution {
*
* Emits a {Drop} event on success.
*/
function drop(bytes32 merkleRoot_, uint256 amount_) external payable returns (uint256 treeId_);
function drop(bytes32 merkleRoot_, uint256 amount_) external returns (uint256 treeId_);

/**
* @notice Claim and transfer tokens
*
* @param treeId_ Tree ID
* @param cid_ Content ID
* @param account_ Address of claim
* @param amount_ Amount of claim
* @param proof_ Merkle proof for (treeId_, cid_, account_, amount_)
* @param share_ Share (percentage with two decimal places to an integer representation, 0-10000) of total amount
* @param proof_ Merkle proof for (treeId_, cid_, account_, share_)
*
* Emits a {Claim} event on success.
*/
function claim(
uint256 treeId_,
string calldata cid_,
address account_,
uint256 amount_,
uint256 share_,
bytes32[] calldata proof_
) external;

Expand Down
82 changes: 43 additions & 39 deletions src/test/Billboard/DistributionTest.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,19 +43,19 @@ contract DistributionTest is DistributionTestBase {

function testDrop() public {
// drop#1
uint256 _amount = 1510000000000000000;
drop(_amount);
uint256 _totalAmount = 1510000000000000000;
drop(_totalAmount);
assertEq(distribution.lastTreeId(), 1);
assertEq(distribution.merkleRoots(1), TREE_1_ROOT);
assertEq(distribution.balances(1), _amount);
assertEq(usdt.balanceOf(address(distribution)), _amount);
assertEq(distribution.balances(1), _totalAmount);
assertEq(usdt.balanceOf(address(distribution)), _totalAmount);

// drop#2
drop(_amount);
drop(_totalAmount);
assertEq(distribution.lastTreeId(), 2);
assertEq(distribution.merkleRoots(2), TREE_1_ROOT);
assertEq(distribution.balances(2), _amount);
assertEq(usdt.balanceOf(address(distribution)), _amount * 2);
assertEq(distribution.balances(2), _totalAmount);
assertEq(usdt.balanceOf(address(distribution)), _totalAmount * 2);
}

function testCannotDropByAttacker() public {
Expand Down Expand Up @@ -99,58 +99,61 @@ contract DistributionTest is DistributionTestBase {
//////////////////////////////
function testClaim() public {
// drop#1
uint256 _amount = 1510000000000000000;
drop(_amount);
uint256 _totalAmount = 1510000000000000000;
drop(_totalAmount);

// claim#Alice
uint256 _amount = distribution.calculateAmount(TREE_1_SHARES[USER_ALICE], _totalAmount);
vm.expectEmit(true, true, false, false);
emit IDistribution.Claim(TREE_1_CIDS[USER_ALICE], USER_ALICE, TREE_1_AMOUNTS[USER_ALICE]);
emit IDistribution.Claim(TREE_1_CIDS[USER_ALICE], USER_ALICE, _amount);
uint256 balanceAlce = address(USER_ALICE).balance;
distribution.claim(
1,
TREE_1_CIDS[USER_ALICE],
USER_ALICE,
TREE_1_AMOUNTS[USER_ALICE],
TREE_1_SHARES[USER_ALICE],
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]);
assertEq(usdt.balanceOf(address(USER_ALICE)), balanceAlce + _amount);
assertEq(usdt.balanceOf(address(distribution)), _totalAmount - _amount);

// claim#Bob
uint256 _amountBob = distribution.calculateAmount(TREE_1_SHARES[USER_BOB], _totalAmount);
vm.expectEmit(true, true, false, false);
emit IDistribution.Claim(TREE_1_CIDS[USER_BOB], USER_BOB, TREE_1_AMOUNTS[USER_BOB]);
emit IDistribution.Claim(TREE_1_CIDS[USER_BOB], USER_BOB, _amountBob);
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]);
distribution.claim(1, TREE_1_CIDS[USER_BOB], USER_BOB, TREE_1_SHARES[USER_BOB], TREE_1_PROOFS[USER_BOB]);
assertEq(usdt.balanceOf(address(USER_BOB)), balanceBob + _amountBob);

// claim#Charlie
uint256 _amountCharlie = distribution.calculateAmount(TREE_1_SHARES[USER_CHARLIE], _totalAmount);
vm.expectEmit(true, true, false, false);
emit IDistribution.Claim(TREE_1_CIDS[USER_CHARLIE], USER_CHARLIE, TREE_1_AMOUNTS[USER_CHARLIE]);
emit IDistribution.Claim(TREE_1_CIDS[USER_CHARLIE], USER_CHARLIE, _amountCharlie);
uint256 balanceCharlie = address(USER_CHARLIE).balance;
distribution.claim(
1,
TREE_1_CIDS[USER_CHARLIE],
USER_CHARLIE,
TREE_1_AMOUNTS[USER_CHARLIE],
TREE_1_SHARES[USER_CHARLIE],
TREE_1_PROOFS[USER_CHARLIE]
);
assertEq(usdt.balanceOf(address(USER_CHARLIE)), balanceCharlie + TREE_1_AMOUNTS[USER_CHARLIE]);
assertEq(usdt.balanceOf(address(USER_CHARLIE)), balanceCharlie + _amountCharlie);

// check balance
assertEq(address(distribution).balance, 0);
}

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

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

Expand All @@ -160,41 +163,42 @@ contract DistributionTest is DistributionTestBase {
1,
TREE_1_CIDS[USER_ALICE],
USER_ALICE,
TREE_1_AMOUNTS[USER_ALICE],
TREE_1_SHARES[USER_ALICE],
TREE_1_PROOFS[USER_ALICE]
);
}

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

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

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

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

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

Expand All @@ -204,7 +208,7 @@ contract DistributionTest is DistributionTestBase {
1,
TREE_1_CIDS[USER_ALICE],
USER_ALICE,
TREE_1_AMOUNTS[USER_ALICE],
TREE_1_SHARES[USER_ALICE],
TREE_1_PROOFS[USER_ALICE]
);
}
Expand All @@ -214,22 +218,22 @@ contract DistributionTest is DistributionTestBase {
//////////////////////////////
function testSweep() public {
// drop
uint256 _amount = 1510000000000000000;
drop(_amount);
uint256 _totalAmount = 1510000000000000000;
drop(_totalAmount);

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

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

// sweep
vm.prank(ATTACKER);
Expand All @@ -239,8 +243,8 @@ contract DistributionTest is DistributionTestBase {

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

// sweep
vm.prank(ADMIN);
Expand Down
Loading

0 comments on commit 8ce7758

Please sign in to comment.