Incorrect rewards logic in `MartenitsaMarketplace::collectRewards`

Severity

High Risk

Relevant GitHub Links

https://github.com/Cyfrin/2024-04-Baba-Marta/blob/5eaab7b51774d1083b926bf5ef116732c5a35cfd/src/MartenitsaMarketplace.sol#L101-L107

Summary

MartenitsaMarketplace::collectRewards contains a logic error that allows users to collect more healthToken rewards than they are entitled to.

Vulnerability Details

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.

Proof of Code

    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].

Impact

The obious and more serious bug lets users to repeatedly claim rewards for the same NFTs without any limitation.
The less obvious flaw results in users being able to claim less rewards (provided they do not take advantage of the more serious bug) than they are entitled to (e.g. if 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).

Tools Used

Manual review, Foundry.

Recommendations

To fix the obvious and more serious bug, update 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.

To adjust the logic completely according to the intention, a comprehensive logic overhaul is needed. In 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);
+    }
+}