From c982ed99e8bd1a01ec637aa1b9cd2c5ae69ddac4 Mon Sep 17 00:00:00 2001
From: royvardhan <royvardhan.official@gmail.com>
Date: Fri, 24 Jan 2025 18:12:56 +0530
Subject: [PATCH] feat: add pause/unpause methods

---
 foundry/src/TychoRouter.sol           | 22 ++++++++++++++++++++--
 foundry/test/Constants.sol            |  2 ++
 foundry/test/TychoRouter.t.sol        | 25 +++++++++++++++++++++++++
 foundry/test/TychoRouterTestSetup.sol |  2 ++
 4 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/foundry/src/TychoRouter.sol b/foundry/src/TychoRouter.sol
index 4383618..7a4df3f 100644
--- a/foundry/src/TychoRouter.sol
+++ b/foundry/src/TychoRouter.sol
@@ -7,6 +7,7 @@ import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@permit2/src/interfaces/IAllowanceTransfer.sol";
 import "./SwapExecutionDispatcher.sol";
 import "./CallbackVerificationDispatcher.sol";
+import "@openzeppelin/contracts/utils/Pausable.sol";
 
 error TychoRouter__WithdrawalFailed();
 error TychoRouter__AddressZero();
@@ -16,7 +17,8 @@ error TychoRouter__NonContractVerifier();
 contract TychoRouter is
     AccessControl,
     SwapExecutionDispatcher,
-    CallbackVerificationDispatcher
+    CallbackVerificationDispatcher,
+    Pausable
 {
     IAllowanceTransfer public immutable permit2;
 
@@ -29,6 +31,8 @@ contract TychoRouter is
         0xe6ad9a47fbda1dc18de1eb5eeb7d935e5e81b4748f3cfc61e233e64f88182060;
     bytes32 public constant PAUSER_ROLE =
         0x65d7a28e3265b37a6474929f336521b332c1681b933f6cb9f3376673440d862a;
+    bytes32 public constant UNPAUSER_ROLE =
+        0x427da25fe773164f88948d3e215c94b6554e2ed5e5f203a821c9f2f6131cf75a;
     bytes32 public constant FUND_RESCUER_ROLE =
         0x912e45d663a6f4cc1d0491d8f046e06c616f40352565ea1cdb86a0e1aaefa41b;
 
@@ -62,6 +66,20 @@ contract TychoRouter is
         // TODO execute generic callback
     }
 
+    /**
+     * @dev Pauses the contract
+     */
+    function pause() external onlyRole(PAUSER_ROLE) {
+        _pause();
+    }
+
+    /**
+     * @dev Unpauses the contract
+     */
+    function unpause() external onlyRole(UNPAUSER_ROLE) {
+        _unpause();
+    }
+
     /**
      * @dev Executes a swap graph supporting internal splits token amount
      *  splits, checking that the user gets more than minUserAmount of buyToken.
@@ -76,7 +94,7 @@ contract TychoRouter is
         bytes calldata swaps,
         IAllowanceTransfer.PermitSingle calldata permitSingle,
         bytes calldata signature
-    ) external returns (uint256 amountOut) {
+    ) external whenNotPaused returns (uint256 amountOut) {
         amountOut = 0;
         // TODO
     }
diff --git a/foundry/test/Constants.sol b/foundry/test/Constants.sol
index 521dff7..fe12d79 100644
--- a/foundry/test/Constants.sol
+++ b/foundry/test/Constants.sol
@@ -11,4 +11,6 @@ contract Constants is Test {
     // dummy contracts
     address DUMMY = makeAddr("dummy");
     address FEE_RECEIVER = makeAddr("feeReceiver");
+    address PAUSER = makeAddr("pauser");
+    address UNPAUSER = makeAddr("unpauser");
 }
diff --git a/foundry/test/TychoRouter.t.sol b/foundry/test/TychoRouter.t.sol
index fb9ad4b..ebfc6c8 100644
--- a/foundry/test/TychoRouter.t.sol
+++ b/foundry/test/TychoRouter.t.sol
@@ -206,4 +206,29 @@ contract TychoRouterTest is TychoRouterTestSetup {
         tychoRouter.setFeeReceiver(FEE_RECEIVER);
         vm.stopPrank();
     }
+
+    function testPause() public {
+        vm.startPrank(PAUSER);
+        assertEq(tychoRouter.paused(), false);
+        tychoRouter.pause();
+        assertEq(tychoRouter.paused(), true);
+        vm.stopPrank();
+
+        vm.startPrank(UNPAUSER);
+        tychoRouter.unpause();
+        assertEq(tychoRouter.paused(), false);
+        vm.stopPrank();
+    }
+
+    function testPauseFailures() public {
+        vm.startPrank(BOB);
+        vm.expectRevert();
+        tychoRouter.pause();
+        vm.stopPrank();
+
+        vm.startPrank(UNPAUSER);
+        vm.expectRevert();
+        tychoRouter.unpause();
+        vm.stopPrank();
+    }
 }
diff --git a/foundry/test/TychoRouterTestSetup.sol b/foundry/test/TychoRouterTestSetup.sol
index 9765d76..31fb266 100644
--- a/foundry/test/TychoRouterTestSetup.sol
+++ b/foundry/test/TychoRouterTestSetup.sol
@@ -17,6 +17,8 @@ contract TychoRouterTestSetup is Test, Constants {
         tychoRouter.grantRole(keccak256("EXECUTOR_SETTER_ROLE"), BOB);
         tychoRouter.grantRole(keccak256("FUND_RESCUER_ROLE"), FUND_RESCUER);
         tychoRouter.grantRole(keccak256("FEE_SETTER_ROLE"), FEE_SETTER);
+        tychoRouter.grantRole(keccak256("PAUSER_ROLE"), PAUSER);
+        tychoRouter.grantRole(keccak256("UNPAUSER_ROLE"), UNPAUSER);
         executorSetter = BOB;
         deployDummyContract();
         vm.stopPrank();