feat: Detect extra rewards at AToken when deposit and withdrawal of Curve gauge tokens.

contracts/adapters/rewards/curve/CurveGaugeRewardsAwareAToken.sol

@@ -224,7 +224,30 @@ contract CurveGaugeRewardsAwareAToken is RewardsAwareAToken {
    */
   function _stake(address token, uint256 amount) internal override returns (uint256) {
     if (token == UNDERLYING_ASSET_ADDRESS()) {
-      ICurveTreasury(CURVE_TREASURY).deposit(token, amount, true);
+      if (_rewardTokens[1] != address(0)) {
+        // Track the pending rewards to distribute at `_retrieveAvailableReward` call
+        uint256[MAX_REWARD_TOKENS] memory priorTokenBalances;
+        for (uint256 index = 1; index < MAX_REWARD_TOKENS; index++) {
+          address rewardToken = _getRewardsTokenAddress(index);
+          if (rewardToken == address(0)) break;
+          if (rewardToken == CRV_TOKEN) continue;
+          priorTokenBalances[index] = IERC20(rewardToken).balanceOf(address(this));
+        }
+        // At deposits it sends extra rewards to aToken
+        ICurveTreasury(CURVE_TREASURY).deposit(token, amount, true);
+
+        for (uint256 index = 1; index < MAX_REWARD_TOKENS; index++) {
+          address rewardToken = _getRewardsTokenAddress(index);
+          if (rewardToken == address(0)) break;
+          if (rewardToken == CRV_TOKEN) continue;
+          uint256 balance = IERC20(rewardToken).balanceOf(address(this));
+          _pendingRewards[rewardToken] = _pendingRewards[rewardToken].add(
+            balance.sub(priorTokenBalances[index])
+          );
+        }
+      } else {
+        ICurveTreasury(CURVE_TREASURY).deposit(token, amount, true);
+      }
     }
     return amount;
   }
@@ -236,7 +259,30 @@ contract CurveGaugeRewardsAwareAToken is RewardsAwareAToken {
    */
   function _unstake(address token, uint256 amount) internal override returns (uint256) {
     if (token == UNDERLYING_ASSET_ADDRESS()) {
-      ICurveTreasury(CURVE_TREASURY).withdraw(token, amount, true);
+      // Claim other Curve gauge tokens, and track the pending rewards to distribute at `_retrieveAvailableReward` call
+      if (_rewardTokens[1] != address(0)) {
+        uint256[MAX_REWARD_TOKENS] memory priorTokenBalances;
+        for (uint256 index = 1; index < MAX_REWARD_TOKENS; index++) {
+          address rewardToken = _getRewardsTokenAddress(index);
+          if (rewardToken == address(0)) break;
+          if (rewardToken == CRV_TOKEN) continue;
+          priorTokenBalances[index] = IERC20(rewardToken).balanceOf(address(this));
+        }
+        // Mint other rewards to aToken
+        ICurveTreasury(CURVE_TREASURY).withdraw(token, amount, true);
+
+        for (uint256 index = 1; index < MAX_REWARD_TOKENS; index++) {
+          address rewardToken = _getRewardsTokenAddress(index);
+          if (rewardToken == address(0)) break;
+          if (rewardToken == CRV_TOKEN) continue;
+          uint256 balance = IERC20(rewardToken).balanceOf(address(this));
+          _pendingRewards[rewardToken] = _pendingRewards[rewardToken].add(
+            balance.sub(priorTokenBalances[index])
+          );
+        }
+      } else {
+        ICurveTreasury(CURVE_TREASURY).withdraw(token, amount, true);
+      }
     }
     return amount;
   }

contracts/adapters/rewards/curve/CurveTreasury.sol

@@ -68,7 +68,6 @@ contract CurveTreasury is ICurveTreasury, VersionedInitializable {
    * @dev Revert if caller and selected token is not a whitelisted entity
    */
   modifier onlyWhitelistedEntity(address token) {
-    console.log(msg.sender, token, _entityTokenWhitelist[msg.sender][token]);
     require(_entityTokenWhitelist[msg.sender][token] == true, 'ENTITY_NOT_WHITELISTED');
     _;
   }
@@ -98,7 +97,35 @@ contract CurveTreasury is ICurveTreasury, VersionedInitializable {
     IERC20(token).safeTransferFrom(msg.sender, address(this), amount);
 
     if (useGauge && _entityTokenGauge[msg.sender][token] != address(0)) {
-      _stakeGauge(_entityTokenGauge[msg.sender][token], amount);
+      address gauge = _entityTokenGauge[msg.sender][token];
+      if (_isGaugeV2Compatible[gauge] == true) {
+        // Claim the extra rewards from Gauge Staking
+        uint256[] memory priorRewardsBalance = new uint256[](MAX_REWARD_TOKENS);
+        uint256[] memory afterRewardsBalance = new uint256[](MAX_REWARD_TOKENS);
+
+        // Calculate balances prior claiming rewards
+        for (uint256 index = 1; index < MAX_REWARD_TOKENS; index++) {
+          address rewardToken = ICurveGaugeView(gauge).reward_tokens(index - 1);
+          if (rewardToken == address(0)) break;
+          priorRewardsBalance[index] = IERC20(rewardToken).balanceOf(address(this));
+        }
+
+        // Claim extra rewards
+        _stakeGauge(_entityTokenGauge[msg.sender][token], amount);
+
+        // Transfer extra rewards to entity
+        for (uint256 index = 1; index < MAX_REWARD_TOKENS; index++) {
+          address rewardToken = ICurveGaugeView(gauge).reward_tokens(index - 1);
+          if (rewardToken == address(0)) break;
+          afterRewardsBalance[index] = IERC20(rewardToken).balanceOf(address(this));
+          uint256 rewardsAmount = afterRewardsBalance[index].sub(priorRewardsBalance[index]);
+          if (rewardsAmount > 0) {
+            IERC20(rewardToken).safeTransfer(msg.sender, rewardsAmount);
+          }
+        }
+      } else {
+        _stakeGauge(_entityTokenGauge[msg.sender][token], amount);
+      }
     }
   }
 
@@ -109,7 +136,35 @@ contract CurveTreasury is ICurveTreasury, VersionedInitializable {
     bool useGauge
   ) external override onlyWhitelistedEntity(token) {
     if (useGauge && _entityTokenGauge[msg.sender][token] != address(0)) {
-      _unstakeGauge(_entityTokenGauge[msg.sender][token], amount);
+      address gauge = _entityTokenGauge[msg.sender][token];
+      if (_isGaugeV2Compatible[gauge] == true) {
+        // Claim the extra rewards from Gauge Staking
+        uint256[] memory priorRewardsBalance = new uint256[](MAX_REWARD_TOKENS);
+        uint256[] memory afterRewardsBalance = new uint256[](MAX_REWARD_TOKENS);
+
+        // Calculate balances prior claiming rewards
+        for (uint256 index = 1; index < MAX_REWARD_TOKENS; index++) {
+          address rewardToken = ICurveGaugeView(gauge).reward_tokens(index - 1);
+          if (rewardToken == address(0)) break;
+          priorRewardsBalance[index] = IERC20(rewardToken).balanceOf(address(this));
+        }
+
+        // Claim extra rewards
+        _unstakeGauge(_entityTokenGauge[msg.sender][token], amount);
+
+        // Transfer extra rewards to entity
+        for (uint256 index = 1; index < MAX_REWARD_TOKENS; index++) {
+          address rewardToken = ICurveGaugeView(gauge).reward_tokens(index - 1);
+          if (rewardToken == address(0)) break;
+          afterRewardsBalance[index] = IERC20(rewardToken).balanceOf(address(this));
+          uint256 rewardsAmount = afterRewardsBalance[index].sub(priorRewardsBalance[index]);
+          if (rewardsAmount > 0) {
+            IERC20(rewardToken).safeTransfer(msg.sender, rewardsAmount);
+          }
+        }
+      } else {
+        _unstakeGauge(_entityTokenGauge[msg.sender][token], amount);
+      }
     }
     IERC20(token).safeTransfer(msg.sender, amount);
   }

test-suites/test-aave/helpers/rewards-distribution/verify.ts

@@ -107,6 +107,4 @@ export const checkRewards = async (
       expect(userRewardsBefore[i]).to.be.eq(userRewardsAfter[i], 'Rewards should stay the same');
     }
   }
-
-  return;
 };

test-suites/test-aave/mainnet/atoken-curve-rewards.main.ts

@@ -333,7 +333,6 @@ makeSuite('Curve Rewards Aware aToken', (testEnv: TestEnv) => {
     curveTreasury = CurveTreasuryFactory.connect(curveTreasuryAddress, testEnv.users[0].signer);
 
     // Enable atoken entities into Curve Treasury
-    console.log(a3POOL.address, aEURS.address, aAAVE3.address, aANKR.address);
     await waitForTx(
       await curveTreasury.setWhitelist(
         [a3POOL.address, aEURS.address, aAAVE3.address, aANKR.address],
@@ -466,11 +465,18 @@ makeSuite('Curve Rewards Aware aToken', (testEnv: TestEnv) => {
     it('Deposit and generate user reward checkpoints', async () => {
       // Deposits
 
-      await depositPoolToken(depositor, GAUGE_EURS, aEURS.address, parseEther('100000'));
+      await depositPoolToken(depositor, GAUGE_EURS, aEURS.address, parseEther('10000'), false);
       const curveATokenBalance = await crvToken.balanceOf(aEURS.address);
       expect(curveATokenBalance).to.be.eq('0', 'CRV should be zero');
     });
 
+    it('Second deposit should add user reward checkpoints', async () => {
+      // Pass time to generate rewards
+      await advanceTimeAndBlock(ONE_DAY * 14);
+
+      await depositPoolToken(depositor, GAUGE_EURS, aEURS.address, parseEther('10000'), true);
+    });
+
     it('Increase time and claim CRV and SNX', async () => {
       // Pass time to generate rewards
       await advanceTimeAndBlock(ONE_DAY * 14);
@@ -484,7 +490,7 @@ makeSuite('Curve Rewards Aware aToken', (testEnv: TestEnv) => {
       await increaseTime(ONE_DAY);
 
       // Withdraw
-      await withdrawPoolToken(depositor, GAUGE_EURS, aEURS.address);
+      await withdrawPoolToken(depositor, GAUGE_EURS, aEURS.address, true);
     });
 
     it('Claim the remaining CRV and SNX', async () => {