fix: use AToken balances to update virtual balances on deposit/withdr…

…aw (#86) * fix: use AToken balances to update virtual balances on deposit/withdraw * fix: Refactor code * fix: Use aTokenBalance directly for virtual balance update * fix: Fix the update virtual balance after withdrawFees --------- Co-authored-by: cedephrase <[email protected]>
aave · Jun 20, 2023 · 23366cc · 23366cc
1 parent 611c20f
commit 23366cc
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/src/ATokenVault.sol b/src/ATokenVault.sol
@@ -422,10 +422,11 @@ contract ATokenVault is ERC4626Upgradeable, OwnableUpgradeable, EIP712Upgradeabl
         require(amount <= _s.accumulatedFees, "INSUFFICIENT_FEES"); // will underflow below anyway, error msg for clarity
 
         _s.accumulatedFees -= uint128(amount);
-        _s.lastVaultBalance -= uint128(amount);
 
         ATOKEN.transfer(to, amount);
 
+        _s.lastVaultBalance = uint128(ATOKEN.balanceOf(address(this)));
+
         emit FeesWithdrawn(to, amount, _s.lastVaultBalance, _s.accumulatedFees);
     }
 
@@ -617,13 +618,11 @@ contract ATokenVault is ERC4626Upgradeable, OwnableUpgradeable, EIP712Upgradeabl
         // Need to transfer before minting or ERC777s could reenter.
         if (asAToken) {
             ATOKEN.transferFrom(depositor, address(this), assets);
-            _s.lastVaultBalance += uint128(assets);
         } else {
             UNDERLYING.safeTransferFrom(depositor, address(this), assets);
-            uint256 aTokenBalanceBefore = ATOKEN.balanceOf(address(this));
             AAVE_POOL.supply(address(UNDERLYING), assets, address(this), REFERRAL_CODE);
-            _s.lastVaultBalance += uint128(ATOKEN.balanceOf(address(this)) - aTokenBalanceBefore);
         }
+        _s.lastVaultBalance = uint128(ATOKEN.balanceOf(address(this)));
 
         _mint(receiver, shares);
 
@@ -647,11 +646,10 @@ contract ATokenVault is ERC4626Upgradeable, OwnableUpgradeable, EIP712Upgradeabl
         // Withdraw assets from Aave v3 and send to receiver
         if (asAToken) {
             ATOKEN.transfer(receiver, assets);
-            _s.lastVaultBalance -= uint128(assets);
         } else {
-            uint256 amountWithdrawn = AAVE_POOL.withdraw(address(UNDERLYING), assets, receiver);
-            _s.lastVaultBalance -= uint128(amountWithdrawn);
+            AAVE_POOL.withdraw(address(UNDERLYING), assets, receiver);
         }
+        _s.lastVaultBalance = uint128(ATOKEN.balanceOf(address(this)));
 
         emit Withdraw(allowanceTarget, receiver, owner, assets, shares);
     }