unioncredit · GeraldHost · Dec 4, 2022 · Nov 24, 2022
diff --git a/contracts/asset/AaveV3Adapter.sol b/contracts/asset/AaveV3Adapter.sol
@@ -201,10 +201,21 @@ contract AaveV3Adapter is Controller, IMoneyMarketAdapter {
      * @dev Deposit tokens into the underlying Aave V3 lending pool
      * @param tokenAddress Token address
      */
-    function deposit(address tokenAddress) external override checkTokenSupported(tokenAddress) {
+    function deposit(address tokenAddress)
+        external
+        override
+        onlyAssetManager
+        checkTokenSupported(tokenAddress)
+        returns (bool)
+    {
         IERC20Upgradeable token = IERC20Upgradeable(tokenAddress);
         uint256 amount = token.balanceOf(address(this));
-        lendingPool.supply(tokenAddress, amount, address(this), 0);
+        try lendingPool.supply(tokenAddress, amount, address(this), 0) {
+            return true;
+        } catch {
+            token.safeTransfer(assetManager, amount);
+            return false;
+        }
     }
 
     /**
@@ -218,9 +229,13 @@ contract AaveV3Adapter is Controller, IMoneyMarketAdapter {
         address tokenAddress,
         address recipient,
         uint256 tokenAmount
-    ) external override onlyAssetManager checkTokenSupported(tokenAddress) {
-        if (_getSupply(tokenAddress) > 0) {
-            lendingPool.withdraw(tokenAddress, tokenAmount, recipient);
+    ) external override onlyAssetManager checkTokenSupported(tokenAddress) returns (bool) {
+        if (_checkBal(tokenAddress)) {
+            try lendingPool.withdraw(tokenAddress, tokenAmount, recipient) {
+                return true;
+            } catch {
+                return false;
+            }
         }
     }
 
@@ -235,9 +250,14 @@ contract AaveV3Adapter is Controller, IMoneyMarketAdapter {
         override
         onlyAssetManager
         checkTokenSupported(tokenAddress)
+        returns (bool)
     {
-        if (_getSupply(tokenAddress) > 0) {
-            lendingPool.withdraw(tokenAddress, type(uint256).max, recipient);
+        if (_checkBal(tokenAddress)) {
+            try lendingPool.withdraw(tokenAddress, type(uint256).max, recipient) {
+                return true;
+            } catch {
+                return false;
+            }
         }
     }
 
@@ -270,4 +290,10 @@ contract AaveV3Adapter is Controller, IMoneyMarketAdapter {
     function _supportsToken(address tokenAddress) internal view returns (bool) {
         return tokenToAToken[tokenAddress] != address(0);
     }
+
+    function _checkBal(address tokenAddress) internal view returns (bool) {
+        address aTokenAddress = tokenToAToken[tokenAddress];
+        IERC20Upgradeable aToken = IERC20Upgradeable(aTokenAddress);
+        return aToken.balanceOf(address(this)) > 0;
+    }
 }
diff --git a/contracts/asset/AssetManager.sol b/contracts/asset/AssetManager.sol
@@ -276,6 +276,7 @@ contract AssetManager is Controller, ReentrancyGuardUpgradeable, IAssetManager {
         }
 
         bool remaining = true;
+        poolToken.safeTransferFrom(msg.sender, address(this), amount);
         if (isMarketSupported(token)) {
             uint256 moneyMarketsLength = moneyMarkets.length;
             // assumption: markets are arranged in order of decreasing liquidity
@@ -287,9 +288,10 @@ contract AssetManager is Controller, ReentrancyGuardUpgradeable, IAssetManager {
                 if (!moneyMarket.supportsToken(token)) continue;
                 if (moneyMarket.floorMap(token) <= moneyMarket.getSupply(token)) continue;
 
-                poolToken.safeTransferFrom(msg.sender, address(moneyMarket), amount);
-                moneyMarket.deposit(token);
-                remaining = false;
+                poolToken.safeTransfer(address(moneyMarket), amount);
+                if (moneyMarket.deposit(token)) {
+                    remaining = false;
+                }
             }
 
             // assumption: less liquid markets provide more yield
@@ -304,18 +306,14 @@ contract AssetManager is Controller, ReentrancyGuardUpgradeable, IAssetManager {
                 uint256 ceiling = moneyMarket.ceilingMap(token);
                 if (supply + amount > ceiling) continue;
 
-                poolToken.safeTransferFrom(msg.sender, address(moneyMarket), amount);
-                moneyMarket.deposit(token);
-                remaining = false;
+                poolToken.safeTransfer(address(moneyMarket), amount);
+                if (moneyMarket.deposit(token)) {
+                    remaining = false;
+                }
             }
         }
 
-        if (remaining) {
-            poolToken.safeTransferFrom(msg.sender, address(this), amount);
-        }
-
         emit LogDeposit(token, msg.sender, amount);
-
         return true;
     }
 
@@ -354,8 +352,9 @@ contract AssetManager is Controller, ReentrancyGuardUpgradeable, IAssetManager {
                 if (supply == 0) continue;
 
                 uint256 withdrawAmount = supply < remaining ? supply : remaining;
-                remaining -= withdrawAmount;
-                moneyMarket.withdraw(token, account, withdrawAmount);
+                if (moneyMarket.withdraw(token, account, withdrawAmount)) {
+                    remaining -= withdrawAmount;
+                }
             }
         }
 
@@ -518,10 +517,9 @@ contract AssetManager is Controller, ReentrancyGuardUpgradeable, IAssetManager {
         for (uint256 i = 0; i < moneyMarketsLength; i++) {
             IMoneyMarketAdapter moneyMarket = moneyMarkets[i];
             if (!moneyMarket.supportsToken(tokenAddress)) continue;
-            moneyMarket.withdrawAll(tokenAddress, address(this));
-
             supportedMoneyMarkets[supportedMoneyMarketsSize] = moneyMarket;
             supportedMoneyMarketsSize++;
+            moneyMarket.withdrawAll(tokenAddress, address(this));
         }
 
         if (percentagesLength + 1 != supportedMoneyMarketsSize) revert NotParity();

diff --git a/contracts/asset/PureTokenAdapter.sol b/contracts/asset/PureTokenAdapter.sol
@@ -140,7 +140,8 @@ contract PureTokenAdapter is Controller, IMoneyMarketAdapter {
     ------------------------------------------------------------------- */
 
     // solhint-disable-next-line no-empty-blocks
-    function deposit(address tokenAddress) external override checkTokenSupported(tokenAddress) {
+    function deposit(address tokenAddress) external override checkTokenSupported(tokenAddress) returns (bool) {
+        return true;
         // Don't have to do anything because AssetManager already transfered tokens here
     }
 
@@ -155,9 +156,10 @@ contract PureTokenAdapter is Controller, IMoneyMarketAdapter {
         address tokenAddress,
         address recipient,
         uint256 tokenAmount
-    ) external override onlyAssetManager checkTokenSupported(tokenAddress) {
+    ) external override onlyAssetManager checkTokenSupported(tokenAddress) returns (bool) {
         IERC20Upgradeable token = IERC20Upgradeable(tokenAddress);
         token.safeTransfer(recipient, tokenAmount);
+        return true;
     }
 
     /**
@@ -171,9 +173,11 @@ contract PureTokenAdapter is Controller, IMoneyMarketAdapter {
         override
         onlyAssetManager
         checkTokenSupported(tokenAddress)
+        returns (bool)
     {
         IERC20Upgradeable token = IERC20Upgradeable(tokenAddress);
         token.safeTransfer(recipient, token.balanceOf(address(this)));
+        return true;
     }
 
     // solhint-disable-next-line no-empty-blocks

diff --git a/contracts/interfaces/IMoneyMarketAdapter.sol b/contracts/interfaces/IMoneyMarketAdapter.sol
@@ -29,7 +29,7 @@ interface IMoneyMarketAdapter {
     /**
      * @dev Deposits the given amount of tokens in the underlying money market.
      */
-    function deposit(address tokenAddress) external;
+    function deposit(address tokenAddress) external returns (bool);
 
     /**
      * @dev Withdraws the given amount of tokens from the underlying money market and transfers them to `recipient`.
@@ -38,12 +38,12 @@ interface IMoneyMarketAdapter {
         address tokenAddress,
         address recipient,
         uint256 amount
-    ) external;
+    ) external returns (bool);
 
     /**
      * @dev Withdraws all the tokens from the underlying money market and transfers them to `recipient`.
      */
-    function withdrawAll(address tokenAddress, address recipient) external;
+    function withdrawAll(address tokenAddress, address recipient) external returns (bool);
 
     /**
      * @dev Returns the supply for the given token, including accrued interest. This function can have side effects.

diff --git a/contracts/mocks/AdapterMock.sol b/contracts/mocks/AdapterMock.sol
@@ -13,6 +13,8 @@ contract AdapterMock is IMoneyMarketAdapter {
 
     uint256 public rate;
 
+    bool public toRevert;
+
     /**
      * @dev Mapping of token address to floor balance
      */
@@ -23,6 +25,10 @@ contract AdapterMock is IMoneyMarketAdapter {
      */
     mapping(address => uint256) public override ceilingMap;
 
+    function setRevert() external {
+        toRevert = !toRevert;
+    }
+
     function setAssetManager(address _assetManager) external {
         assetManager = _assetManager;
     }
@@ -55,20 +61,32 @@ contract AdapterMock is IMoneyMarketAdapter {
         return true;
     }
 
-    function deposit(address) external override {}
+    function deposit(address tokenAddress) external override returns (bool) {
+        if (toRevert) {
+            IERC20Upgradeable token = IERC20Upgradeable(tokenAddress);
+            uint256 tokenAmount = token.balanceOf(address(this));
+            token.safeTransfer(msg.sender, tokenAmount);
+            return false;
+        }
+        return true;
+    }
 
     function withdraw(
         address tokenAddress,
         address recipient,
         uint256 tokenAmount
-    ) external override {
+    ) external override returns (bool) {
+        if (toRevert) return false;
         IERC20Upgradeable token = IERC20Upgradeable(tokenAddress);
         token.safeTransfer(recipient, tokenAmount);
+        return true;
     }
 
-    function withdrawAll(address tokenAddress, address recipient) external override {
+    function withdrawAll(address tokenAddress, address recipient) external override returns (bool) {
+        if (toRevert) return false;
         IERC20Upgradeable token = IERC20Upgradeable(tokenAddress);
         token.safeTransfer(recipient, token.balanceOf(address(this)));
+        return true;
     }
 
     function claimRewards(address, address) external override {}

diff --git a/contracts/mocks/CompoundAdapterMock.sol b/contracts/mocks/CompoundAdapterMock.sol
@@ -45,19 +45,23 @@ contract CompoundAdapterMock is Initializable, IMoneyMarketAdapter {
         return token.balanceOf(address(this));
     }
 
-    function deposit(address) external {}
+    function deposit(address) external returns (bool) {
+        return true;
+    }
 
     function withdraw(
         address tokenAddress,
         address recipient,
         uint256 tokenAmount
-    ) external {
+    ) external returns (bool) {
         IERC20Upgradeable token = IERC20Upgradeable(tokenAddress);
         token.transfer(recipient, tokenAmount);
+        return true;
     }
 
-    function withdrawAll(address tokenAddress, address recipient) external {
+    function withdrawAll(address tokenAddress, address recipient) external returns (bool) {
         IERC20Upgradeable token = IERC20Upgradeable(tokenAddress);
         token.transfer(recipient, token.balanceOf(address(this)));
+        return true;
     }
 }
diff --git a/test/foundry/assetManager/TestDepositWithdraw.t.sol b/test/foundry/assetManager/TestDepositWithdraw.t.sol
@@ -11,6 +11,8 @@ contract TestDepositWithdraw is TestAssetManagerBase {
 
         daiMock.mint(address(this), daiAmount);
         daiMock.approve(address(assetManager), daiAmount);
+        assetManager.addToken(address(daiMock));
+        assetManager.addAdapter(address(adapterMock));
     }
 
     function setTokens(address a, address b) public {
@@ -20,10 +22,24 @@ contract TestDepositWithdraw is TestAssetManagerBase {
 
     function testDeposit(uint256 amount) public {
         vm.assume(amount != 0 && amount < daiAmount);
+        adapterMock.setCeiling(address(daiMock), amount);
+
+        setTokens(address(this), address(123));
+        assetManager.deposit(address(daiMock), amount);
+        assertEq(assetManager.totalPrincipal(address(daiMock)), amount);
+        assertEq(assetManager.balances(address(this), address(daiMock)), amount);
+        assertEq(daiMock.balanceOf(address(assetManager)), 0);
+    }
+
+    function testDepositWhenAdapterRevert(uint256 amount) public {
+        vm.assume(amount != 0 && amount < daiAmount);
+        adapterMock.setRevert();
+        adapterMock.setCeiling(address(daiMock), amount);
         setTokens(address(this), address(123));
         assetManager.deposit(address(daiMock), amount);
         assertEq(assetManager.totalPrincipal(address(daiMock)), amount);
         assertEq(assetManager.balances(address(this), address(daiMock)), amount);
+        assertEq(daiMock.balanceOf(address(assetManager)), amount);
     }
 
     function testDepositAsUToken(uint256 amount) public {

diff --git a/test/foundry/assetManager/TestRebalance.t.sol b/test/foundry/assetManager/TestRebalance.t.sol
@@ -10,12 +10,14 @@ contract FakeAdapter {
         return true;
     }
 
-    function withdrawAll(address token, address to) public {
+    function withdrawAll(address token, address to) public returns (bool) {
         uint256 balance = IERC20(token).balanceOf(address(this));
         IERC20(token).transfer(to, balance);
     }
 
-    function deposit(address) public {}
+    function deposit(address) public returns (bool) {
+        return true;
+    }
 }
 
 contract TestRebalance is TestAssetManagerBase {