fix: optional forwarding

synapsecns · Jan 2, 2025 · 93c28e6 · 93c28e6
1 parent 2414e3c
commit 93c28e6
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 1 deletion.
diff --git a/packages/contracts-rfq/contracts/zaps/TokenZapV1.sol b/packages/contracts-rfq/contracts/zaps/TokenZapV1.sol
@@ -176,9 +176,11 @@ contract TokenZapV1 is IZapRecipient {
         }
         // Check the token address and its balance to be safely forwarded.
         if (token == address(0)) revert TokenZapV1__FinalTokenNotSpecified();
+        // Perform the balance check before forwarding.
         uint256 amount = token == NATIVE_GAS_TOKEN ? address(this).balance : IERC20(token).balanceOf(address(this));
         if (amount == 0 || amount < minFinalBalance) revert TokenZapV1__FinalBalanceBelowMin();
-        // Forward the full balance of the final token to the specified recipient.
+        // Forward the full balance of the final token to the specified recipient, if required.
+        if (forwardTo == address(0)) return;
         if (token == NATIVE_GAS_TOKEN) {
             Address.sendValue({recipient: payable(forwardTo), amount: amount});
         } else {

diff --git a/packages/contracts-rfq/test/zaps/TokenZapV1.t.sol b/packages/contracts-rfq/test/zaps/TokenZapV1.t.sol
@@ -367,6 +367,24 @@ contract TokenZapV1Test is Test {
         assertEq(payableMock.balance, 2 * AMOUNT);
     }
 
+    function test_zap_unwrap_zeroForwardTo_withMinFinalBalance() public {
+        bytes memory zapData = getZapDataUnwrapAndForward(0, nativeGasToken, address(0), AMOUNT);
+        weth.transfer(address(tokenZap), AMOUNT);
+        bytes4 returnValue = tokenZap.zap(address(weth), AMOUNT, zapData);
+        assertEq(returnValue, tokenZap.zap.selector);
+        // Check that the unwrapped tokens remain in the TokenZap
+        assertEq(address(tokenZap).balance, AMOUNT);
+    }
+
+    function test_zap_unwrap_zeroForwardTo_zeroMinFinalBalance() public {
+        bytes memory zapData = getZapDataUnwrapAndForward(0, nativeGasToken, address(0), 0);
+        weth.transfer(address(tokenZap), AMOUNT);
+        bytes4 returnValue = tokenZap.zap(address(weth), AMOUNT, zapData);
+        assertEq(returnValue, tokenZap.zap.selector);
+        // Check that the unwrapped tokens remain in the TokenZap
+        assertEq(address(tokenZap).balance, AMOUNT);
+    }
+
     function test_zap_wrap_depositWETH_placeholderZero() public {
         bytes memory zapDataWrap = getZapDataWrap();
         bytes memory zapDataDeposit = getZapDataDeposit(getVaultPayload(address(weth), 0));
@@ -429,6 +447,22 @@ contract TokenZapV1Test is Test {
         assertEq(weth.balanceOf(user), 2 * AMOUNT);
     }
 
+    function test_zap_wrap_zeroForwardTo_withMinFinalBalance() public {
+        bytes memory zapData = getZapDataWrapAndForward(address(weth), address(0), AMOUNT);
+        bytes4 returnValue = tokenZap.zap{value: AMOUNT}(nativeGasToken, AMOUNT, zapData);
+        assertEq(returnValue, tokenZap.zap.selector);
+        // Check that the wrapped tokens remain in the TokenZap
+        assertEq(weth.balanceOf(address(tokenZap)), AMOUNT);
+    }
+
+    function test_zap_wrap_zeroForwardTo_zeroMinFinalBalance() public {
+        bytes memory zapData = getZapDataWrapAndForward(address(weth), address(0), 0);
+        bytes4 returnValue = tokenZap.zap{value: AMOUNT}(nativeGasToken, AMOUNT, zapData);
+        assertEq(returnValue, tokenZap.zap.selector);
+        // Check that the wrapped tokens remain in the TokenZap
+        assertEq(weth.balanceOf(address(tokenZap)), AMOUNT);
+    }
+
     function getZapDataTransferNative(address target) public view returns (bytes memory) {
         return tokenZap.encodeZapData({
             target: target,