From 530add7ee332632b4ea48d5ddfcb6138efa1e86c Mon Sep 17 00:00:00 2001
From: Shawn <44221603+shaspitz@users.noreply.github.com>
Date: Wed, 6 Nov 2024 00:28:00 -0800
Subject: [PATCH] fix: issue 3.2.4 from audit report, introduce setRelayer
 function for gateway.sol

---
 contracts/contracts/interfaces/IGateway.sol   |  4 ++
 .../contracts/standard-bridge/Gateway.sol     | 10 ++--
 .../test/standard-bridge/L1GatewayTest.sol    | 18 +++++++
 .../standard-bridge/SettlementGatewayTest.sol | 50 +++++++++++++++++++
 4 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/contracts/contracts/interfaces/IGateway.sol b/contracts/contracts/interfaces/IGateway.sol
index a8d1e7f07..808e93aff 100644
--- a/contracts/contracts/interfaces/IGateway.sol
+++ b/contracts/contracts/interfaces/IGateway.sol
@@ -34,10 +34,14 @@ interface IGateway {
 
     event FinalizationFeeSet(uint256 finalizationFee);
     event CounterpartyFeeSet(uint256 counterpartyFee);
+    event RelayerSet(address indexed relayer);
 
     error SenderNotRelayer(address sender, address relayer);
     error AmountTooSmall(uint256 amount, uint256 counterpartyFee);
     error InvalidCounterpartyIndex(uint256 counterpartyIdx, uint256 transferFinalizedIdx);
+    error FinalizationFeeTooSmall(uint256 _finalizationFee);
+    error CounterpartyFeeTooSmall(uint256 _counterpartyFee);
+    error RelayerCannotBeZeroAddress();
 
     /**
      * @notice Initiates a cross-chain transfer.
diff --git a/contracts/contracts/standard-bridge/Gateway.sol b/contracts/contracts/standard-bridge/Gateway.sol
index 70632d404..f4b991551 100644
--- a/contracts/contracts/standard-bridge/Gateway.sol
+++ b/contracts/contracts/standard-bridge/Gateway.sol
@@ -11,9 +11,6 @@ import {GatewayStorage} from "./GatewayStorage.sol";
 abstract contract Gateway is IGateway, GatewayStorage,
     Ownable2StepUpgradeable, UUPSUpgradeable, PausableUpgradeable, ReentrancyGuardUpgradeable {   
 
-    error FinalizationFeeTooSmall(uint256 _finalizationFee);
-    error CounterpartyFeeTooSmall(uint256 _counterpartyFee);
-
     modifier onlyRelayer() {
         require(msg.sender == relayer, SenderNotRelayer(msg.sender, relayer));
         _;
@@ -53,6 +50,13 @@ abstract contract Gateway is IGateway, GatewayStorage,
     /// @dev Allows owner to unpause the contract.
     function unpause() external onlyOwner { _unpause(); }
 
+    /// @dev Allows owner to set a new relayer account.
+    function setRelayer(address _relayer) external onlyOwner {
+        require(_relayer != address(0), RelayerCannotBeZeroAddress());
+        relayer = _relayer;
+        emit RelayerSet(_relayer);
+    }
+
     /// @dev Allows owner to set a new finalization fee.
     /// @notice If using this function, ensure the same value is set as the `counterpartyFee` in the counterparty contract.
     function setFinalizationFee(uint256 _finalizationFee) external onlyOwner {
diff --git a/contracts/test/standard-bridge/L1GatewayTest.sol b/contracts/test/standard-bridge/L1GatewayTest.sol
index 9ab923cb5..d2cecc0d2 100644
--- a/contracts/test/standard-bridge/L1GatewayTest.sol
+++ b/contracts/test/standard-bridge/L1GatewayTest.sol
@@ -71,6 +71,24 @@ contract L1GatewayTest is Test {
         assertEq(l1Gateway.counterpartyFee(), 0.0005 ether);
     }
 
+    event RelayerSet(address indexed relayer);
+
+    function test_SetRelayer() public {
+
+        vm.expectRevert(
+            abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, vm.addr(888))
+        );
+        vm.prank(vm.addr(888));
+        l1Gateway.setRelayer(address(0x123));
+
+        assertEq(l1Gateway.relayer(), address(0x78));
+
+        vm.expectEmit(true, true, true, true);
+        emit RelayerSet(address(0x12345));
+        l1Gateway.setRelayer(address(0x12345));
+        assertEq(l1Gateway.relayer(), address(0x12345));
+    }
+
     // Expected event signature emitted in initiateTransfer()
     event TransferInitiated(
         address indexed sender, address indexed recipient, uint256 amount, uint256 indexed transferIdx);
diff --git a/contracts/test/standard-bridge/SettlementGatewayTest.sol b/contracts/test/standard-bridge/SettlementGatewayTest.sol
index 7ef8769be..b0b89c283 100644
--- a/contracts/test/standard-bridge/SettlementGatewayTest.sol
+++ b/contracts/test/standard-bridge/SettlementGatewayTest.sol
@@ -9,6 +9,7 @@ import {IGateway} from "../../contracts/interfaces/IGateway.sol";
 import {IAllocator} from "../../contracts/interfaces/IAllocator.sol";
 import {RevertingReceiver} from "./RevertingReceiver.sol";
 import {EventReceiver} from "./EventReceiver.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 
 contract SettlementGatewayTest is Test {
 
@@ -65,6 +66,55 @@ contract SettlementGatewayTest is Test {
     event TransferNeedsWithdrawal(address indexed recipient, uint256 amount);
     event TransferSuccess(address indexed recipient, uint256 amount);
 
+    event FinalizationFeeSet(uint256 finalizationFee);
+    event CounterpartyFeeSet(uint256 counterpartyFee);
+
+    function test_SetFinalizationFee() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, vm.addr(888))
+        );
+        vm.prank(vm.addr(888));
+        settlementGateway.setFinalizationFee(0.0015 ether);
+
+        assertEq(settlementGateway.finalizationFee(), 0.05 ether);
+        vm.expectEmit(true, true, true, true);
+        emit FinalizationFeeSet(0.0015 ether);
+        settlementGateway.setFinalizationFee(0.0015 ether);
+        assertEq(settlementGateway.finalizationFee(), 0.0015 ether);
+    }
+
+    function test_SetCounterpartyFee() public {
+        vm.expectRevert(
+            abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, vm.addr(888))
+        );
+        vm.prank(vm.addr(888));
+        settlementGateway.setCounterpartyFee(0.0005 ether);
+
+        assertEq(settlementGateway.counterpartyFee(), 0.1 ether);
+        vm.expectEmit(true, true, true, true);
+        emit CounterpartyFeeSet(0.0005 ether);
+        settlementGateway.setCounterpartyFee(0.0005 ether);
+        assertEq(settlementGateway.counterpartyFee(), 0.0005 ether);
+    }
+
+    event RelayerSet(address indexed relayer);
+
+    function test_SetRelayer() public {
+
+        vm.expectRevert(
+            abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, vm.addr(888))
+        );
+        vm.prank(vm.addr(888));
+        settlementGateway.setRelayer(address(0x123));
+
+        assertEq(settlementGateway.relayer(), address(0x78));
+
+        vm.expectEmit(true, true, true, true);
+        emit RelayerSet(address(0x12345));
+        settlementGateway.setRelayer(address(0x12345));
+        assertEq(settlementGateway.relayer(), address(0x12345));
+    }
+
     function test_InitiateTransferSuccess() public {
         vm.deal(bridgeUser, 100 ether);
         uint256 amount = 7 ether;