feat: Sketch for OneTransferFromOnly.sol

Took 49 seconds

foundry/src/OneTransferFromOnly.sol

@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.26;
+
+import "@interfaces/IExecutor.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import "@permit2/src/interfaces/IAllowanceTransfer.sol";
+
+error TokenTransfer__AddressZero();
+
+contract OneTransferFromOnly {
+    using SafeERC20 for IERC20;
+
+    // this is a stupid name but the compiler was complaining that we already had a permit2 variable in TychoRouter
+    IAllowanceTransfer public immutable permit2lal;
+    uint256 private constant _TOKEN_IN_SLOT = 123;
+    uint256 private constant _AMOUNT_IN_SLOT = 124;
+    uint256 private constant _IS_PERMIT2_SLOT = 125;
+    uint256 private constant _SENDER_SLOT = 126;
+    uint256 private constant _IS_TRANSFER_EXECUTED_SLOT = 127;
+
+    constructor(address _permit2) {
+        if (_permit2 == address(0)) {
+            revert TokenTransfer__AddressZero();
+        }
+        permit2lal = IAllowanceTransfer(_permit2);
+    }
+
+    function tstoreTransferFromInfo(
+        address tokenIn,
+        address amountIn,
+        bool isPermit2,
+        address sender
+    ) internal {
+        assembly {
+            tstore(_TOKEN_IN_SLOT, tokenIn)
+            tstore(_AMOUNT_IN_SLOT, amountIn)
+            tstore(_IS_PERMIT2_SLOT, isPermit2)
+            tstore(_SENDER_SLOT, sender)
+            tstore(_IS_TRANSFER_EXECUTED_SLOT, false)
+        }
+    }
+
+    function _transfer(address receiver)
+        // we could pass the amount and address too and compare to what is in the slots?
+        internal
+    {
+        address tokenIn;
+        uint256 amount;
+        bool isPermit2;
+        address sender;
+        bool isTransferExecuted;
+        assembly {
+            tokenIn := tload(_TOKEN_IN_SLOT)
+            amount := tload(_AMOUNT_IN_SLOT)
+            isPermit2 := tload(_IS_PERMIT2_SLOT)
+            sender := tload(_SENDER_SLOT)
+            isTransferExecuted := tload(_IS_TRANSFER_EXECUTED_SLOT)
+        }
+        if (isTransferExecuted) {
+            return; // or revert?
+        }
+
+        if (isPermit2) {
+            // Permit2.permit is already called from the TychoRouter
+            permit2lal.transferFrom(sender, receiver, uint160(amount), tokenIn);
+            assembly {
+                tstore(_IS_TRANSFER_EXECUTED_SLOT, true)
+            }
+        } else {
+            // slither-disable-next-line arbitrary-send-erc20
+            IERC20(tokenIn).safeTransferFrom(sender, receiver, amount);
+            assembly {
+                tstore(_IS_TRANSFER_EXECUTED_SLOT, true)
+            }
+        }
+    }
+}

foundry/src/TychoRouter.sol

@@ -14,6 +14,7 @@ import "@permit2/src/interfaces/IAllowanceTransfer.sol";
 import "./Dispatcher.sol";
 import {LibSwap} from "../lib/LibSwap.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {OneTransferFromOnly} from "./OneTransferFromOnly.sol";
 
 //                                         ✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷
 //                                   ✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷✷
@@ -65,7 +66,13 @@ error TychoRouter__MessageValueMismatch(uint256 value, uint256 amount);
 error TychoRouter__InvalidDataLength();
 error TychoRouter__UndefinedMinAmountOut();
 
-contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
+contract TychoRouter is
+    AccessControl,
+    Dispatcher,
+    Pausable,
+    ReentrancyGuard,
+    OneTransferFromOnly
+{
     IAllowanceTransfer public immutable permit2;
     IWETH private immutable _weth;
 
@@ -87,7 +94,7 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         address indexed token, uint256 amount, address indexed receiver
     );
 
