From 0d3b155eeb39a0a8efd05b368ca30a7c109d8384 Mon Sep 17 00:00:00 2001
From: David Racero <canillas.mail@gmail.com>
Date: Thu, 3 Jun 2021 14:56:42 +0200
Subject: [PATCH] feat: RewardsAwareAToken replace transfer to safeTransfer at
 rewards, add early return at updateDistribution if address is 0x0, remove
 unneeded updateDistribution calls at mint and burn.

---
 .../tokenization/RewardsAwareAToken.sol       | 12 ++++++------
 .../test-aave/rewardAwareAToken.spec.ts       | 19 +++++++++++++++++--
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/contracts/protocol/tokenization/RewardsAwareAToken.sol b/contracts/protocol/tokenization/RewardsAwareAToken.sol
index 32775d35..6f5f60ab 100644
--- a/contracts/protocol/tokenization/RewardsAwareAToken.sol
+++ b/contracts/protocol/tokenization/RewardsAwareAToken.sol
@@ -10,7 +10,6 @@ import {SafeERC20} from '../../dependencies/openzeppelin/contracts/SafeERC20.sol
 import {IERC20} from '../../dependencies/openzeppelin/contracts/IERC20.sol';
 import {IRewardsAwareAToken} from '../../interfaces/IRewardsAwareAToken.sol';
 import {IAToken} from '../../interfaces/IRewardsAwareAToken.sol';
-import {WadRayMath} from '../../protocol/libraries/math/WadRayMath.sol';
 import {PercentageMath} from '../../protocol/libraries/math/PercentageMath.sol';
 
 /**
@@ -20,7 +19,6 @@ import {PercentageMath} from '../../protocol/libraries/math/PercentageMath.sol';
  */
 abstract contract RewardsAwareAToken is AToken, IRewardsAwareAToken {
   using SafeERC20 for IERC20;
-  using WadRayMath for uint256;
   using PercentageMath for uint256;
 
   // Precision of the multiplier for calculating distribution percentages
@@ -172,7 +170,7 @@ abstract contract RewardsAwareAToken is AToken, IRewardsAwareAToken {
     _unstake(UNDERLYING_ASSET_ADDRESS(), amount);
 
     // Update distribution of rewards
-    _updateDistribution(user);
+    // _updateDistribution(user);
 
     // burns aTokens
     return super.burn(user, receiverOfUnderlying, amount, index);
@@ -195,7 +193,7 @@ abstract contract RewardsAwareAToken is AToken, IRewardsAwareAToken {
     _stake(UNDERLYING_ASSET_ADDRESS(), amount);
 
     // Update distribution of rewards
-    _updateDistribution(user);
+    // _updateDistribution(user);
 
     // mint aTokens
     return super.mint(user, amount, index);
@@ -384,6 +382,8 @@ abstract contract RewardsAwareAToken is AToken, IRewardsAwareAToken {
    * @param user The `user` address to update the rewards index and retrieve latest reward distribution
    */
   function _updateDistribution(address user) internal {
+    if (user == address(0)) return;
+
     uint256 aTokenBalance = balanceOf(user);
     if (aTokenBalance == 0) {
       return;
@@ -444,12 +444,12 @@ abstract contract RewardsAwareAToken is AToken, IRewardsAwareAToken {
         uint256 reserveRewards = unstaked.percentMul(rewardsReserveFactor);
         if (reserveRewards > 0) {
           userRewards = unstaked.sub(reserveRewards);
-          IERC20(token).transfer(_treasury, reserveRewards);
+          IERC20(token).safeTransfer(_treasury, reserveRewards);
         }
       }
 
       // Transfer rewards to user
-      IERC20(token).transfer(msg.sender, userRewards);
+      IERC20(token).safeTransfer(msg.sender, userRewards);
       emit Claim(msg.sender, token, accruedRewards, userRewards);
     }
   }
diff --git a/test-suites/test-aave/rewardAwareAToken.spec.ts b/test-suites/test-aave/rewardAwareAToken.spec.ts
index 75f6be2f..d1c0880a 100644
--- a/test-suites/test-aave/rewardAwareAToken.spec.ts
+++ b/test-suites/test-aave/rewardAwareAToken.spec.ts
@@ -259,8 +259,13 @@ makeSuite('Reward Aware AToken', (testEnv: TestEnv) => {
           `\taRew Total supply: ${formatEther(
             await aRew.totalSupply()
           )}. Rewards in contract: ${formatEther(
-            await rew.balanceOf(aRew.address)
-          )}. Summed claimable: ${formatEther(summedClaimable)}`
+            await (await rew.balanceOf(aRew.address)).sub(await aRew.totalSupply())
+          )}. Summed claimable: ${formatEther(summedClaimable)}
+          . Summed aRew REW balance + claimable REW: ${formatEther(
+            await (await rew.balanceOf(aRew.address))
+              .sub(await aRew.totalSupply())
+              .add(await rew.getClaimableRewards(aRew.address))
+          )}`
         );
       };
 
@@ -289,15 +294,25 @@ makeSuite('Reward Aware AToken', (testEnv: TestEnv) => {
 
       console.log(`Withdraws`);
       for (let i = 0; i < users.length; i++) {
+        console.log('--p');
+        await printState();
+        console.log('--p');
         await pool
           .connect(users[i].signer)
           .withdraw(rew.address, MAX_UINT_AMOUNT, users[i].address);
+        console.log('--a');
+        await printState();
+        console.log('--a');
       }
 
       // Claims and check rewards
       console.log(`Claim`);
       for (let i = 0; i < users.length; i++) {
+        console.log('--c');
+        await printState();
+        console.log('--c');
         await claim(users[i]);
+        console.log('claimed', users[i].address);
       }
 
       await printState();