From 3d469a46b670dd02167bc0c5d395b092b7db1229 Mon Sep 17 00:00:00 2001
From: Rooh Afza <96720500+r0ohafza@users.noreply.github.com>
Date: Tue, 7 Feb 2023 11:45:46 -0800
Subject: [PATCH] fix: only add assets when balance > 0 (#243) (#244)

* fix: only add assets when balance > 0

* fix: address comments
---
 src/core/AccountManager.sol          |  9 +++++----
 src/test/account/Account.t.sol       | 10 ++++++++++
 src/test/integrations/Balancer.t.sol |  2 +-
 src/test/integrations/Weth.t.sol     |  1 +
 4 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/core/AccountManager.sol b/src/core/AccountManager.sol
index 4535ead..1b2cce4 100644
--- 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
index 9327963..fd45742 100644
--- 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
index eebc49b..d0c106a 100644
--- 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
index e72a306..dd30b53 100644
--- 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()");