From b02f9d258461697b354434fe9b9f162aec0dce7f Mon Sep 17 00:00:00 2001
From: Konstantin Fastov <kfastov@gmail.com>
Date: Mon, 16 Sep 2024 12:11:58 +0300
Subject: [PATCH] Fix issues from audit (#61)

* fix: reject on incorrect fee, fix comment

* test: fix tests after adding fee check
---
 src/CrossDomainOwnableLinea.sol |  4 ++--
 src/LineaStateBridge.sol        | 23 +++++++++++++++++++++++
 test/LineaStateBridge.t.sol     |  6 ++++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/CrossDomainOwnableLinea.sol b/src/CrossDomainOwnableLinea.sol
index 4961c4f..0264f7f 100644
--- a/src/CrossDomainOwnableLinea.sol
+++ b/src/CrossDomainOwnableLinea.sol
@@ -11,8 +11,8 @@ import { IMessageService } from "linea-contracts/interfaces/IMessageService.sol"
 abstract contract CrossDomainOwnableLinea is Ownable {
     IMessageService public messageService;
 
-    /// @notice If true, the contract uses the cross domain _checkOwner function override.
-    ///         If false it uses the standard Ownable _checkOwner function.
+    /// @notice If true, the owner is on the same layer as this contract (L2).
+    ///         If false, the owner is on a different layer (L1) and ownership checks use cross-layer messaging.
     bool public isLocal = true;
 
     /// @notice Emits when ownership of the contract is transferred. Includes the
diff --git a/src/LineaStateBridge.sol b/src/LineaStateBridge.sol
index 415347b..6f41b5f 100644
--- a/src/LineaStateBridge.sol
+++ b/src/LineaStateBridge.sol
@@ -91,6 +91,9 @@ contract LineaStateBridge is Ownable2Step {
     /// @notice Emitted when an attempt is made to set an address to zero
     error AddressZero();
 
+    /// @notice Emitted when the incorrect message fee is sent
+    error IncorrectMessageFeeSent();
+
     ///////////////////////////////////////////////////////////////////
     ///                         CONSTRUCTOR                         ///
     ///////////////////////////////////////////////////////////////////
@@ -126,6 +129,11 @@ contract LineaStateBridge is Ownable2Step {
     function propagateRoot() external payable {
         uint256 latestRoot = IWorldIDIdentityManager(worldIDAddress).latestRoot();
 
+        // Ensure that the correct fee is sent with the transaction
+        if (msg.value != _feePropagateRoot) {
+            revert IncorrectMessageFeeSent();
+        }
+
         // The `encodeCall` function is strongly typed, so this checks that we are passing the
         // correct data to the Linea message service.
         bytes memory message = abi.encodeCall(ILineaWorldID.receiveRoot, (latestRoot));
@@ -149,6 +157,11 @@ contract LineaStateBridge is Ownable2Step {
             revert AddressZero();
         }
 
+        // Ensure that the correct fee is sent with the transaction
+        if (msg.value != _feeTransferOwnership) {
+            revert IncorrectMessageFeeSent();
+        }
+
         // Encoding the call to transferOwnership on ICrossDomainOwnableLinea
         bytes memory message = abi.encodeCall(ICrossDomainOwnableLinea.transferOwnership, (_owner, _isLocal));
 
@@ -167,6 +180,11 @@ contract LineaStateBridge is Ownable2Step {
             revert AddressZero();
         }
 
+        // Ensure that the correct fee is sent with the transaction
+        if (msg.value != _feeSetMessageService) {
+            revert IncorrectMessageFeeSent();
+        }
+
         // Encoding the call to setMessageService on ICrossDomainOwnableLinea
         bytes memory message = abi.encodeCall(ICrossDomainOwnableLinea.setMessageService, (_messageService));
 
@@ -181,6 +199,11 @@ contract LineaStateBridge is Ownable2Step {
     /// @notice Adds functionality to the StateBridge to set the root history expiry on LineaWorldID
     /// @param _rootHistoryExpiry new root history expiry
     function setRootHistoryExpiry(uint256 _rootHistoryExpiry) external payable onlyOwner {
+        // Ensure that the correct fee is sent with the transaction
+        if (msg.value != _feeSetRootHistoryExpiry) {
+            revert IncorrectMessageFeeSent();
+        }
+
         // The `encodeCall` function is strongly typed, so this checks that we are passing the
         // correct data to the linea bridge.
         bytes memory message = abi.encodeCall(IRootHistory.setRootHistoryExpiry, (_rootHistoryExpiry));
diff --git a/test/LineaStateBridge.t.sol b/test/LineaStateBridge.t.sol
index 82ab290..cf0fe40 100644
--- a/test/LineaStateBridge.t.sol
+++ b/test/LineaStateBridge.t.sol
@@ -101,6 +101,12 @@ contract LineaStateBridgeTest is PRBTest, StdCheats {
         lineaStateBridge =
             new LineaStateBridge(mockWorldIDAddress, lineaWorldIDAddress, lineaCrossDomainMessengerAddress);
 
+        // Set fee for every action to 1
+        lineaStateBridge.setFeePropagateRoot(1);
+        lineaStateBridge.setFeeSetRootHistoryExpiry(1);
+        lineaStateBridge.setFeeTransferOwnershipLinea(1);
+        lineaStateBridge.setFeeSetMessageServiceLinea(1);
+
         owner = lineaStateBridge.owner();
     }