From 70230bf05f4bdfdd54f62799b72bd998d351983e Mon Sep 17 00:00:00 2001 From: Diana Carvalho Date: Tue, 13 May 2025 11:41:11 +0100 Subject: [PATCH 1/3] feat: Verify the amount out was received correctly for arbitrage swaps Took 25 minutes Took 3 minutes Took 23 seconds --- foundry/src/TychoRouter.sol | 75 +++++++++++++++++++++++-------------- 1 file changed, 47 insertions(+), 28 deletions(-) diff --git a/foundry/src/TychoRouter.sol b/foundry/src/TychoRouter.sol index c237d64..5c41e9f 100644 --- a/foundry/src/TychoRouter.sol +++ b/foundry/src/TychoRouter.sol @@ -436,15 +436,14 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard { Address.sendValue(payable(receiver), amountOut); } - if (tokenIn != tokenOut) { - uint256 currentBalanceTokenOut = _balanceOf(tokenOut, receiver); - uint256 userAmount = currentBalanceTokenOut - initialBalanceTokenOut; - if (userAmount != amountOut) { - revert TychoRouter__AmountOutNotFullyReceived( - userAmount, amountOut - ); - } - } + _verifyAmountOutWasReceived( + tokenIn, + tokenOut, + initialBalanceTokenOut, + amountOut, + receiver, + amountIn + ); } /** @@ -493,15 +492,14 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard { Address.sendValue(payable(receiver), amountOut); } - if (tokenIn != tokenOut) { - uint256 currentBalanceTokenOut = _balanceOf(tokenOut, receiver); - uint256 userAmount = currentBalanceTokenOut - initialBalanceTokenOut; - if (userAmount != amountOut) { - revert TychoRouter__AmountOutNotFullyReceived( - userAmount, amountOut - ); - } - } + _verifyAmountOutWasReceived( + tokenIn, + tokenOut, + initialBalanceTokenOut, + amountOut, + receiver, + amountIn + ); } /** @@ -546,16 +544,14 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard { _unwrapETH(amountOut); Address.sendValue(payable(receiver), amountOut); } - - if (tokenIn != tokenOut) { - uint256 currentBalanceTokenOut = _balanceOf(tokenOut, receiver); - uint256 userAmount = currentBalanceTokenOut - initialBalanceTokenOut; - if (userAmount != amountOut) { - revert TychoRouter__AmountOutNotFullyReceived( - userAmount, amountOut - ); - } - } + _verifyAmountOutWasReceived( + tokenIn, + tokenOut, + initialBalanceTokenOut, + amountOut, + receiver, + amountIn + ); } /** @@ -784,4 +780,27 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard { return token == address(0) ? owner.balance : IERC20(token).balanceOf(owner); } + + /** + * @dev Verifies that the expected amount of output tokens was received by the receiver. + * It also handles the case of arbitrage swaps where the input and output tokens are the same. + */ + function _verifyAmountOutWasReceived( + address tokenIn, + address tokenOut, + uint256 initialBalanceTokenOut, + uint256 amountOut, + address receiver, + uint256 amountIn + ) internal view { + uint256 currentBalanceTokenOut = _balanceOf(tokenOut, receiver); + if (tokenIn == tokenOut) { + // If it is an arbitrage, we need to remove the amountIn from the initial balance to get a correct userAmount + initialBalanceTokenOut -= amountIn; + } + uint256 userAmount = currentBalanceTokenOut - initialBalanceTokenOut; + if (userAmount != amountOut) { + revert TychoRouter__AmountOutNotFullyReceived(userAmount, amountOut); + } + } } From 65bd0d07499b79c261efdb00debb19487d6af543 Mon Sep 17 00:00:00 2001 From: Diana Carvalho Date: Tue, 13 May 2025 12:31:44 +0100 Subject: [PATCH 2/3] feat: Explicitly handle the TransferType.NONE case Took 8 minutes --- foundry/src/executors/TokenTransfer.sol | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/foundry/src/executors/TokenTransfer.sol b/foundry/src/executors/TokenTransfer.sol index b5f8629..05e296a 100644 --- a/foundry/src/executors/TokenTransfer.sol +++ b/foundry/src/executors/TokenTransfer.sol @@ -45,7 +45,9 @@ contract TokenTransfer { uint256 amount, TransferType transferType ) internal { - if (transferType == TransferType.TRANSFER_TO_PROTOCOL) { + if (transferType == TransferType.NONE) { + return; + } else if (transferType == TransferType.TRANSFER_TO_PROTOCOL) { if (tokenIn == address(0)) { payable(receiver).transfer(amount); } else { From b0c254add44e56d10f203eae30301c861fd9d8ff Mon Sep 17 00:00:00 2001 From: Diana Carvalho Date: Wed, 14 May 2025 09:02:55 +0100 Subject: [PATCH 3/3] fix: Revert if the TransferType is not valid Took 8 minutes --- foundry/src/executors/TokenTransfer.sol | 3 +++ 1 file changed, 3 insertions(+) diff --git a/foundry/src/executors/TokenTransfer.sol b/foundry/src/executors/TokenTransfer.sol index 05e296a..6eac1c0 100644 --- a/foundry/src/executors/TokenTransfer.sol +++ b/foundry/src/executors/TokenTransfer.sol @@ -6,6 +6,7 @@ import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; import "@permit2/src/interfaces/IAllowanceTransfer.sol"; error TokenTransfer__AddressZero(); +error TokenTransfer__InvalidTransferType(); contract TokenTransfer { using SafeERC20 for IERC20; @@ -67,6 +68,8 @@ contract TokenTransfer { permit2.transferFrom( sender, address(this), uint160(amount), tokenIn ); + } else { + revert TokenTransfer__InvalidTransferType(); } } }