From 08d15c55ffba27dc64bd17aca2a7fe88f6b5e061 Mon Sep 17 00:00:00 2001 From: Denis Fadeev Date: Fri, 6 Dec 2024 13:01:14 +0300 Subject: [PATCH] Add revert handling to the Swap example (#219) --- examples/swap/contracts/Swap.sol | 93 ++++++++--- examples/swap/contracts/SwapToAnyToken.sol | 176 ++++++++++++++------- examples/swap/scripts/localnet.sh | 53 ++++++- examples/swap/tasks/deploy.ts | 11 +- 4 files changed, 255 insertions(+), 78 deletions(-) diff --git a/examples/swap/contracts/Swap.sol b/examples/swap/contracts/Swap.sol index d6c8ade1..1c112c63 100644 --- a/examples/swap/contracts/Swap.sol +++ b/examples/swap/contracts/Swap.sol @@ -15,20 +15,27 @@ contract Swap is UniversalContract { address public immutable uniswapRouter; GatewayZEVM public gateway; uint256 constant BITCOIN = 18332; + uint256 public immutable gasLimit; error InvalidAddress(); error Unauthorized(); + error ApprovalFailed(); modifier onlyGateway() { if (msg.sender != address(gateway)) revert Unauthorized(); _; } - constructor(address payable gatewayAddress, address uniswapRouterAddress) { + constructor( + address payable gatewayAddress, + address uniswapRouterAddress, + uint256 gasLimitAmount + ) { if (gatewayAddress == address(0) || uniswapRouterAddress == address(0)) revert InvalidAddress(); uniswapRouter = uniswapRouterAddress; gateway = GatewayZEVM(gatewayAddress); + gasLimit = gasLimitAmount; } struct Params { @@ -57,15 +64,19 @@ contract Swap is UniversalContract { params.to = recipient; } - swapAndWithdraw(zrc20, amount, params.target, params.to); + (uint256 out, address gasZRC20, uint256 gasFee) = handleGasAndSwap( + zrc20, + amount, + params.target + ); + withdraw(params, context.sender, gasFee, gasZRC20, out, zrc20); } - function swapAndWithdraw( + function handleGasAndSwap( address inputToken, uint256 amount, - address targetToken, - bytes memory recipient - ) internal { + address targetToken + ) internal returns (uint256, address, uint256) { uint256 inputForGas; address gasZRC20; uint256 gasFee; @@ -93,29 +104,75 @@ contract Swap is UniversalContract { targetToken, 0 ); + return (outputAmount, gasZRC20, gasFee); + } - if (gasZRC20 == targetToken) { - IZRC20(gasZRC20).approve(address(gateway), outputAmount + gasFee); + function withdraw( + Params memory params, + address sender, + uint256 gasFee, + address gasZRC20, + uint256 outputAmount, + address inputToken + ) public { + if (gasZRC20 == params.target) { + if ( + !IZRC20(gasZRC20).approve( + address(gateway), + outputAmount + gasFee + ) + ) { + revert ApprovalFailed(); + } } else { - IZRC20(gasZRC20).approve(address(gateway), gasFee); - IZRC20(targetToken).approve(address(gateway), outputAmount); + if (!IZRC20(gasZRC20).approve(address(gateway), gasFee)) { + revert ApprovalFailed(); + } + if ( + !IZRC20(params.target).approve(address(gateway), outputAmount) + ) { + revert ApprovalFailed(); + } } - gateway.withdraw( - recipient, + params.to, outputAmount, - targetToken, + params.target, + RevertOptions({ + revertAddress: address(this), + callOnRevert: true, + abortAddress: address(0), + revertMessage: abi.encode(sender, inputToken), + onRevertGasLimit: gasLimit + }) + ); + } + + function onRevert(RevertContext calldata context) external onlyGateway { + (address sender, address zrc20) = abi.decode( + context.revertMessage, + (address, address) + ); + (uint256 out, , ) = handleGasAndSwap( + context.asset, + context.amount, + zrc20 + ); + gateway.withdraw( + abi.encodePacked(sender), + out, + zrc20, RevertOptions({ - revertAddress: address(0), + revertAddress: sender, callOnRevert: false, abortAddress: address(0), revertMessage: "", - onRevertGasLimit: 0 + onRevertGasLimit: gasLimit }) ); } - function onRevert( - RevertContext calldata revertContext - ) external onlyGateway {} + fallback() external payable {} + + receive() external payable {} } diff --git a/examples/swap/contracts/SwapToAnyToken.sol b/examples/swap/contracts/SwapToAnyToken.sol index d3022c99..4c25b08d 100644 --- a/examples/swap/contracts/SwapToAnyToken.sol +++ b/examples/swap/contracts/SwapToAnyToken.sol @@ -16,20 +16,28 @@ contract SwapToAnyToken is UniversalContract { address public immutable uniswapRouter; GatewayZEVM public gateway; uint256 constant BITCOIN = 18332; + uint256 public immutable gasLimit; error InvalidAddress(); error Unauthorized(); + error ApprovalFailed(); + error TransferFailed(); modifier onlyGateway() { if (msg.sender != address(gateway)) revert Unauthorized(); _; } - constructor(address payable gatewayAddress, address uniswapRouterAddress) { + constructor( + address payable gatewayAddress, + address uniswapRouterAddress, + uint256 gasLimitAmount + ) { if (gatewayAddress == address(0) || uniswapRouterAddress == address(0)) revert InvalidAddress(); uniswapRouter = uniswapRouterAddress; gateway = GatewayZEVM(gatewayAddress); + gasLimit = gasLimitAmount; } struct Params { @@ -69,95 +77,155 @@ contract SwapToAnyToken is UniversalContract { params.withdraw = withdrawFlag; } - swapAndWithdraw( + (uint256 out, address gasZRC20, uint256 gasFee) = handleGasAndSwap( zrc20, amount, - params.target, - params.to, - params.withdraw + params.target ); + withdraw(params, context.sender, gasFee, gasZRC20, out, zrc20); } - function swapAndWithdraw( + function swap( address inputToken, uint256 amount, address targetToken, bytes memory recipient, - bool withdraw - ) internal { + bool withdrawFlag + ) public { + bool success = IZRC20(inputToken).transferFrom( + msg.sender, + address(this), + amount + ); + if (!success) { + revert TransferFailed(); + } + + (uint256 out, address gasZRC20, uint256 gasFee) = handleGasAndSwap( + inputToken, + amount, + targetToken + ); + + withdraw( + Params({ + target: targetToken, + to: recipient, + withdraw: withdrawFlag + }), + msg.sender, + gasFee, + gasZRC20, + out, + inputToken + ); + } + + function handleGasAndSwap( + address inputToken, + uint256 amount, + address targetToken + ) internal returns (uint256, address, uint256) { uint256 inputForGas; address gasZRC20; uint256 gasFee; - uint256 swapAmount = amount; + uint256 swapAmount; - if (withdraw) { - (gasZRC20, gasFee) = IZRC20(targetToken).withdrawGasFee(); + (gasZRC20, gasFee) = IZRC20(targetToken).withdrawGasFee(); - if (gasZRC20 == inputToken) { - swapAmount = amount - gasFee; - } else { - inputForGas = SwapHelperLib.swapTokensForExactTokens( - uniswapRouter, - inputToken, - gasFee, - gasZRC20, - amount - ); - swapAmount = amount - inputForGas; - } + if (gasZRC20 == inputToken) { + swapAmount = amount - gasFee; + } else { + inputForGas = SwapHelperLib.swapTokensForExactTokens( + uniswapRouter, + inputToken, + gasFee, + gasZRC20, + amount + ); + swapAmount = amount - inputForGas; } - uint256 outputAmount = SwapHelperLib.swapExactTokensForTokens( + uint256 out = SwapHelperLib.swapExactTokensForTokens( uniswapRouter, inputToken, swapAmount, targetToken, 0 ); + return (out, gasZRC20, gasFee); + } - if (withdraw) { - if (gasZRC20 == targetToken) { - IZRC20(gasZRC20).approve( - address(gateway), - outputAmount + gasFee - ); + function withdraw( + Params memory params, + address sender, + uint256 gasFee, + address gasZRC20, + uint256 out, + address inputToken + ) public { + if (params.withdraw) { + if (gasZRC20 == params.target) { + if (!IZRC20(gasZRC20).approve(address(gateway), out + gasFee)) { + revert ApprovalFailed(); + } } else { - IZRC20(gasZRC20).approve(address(gateway), gasFee); - IZRC20(targetToken).approve(address(gateway), outputAmount); + if (!IZRC20(gasZRC20).approve(address(gateway), gasFee)) { + revert ApprovalFailed(); + } + if (!IZRC20(params.target).approve(address(gateway), out)) { + revert ApprovalFailed(); + } } gateway.withdraw( - recipient, - outputAmount, - targetToken, + abi.encodePacked(params.to), + out, + params.target, RevertOptions({ - revertAddress: address(0), - callOnRevert: false, + revertAddress: address(this), + callOnRevert: true, abortAddress: address(0), - revertMessage: "", - onRevertGasLimit: 0 + revertMessage: abi.encode(sender, inputToken), + onRevertGasLimit: gasLimit }) ); } else { - IWETH9(targetToken).transfer( - address(uint160(bytes20(recipient))), - outputAmount + bool success = IWETH9(params.target).transfer( + address(uint160(bytes20(params.to))), + out ); + if (!success) { + revert TransferFailed(); + } } } - function swap( - address inputToken, - uint256 amount, - address targetToken, - bytes memory recipient, - bool withdraw - ) public { - IZRC20(inputToken).transferFrom(msg.sender, address(this), amount); + function onRevert(RevertContext calldata context) external onlyGateway { + (address sender, address zrc20) = abi.decode( + context.revertMessage, + (address, address) + ); + (uint256 out, , ) = handleGasAndSwap( + context.asset, + context.amount, + zrc20 + ); - swapAndWithdraw(inputToken, amount, targetToken, recipient, withdraw); + gateway.withdraw( + abi.encodePacked(sender), + out, + zrc20, + RevertOptions({ + revertAddress: sender, + callOnRevert: false, + abortAddress: address(0), + revertMessage: "", + onRevertGasLimit: gasLimit + }) + ); } - function onRevert( - RevertContext calldata revertContext - ) external onlyGateway {} + fallback() external payable {} + + receive() external payable {} } diff --git a/examples/swap/scripts/localnet.sh b/examples/swap/scripts/localnet.sh index 9867ddef..0ea56366 100755 --- a/examples/swap/scripts/localnet.sh +++ b/examples/swap/scripts/localnet.sh @@ -9,19 +9,64 @@ npx hardhat compile --force --quiet ZRC20_ETHEREUM=$(jq -r '.addresses[] | select(.type=="ZRC-20 ETH on 5") | .address' localnet.json) ZRC20_USDC=$(jq -r '.addresses[] | select(.type=="ZRC-20 USDC on 97") | .address' localnet.json) +ZRC20_BNB=$(jq -r '.addresses[] | select(.type=="ZRC-20 BNB on 97") | .address' localnet.json) GATEWAY_ETHEREUM=$(jq -r '.addresses[] | select(.type=="gatewayEVM" and .chain=="ethereum") | .address' localnet.json) GATEWAY_ZETACHAIN=$(jq -r '.addresses[] | select(.type=="gatewayZEVM" and .chain=="zetachain") | .address' localnet.json) UNISWAP_ROUTER=$(jq -r '.addresses[] | select(.type=="uniswapRouterInstance" and .chain=="zetachain") | .address' localnet.json) SENDER=0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266 -CONTRACT_ZETACHAIN=$(npx hardhat deploy --name Swap --network localhost --gateway "$GATEWAY_ZETACHAIN" --uniswap-router "$UNISWAP_ROUTER" --json | jq -r '.contractAddress') -echo -e "\nšŸš€ Deployed contract on ZetaChain: $CONTRACT_ZETACHAIN" +CONTRACT_SWAP=$(npx hardhat deploy --name Swap --network localhost --gateway "$GATEWAY_ZETACHAIN" --uniswap-router "$UNISWAP_ROUTER" --json | jq -r '.contractAddress') +echo -e "\nšŸš€ Deployed Swap contract on ZetaChain: $CONTRACT_SWAP" + +CONTRACT_SWAPTOANY=$(npx hardhat deploy --name SwapToAnyToken --network localhost --gateway "$GATEWAY_ZETACHAIN" --uniswap-router "$UNISWAP_ROUTER" --json | jq -r '.contractAddress') +echo -e "\nšŸš€ Deployed SwapToAny contract on ZetaChain: $CONTRACT_SWAPTOANY" + +npx hardhat swap-from-evm \ + --network localhost \ + --receiver "$CONTRACT_SWAP" \ + --amount 1 \ + --target "$ZRC20_BNB" \ + --recipient "$SENDER" + +npx hardhat localnet-check + +npx hardhat swap-from-evm \ + --network localhost \ + --receiver "$CONTRACT_SWAP" \ + --amount 1 \ + --target "$ZRC20_BNB" \ + --recipient "$SENDER" \ + --withdraw false + +npx hardhat localnet-check npx hardhat swap-from-evm \ --network localhost \ - --receiver "$CONTRACT_ZETACHAIN" \ + --receiver "$CONTRACT_SWAPTOANY" \ --amount 1 \ - --target "$ZRC20_USDC" \ + --target "$ZRC20_BNB" \ --recipient "$SENDER" +npx hardhat localnet-check + +npx hardhat swap-from-evm \ + --network localhost \ + --receiver "$CONTRACT_SWAPTOANY" \ + --amount 1 \ + --target "$ZRC20_BNB" \ + --recipient "$SENDER" \ + --withdraw false + +npx hardhat localnet-check + +npx hardhat swap-from-zetachain \ + --network localhost \ + --contract "$CONTRACT_SWAPTOANY" \ + --amount 1 \ + --zrc20 "$ZRC20_BNB" \ + --target "$ZRC20_ETHEREUM" \ + --recipient "$SENDER" + +npx hardhat localnet-check + if [ "$1" = "start" ]; then npx hardhat localnet-stop; fi \ No newline at end of file diff --git a/examples/swap/tasks/deploy.ts b/examples/swap/tasks/deploy.ts index 8475d31e..8d78f17c 100644 --- a/examples/swap/tasks/deploy.ts +++ b/examples/swap/tasks/deploy.ts @@ -1,4 +1,4 @@ -import { task } from "hardhat/config"; +import { task, types } from "hardhat/config"; import { HardhatRuntimeEnvironment } from "hardhat/types"; const main = async (args: any, hre: HardhatRuntimeEnvironment) => { @@ -14,7 +14,8 @@ const main = async (args: any, hre: HardhatRuntimeEnvironment) => { const factory = await hre.ethers.getContractFactory(args.name); const contract = await (factory as any).deploy( args.gateway, - args.uniswapRouter + args.uniswapRouter, + args.gasLimit ); await contract.deployed(); @@ -44,4 +45,10 @@ task("deploy", "Deploy the contract", main) "gateway", "Gateway address (default: ZetaChain Gateway)", "0x6c533f7fe93fae114d0954697069df33c9b74fd7" + ) + .addOptionalParam( + "gasLimit", + "Gas limit for the transaction", + 1000000, + types.int );