fix: ATokenVaultFactory safe approval (#106)

donosonaumczuk · web-flow · commit c58b21d34350 · 2025-09-01T22:08:11.000+01:00
* test: Non-standard ERC-20 USDT mock added

* test: Use SafeERC20 for MockAToken

* test: MockAavePool - Allow to mock multiple assets and use SafeERC20

* test: Adapt old tests to the new Pool mock

* fix: use safeApprove for underlying approval in the factory

* test: Refactor ATokenVaultFactory test to allow overriding the underlying

* test: Factory test created to run all against USDT non-standard ERC-20
diff --git a/src/ATokenVaultFactory.sol b/src/ATokenVaultFactory.sol
@@ -133,7 +133,7 @@ contract ATokenVaultFactory {
             ""
         ));
 
-        IERC20(params.underlying).approve(vault, params.initialLockDeposit);
+        IERC20(params.underlying).safeApprove(vault, params.initialLockDeposit);
 
         ATokenVault(vault).initialize(
             params.owner,
diff --git a/test/ATokenVaultFactory.t.sol b/test/ATokenVaultFactory.t.sol
diff --git a/test/ATokenVaultFactoryNonStandardERC20.t.sol b/test/ATokenVaultFactoryNonStandardERC20.t.sol
@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.10;
+
+import {MockUSDT} from "./mocks/MockUSDT.sol";
+
+contract ATokenVaultFactoryNonStandardERC20Test {
+    
+    function _deployUnderlying() internal virtual returns (address) {
+        return address(new MockUSDT());
+    }
+}
diff --git a/test/ATokenVaultFork.t.sol b/test/ATokenVaultFork.t.sol
@@ -506,7 +506,8 @@ contract ATokenVaultForkTest is ATokenVaultForkBaseTest {
 
     function testDepositFailsWithExceedsMax() public {
         // mock call to Aave Pool
-        MockAavePool mp = new MockAavePool(new MockAToken(address(dai)));
+        MockAavePool mp = new MockAavePool();
+        mp.mockReserve(address(dai), new MockAToken(address(dai)));
         mp.setReserveConfigMap(RESERVE_CONFIG_MAP_INACTIVE);
         vm.mockCall(
             address(vault.AAVE_POOL()),
@@ -564,7 +565,8 @@ contract ATokenVaultForkTest is ATokenVaultForkBaseTest {
 
     function testDepositATokensWithExceedsMax() public {
         // mock call to Aave Pool
-        MockAavePool mp = new MockAavePool(new MockAToken(address(dai)));
+        MockAavePool mp = new MockAavePool();
+        mp.mockReserve(address(dai), new MockAToken(address(dai)));
         mp.setReserveConfigMap(RESERVE_CONFIG_MAP_INACTIVE);
         vm.mockCall(
             address(vault.AAVE_POOL()),
@@ -619,7 +621,8 @@ contract ATokenVaultForkTest is ATokenVaultForkBaseTest {
 
     function testMintFailsWithExceedsMax() public {
         // mock call to Aave Pool
-        MockAavePool mp = new MockAavePool(new MockAToken(address(dai)));
+        MockAavePool mp = new MockAavePool();
+        mp.mockReserve(address(dai), new MockAToken(address(dai)));
         mp.setReserveConfigMap(RESERVE_CONFIG_MAP_INACTIVE);
         vm.mockCall(
             address(vault.AAVE_POOL()),
@@ -658,7 +661,8 @@ contract ATokenVaultForkTest is ATokenVaultForkBaseTest {
 
     function testMintATokensWithExceedsMax() public {
         // mock call to Aave Pool
-        MockAavePool mp = new MockAavePool(new MockAToken(address(dai)));
+        MockAavePool mp = new MockAavePool();
+        mp.mockReserve(address(dai), new MockAToken(address(dai)));
         mp.setReserveConfigMap(RESERVE_CONFIG_MAP_INACTIVE);
         vm.mockCall(
             address(vault.AAVE_POOL()),
diff --git a/test/ATokenVaultMocks.t.sol b/test/ATokenVaultMocks.t.sol
@@ -21,7 +21,8 @@ contract ATokenVaultMocksTest is ATokenVaultBaseTest {
     function setUp() public override {
         dai = new MockDAI();
         aDai = new MockAToken(address(dai));
-        pool = new MockAavePool(aDai);
+        pool = new MockAavePool();
+        pool.mockReserve(address(dai), aDai);
         poolAddrProvider = new MockAavePoolAddressesProvider(address(pool));
 
         pool.setReserveConfigMap(RESERVE_CONFIG_MAP_UNCAPPED_ACTIVE);
diff --git a/test/ATokenVaultProperties.t.sol b/test/ATokenVaultProperties.t.sol
@@ -22,7 +22,8 @@ contract ATokenVaultPropertiesTest is ERC4626Test, ATokenVaultBaseTest {
     function setUp() public override(ERC4626Test, ATokenVaultBaseTest) {
         dai = new MockDAI();
         aDai = new MockAToken(address(dai));
-        pool = new MockAavePool(aDai);
+        pool = new MockAavePool();
+        pool.mockReserve(address(dai), aDai);
         poolAddrProvider = new MockAavePoolAddressesProvider(address(pool));
 
         pool.setReserveConfigMap(RESERVE_CONFIG_MAP_UNCAPPED_ACTIVE);
diff --git a/test/ATokenVaultWithSig.t.sol b/test/ATokenVaultWithSig.t.sol
@@ -28,7 +28,8 @@ contract ATokenVaultWithSigTest is ATokenVaultBaseTest {
         dai = new MockDAI();
 
         aDai = new MockAToken(address(dai));
-        pool = new MockAavePool(aDai);
+        pool = new MockAavePool();
+        pool.mockReserve(address(dai), aDai);
         poolAddrProvider = new MockAavePoolAddressesProvider(address(pool));
 
         vaultAssetAddress = address(aDai);
diff --git a/test/ImmutableATokenVaultFork.t.sol b/test/ImmutableATokenVaultFork.t.sol
@@ -509,7 +509,8 @@ contract ImmutableATokenVaultForkTest is ImmutableATokenVaultForkBaseTest {
 
     function testDepositFailsWithExceedsMax() public {
         // mock call to Aave Pool
-        MockAavePool mp = new MockAavePool(new MockAToken(address(dai)));
+        MockAavePool mp = new MockAavePool();
+        mp.mockReserve(address(dai), new MockAToken(address(dai)));
         mp.setReserveConfigMap(RESERVE_CONFIG_MAP_INACTIVE);
         vm.mockCall(
             address(vault.AAVE_POOL()),
@@ -567,7 +568,8 @@ contract ImmutableATokenVaultForkTest is ImmutableATokenVaultForkBaseTest {
 
     function testDepositATokensWithExceedsMax() public {
         // mock call to Aave Pool
-        MockAavePool mp = new MockAavePool(new MockAToken(address(dai)));
+        MockAavePool mp = new MockAavePool();
+        mp.mockReserve(address(dai), new MockAToken(address(dai)));
         mp.setReserveConfigMap(RESERVE_CONFIG_MAP_INACTIVE);
         vm.mockCall(
             address(vault.AAVE_POOL()),
@@ -622,7 +624,8 @@ contract ImmutableATokenVaultForkTest is ImmutableATokenVaultForkBaseTest {
 
     function testMintFailsWithExceedsMax() public {
         // mock call to Aave Pool
-        MockAavePool mp = new MockAavePool(new MockAToken(address(dai)));
+        MockAavePool mp = new MockAavePool();
+        mp.mockReserve(address(dai), new MockAToken(address(dai)));
         mp.setReserveConfigMap(RESERVE_CONFIG_MAP_INACTIVE);
         vm.mockCall(
             address(vault.AAVE_POOL()),
@@ -661,7 +664,8 @@ contract ImmutableATokenVaultForkTest is ImmutableATokenVaultForkBaseTest {
 
     function testMintATokensWithExceedsMax() public {
         // mock call to Aave Pool
-        MockAavePool mp = new MockAavePool(new MockAToken(address(dai)));
+        MockAavePool mp = new MockAavePool();
+        mp.mockReserve(address(dai), new MockAToken(address(dai)));
         mp.setReserveConfigMap(RESERVE_CONFIG_MAP_INACTIVE);
         vm.mockCall(
             address(vault.AAVE_POOL()),
diff --git a/test/ImmutableATokenVaultMocks.t.sol b/test/ImmutableATokenVaultMocks.t.sol
@@ -18,7 +18,8 @@ contract ImmutableATokenVaultMocksTest is ImmutableATokenVaultBaseTest {
     function setUp() public override {
         dai = new MockDAI();
         aDai = new MockAToken(address(dai));
-        pool = new MockAavePool(aDai);
+        pool = new MockAavePool();
+        pool.mockReserve(address(dai), aDai);
         poolAddrProvider = new MockAavePoolAddressesProvider(address(pool));
 
         pool.setReserveConfigMap(RESERVE_CONFIG_MAP_UNCAPPED_ACTIVE);
diff --git a/test/ImmutableATokenVaultProperties.t.sol b/test/ImmutableATokenVaultProperties.t.sol
@@ -20,7 +20,8 @@ contract ImmutableATokenVaultPropertiesTest is ERC4626Test, ImmutableATokenVault
     function setUp() public override(ERC4626Test, ImmutableATokenVaultBaseTest) {
         dai = new MockDAI();
         aDai = new MockAToken(address(dai));
-        pool = new MockAavePool(aDai);
+        pool = new MockAavePool();
+        pool.mockReserve(address(dai), aDai);
         poolAddrProvider = new MockAavePoolAddressesProvider(address(pool));
 
         pool.setReserveConfigMap(RESERVE_CONFIG_MAP_UNCAPPED_ACTIVE);
diff --git a/test/ImmutableATokenVaultWithSig.t.sol b/test/ImmutableATokenVaultWithSig.t.sol
@@ -27,7 +27,8 @@ contract ImmutableATokenVaultWithSigTest is ImmutableATokenVaultBaseTest {
         dai = new MockDAI();
 
         aDai = new MockAToken(address(dai));
-        pool = new MockAavePool(aDai);
+        pool = new MockAavePool();
+        pool.mockReserve(address(dai), aDai);
         poolAddrProvider = new MockAavePoolAddressesProvider(address(pool));
 
         vaultAssetAddress = address(aDai);
diff --git a/test/mocks/MockAToken.sol b/test/mocks/MockAToken.sol
@@ -3,8 +3,11 @@
 pragma solidity ^0.8.10;
 
 import {ERC20} from "@openzeppelin/token/ERC20/ERC20.sol";
+import {SafeERC20} from "@openzeppelin/token/ERC20/utils/SafeERC20.sol";
 
 contract MockAToken is ERC20 {
+    using SafeERC20 for ERC20;
+
     address internal _underlyingAsset;
 
     constructor(address underlyingAsset) ERC20("Mock aToken", "MAT") {
@@ -19,7 +22,7 @@ contract MockAToken is ERC20 {
     function burn(address from, address receiverOfUnderlying, uint256 amount, uint256 /* index */) external {
         _burn(from, amount);
         if (receiverOfUnderlying != address(this)) {
-            ERC20(_underlyingAsset).transfer(receiverOfUnderlying, amount);
+            ERC20(_underlyingAsset).safeTransfer(receiverOfUnderlying, amount);
         }
     }
 
diff --git a/test/mocks/MockAavePool.sol b/test/mocks/MockAavePool.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.10;
 
 import {ERC20} from "@openzeppelin/token/ERC20/ERC20.sol";
 import {DataTypes} from "@aave-v3-core/protocol/libraries/types/DataTypes.sol";
+import {SafeERC20} from "@openzeppelin/token/ERC20/utils/SafeERC20.sol";
 import {MockAToken} from "./MockAToken.sol";
 
 // NOTE: Yield Simulation Design
@@ -14,22 +15,24 @@ import {MockAToken} from "./MockAToken.sol";
 // implement utility functions which retain 1:1 asset:aToken parity
 
 contract MockAavePool {
+    using SafeERC20 for ERC20;
+
     uint256 constant SCALE = 1e18;
 
-    MockAToken public aToken;
+    mapping(address => MockAToken) internal _reserves;
 
     uint256 public reserveConfigMap;
 
-    constructor(MockAToken _aToken) {
-        aToken = _aToken;
+    function mockReserve(address asset, MockAToken aToken) public {
+        _reserves[asset] = aToken;
     }
 
     // For mock test purposes
     function setReserveConfigMap(uint256 _reserveConfigMap) public {
         reserveConfigMap = _reserveConfigMap;
     }
 
-    function getReserveData(address) public view returns (DataTypes.ReserveData memory) {
+    function getReserveData(address asset) public view returns (DataTypes.ReserveData memory) {
         return
             DataTypes.ReserveData({
                 //stores the reserve configuration
@@ -49,7 +52,7 @@ contract MockAavePool {
                 //the id of the reserve. Represents the position in the list of the active reserves
                 id: 0,
                 //aToken address
-                aTokenAddress: address(aToken),
+                aTokenAddress: address(_reserves[asset]),
                 //stableDebtToken address
                 stableDebtTokenAddress: address(0),
                 //variableDebtToken address
@@ -66,23 +69,23 @@ contract MockAavePool {
     }
 
     function supply(address _asset, uint256 _amount, address _onBehalfOf, uint16) public {
-        ERC20(_asset).transferFrom(msg.sender, address(this), _amount);
-        aToken.mint(address(this), _onBehalfOf, _amount, 0);
-        ERC20(_asset).transfer(address(aToken), _amount);
+        ERC20(_asset).safeTransferFrom(msg.sender, address(this), _amount);
+        _reserves[_asset].mint(address(this), _onBehalfOf, _amount, 0);
+        ERC20(_asset).safeTransfer(address(_reserves[_asset]), _amount);
     }
 
-    function withdraw(address /* _asset */, uint256 _amount, address _receiver) public returns (uint256) {
-        aToken.burn(msg.sender, _receiver, _amount, 0);
+    function withdraw(address _asset, uint256 _amount, address _receiver) public returns (uint256) {
+        _reserves[_asset].burn(msg.sender, _receiver, _amount, 0);
         return _amount;
     }
 
     // Mints recipient new tokens based on current aToken balance
     // to simulate new yield accrued
     // _yield is fraction > 1 using SCALE as denominator
-    function simulateYield(address _recipient, uint256 _yield) public {
-        uint256 balanceBefore = aToken.balanceOf(_recipient);
+    function simulateYield(address _asset, address _recipient, uint256 _yield) public {
+        uint256 balanceBefore = _reserves[_asset].balanceOf(_recipient);
         uint256 balanceAfter = (balanceBefore * _yield) / SCALE;
 
-        aToken.mint(address(this), _recipient, balanceAfter - balanceBefore, 0);
+        _reserves[_asset].mint(address(this), _recipient, balanceAfter - balanceBefore, 0);
     }
 }
diff --git a/test/mocks/MockUSDT.sol b/test/mocks/MockUSDT.sol