-    constructor(address _permit2, address weth) {
+    constructor(address _permit2, address weth) OneTransferFromOnly(_permit2) {
         if (_permit2 == address(0) || weth == address(0)) {
             revert TychoRouter__AddressZero();
         }
@@ -130,6 +137,7 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         address receiver,
         bytes calldata swaps
     ) public payable whenNotPaused nonReentrant returns (uint256 amountOut) {
+        tstoreTransferFromInfo(tokenIn, amountIn, false, msg.sender);
         return _splitSwapChecked(
             amountIn,
             tokenIn,
@@ -187,6 +195,8 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
             permit2.permit(msg.sender, permitSingle, signature);
         }
 
+        tstoreTransferFromInfo(tokenIn, amountIn, true, msg.sender);
+
         return _splitSwapChecked(
             amountIn,
             tokenIn,
@@ -232,6 +242,7 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         address receiver,
         bytes calldata swaps
     ) public payable whenNotPaused nonReentrant returns (uint256 amountOut) {
+        tstoreTransferFromInfo(tokenIn, amountIn, false, msg.sender);
         return _sequentialSwapChecked(
             amountIn,
             tokenIn,
@@ -285,6 +296,7 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
             permit2.permit(msg.sender, permitSingle, signature);
         }
 
+        tstoreTransferFromInfo(tokenIn, amountIn, true, msg.sender);
         return _sequentialSwapChecked(
             amountIn,
             tokenIn,
@@ -325,8 +337,14 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         bool wrapEth,
         bool unwrapEth,
         address receiver,
+        bool inTransferNeeded,
+        address fundsReceiver,
         bytes calldata swapData
     ) public payable whenNotPaused nonReentrant returns (uint256 amountOut) {
+        tstoreTransferFromInfo(tokenIn, amountIn, false, msg.sender);
+        if (inTransferNeeded) {
+            _transfer(fundsReceiver);
+        }
         return _singleSwap(
             amountIn,
             tokenIn,
@@ -379,7 +397,7 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
         if (tokenIn != address(0)) {
             permit2.permit(msg.sender, permitSingle, signature);
         }
-
+        tstoreTransferFromInfo(tokenIn, amountIn, true, msg.sender);
         return _singleSwap(
             amountIn,
             tokenIn,

foundry/src/executors/UniswapV2Executor.sol

@@ -4,7 +4,6 @@ pragma solidity ^0.8.26;
 import "@interfaces/IExecutor.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@uniswap-v2/contracts/interfaces/IUniswapV2Pair.sol";
-import "./TokenTransfer.sol";
 
 error UniswapV2Executor__InvalidDataLength();
 error UniswapV2Executor__InvalidTarget();
@@ -12,7 +11,7 @@ error UniswapV2Executor__InvalidFactory();
 error UniswapV2Executor__InvalidInitCode();
 error UniswapV2Executor__InvalidFee();
 
-contract UniswapV2Executor is IExecutor, TokenTransfer {
+contract UniswapV2Executor is IExecutor {
     using SafeERC20 for IERC20;
 
     address public immutable factory;
@@ -25,7 +24,7 @@ contract UniswapV2Executor is IExecutor, TokenTransfer {
         bytes32 _initCode,
         address _permit2,
         uint256 _feeBps
-    ) TokenTransfer(_permit2) {
+    ) {
         if (_factory == address(0)) {
             revert UniswapV2Executor__InvalidFactory();
         }
@@ -51,17 +50,23 @@ contract UniswapV2Executor is IExecutor, TokenTransfer {
         address target;
         address receiver;
         bool zeroForOne;
-        TransferType transferType;
+        bool transferNeeded;
 
-        (tokenIn, target, receiver, zeroForOne, transferType) =
+        (tokenIn, target, receiver, zeroForOne, transferNeeded) =
             _decodeData(data);
 
         _verifyPairAddress(target);
 
         calculatedAmount = _getAmountOut(target, givenAmount, zeroForOne);
-        _transfer(
+
-            address(tokenIn), msg.sender, target, givenAmount, transferType
+        if (transferNeeded){
-        );
+            if (tokenIn == address(0)) {
+                payable(target).transfer(givenAmount);
+            } else {
+                // slither-disable-next-line arbitrary-send-erc20
+                tokenIn.safeTransferFrom(msg.sender, target, givenAmount);
+            }
+        }
 
         IUniswapV2Pair pool = IUniswapV2Pair(target);
         if (zeroForOne) {
@@ -79,7 +84,7 @@ contract UniswapV2Executor is IExecutor, TokenTransfer {
             address target,
             address receiver,
             bool zeroForOne,
-            TransferType transferType
+            bool transferNeeded
         )
     {
         if (data.length != 62) {
@@ -89,7 +94,7 @@ contract UniswapV2Executor is IExecutor, TokenTransfer {
         target = address(bytes20(data[20:40]));
         receiver = address(bytes20(data[40:60]));
         zeroForOne = uint8(data[60]) > 0;
-        transferType = TransferType(uint8(data[61]));
+        transferNeeded = bool(data[61]);
     }
 
     function _getAmountOut(address target, uint256 amountIn, bool zeroForOne)

foundry/src/executors/UniswapV3Executor.sol

@@ -6,14 +6,14 @@ import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@uniswap/v3-core/contracts/interfaces/IUniswapV3Pool.sol";
 import "@interfaces/ICallback.sol";
 import {TokenTransfer} from "./TokenTransfer.sol";
+import {OneTransferFromOnly} from "../OneTransferFromOnly.sol";
 
 error UniswapV3Executor__InvalidDataLength();
 error UniswapV3Executor__InvalidFactory();
 error UniswapV3Executor__InvalidTarget();
 error UniswapV3Executor__InvalidInitCode();
-error UniswapV3Executor__InvalidTransferType(uint8 transferType);
 
-contract UniswapV3Executor is IExecutor, ICallback, TokenTransfer {
+contract UniswapV3Executor is IExecutor, ICallback, OneTransferFromOnly {
     using SafeERC20 for IERC20;
 
     uint160 private constant MIN_SQRT_RATIO = 4295128739;
@@ -25,7 +25,7 @@ contract UniswapV3Executor is IExecutor, ICallback, TokenTransfer {
     address private immutable self;
 
     constructor(address _factory, bytes32 _initCode, address _permit2)
-        TokenTransfer(_permit2)
+        OneTransferFromOnly(_permit2)
     {
         if (_factory == address(0)) {
             revert UniswapV3Executor__InvalidFactory();
@@ -51,7 +51,8 @@ contract UniswapV3Executor is IExecutor, ICallback, TokenTransfer {
             address receiver,
             address target,
             bool zeroForOne,
-            TransferType transferType
+            bool inTransferNeeded,
+            bool inBetweenSwapsTransferNeeded
         ) = _decodeData(data);
 
         _verifyPairAddress(tokenIn, tokenOut, fee, target);
@@ -60,8 +61,13 @@ contract UniswapV3Executor is IExecutor, ICallback, TokenTransfer {
         int256 amount1;
         IUniswapV3Pool pool = IUniswapV3Pool(target);
 
-        bytes memory callbackData =
+        bytes memory callbackData = _makeV3CallbackData(
-            _makeV3CallbackData(tokenIn, tokenOut, fee, transferType);
+            tokenIn,
+            tokenOut,
+            fee,
+            inTransferNeeded,
+            inBetweenSwapsTransferNeeded
+        );
 
         {
             (amount0, amount1) = pool.swap(
@@ -98,12 +104,8 @@ contract UniswapV3Executor is IExecutor, ICallback, TokenTransfer {
 
         address tokenIn = address(bytes20(msgData[132:152]));
 
-        // Transfer type does not exist
+        bool inTransferNeeded = bool(msgData[175]);
-        if (uint8(msgData[175]) > uint8(TransferType.NONE)) {
+        bool inBetweenSwapsTransferNeeded = bool(msgData[176]);
-            revert UniswapV3Executor__InvalidTransferType(uint8(msgData[175]));
-        }
-
-        TransferType transferType = TransferType(uint8(msgData[175]));
         address sender = address(bytes20(msgData[176:196]));
 
         verifyCallback(msgData[132:]);
@@ -111,7 +113,15 @@ contract UniswapV3Executor is IExecutor, ICallback, TokenTransfer {
         uint256 amountOwed =
             amount0Delta > 0 ? uint256(amount0Delta) : uint256(amount1Delta);
 
-        _transfer(tokenIn, sender, msg.sender, amountOwed, transferType);
+        if (inTransferNeeded) {
+            _transfer(msg.sender);
+        } else if (inBetweenSwapsTransferNeeded) {
+            if (tokenIn == address(0)) {
+                payable(msg.sender).transfer(amountOwed);
+            } else {
+                IERC20(tokenIn).safeTransfer(msg.sender, amountOwed);
+            }
+        }
 
         return abi.encode(amountOwed, tokenIn);
     }
@@ -142,7 +152,8 @@ contract UniswapV3Executor is IExecutor, ICallback, TokenTransfer {
             address receiver,
             address target,
             bool zeroForOne,
-            TransferType transferType
+            bool inTransferNeeded,
+            bool inBetweenSwapsTransferNeeded
         )
     {
         if (data.length != 85) {
@@ -154,17 +165,24 @@ contract UniswapV3Executor is IExecutor, ICallback, TokenTransfer {
         receiver = address(bytes20(data[43:63]));
         target = address(bytes20(data[63:83]));
         zeroForOne = uint8(data[83]) > 0;
-        transferType = TransferType(uint8(data[84]));
+        inTransferNeeded = bool(data[84]);
+        inBetweenSwapsTransferNeeded = bool(data[85]);
     }
 
     function _makeV3CallbackData(
         address tokenIn,
         address tokenOut,
         uint24 fee,
-        TransferType transferType
+        bool inTransferNeeded,
+        bool inBetweenSwapsTransferNeeded
     ) internal view returns (bytes memory) {
         return abi.encodePacked(
-            tokenIn, tokenOut, fee, uint8(transferType), msg.sender
+            tokenIn,
+            tokenOut,
+            fee,
+            inTransferNeeded,
+            inBetweenSwapsTransferNeeded,
+            msg.sender
         );
     }