From 0be69c9aea78f486338ab50ba803ac4869df93ac Mon Sep 17 00:00:00 2001 From: royvardhan Date: Wed, 26 Feb 2025 19:04:26 +0530 Subject: [PATCH] refactor: check if amountIn was fully consumed based on balance changes --- foundry/src/TychoRouter.sol | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/foundry/src/TychoRouter.sol b/foundry/src/TychoRouter.sol index 0896279..6f10fed 100644 --- a/foundry/src/TychoRouter.sol +++ b/foundry/src/TychoRouter.sol @@ -251,8 +251,23 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard { address receiver, bytes calldata swaps ) internal returns (uint256 amountOut) { + uint256 initialBalance = tokenIn == address(0) + ? address(this).balance + : IERC20(tokenIn).balanceOf(address(this)); + amountOut = _swap(amountIn, nTokens, swaps); + uint256 currentBalance = tokenIn == address(0) + ? address(this).balance + : IERC20(tokenIn).balanceOf(address(this)); + + uint256 amountConsumed = initialBalance - currentBalance; + + if (amountConsumed < amountIn) { + uint256 leftoverAmount = amountIn - amountConsumed; + revert TychoRouter__AmountInNotFullySpent(leftoverAmount); + } + if (fee > 0) { uint256 feeAmount = (amountOut * fee) / 10000; amountOut -= feeAmount; @@ -263,17 +278,6 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard { revert TychoRouter__NegativeSlippage(amountOut, minAmountOut); } - uint256 leftoverAmountIn; - if (tokenIn == address(0)) { - leftoverAmountIn = address(this).balance; - } else { - leftoverAmountIn = IERC20(tokenIn).balanceOf(address(this)); - } - - if (leftoverAmountIn > 0) { - revert TychoRouter__AmountInNotFullySpent(leftoverAmountIn); - } - if (unwrapEth) { _unwrapETH(amountOut); }