Added transfer of the funds to aTokens

contracts/flashloan/base/FlashLoanReceiverBase.sol

@@ -21,14 +21,13 @@ abstract contract FlashLoanReceiverBase is IFlashLoanReceiver {
 
     receive() external payable {}
 
-    function transferFundsBackToPoolInternal(address _reserve, uint256 _amount) internal {
+    function transferFundsBackInternal(address _reserve, address _destination, uint256 _amount) internal {
 
-        address payable pool = payable(addressesProvider.getLendingPool());
-
-        transferInternal(pool,_reserve, _amount);
+        transferInternal(payable(_destination),_reserve, _amount);
     }
 
     function transferInternal(address payable _destination, address _reserve, uint256  _amount) internal {
+        
         if(_reserve == EthAddressLib.ethAddress()) {
             //solium-disable-next-line
             _destination.call{value: _amount}("");

contracts/flashloan/interfaces/IFlashLoanReceiver.sol

@@ -9,5 +9,5 @@ pragma solidity ^0.6.8;
 **/
 interface IFlashLoanReceiver {
 
-    function executeOperation(address _reserve, uint256 _amount, uint256 _fee, bytes calldata _params) external;
+    function executeOperation(address _reserve, address _destination, uint256 _amount, uint256 _fee, bytes calldata _params) external;
 }

contracts/lendingpool/LendingPool.sol

@@ -261,12 +261,13 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
     uint256 _amount,
     uint16 _referralCode
   ) external payable nonReentrant {
+
     ReserveLogic.ReserveData storage reserve = reserves[_reserve];
     UserLogic.UserReserveData storage user = usersReserveData[msg.sender][_reserve];
 
     ValidationLogic.validateDeposit(reserve, _amount);
 
-    AToken aToken = AToken(reserve.aTokenAddress);
+    AToken aToken = AToken(payable(reserve.aTokenAddress));
 
     bool isFirstDeposit = aToken.balanceOf(msg.sender) == 0;
 
@@ -280,8 +281,8 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
     //minting AToken to user 1:1 with the specific exchange rate
     aToken.mintOnDeposit(msg.sender, _amount);
 
-    //transfer to the core contract
-    IERC20(_reserve).universalTransferFromSenderToThis(_amount, true);
+    //transfer to the aToken contract
+    IERC20(_reserve).universalTransferFrom(msg.sender, address(aToken), _amount, true);
 
     //solium-disable-next-line
     emit Deposit(_reserve, msg.sender, _amount, _referralCode, block.timestamp);
@@ -303,16 +304,20 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
     ReserveLogic.ReserveData storage reserve = reserves[_reserve];
     UserLogic.UserReserveData storage user = usersReserveData[_user][_reserve];
 
+    AToken aToken = AToken(payable(reserve.aTokenAddress));
+
     ValidationLogic.validateRedeem(reserve, _reserve, _amount);
 
+    reserve.updateCumulativeIndexesAndTimestamp();
+
+    reserve.updateInterestRates(_reserve, 0, _amount);
+
     if (_aTokenBalanceAfterRedeem == 0) {
       user.useAsCollateral = false;
     }
 
-    reserve.updateCumulativeIndexesAndTimestamp();
-    reserve.updateInterestRates(_reserve, 0, _amount);
-
-    IERC20(_reserve).universalTransfer(_user, _amount);
+    
+    AToken(reserve.aTokenAddress).transferUnderlyingTo(_user, _amount);
 
     //solium-disable-next-line
     emit RedeemUnderlying(_reserve, _user, _amount, block.timestamp);
@@ -372,7 +377,7 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
     reserve.updateInterestRates(_reserve, 0, _amount);
 
     //if we reached this point, we can transfer
-    IERC20(_reserve).universalTransfer(msg.sender, _amount);
+    AToken(reserve.aTokenAddress).transferUnderlyingTo(msg.sender, _amount);
 
     emit Borrow(
       _reserve,
@@ -453,7 +458,7 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
 
     reserve.updateInterestRates(_reserve, vars.paybackAmount, 0);
 
-    IERC20(_reserve).universalTransferFromSenderToThis(vars.paybackAmount, false);
+    IERC20(_reserve).universalTransferFrom(msg.sender, reserve.aTokenAddress, vars.paybackAmount, false);
 
     if (IERC20(_reserve).isETH()) {
       //send excess ETH back to the caller if needed
@@ -651,6 +656,7 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
     uint256 protocolFeeBips;
     uint256 amountFee;
     uint256 protocolFee;
+    address payable aTokenAddress;
   }
 
   /**
@@ -667,12 +673,15 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
     uint256 _amount,
     bytes memory _params
   ) public nonReentrant {
+
     FlashLoanLocalVars memory vars;
 
     ReserveLogic.ReserveData storage reserve = reserves[_reserve];
 
+    vars.aTokenAddress = payable(reserve.aTokenAddress);
+
     //check that the reserve has enough available liquidity
-    vars.availableLiquidityBefore = IERC20(_reserve).universalBalanceOf(address(this));
+    vars.availableLiquidityBefore = IERC20(_reserve).universalBalanceOf(vars.aTokenAddress);
 
     //calculate amount fee
     vars.amountFee = _amount.mul(FLASHLOAN_FEE_TOTAL).div(10000);
@@ -680,6 +689,7 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
     //protocol fee is the part of the amountFee reserved for the protocol - the rest goes to depositors
     vars.protocolFee = vars.amountFee.mul(FLASHLOAN_FEE_PROTOCOL).div(10000);
 
+
     require(
       vars.availableLiquidityBefore >= _amount,
       'There is not enough liquidity available to borrow'
@@ -695,13 +705,13 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
     address payable userPayable = address(uint160(_receiver));
 
     //transfer funds to the receiver
-    IERC20(_reserve).universalTransfer(userPayable, _amount);
+    AToken(vars.aTokenAddress).transferUnderlyingTo(userPayable, _amount);
 
     //execute action of the receiver
-    receiver.executeOperation(_reserve, _amount, vars.amountFee, _params);
+    receiver.executeOperation(_reserve, vars.aTokenAddress, _amount, vars.amountFee, _params);
 
     //check that the actual balance of the core contract includes the returned amount
-    uint256 availableLiquidityAfter = IERC20(_reserve).universalBalanceOf(address(this));
+    uint256 availableLiquidityAfter = IERC20(_reserve).universalBalanceOf(vars.aTokenAddress);
 
     require(
       availableLiquidityAfter == vars.availableLiquidityBefore.add(vars.amountFee),
@@ -791,7 +801,7 @@ contract LendingPool is ReentrancyGuard, VersionedInitializable {
   {
     ReserveLogic.ReserveData memory reserve = reserves[_reserve];
     return (
-      IERC20(_reserve).universalBalanceOf(address(this)),
+      IERC20(_reserve).universalBalanceOf(reserve.aTokenAddress),
       IERC20(reserve.stableDebtTokenAddress).totalSupply(),
       IERC20(reserve.variableDebtTokenAddress).totalSupply(),
       reserve.currentLiquidityRate,

contracts/lendingpool/LendingPoolLiquidationManager.sol

@@ -192,6 +192,9 @@ contract LendingPoolLiquidationManager is ReentrancyGuard, VersionedInitializabl
             vars.userCollateralBalance
         );
 
+        vars.collateralAtoken = AToken(payable(collateralReserve.aTokenAddress));
+
+
         //if principalAmountNeeded < vars.ActualAmountToLiquidate, there isn't enough
         //of _collateral to cover the actual amount that is being liquidated, hence we liquidate
         //a smaller amount
@@ -202,7 +205,7 @@ contract LendingPoolLiquidationManager is ReentrancyGuard, VersionedInitializabl
 
         //if liquidator reclaims the underlying asset, we make sure there is enough available collateral in the reserve
         if (!_receiveAToken) {
-            uint256 currentAvailableCollateral = IERC20(_collateral).universalBalanceOf(address(this));
+            uint256 currentAvailableCollateral = IERC20(_collateral).universalBalanceOf(address(vars.collateralAtoken));
             if (currentAvailableCollateral < vars.maxCollateralToLiquidate) {
                 return (
                     uint256(LiquidationErrors.NOT_ENOUGH_LIQUIDITY),
@@ -219,9 +222,7 @@ contract LendingPoolLiquidationManager is ReentrancyGuard, VersionedInitializabl
             IVariableDebtToken(principalReserve.variableDebtTokenAddress).burn(_user, vars.userVariableDebt);
             IStableDebtToken(principalReserve.stableDebtTokenAddress).burn(_user, vars.actualAmountToLiquidate.sub(vars.userVariableDebt));
         }
-
-        vars.collateralAtoken = AToken(collateralReserve.aTokenAddress);
-
+  
         //if liquidator reclaims the aToken, he receives the equivalent atoken amount
         if (_receiveAToken) {
             vars.collateralAtoken.transferOnLiquidation(
@@ -233,13 +234,11 @@ contract LendingPoolLiquidationManager is ReentrancyGuard, VersionedInitializabl
             //otherwise receives the underlying asset
             //burn the equivalent amount of atoken
             vars.collateralAtoken.burnOnLiquidation(_user, vars.maxCollateralToLiquidate);
-   
-            IERC20(_collateral).universalTransfer(msg.sender, vars.maxCollateralToLiquidate);
+            vars.collateralAtoken.transferUnderlyingTo(msg.sender, vars.maxCollateralToLiquidate);
         }
 
-        //transfers the principal currency to the pool
-        IERC20(_reserve).universalTransferFromSenderToThis(vars.actualAmountToLiquidate, true);
-
+        //transfers the principal currency to the aToken
+        IERC20(_reserve).universalTransferFrom(msg.sender, principalReserve.aTokenAddress, vars.actualAmountToLiquidate, true);
 
         emit LiquidationCall(
             _collateral,

contracts/libraries/ReserveLogic.sol

@@ -15,6 +15,7 @@ import '../interfaces/ILendingRateOracle.sol';
 import '../interfaces/IReserveInterestRateStrategy.sol';
 import '../tokenization/AToken.sol';
 import './WadRayMath.sol';
+import '@nomiclabs/buidler/console.sol';
 
 /**
  * @title ReserveLogic library
@@ -56,7 +57,7 @@ library ReserveLogic {
     /**
      * @dev address of the aToken representing the asset
      **/
-    address aTokenAddress;
+    address payable aTokenAddress;
     address stableDebtTokenAddress;
     address variableDebtTokenAddress;
     /**
@@ -193,7 +194,7 @@ library ReserveLogic {
       _self.lastVariableBorrowCumulativeIndex = WadRayMath.ray();
     }
 
-    _self.aTokenAddress = _aTokenAddress;
+    _self.aTokenAddress = payable(_aTokenAddress);
     _self.stableDebtTokenAddress = _stableDebtAddress;
     _self.variableDebtTokenAddress = _variableDebtAddress;
     _self.decimals = _decimals;
@@ -353,15 +354,11 @@ library ReserveLogic {
     uint256 _liquidityAdded,
     uint256 _liquidityTaken
   ) internal {
+
     uint256 currentAvgStableRate = IStableDebtToken(_reserve.stableDebtTokenAddress)
       .getAverageStableRate();
 
-    uint256 balance = IERC20(_reserveAddress).universalBalanceOf(address(this));
-
-    //if the reserve is ETH, the msg.value has already been cumulated to the balance of the reserve
-    if (IERC20(_reserveAddress).isETH()) {
-      balance = balance.sub(msg.value);
-    }
+    uint256 balance = IERC20(_reserveAddress).universalBalanceOf(_reserve.aTokenAddress);
 
     (
       uint256 newLiquidityRate,

contracts/libraries/ValidationLogic.sol

@@ -53,7 +53,8 @@ library ValidationLogic {
 
     require(msg.sender == _reserve.aTokenAddress, '31');
 
-    uint256 currentAvailableLiquidity = IERC20(_reserveAddress).universalBalanceOf(address(this));
+    uint256 currentAvailableLiquidity = IERC20(_reserveAddress).universalBalanceOf(address(_reserve.aTokenAddress));
+
     require(currentAvailableLiquidity >= _amount, '4');
   }
 
@@ -117,7 +118,7 @@ library ValidationLogic {
     );
 
     //check that the amount is available in the reserve
-    vars.availableLiquidity = IERC20(_reserveAddress).universalBalanceOf(address(this));
+    vars.availableLiquidity = IERC20(_reserveAddress).universalBalanceOf(address(_reserve.aTokenAddress));
 
     require(vars.availableLiquidity >= _amount, '7');
 

contracts/misc/AaveProtocolTestHelpers.sol

@@ -38,7 +38,7 @@ contract AaveProtocolTestHelpers {
         for (uint256 i = 0; i < reserves.length; i++) {
             (address aTokenAddress,,) = pool.getReserveTokensAddresses(reserves[i]);
             aTokens[i] = TokenData({
-                symbol: AToken(aTokenAddress).symbol(),
+                symbol: AToken(payable(aTokenAddress)).symbol(),
                 tokenAddress: aTokenAddress
             });
         }

contracts/mocks/flashloan/MockFlashLoanReceiver.sol

@@ -24,6 +24,7 @@ contract MockFlashLoanReceiver is FlashLoanReceiverBase {
 
     function executeOperation(
         address _reserve,
+        address _destination,
         uint256 _amount,
         uint256 _fee,
         bytes memory _params) public override {
@@ -46,7 +47,7 @@ contract MockFlashLoanReceiver is FlashLoanReceiverBase {
             token.mint(_fee);
         }
         //returning amount + fee to the destination
-        transferFundsBackToPoolInternal(_reserve, _amount.add(_fee));
+        transferFundsBackInternal(_reserve, _destination, _amount.add(_fee));
         emit ExecutedWithSuccess(_reserve, _amount, _fee);
     }
 }