feat: Support returning values from the callback

Rollback some of the Ekubo's Executor changes to a previous version to use the generic callback logic using transient storage

Took 1 hour 25 minutes


Took 13 seconds

foundry/src/Dispatcher.sol

@@ -85,7 +85,10 @@ contract Dispatcher {
     }
 
     // slither-disable-next-line assembly
-    function _handleCallback(bytes calldata data) internal {
+    function _handleCallback(bytes calldata data)
+        internal
+        returns (bytes memory)
+    {
         address executor;
         assembly {
             executor := tload(0)
@@ -114,5 +117,9 @@ contract Dispatcher {
         assembly {
             tstore(0, 0)
         }
+
+        // this is necessary because the delegatecall will prepend extra bytes we don't want like the length and prefix
+        bytes memory decodedResult = abi.decode(result, (bytes));
+        return decodedResult;
     }
 }

foundry/src/TychoRouter.sol

@@ -367,7 +367,12 @@ contract TychoRouter is AccessControl, Dispatcher, Pausable, ReentrancyGuard {
      * @dev We use the fallback function to allow flexibility on callback.
      */
     fallback() external {
-        _handleCallback(msg.data);
+        bytes memory result = _handleCallback(msg.data);
+        // slither-disable-next-line assembly
+        assembly ("memory-safe") {
+            // Propagate the calculatedAmount
+            return(add(result, 32), 16)
+        }
     }
 
     /**

foundry/src/executors/EkuboExecutor.sol

@@ -3,6 +3,7 @@ pragma solidity ^0.8.26;
 
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {IExecutor} from "@interfaces/IExecutor.sol";
+import {ICallback} from "@interfaces/ICallback.sol";
 import {ICore} from "@ekubo/interfaces/ICore.sol";
 import {ILocker, IPayer} from "@ekubo/interfaces/IFlashAccountant.sol";
 import {NATIVE_TOKEN_ADDRESS} from "@ekubo/math/constants.sol";
@@ -11,16 +12,19 @@ import {LibBytes} from "@solady/utils/LibBytes.sol";
 import {Config, EkuboPoolKey} from "@ekubo/types/poolKey.sol";
 import {MAX_SQRT_RATIO, MIN_SQRT_RATIO} from "@ekubo/types/sqrtRatio.sol";
 
-contract EkuboExecutor is IExecutor, ILocker, IPayer {
+contract EkuboExecutor is IExecutor, ILocker, IPayer, ICallback {
     error EkuboExecutor__InvalidDataLength();
     error EkuboExecutor__CoreOnly();
     error EkuboExecutor__UnknownCallback();
 
     ICore immutable core;
 
-    uint256 constant POOL_DATA_OFFSET = 92;
+    uint256 constant POOL_DATA_OFFSET = 56;
     uint256 constant HOP_BYTE_LEN = 52;
 
+    bytes4 constant LOCKED_SELECTOR = 0xb45a3c0e; // locked(uint256)
+    bytes4 constant PAY_CALLBACK_SELECTOR = 0x599d0714; // payCallback(uint256,address)
+
     constructor(address _core) {
         core = ICore(_core);
     }
@@ -37,60 +41,45 @@ contract EkuboExecutor is IExecutor, ILocker, IPayer {
             uint256(_lock(bytes.concat(bytes16(uint128(amountIn)), data)));
     }
 
-    function locked(uint256) external coreOnly {
+    function handleCallback(bytes calldata raw)
-        int128 nextAmountIn = int128(uint128(bytes16(msg.data[36:52])));
+        external
-        uint128 tokenInDebtAmount = uint128(nextAmountIn);
+        returns (bytes memory)
+    {
+        verifyCallback(raw);
 
-        address receiver = address(bytes20(msg.data[52:72]));
+        // Without selector and locker id
-        address tokenIn = address(bytes20(msg.data[72:POOL_DATA_OFFSET]));
+        bytes calldata stripped = raw[36:];
 
-        address nextTokenIn = tokenIn;
+        bytes4 selector = bytes4(raw[:4]);
 
-        uint256 hopsLength = (msg.data.length - POOL_DATA_OFFSET) / HOP_BYTE_LEN;
+        bytes memory result = "";
-
+        if (selector == LOCKED_SELECTOR) {
-        uint256 offset = POOL_DATA_OFFSET;
+            int128 calculatedAmount = _locked(stripped);
-
+            result = abi.encodePacked(calculatedAmount);
-        for (uint256 i = 0; i < hopsLength; i++) {
+        } else if (selector == PAY_CALLBACK_SELECTOR) {
-            address nextTokenOut =
+            _payCallback(stripped);
-                address(bytes20(LibBytes.loadCalldata(msg.data, offset)));
+        } else {
-            Config poolConfig =
+            revert EkuboExecutor__UnknownCallback();
-                Config.wrap(LibBytes.loadCalldata(msg.data, offset + 20));
-
-            (address token0, address token1, bool isToken1) = nextTokenIn
-                > nextTokenOut
-                ? (nextTokenOut, nextTokenIn, true)
-                : (nextTokenIn, nextTokenOut, false);
-
-            // slither-disable-next-line calls-loop
-            (int128 delta0, int128 delta1) = core.swap_611415377(
-                EkuboPoolKey(token0, token1, poolConfig),
-                nextAmountIn,
-                isToken1,
-                isToken1 ? MAX_SQRT_RATIO : MIN_SQRT_RATIO,
-                0
-            );
-
-            nextTokenIn = nextTokenOut;
-            nextAmountIn = -(isToken1 ? delta0 : delta1);
-
-            offset += HOP_BYTE_LEN;
         }
 
-        _pay(tokenIn, tokenInDebtAmount);
+        return result;
+    }
 
-        core.withdraw(nextTokenIn, receiver, uint128(nextAmountIn));
+    function verifyCallback(bytes calldata) public view coreOnly {}
 
+    function locked(uint256) external coreOnly {
+        // Without selector and locker id
+        int128 calculatedAmount = _locked(msg.data[36:]);
         // slither-disable-next-line assembly
         assembly ("memory-safe") {
-            mstore(0, nextAmountIn)
+            mstore(0, calculatedAmount)
             return(0x10, 16)
         }
     }
 
-    function payCallback(uint256, address token) external coreOnly {
+    function payCallback(uint256, address /*token*/ ) external coreOnly {
-        uint128 amount = uint128(bytes16(msg.data[68:84]));
+        // Without selector and locker id
-
+        _payCallback(msg.data[36:]);
-        SafeTransferLib.safeTransfer(token, address(core), amount);
     }
 
     function _lock(bytes memory data)
@@ -121,6 +110,52 @@ contract EkuboExecutor is IExecutor, ILocker, IPayer {
         }
     }
 
+    function _locked(bytes calldata swapData) internal returns (int128) {
+        int128 nextAmountIn = int128(uint128(bytes16(swapData[0:16])));
+        uint128 tokenInDebtAmount = uint128(nextAmountIn);
+
+        address receiver = address(bytes20(swapData[16:36]));
+        address tokenIn = address(bytes20(swapData[36:POOL_DATA_OFFSET]));
+
+        address nextTokenIn = tokenIn;
+
+        uint256 hopsLength = (swapData.length - POOL_DATA_OFFSET) / HOP_BYTE_LEN;
+
+        uint256 offset = POOL_DATA_OFFSET;
+
+        for (uint256 i = 0; i < hopsLength; i++) {
+            address nextTokenOut =
+                address(bytes20(LibBytes.loadCalldata(swapData, offset)));
+            Config poolConfig =
+                Config.wrap(LibBytes.loadCalldata(swapData, offset + 20));
+
+            (address token0, address token1, bool isToken1) = nextTokenIn
+                > nextTokenOut
+                ? (nextTokenOut, nextTokenIn, true)
+                : (nextTokenIn, nextTokenOut, false);
+
+            // slither-disable-next-line calls-loop
+            (int128 delta0, int128 delta1) = core.swap_611415377(
+                EkuboPoolKey(token0, token1, poolConfig),
+                nextAmountIn,
+                isToken1,
+                isToken1 ? MAX_SQRT_RATIO : MIN_SQRT_RATIO,
+                0
+            );
+
+            nextTokenIn = nextTokenOut;
+            nextAmountIn = -(isToken1 ? delta0 : delta1);
+
+            offset += HOP_BYTE_LEN;
+        }
+
+        _pay(tokenIn, tokenInDebtAmount);
+
+        core.withdraw(nextTokenIn, receiver, uint128(nextAmountIn));
+
+        return nextAmountIn;
+    }
+
     function _pay(address token, uint128 amount) internal {
         address target = address(core);
 
@@ -144,6 +179,13 @@ contract EkuboExecutor is IExecutor, ILocker, IPayer {
         }
     }
 
+    function _payCallback(bytes calldata payData) internal {
+        address token = address(bytes20(payData[12:32])); // This arg is abi-encoded
+        uint128 amount = uint128(bytes16(payData[32:48]));
+
+        SafeTransferLib.safeTransfer(token, address(core), amount);
+    }
+
     // To receive withdrawals from Core
     receive() external payable {}
 

foundry/test/TychoRouter.t.sol

@@ -1023,7 +1023,7 @@ contract TychoRouterTest is TychoRouterTestSetup {
         vm.startPrank(ALICE);
         // Encoded solution generated using `test_split_encoding_strategy_ekubo`
         (bool success,) = address(tychoRouter).call{value: 1 ether}(
-            hex"0a83cb080000000000000000000000000000000000000000000000000de0b6b3a76400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000a0b86991c6218b36c1d19d4a2e9eb0ce3606eb480000000000000000000000000000000000000000000000000000000000000001000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000002000000000000000000000000cd09f75e2bf2a4d11f3ab23f1389fcc1621c0cc200000000000000000000000000000000000000000000000000000000000001200000000000000000000000000000000000000000000000000000000000000077007500010000002a07706473244bc757e10f2a9e86fb532828afe31d1499e622d69689cdf9004d05ec547d650ff2110000000000000000000000000000000000000000a0b86991c6218b36c1d19d4a2e9eb0ce3606eb4851d02a5948496a67827242eabc5725531342527c000000000000000000000000000000000000000000"
+            hex"0a83cb080000000000000000000000000000000000000000000000000de0b6b3a76400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000a0b86991c6218b36c1d19d4a2e9eb0ce3606eb480000000000000000000000000000000000000000000000000000000000000001000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000002000000000000000000000000cd09f75e2bf2a4d11f3ab23f1389fcc1621c0cc200000000000000000000000000000000000000000000000000000000000001200000000000000000000000000000000000000000000000000000000000000077007500010000003d7ebc40af7092e3f1c81f2e996cba5cae2090d7a4ad4f68d0b91cfd19687c881e50f3a00242828c0000000000000000000000000000000000000000a0b86991c6218b36c1d19d4a2e9eb0ce3606eb4851d02a5948496a67827242eabc5725531342527c000000000000000000000000000000000000000000"
         );
 
         uint256 balanceAfter = IERC20(USDC_ADDR).balanceOf(ALICE);