MartenitsaMarketplace::collectRewards
High Risk
https://github.com/Cyfrin/2024-04-Baba-Marta/blob/5eaab7b51774d1083b926bf5ef116732c5a35cfd/src/MartenitsaMarketplace.sol#L101-L107
MartenitsaMarketplace::collectRewards
contains a logic error that allows users to collect more healthToken
rewards than they are entitled to.
MartentisaToken::collectReward
is supposed to allow non-producer users to collect healthToken
rewards for every three different Martenitsa NFTs they own. However, the function logic does not follow this intention. The function does not properly account for rewards that have already been collected, nor does it adequately update the internal state to reflect the number of rewards that should be available based on current NFT ownership.
In the provided test, a user acquires five NFTs and collects their corresponding rewards (1 healthToken
). However, after receiving a sixth NFT, they are able to repeatedly collect rewards for the same NFTs without any limitations.
function testRewardsLogicIsIncorrect() public {
address user = makeAddr("user");
// jack gives 5 NFTs to user as a gift
vm.startPrank(jack);
martenitsaToken.createMartenitsa("bracelet");
martenitsaToken.createMartenitsa("bracelet");
martenitsaToken.createMartenitsa("bracelet");
martenitsaToken.createMartenitsa("bracelet");
martenitsaToken.createMartenitsa("bracelet");
martenitsaToken.approve(address(marketplace), 0);
martenitsaToken.approve(address(marketplace), 1);
martenitsaToken.approve(address(marketplace), 2);
martenitsaToken.approve(address(marketplace), 3);
martenitsaToken.approve(address(marketplace), 4);
marketplace.makePresent(user, 0);
marketplace.makePresent(user, 1);
marketplace.makePresent(user, 2);
marketplace.makePresent(user, 3);
marketplace.makePresent(user, 4);
vm.stopPrank();
assert(martenitsaToken.ownerOf(0) == user);
// user collects rewards for 3 NFTs
vm.prank(user);
marketplace.collectReward();
// jack gives a 6th NFT to the user as present
vm.startPrank(jack);
martenitsaToken.createMartenitsa("bracelet");
martenitsaToken.approve(address(marketplace), 5);
marketplace.makePresent(user, 5);
vm.stopPrank();
// user can collect more rewards than eligible for
for (uint256 i = 0; i < 100; i++) {
vm.prank(user);
marketplace.collectReward();
}
assert(healthToken.balanceOf(user) > 2);
console.log("User's health token balance: %e", healthToken.balanceOf(user));
}
Note that there is an additional, lot less obvious flaw. Confirmed by the protocol owner, the real goal is to ensure that rewards are granted based on unique sets of three NFTs, where reusing any NFT from a previously rewarded set disqualifies the new set from earning another reward (for the same user). I.e. all of the following must be satisfied for new users (who has not claimed any rewards before):
userA
can claim 1 healthToken
for any set of 3 unique NFTs, e.g. for tokenIdSet_1=[1,2,3]
, and then-- userA
cannot claim another healthToken
with any set of 3 unique NFTs that share any elements with tokenIdSet_1
. E.g. cannot claim another reward with tokenIdSet_2=[1,4,5]
or tokenIdSet_3=[4,5,100]
.
-- userA
can claim another healthToken
with any set of 3 unique NFTs that DO NOT share any elements with tokenIdSet_1
. E.g. they can claim with tokenIdSet_4=[4,5,6]
or tokenIdSet_5=[x,y,z]
where x!=1
, x!=2, x!=3
, y!=1, y!=2
, y!=3, z!=1
, z!=2, z!=3
.
-- userB
or any other user can claim 1 healthToken
with tokenIdSet_1=[1,2,3]
.
userA
claims rewards for tokenIdSet_1=[1,2,3]
, then sells or gives away these 3 NFTs and acquires 3 new ones and tries to claim once again).Manual review, Foundry.
MartenitsaMarketplace
by correcting the reward tracking logic to accurately count already paid rewards as follows: function collectReward() external {
require(!martenitsaToken.isProducer(msg.sender), "You are producer and not eligible for a reward!");
uint256 count = martenitsaToken.getCountMartenitsaTokensOwner(msg.sender);
uint256 amountRewards = (count / requiredMartenitsaTokens) - _collectedRewards[msg.sender];
if (amountRewards > 0) {
- _collectedRewards[msg.sender] = amountRewards;
+ _collectedRewards[msg.sender] += amountRewards;
healthToken.distributeHealthToken(msg.sender, amountRewards);
}
}
With this partial fix users are rewarded based on the net increase in the number of complete sets (3 NFTs) they hold. Rewards can only be claimed for new sets that increase the user's holdings beyond previously rewarded levels. It effectively prevents exploiting the rewards system by recalculating eligible sets from scratch each time, thus aligning reward claims with actual holdings growth.
MartenitsaMarketplace
, implement set-based tracking to ensure no NFT is reused in multiple reward claims:+ import "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
...
contract MartenitsaMarketplace {
+ using EnumerableSet for EnumerableSet.UintSet;
+ mapping(address => EnumerableSet.UintSet) private claimedTokens;
+ function claimReward(uint256[] calldata tokenIds) external {
+ require(tokenIds.length == 3, "Must provide exactly three tokens");
+ require(areTokensUnique(tokenIds), "Tokens must be unique and unclaimed");
// Mark tokens as claimed
+ for (uint256 i = 0; i < tokenIds.length; i++) {
+ claimedTokens[msg.sender].add(tokenIds[i]);
+ }
// Logic to distribute rewards
+ distributeReward(msg.sender);
+ }
+ function areTokensUnique(uint256[] memory tokenIds) private view returns (bool) {
+ for (uint256 i = 0; i < tokenIds.length; i++) {
+ if (claimedTokens[msg.sender].contains(tokenIds[i])) {
+ return false;
+ }
+ }
+ return true;
+ }
+ function distributeReward(address recipient) private {
+ healthToken.distributeHealthToken(msg.sender, 1);
+ }
+}