Merge pull request #3 from euler-xyz/cantina-231

Cantina 231

README.md

@@ -46,8 +46,6 @@ Reward Streams operates in two modes of rewards distribution: staking and balanc
 
 The balance-tracking `TrackingRewardStreams` implementation inherits from the `BaseRewardStreams` contract. It defines the `IBalanceTracker.balanceTrackerHook` function, which is required to be called on every transfer of the rewarded token if a user opted in for the hook to be called. 
 
-In this mode, the rewarded token contract not only calls the `balanceTrackerHook` function whenever a given account balance changes, but also implements the `IBalanceForwarder` interface. This interface defines two functions: `enableBalanceForwarding` and `disableBalanceForwarding`, which are used to opt in and out of the hook being called.
-
 ### Staking Reward Distribution
 
 The staking `StakingRewardStreams` implementation also inherits from the `BaseRewardStreams` contract. It defines two functions: `stake` and `unstake`, which are used to stake and unstake the rewarded token.

src/TrackingRewardStreams.sol

@@ -27,7 +27,8 @@ contract TrackingRewardStreams is BaseRewardStreams, ITrackingRewardStreams {
     /// @notice Executes the balance tracking hook for an account
     /// @param account The account address to execute the hook for
     /// @param newAccountBalance The new balance of the account
-    /// @param forfeitRecentReward Whether to forfeit the most recent reward and not update the accumulator
+    /// @param forfeitRecentReward Whether to forfeit the most recent reward and not update the accumulator. Ignored
+    /// when the new balance is greater than the current balance.
     function balanceTrackerHook(
         address account,
         uint256 newAccountBalance,
@@ -38,6 +39,8 @@ contract TrackingRewardStreams is BaseRewardStreams, ITrackingRewardStreams {
         uint256 currentAccountBalance = accountStorage.balance;
         address[] memory rewards = accountStorage.enabledRewards.get();
 
+        if (newAccountBalance > currentAccountBalance) forfeitRecentReward = false;
+
         for (uint256 i = 0; i < rewards.length; ++i) {
             address reward = rewards[i];
             DistributionStorage storage distributionStorage = distributions[rewarded][reward];

src/interfaces/IBalanceTracker.sol

@@ -8,11 +8,10 @@ pragma solidity >=0.8.0;
 /// @notice Provides an interface for tracking the balance of accounts.
 interface IBalanceTracker {
     /// @notice Executes the balance tracking hook for an account.
-    /// @dev This function is called by the Balance Forwarder contract which was enabled for the account. This function
-    /// must be called with the current balance of the account when enabling the balance forwarding for it. This
-    /// function must be called with 0 balance of the account when disabling the balance forwarding for it. This
-    /// function allows to be called on zero balance transfers, when the newAccountBalance is the same as the previous
-    /// one. To prevent DOS attacks, forfeitRecentReward should be used appropriately.
+    /// @dev This function must be called with the current balance of the account when enabling the balance forwarding
+    /// for it. This function must be called with 0 balance of the account when disabling the balance forwarding for it.
+    /// This function allows to be called on zero balance transfers, when the newAccountBalance is the same as the
+    /// previous one. To prevent DOS attacks, forfeitRecentReward should be used appropriately.
     /// @param account The account address to execute the hook for.
     /// @param newAccountBalance The new balance of the account.
     /// @param forfeitRecentReward Whether to forfeit the most recent reward and not update the accumulator.

test/unit/Scenarios.t.sol

@@ -6,7 +6,7 @@ import "forge-std/Test.sol";
 import "evc/EthereumVaultConnector.sol";
 import "../harness/StakingRewardStreamsHarness.sol";
 import "../harness/TrackingRewardStreamsHarness.sol";
-import {MockERC20, MockERC20BalanceForwarder} from "../utils/MockERC20.sol";
+import {MockERC20, MockERC20BalanceForwarder, MockERC20BalanceForwarderMessedUp} from "../utils/MockERC20.sol";
 import {MockController} from "../utils/MockController.sol";
 import {boundAddr} from "../utils/TestUtils.sol";
 
@@ -1920,7 +1920,7 @@ contract ScenarioTest is Test {
         vm.stopPrank();
 
         // forward the time to the middle of the second epoch of the distribution scheme
-        vm.warp(stakingDistributor.getEpochStartTimestamp(stakingDistributor.currentEpoch() + 1) + 15 days);
+        vm.warp(trackingDistributor.getEpochStartTimestamp(trackingDistributor.currentEpoch() + 1) + 15 days);
 
         // verify earnings
         assertApproxEqRel(
@@ -2237,6 +2237,88 @@ contract ScenarioTest is Test {
         assertEq(MockERC20(trackingRewarded).balanceOf(PARTICIPANT_1), preBalance - 10e18);
     }
 
+    // reward and rewarded are the same
+    function test_Scenario_MaliciousBalanceForwarder(uint48 blockTimestamp) external {
+        blockTimestamp = uint48(bound(blockTimestamp, 1, type(uint48).max - 365 days));
+
+        uint256 ALLOWED_DELTA = 1e12; // 0.0001%
+
+        address trackingRewardedMessedUp = address(
+            new MockERC20BalanceForwarderMessedUp(evc, trackingDistributor, "Tracking Rewarded Malicious", "SFRWDDM")
+        );
+
+        // mint the tracking rewarded token to participants
+        MockERC20(trackingRewardedMessedUp).mint(PARTICIPANT_1, 10e18);
+        MockERC20(trackingRewardedMessedUp).mint(PARTICIPANT_2, 10e18);
+        MockERC20(trackingRewardedMessedUp).mint(PARTICIPANT_3, 10e18);
+
+        vm.warp(blockTimestamp);
+
+        // prepare the amounts; 3 epochs
+        uint128[] memory amounts = new uint128[](3);
+        amounts[0] = 10e18;
+        amounts[1] = 10e18;
+        amounts[2] = 10e18;
+
+        // register the distribution scheme
+        vm.startPrank(seeder);
+        trackingDistributor.registerReward(trackingRewardedMessedUp, reward, 0, amounts);
+        vm.stopPrank();
+
+        // enable reward and balance forwarding for participant 1 and participant 2
+        vm.startPrank(PARTICIPANT_1);
+        MockERC20BalanceForwarder(trackingRewardedMessedUp).enableBalanceForwarding();
+        trackingDistributor.enableReward(trackingRewardedMessedUp, reward);
+        vm.stopPrank();
+
+        vm.startPrank(PARTICIPANT_2);
+        MockERC20BalanceForwarder(trackingRewardedMessedUp).enableBalanceForwarding();
+        trackingDistributor.enableReward(trackingRewardedMessedUp, reward);
+        vm.stopPrank();
+
+        // forward the time to the end of the first epoch of the distribution scheme
+        vm.warp(trackingDistributor.getEpochStartTimestamp(trackingDistributor.currentEpoch() + 1) + 10 days);
+
+        // lock in current earnings
+        trackingDistributor.updateReward(trackingRewardedMessedUp, reward, address(0));
+
+        // verify earnings
+        assertApproxEqRel(
+            trackingDistributor.earnedReward(PARTICIPANT_1, trackingRewardedMessedUp, reward, false),
+            5e18,
+            ALLOWED_DELTA
+        );
+        assertApproxEqRel(
+            trackingDistributor.earnedReward(PARTICIPANT_2, trackingRewardedMessedUp, reward, false),
+            5e18,
+            ALLOWED_DELTA
+        );
+        assertEq(trackingDistributor.earnedReward(address(0), trackingRewardedMessedUp, reward, false), 0);
+
+        // forward the time
+        vm.warp(block.timestamp + 20 days);
+
+        // participant 3 (who doesn't have balance forwarding enabled) transfers tokens to participant 1 (who has
+        // balance forwarding enabled). despite messed up balance forwarder integration, the accounting will be correct
+        // because the forfeitRecentReward flag setting will be overriden
+        vm.prank(PARTICIPANT_3);
+        MockERC20(trackingRewardedMessedUp).transfer(PARTICIPANT_1, 10e18);
+
+        // verify earnings. both participants should earn the same amount of rewards, despite messed up balance
+        // forwarder integrations
+        assertApproxEqRel(
+            trackingDistributor.earnedReward(PARTICIPANT_1, trackingRewardedMessedUp, reward, false),
+            15e18,
+            ALLOWED_DELTA
+        );
+        assertApproxEqRel(
+            trackingDistributor.earnedReward(PARTICIPANT_2, trackingRewardedMessedUp, reward, false),
+            15e18,
+            ALLOWED_DELTA
+        );
+        assertEq(trackingDistributor.earnedReward(address(0), trackingRewardedMessedUp, reward, false), 0);
+    }
+
     function test_AssertionTrigger(
         address _account,
         address _rewarded,

test/utils/MockERC20.sol

@@ -94,3 +94,25 @@ contract MockERC20BalanceForwarder is MockERC20, IBalanceForwarder {
         balanceTracker.balanceTrackerHook(account, newAccountBalance, forfeitRecentReward);
     }
 }
+
+contract MockERC20BalanceForwarderMessedUp is MockERC20BalanceForwarder {
+    constructor(
+        IEVC _evc,
+        IBalanceTracker _balanceTracker,
+        string memory _name,
+        string memory _symbol
+    ) MockERC20BalanceForwarder(_evc, _balanceTracker, _name, _symbol) {}
+
+    function _update(address from, address to, uint256 value) internal virtual override {
+        ERC20._update(from, to, value);
+
+        if (forwardingEnabled[from]) {
+            balanceTracker.balanceTrackerHook(from, balanceOf(from), false);
+        }
+
+        // always pass forfeitRecentReward = true when increasing balance to mess up the accounting
+        if (from != to && forwardingEnabled[to]) {
+            balanceTracker.balanceTrackerHook(to, balanceOf(to), true);
+        }
+    }
+}