fix: only add assets when balance > 0 (#243) (#244)

* fix: only add assets when balance > 0 * fix: address comments
sentimentxyz · Feb 7, 2023 · 3d469a4 · 3d469a4
1 parent 06c35d2
commit 3d469a4
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 5 deletions.
diff --git a/src/core/AccountManager.sol b/src/core/AccountManager.sol
@@ -314,10 +314,10 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
             data
         );
         if (!isAllowed) revert Errors.FunctionCallRestricted();
-        _updateTokensIn(account, tokensIn);
         (bool success, ) = IAccount(account).exec(target, amt, data);
         if (!success)
             revert Errors.AccountInteractionFailure(account, target, amt, data);
+        _updateTokensIn(account, tokensIn);
         _updateTokensOut(account, tokensOut);
         if (IAccount(account).getAssets().length > assetCap + 1)
             revert Errors.MaxAssetCap();
@@ -358,8 +358,9 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
     {
         uint256 tokensInLen = tokensIn.length;
         for (uint256 i; i < tokensInLen; ++i) {
-            if (IAccount(account).hasAsset(tokensIn[i]) == false)
-                IAccount(account).addAsset(tokensIn[i]);
+            address token = tokensIn[i];
+            if (IAccount(account).hasAsset(token) == false && IERC20(token).balanceOf(account) > 0)
+                IAccount(account).addAsset(token);
         }
     }
 
@@ -368,7 +369,7 @@ contract AccountManager is ReentrancyGuard, Pausable, IAccountManager {
     {
         uint256 tokensOutLen = tokensOut.length;
         for (uint256 i; i < tokensOutLen; ++i) {
-            if (tokensOut[i].balanceOf(account) == 0)
+            if (IAccount(account).hasAsset(tokensOut[i]) == true && tokensOut[i].balanceOf(account) == 0)
                 IAccount(account).removeAsset(tokensOut[i]);
         }
     }

diff --git a/src/test/account/Account.t.sol b/src/test/account/Account.t.sol
@@ -65,6 +65,16 @@ contract AccountTest is BaseTest {
         assertFalse(account.hasAsset(token));
     }
 
+    function testRemoveNonExistingAsset(address token) public {
+        // Test
+        cheats.prank(address(accountManager));
+        account.removeAsset(token);
+
+        // Assert
+        assertEq(0, account.getAssets().length);
+        assertFalse(account.hasAsset(token));
+    }
+
     function testRemoveAssetError(address token) public {
         // Test
         cheats.expectRevert(Errors.AccountManagerOnly.selector);

diff --git a/src/test/integrations/Balancer.t.sol b/src/test/integrations/Balancer.t.sol
@@ -170,7 +170,7 @@ contract BalancerIntegrationTest is IntegrationBaseTest {
         // Assert
         assertEq(IERC20(balancerWeightedPool).balanceOf(account), 0);
         assertGt(account.balance, 0);
-        assertEq(IAccount(account).getAssets().length, 1);
+        assertEq(IAccount(account).getAssets().length, 0);
     }
 
     function testVaultExitStablePool() public {

diff --git a/src/test/integrations/Weth.t.sol b/src/test/integrations/Weth.t.sol
@@ -20,6 +20,7 @@ contract WethIntegrationTest is IntegrationBaseTest {
 
     function testWrapEth(uint96 amt) public {
         // Setup
+        cheats.assume(amt > 0);
         deposit(user, account, address(0), amt);
         bytes memory data = abi.encodeWithSignature("deposit()");