From ff83693b2872743f1673091fd0906299e8dd1579 Mon Sep 17 00:00:00 2001
From: royvardhan <royvardhan.official@gmail.com>
Date: Wed, 26 Feb 2025 20:44:29 +0530
Subject: [PATCH] refactor: move wrap operation and receiver check to
 _executeSwap

---
 foundry/src/TychoRouter.sol | 33 ++++++++++++++-------------------
 1 file changed, 14 insertions(+), 19 deletions(-)

diff --git a/foundry/src/TychoRouter.sol b/foundry/src/TychoRouter.sol
index 0896279..ac50628 100644
--- a/foundry/src/TychoRouter.sol
+++ b/foundry/src/TychoRouter.sol
@@ -141,21 +141,12 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         address receiver,
         bytes calldata swaps
     ) external payable whenNotPaused nonReentrant returns (uint256 amountOut) {
-        if (receiver == address(0)) {
-            revert TychoRouter__AddressZero();
-        }
-
-        // Assume funds already in our router.
-        if (wrapEth) {
-            _wrapETH(amountIn);
-            tokenIn = address(_weth);
-        }
-
         return _executeSwap(
             amountIn,
             tokenIn,
             tokenOut,
             minAmountOut,
+            wrapEth,
             unwrapEth,
             nTokens,
             receiver,
@@ -203,15 +194,8 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         bytes calldata signature,
         bytes calldata swaps
     ) external payable whenNotPaused nonReentrant returns (uint256 amountOut) {
-        if (receiver == address(0)) {
-            revert TychoRouter__AddressZero();
-        }
-
-        // For native ETH, assume funds already in our router. Else, transfer and handle approval.
-        if (wrapEth) {
-            _wrapETH(amountIn);
-            tokenIn = address(_weth);
-        } else if (tokenIn != address(0)) {
+        // For non native ETH, transfer and handle approval.
+        if (tokenIn != address(0)) {
             permit2.permit(msg.sender, permitSingle, signature);
             permit2.transferFrom(
                 msg.sender,
@@ -226,6 +210,7 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
             tokenIn,
             tokenOut,
             minAmountOut,
+            wrapEth,
             unwrapEth,
             nTokens,
             receiver,
@@ -246,11 +231,21 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         address tokenIn,
         address tokenOut,
         uint256 minAmountOut,
+        bool wrapEth,
         bool unwrapEth,
         uint256 nTokens,
         address receiver,
         bytes calldata swaps
     ) internal returns (uint256 amountOut) {
+        if (receiver == address(0)) {
+            revert TychoRouter__AddressZero();
+        }
+        // For native ETH, assume funds are already in the router.
+        if (wrapEth) {
+            _wrapETH(amountIn);
+            tokenIn = address(_weth);
+        }
+
         amountOut = _swap(amountIn, nTokens, swaps);
 
         if (fee > 0) {