From 0964bb1fdf6bbc1ce8b14e3a2ffe8d2785a4bdef Mon Sep 17 00:00:00 2001
From: r0ohafza <coldroohafza@gmail.com>
Date: Thu, 10 Nov 2022 10:56:09 -0800
Subject: [PATCH 1/2] fix: deposit/withdraw/claim accounting and sig

---
 .github/workflows/ci-jobs.yml                |  2 +-
 src/balancer/BalancerLPStakingController.sol | 25 +++++++++---
 src/tests/BalancerLPStaking.t.sol            | 41 ++++++++++++--------
 3 files changed, 45 insertions(+), 23 deletions(-)

diff --git a/.github/workflows/ci-jobs.yml b/.github/workflows/ci-jobs.yml
index 1d0dd96..4d823d8 100644
--- a/.github/workflows/ci-jobs.yml
+++ b/.github/workflows/ci-jobs.yml
@@ -24,4 +24,4 @@ jobs:
         run: forge test --no-match-contract Arbi --fork-url https://rpc.ankr.com/eth
 
       - name: Run arbitrum tests
-        run: forge test --mc Arbi --fork-url https://arb-mainnet.g.alchemy.com/v2/WyhayD-L2Ox60Yngd6nzQvDKWtsCoCne
\ No newline at end of file
+        run: forge test --mc Arbi --fork-url  https://arb1.arbitrum.io/rpc
\ No newline at end of file
diff --git a/src/balancer/BalancerLPStakingController.sol b/src/balancer/BalancerLPStakingController.sol
index 585d070..469f757 100644
--- a/src/balancer/BalancerLPStakingController.sol
+++ b/src/balancer/BalancerLPStakingController.sol
@@ -27,8 +27,8 @@ contract BalancerLPStakingController is IController {
     /// @notice withdraw(uint256)
     bytes4 constant WITHDRAW = 0x2e1a7d4d;
 
-    /// @notice withdraw(uint256,address,bool)
-    bytes4 constant WITHDRAWCLAIM = 0x00ebf5dd;
+    /// @notice withdraw(uint256,bool)
+    bytes4 constant WITHDRAWCLAIM = 0x38d07436;
 
     /// @notice claim_rewards()
     bytes4 constant CLAIM = 0xe6f1daf2;
@@ -49,9 +49,9 @@ contract BalancerLPStakingController is IController {
         bytes4 sig = bytes4(data);
 
         if (sig == DEPOSIT) return canDeposit(target);
-        if (sig == DEPOSITCLAIM) return canDepositAndClaim(target);
+        if (sig == DEPOSITCLAIM) return canDepositAndClaim(target, data);
         if (sig == WITHDRAW) return canWithdraw(target);
-        if (sig == WITHDRAWCLAIM) return canWithdrawAndClaim(target);
+        if (sig == WITHDRAWCLAIM) return canWithdrawAndClaim(target, data);
         if (sig == CLAIM) return canClaim(target);
 
         return (false, new address[](0), new address[](0));
@@ -69,11 +69,17 @@ contract BalancerLPStakingController is IController {
         return (true, tokensIn, tokensOut);
     }
 
-    function canDepositAndClaim(address target)
+    function canDepositAndClaim(address target, bytes calldata data)
         internal
         view
         returns (bool, address[] memory, address[] memory)
     {
+        (,,bool claim) = abi.decode(
+            data[4:], (uint256, address, bool)
+        );
+
+        if (!claim) return canDeposit(target);
+
         address[] memory tokensIn = new address[](rewardsCount + 1);
         address reward; uint i;
         for (; i<rewardsCount; i++) {
@@ -103,11 +109,18 @@ contract BalancerLPStakingController is IController {
         return (true, tokensIn, tokensOut);
     }
 
-    function canWithdrawAndClaim(address target)
+    function canWithdrawAndClaim(address target, bytes calldata data)
         internal
         view
         returns (bool, address[] memory, address[] memory)
     {
+
+        (,bool claim) = abi.decode(
+            data[4:], (uint256, bool)
+        );
+
+        if (!claim) return canWithdraw(target);
+
         address[] memory tokensIn = new address[](rewardsCount + 1);
         address reward;
         uint i;
diff --git a/src/tests/BalancerLPStaking.t.sol b/src/tests/BalancerLPStaking.t.sol
index b2e3efd..0fc1d68 100644
--- a/src/tests/BalancerLPStaking.t.sol
+++ b/src/tests/BalancerLPStaking.t.sol
@@ -50,7 +50,7 @@ contract TestBalancerStakingArbi is TestBase {
         assertTrue(!canCall);
     }
 
-    function testCanDepositAndClaim() public {
+    function testCanDepositAndClaim(bool claim) public {
         // Setup
         controllerFacade.toggleTokenAllowance(gauge);
         controllerFacade.toggleTokenAllowance(r1);
@@ -59,7 +59,7 @@ contract TestBalancerStakingArbi is TestBase {
         bytes memory data = abi.encodeWithSelector(0x83df6747,
             0,
             address(0),
-            true
+            claim
         );
 
         // Test
@@ -68,11 +68,16 @@ contract TestBalancerStakingArbi is TestBase {
 
         // Assert
         assertTrue(canCall);
-        assertEq(tokensIn[0], r1);
-        assertEq(tokensIn[1], r2);
-        assertEq(tokensIn[2], gauge);
-        assertEq(tokensIn.length, 3);
         assertEq(tokensOut[0], lp);
+        if (claim) {
+            assertEq(tokensIn[0], r1);
+            assertEq(tokensIn[1], r2);
+            assertEq(tokensIn[2], gauge);
+            assertEq(tokensIn.length, 3);
+        } else {
+            assertEq(tokensIn[0], gauge);
+            assertEq(tokensIn.length, 1);
+        }
     }
 
     function testCannotDepositAndClaim() public {
@@ -90,16 +95,15 @@ contract TestBalancerStakingArbi is TestBase {
         assertTrue(!canCall);
     }
 
-    function testCanWithdrawAndClaim() public {
+    function testCanWithdrawAndClaim(bool claim) public {
         // Setup
         controllerFacade.toggleTokenAllowance(lp);
         controllerFacade.toggleTokenAllowance(r1);
         controllerFacade.toggleTokenAllowance(r2);
 
-        bytes memory data = abi.encodeWithSelector(0x00ebf5dd,
+        bytes memory data = abi.encodeWithSelector(0x38d07436,
             0,
-            address(0),
-            true
+            claim
         );
 
         // Test
@@ -108,17 +112,22 @@ contract TestBalancerStakingArbi is TestBase {
 
         // Assert
         assertTrue(canCall);
-        assertEq(tokensIn[0], r1);
-        assertEq(tokensIn[1], r2);
-        assertEq(tokensIn[2], lp);
-        assertEq(tokensIn.length, 3);
         assertEq(tokensOut[0], gauge);
+
+        if (claim) {
+            assertEq(tokensIn[0], r1);
+            assertEq(tokensIn[1], r2);
+            assertEq(tokensIn[2], lp);
+            assertEq(tokensIn.length, 3);
+        } else {
+            assertEq(tokensIn[0], lp);
+            assertEq(tokensIn.length, 1);
+        }
     }
 
     function testCannotWithdrawAndClaim() public {
-        bytes memory data = abi.encodeWithSelector(0x00ebf5dd,
+        bytes memory data = abi.encodeWithSelector(0x38d07436,
             0,
-            address(0),
             true
         );
 

From 6091affaf3f2d16b3c2bb6d6cbf3cfc1ab072953 Mon Sep 17 00:00:00 2001
From: r0ohafza <coldroohafza@gmail.com>
Date: Thu, 10 Nov 2022 14:43:11 -0800
Subject: [PATCH 2/2] fix: curve reward tokens accounting

---
 src/curve/CurveLPStakingController.sol | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/curve/CurveLPStakingController.sol b/src/curve/CurveLPStakingController.sol
index 6cd9413..bdea423 100644
--- a/src/curve/CurveLPStakingController.sol
+++ b/src/curve/CurveLPStakingController.sol
@@ -47,9 +47,9 @@ contract CurveLPStakingController is IController {
         bytes4 sig = bytes4(data);
 
         if (sig == DEPOSIT) return canDeposit(target);
-        if (sig == DEPOSITCLAIM) return canDepositAndClaim(target);
+        if (sig == DEPOSITCLAIM) return canDepositAndClaim(target, data);
         if (sig == WITHDRAW) return canWithdraw(target);
-        if (sig == WITHDRAWCLAIM) return canWithdrawAndClaim(target);
+        if (sig == WITHDRAWCLAIM) return canWithdrawAndClaim(target, data);
         if (sig == CLAIM) return canClaim(target);
 
         return (false, new address[](0), new address[](0));
@@ -67,11 +67,16 @@ contract CurveLPStakingController is IController {
         return (true, tokensIn, tokensOut);
     }
 
-    function canDepositAndClaim(address target)
+    function canDepositAndClaim(address target, bytes calldata data)
         internal
         view
         returns (bool, address[] memory, address[] memory)
     {
+        (,,bool claim) = abi.decode(
+            data[4:], (uint256, address, bool)
+        );
+        if (!claim) return canDeposit(target);
+
         uint count = IChildGauge(target).reward_count();
 
         address[] memory tokensIn = new address[](count + 1);
@@ -98,11 +103,17 @@ contract CurveLPStakingController is IController {
         return (true, tokensIn, tokensOut);
     }
 
-    function canWithdrawAndClaim(address target)
+    function canWithdrawAndClaim(address target, bytes calldata data)
         internal
         view
         returns (bool, address[] memory, address[] memory)
     {
+        (,,bool claim) = abi.decode(
+            data[4:], (uint256, address, bool)
+        );
+
+        if (!claim) return canWithdraw(target);
+
         uint count = IChildGauge(target).reward_count();
 
         address[] memory tokensIn = new address[](count + 1);