Update ZkBobPay contract (#85)

zkBob · Aug 18, 2023 · 4777adf · 4777adf
1 parent e4cf6d8
commit 4777adf
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 22 deletions.
diff --git a/script/scripts/ZkBobPay.s.sol b/script/scripts/ZkBobPay.s.sol
@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: CC0-1.0
+
+pragma solidity 0.8.15;
+
+import "forge-std/Script.sol";
+import "./Env.s.sol";
+import "../../src/zkbob/periphery/ZkBobPay.sol";
+import "../../src/proxy/EIP1967Proxy.sol";
+
+contract DeployZkBobPay is Script {
+    function run() external {
+        vm.startBroadcast();
+
+        address token = address(0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174);
+        address queue = address(0x668c5286eAD26fAC5fa944887F9D2F20f7DDF289);
+        address permit2 = address(0x000000000022D473030F116dDEE9F6B43aC78BA3);
+        address lifiRouter1 = address(0x1231DEB6f5749EF6cE6943a275A1D3E7486F4EaE);
+        address lifiRouter2 = address(0x9b11bc9FAc17c058CAB6286b0c785bE6a65492EF);
+        address feeReceiver = address(0x39F0bD56c1439a22Ee90b4972c16b7868D161981);
+
+        ZkBobPay pay = new ZkBobPay(token, queue, permit2);
+        EIP1967Proxy proxy =
+            new EIP1967Proxy(tx.origin, address(pay), abi.encodeWithSelector(ZkBobPay.initialize.selector));
+        pay = ZkBobPay(address(proxy));
+
+        pay.updateFeeReceiver(feeReceiver);
+        bytes4[] memory selectors = new bytes4[](1);
+        selectors[0] = 0x4630a0d8;
+        pay.updateRouter(lifiRouter1, selectors, true);
+        pay.updateRouter(lifiRouter2, selectors, true);
+
+        vm.stopBroadcast();
+
+        console2.log("ZkBobPay:", address(pay));
+    }
+}
diff --git a/src/zkbob/periphery/ZkBobPay.sol b/src/zkbob/periphery/ZkBobPay.sol
@@ -9,14 +9,16 @@ import "../../interfaces/IERC677.sol";
 import "../../interfaces/IERC20Permit.sol";
 import "../../interfaces/IUSDCPermit.sol";
 import "../../interfaces/IPermit2.sol";
+import "../../proxy/EIP1967Admin.sol";
 
 /**
  * @title ZkBobPay
  */
-contract ZkBobPay {
+contract ZkBobPay is EIP1967Admin {
     using SafeERC20 for IERC20;
 
     event UpdateFeeReceiver(address receiver);
+    event UpdateRouter(address router, bytes4[] selectors, bool enabled);
     event Pay(uint256 indexed id, address indexed sender, bytes receiver, uint256 amount, address inToken, bytes note);
 
     error InvalidToken();
@@ -28,17 +30,22 @@ contract ZkBobPay {
     address public immutable token;
     IZkBobDirectDeposits public immutable queue;
     IPermit2 public immutable permit2;
-    address public immutable oneInchRouter;
+
+    mapping(address => mapping(bytes4 => bool)) public enabledRouter;
     address public feeReceiver;
 
-    constructor(address _token, address _queue, address _permit2, address _oneInchRouter, address _feeReceiver) {
+    constructor(address _token, address _queue, address _permit2) {
         token = _token;
         queue = IZkBobDirectDeposits(_queue);
         permit2 = IPermit2(_permit2);
-        oneInchRouter = _oneInchRouter;
-        feeReceiver = _feeReceiver;
+    }
+
+    function initialize() external {
+        if (msg.sender != address(this)) {
+            revert Unauthorized();
+        }
 
-        IERC20(token).approve(_queue, type(uint256).max);
+        IERC20(token).approve(address(queue), type(uint256).max);
     }
 
     function onTokenTransfer(address _from, uint256 _amount, bytes memory _data) external returns (bool) {
@@ -61,17 +68,19 @@ contract ZkBobPay {
      * @param _inAmount input token amount.
      * @param _depositAmount zkBob deposit amount, inclusive of direct deposit fee.
      * @param _permit input token approval permit, in one of the supported formats.
-     * @param _oneInchData 1inch swap calldata.
+     * @param _router router contract address, should be whitelisted by contract admin first.
+     * @param _routerData router swap calldata.
      * @param _note optional payment-specific note for the receiver.
      */
     function pay(
-        bytes memory _zkAddress,
+        bytes calldata _zkAddress,
         address _inToken,
         uint256 _inAmount,
         uint256 _depositAmount,
         bytes memory _permit,
-        bytes memory _oneInchData,
-        bytes memory _note
+        address _router,
+        bytes calldata _routerData,
+        bytes calldata _note
     )
         external
         payable
@@ -89,13 +98,17 @@ contract ZkBobPay {
                 revert InsufficientAmount();
             }
         } else {
+            if (!enabledRouter[_router][bytes4(_routerData[:4])]) {
+                revert Unauthorized();
+            }
+
             uint256 balance = IERC20(token).balanceOf(address(this));
 
             if (_inToken != address(0)) {
-                IERC20(_inToken).approve(oneInchRouter, _inAmount);
+                IERC20(_inToken).approve(_router, _inAmount);
             }
 
-            (bool status,) = oneInchRouter.call{value: msg.value}(_oneInchData);
+            (bool status,) = _router.call{value: msg.value}(_routerData);
             if (!status) {
                 revert SwapFailed();
             }
@@ -115,7 +128,7 @@ contract ZkBobPay {
             revert Unauthorized();
         }
 
-        for (uint256 i = 0; i < _tokens.length; i++) {
+        for (uint256 i = 0; i < _tokens.length; ++i) {
             if (_tokens[i] == address(0)) {
                 payable(msg.sender).transfer(address(this).balance);
             } else {
@@ -124,8 +137,20 @@ contract ZkBobPay {
         }
     }
 
+    function updateRouter(address _router, bytes4[] calldata _selectors, bool _enabled) external {
+        if (msg.sender != _admin()) {
+            revert Unauthorized();
+        }
+
+        for (uint256 i = 0; i < _selectors.length; ++i) {
+            enabledRouter[_router][_selectors[i]] = _enabled;
+        }
+
+        emit UpdateRouter(_router, _selectors, _enabled);
+    }
+
     function updateFeeReceiver(address _receiver) external {
-        if (msg.sender != feeReceiver) {
+        if (msg.sender != feeReceiver && msg.sender != _admin()) {
             revert Unauthorized();
         }
 

diff --git a/test/zkbob/ZkBobPay.t.sol b/test/zkbob/ZkBobPay.t.sol
@@ -7,6 +7,7 @@ import "forge-std/Test.sol";
 import "../shared/Env.t.sol";
 import "../shared/ForkTests.t.sol";
 import "../../src/zkbob/periphery/ZkBobPay.sol";
+import "../../src/proxy/EIP1967Proxy.sol";
 
 contract ZkBobPayTest is AbstractPolygonForkTest {
     ZkBobPay pay;
@@ -33,7 +34,16 @@ contract ZkBobPayTest is AbstractPolygonForkTest {
     function setUp() public {
         vm.createSelectFork(forkRpcUrl, 45500000);
 
-        pay = new ZkBobPay(usdc, queue, permit2, oneInchRouter, user2);
+        pay = new ZkBobPay(usdc, queue, permit2);
+        EIP1967Proxy proxy =
+            new EIP1967Proxy(address(this), address(pay), abi.encodeWithSelector(ZkBobPay.initialize.selector));
+        pay = ZkBobPay(address(proxy));
+
+        pay.updateFeeReceiver(user2);
+        bytes4[] memory selectors = new bytes4[](2);
+        selectors[0] = 0xe449022e;
+        selectors[1] = 0x12aa3caf;
+        pay.updateRouter(oneInchRouter, selectors, true);
     }
 
     function testPaymentWithUSDC() public {
@@ -44,7 +54,7 @@ contract ZkBobPayTest is AbstractPolygonForkTest {
         vm.prank(user1);
         vm.expectEmit(false, true, false, true);
         emit Pay(0, user1, zkAddress, 100e6, usdc, note);
-        pay.pay(zkAddress, usdc, 105e6, 100e6, "", "", note);
+        pay.pay(zkAddress, usdc, 105e6, 100e6, "", address(0), "", note);
 
         assertGt(IERC20(usdc).balanceOf(address(pay)), 1e6);
         assertLt(IERC20(usdc).balanceOf(address(pay)), 10e6);
@@ -61,7 +71,7 @@ contract ZkBobPayTest is AbstractPolygonForkTest {
         vm.prank(user1);
         vm.expectEmit(false, true, false, true);
         emit Pay(0, user1, zkAddress, 100e6, usdc, note);
-        pay.pay(zkAddress, usdc, 105e6, 100e6, permit, "", note);
+        pay.pay(zkAddress, usdc, 105e6, 100e6, permit, address(0), "", note);
 
         assertGt(IERC20(usdc).balanceOf(address(pay)), 1e6);
         assertLt(IERC20(usdc).balanceOf(address(pay)), 10e6);
@@ -79,7 +89,7 @@ contract ZkBobPayTest is AbstractPolygonForkTest {
         vm.prank(user1);
         vm.expectEmit(false, true, false, true);
         emit Pay(0, user1, zkAddress, 100e6, frax, note);
-        pay.pay(zkAddress, frax, 105 ether, 100e6, permit, oneInchData, note);
+        pay.pay(zkAddress, frax, 105 ether, 100e6, permit, oneInchRouter, oneInchData, note);
 
         assertGt(IERC20(usdc).balanceOf(address(pay)), 1e6);
         assertLt(IERC20(usdc).balanceOf(address(pay)), 10e6);
@@ -92,7 +102,7 @@ contract ZkBobPayTest is AbstractPolygonForkTest {
         IERC20(wmatic).approve(permit2, type(uint256).max);
 
         bytes memory oneInchData =
-            hex"12aa3caf000000000000000000000000cfd674f8731e801a4a15c1ae31770960e1afded10000000000000000000000000d500b1d8e8ef31e21c99d1db9a6444d3adf12700000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa841740000000000000000000000001093ced81987bf532c2b7907b2a8525cd0c172950000000000000000000000005615deb798bb3e4dfa0139dfa1b3d433cc23b72f000000000000000000000000000000000000000000000007dc477bc1cfa4000000000000000000000000000000000000000000000000000000000000063023e70000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000014000000000000000000000000000000000000000000000000000000000000001600000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000005800000000000000000000000000000000000000000000000000000000003a40201093ced81987bf532c2b7907b2a8525cd0c17295bd6015b40000000000000000000000001111111254eeb25477b68fb85ed929f73a9605820000000000000000cfee7c08";
+            hex"12aa3caf000000000000000000000000cfd674f8731e801a4a15c1ae31770960e1afded10000000000000000000000000d500b1d8e8ef31e21c99d1db9a6444d3adf12700000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa841740000000000000000000000001093ced81987bf532c2b7907b2a8525cd0c172950000000000000000000000002e234dae75c793f67a35089c9d99245e1c58470b000000000000000000000000000000000000000000000007dc477bc1cfa4000000000000000000000000000000000000000000000000000000000000063023e70000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000014000000000000000000000000000000000000000000000000000000000000001600000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000005800000000000000000000000000000000000000000000000000000000003a40201093ced81987bf532c2b7907b2a8525cd0c17295bd6015b40000000000000000000000001111111254eeb25477b68fb85ed929f73a9605820000000000000000cfee7c08";
         uint256 deadline = block.timestamp + 1 hours;
         uint256 nonce = 0x11223344 + 2 ** 248;
         bytes32 digest = _digestPermit2(wmatic, address(pay), 145 ether, nonce, deadline);
@@ -101,7 +111,7 @@ contract ZkBobPayTest is AbstractPolygonForkTest {
         vm.prank(user1);
         vm.expectEmit(false, true, false, true);
         emit Pay(0, user1, zkAddress, 100e6, wmatic, note);
-        pay.pay(zkAddress, wmatic, 145 ether, 100e6, permit, oneInchData, note);
+        pay.pay(zkAddress, wmatic, 145 ether, 100e6, permit, oneInchRouter, oneInchData, note);
 
         assertGt(IERC20(usdc).balanceOf(address(pay)), 1e6);
         assertLt(IERC20(usdc).balanceOf(address(pay)), 10e6);
@@ -111,11 +121,11 @@ contract ZkBobPayTest is AbstractPolygonForkTest {
         deal(address(user1), 145 ether);
 
         bytes memory oneInchData =
-            hex"12aa3caf000000000000000000000000cfd674f8731e801a4a15c1ae31770960e1afded1000000000000000000000000eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee0000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa84174000000000000000000000000cfd674f8731e801a4a15c1ae31770960e1afded10000000000000000000000005615deb798bb3e4dfa0139dfa1b3d433cc23b72f000000000000000000000000000000000000000000000007dc477bc1cfa4000000000000000000000000000000000000000000000000000000000000063023e70000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000014000000000000000000000000000000000000000000000000000000000000001600000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000008600000000000000000000000000000000000000000000000000006800001a40410d500b1d8e8ef31e21c99d1db9a6444d3adf1270d0e30db048201093ced81987bf532c2b7907b2a8525cd0c172950d500b1d8e8ef31e21c99d1db9a6444d3adf1270bd6015b40000000000000000000000001111111254eeb25477b68fb85ed929f73a9605820000000000000000000000000000000000000000000000000000cfee7c08";
+            hex"12aa3caf000000000000000000000000cfd674f8731e801a4a15c1ae31770960e1afded1000000000000000000000000eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee0000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa84174000000000000000000000000cfd674f8731e801a4a15c1ae31770960e1afded10000000000000000000000002e234dae75c793f67a35089c9d99245e1c58470b000000000000000000000000000000000000000000000007dc477bc1cfa4000000000000000000000000000000000000000000000000000000000000063023e70000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000014000000000000000000000000000000000000000000000000000000000000001600000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000008600000000000000000000000000000000000000000000000000006800001a40410d500b1d8e8ef31e21c99d1db9a6444d3adf1270d0e30db048201093ced81987bf532c2b7907b2a8525cd0c172950d500b1d8e8ef31e21c99d1db9a6444d3adf1270bd6015b40000000000000000000000001111111254eeb25477b68fb85ed929f73a9605820000000000000000000000000000000000000000000000000000cfee7c08";
         vm.prank(user1);
         vm.expectEmit(false, true, false, true);
         emit Pay(0, user1, zkAddress, 100e6, address(0), note);
-        pay.pay{value: 145 ether}(zkAddress, address(0), 0, 100e6, "", oneInchData, note);
+        pay.pay{value: 145 ether}(zkAddress, address(0), 0, 100e6, "", oneInchRouter, oneInchData, note);
 
         assertGt(IERC20(usdc).balanceOf(address(pay)), 1e6);
         assertLt(IERC20(usdc).balanceOf(address(pay)), 10e6);