Refactor message bus hierarchy

contracts/src/cross_chain_messaging/L1/MerkleTreeMessageBus.sol

@@ -3,7 +3,7 @@ pragma solidity >=0.7.0 <0.9.0;
 
 import "../../common/Structs.sol";
 import "./IMerkleTreeMessageBus.sol";
-import "../common/MessageBus.sol";
+import "../common/BaseMessageBus.sol";
 import "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol";
 import "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol";
 
@@ -13,7 +13,7 @@ import "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol"
  * This contract manages state roots and verifies message inclusion through Merkle proofs.
  * It implements a role-based access control system for state root and withdrawal management.
  */
-contract MerkleTreeMessageBus is IMerkleTreeMessageBus, MessageBus, AccessControlUpgradeable {
+contract MerkleTreeMessageBus is IMerkleTreeMessageBus, BaseMessageBus, AccessControlUpgradeable {
 
     /**
      * @dev Role identifier for accounts that can manage state roots
@@ -32,15 +32,15 @@ contract MerkleTreeMessageBus is IMerkleTreeMessageBus, MessageBus, AccessContro
     mapping(bytes32 stateRoot => uint256 activationTime) rootValidAfter;
 
     /// @custom:oz-upgrades-unsafe-allow constructor
-    constructor() MessageBus() {
+    constructor() BaseMessageBus() {
         // Constructor intentionally left empty
     }
 
     /**
      * @dev Initializes the contract with provided owner
      * @param initialOwner Address that will be granted the DEFAULT_ADMIN_ROLE and STATE_ROOT_MANAGER_ROLE
      */
-    function initialize(address initialOwner, address withdrawalManager) public override(IMerkleTreeMessageBus, MessageBus) initializer {
+    function initialize(address initialOwner, address withdrawalManager) public override(IMerkleTreeMessageBus, BaseMessageBus) initializer {
         // Initialize parent contracts
         //super.initialize(initialOwner, address(0));
         __Ownable_init(initialOwner);

contracts/src/cross_chain_messaging/common/BaseMessageBus.sol

@@ -0,0 +1,68 @@
+// SPDX-License-Identifier: Apache 2
+pragma solidity >=0.7.0 <0.9.0;
+
+import "../../common/Structs.sol";
+import "../../system/interfaces/IFees.sol";
+import "./IMessageBus.sol";
+import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import "../../common/UnrenouncableOwnable2Step.sol";
+
+/// @title BaseMessageBus
+/// @dev Implements the chain-agnostic publish API shared by L1 and L2 buses.
+abstract contract BaseMessageBus is IMessageBus, Initializable, UnrenouncableOwnable2Step {
+    IFees internal fees;
+    mapping(address sender => uint64 sequence) internal addressSequences;
+
+    constructor() {
+        _transferOwnership(msg.sender);
+    }
+
+    /// @dev Initializes the contract with an owner and optional fee contract.
+    function initialize(address caller, address feesAddress) public virtual initializer {
+        __Ownable_init(caller);
+        fees = IFees(feesAddress);
+    }
+
+    /// @dev Increments and returns the sequence number for a sender.
+    function incrementSequence(address sender) internal returns (uint64 sequence) {
+        sequence = addressSequences[sender];
+        addressSequences[sender] += 1;
+    }
+
+    /// @inheritdoc IMessageBus
+    function getPublishFee() public view virtual override returns (uint256) {
+        return fees.messageFee();
+    }
+
+    /// @inheritdoc IMessageBus
+    function publishMessage(
+        uint32 nonce,
+        uint32 topic,
+        bytes calldata payload,
+        uint8 consistencyLevel
+    ) external payable virtual override returns (uint64 sequence) {
+        if (address(fees) != address(0)) {
+            uint256 fee = getPublishFee();
+            require(msg.value >= fee, "Insufficient funds to publish message");
+            (bool ok, ) = address(fees).call{value: fee}("");
+            require(ok, "Failed to send fees to fees contract");
+        }
+
+        sequence = incrementSequence(msg.sender);
+        emit LogMessagePublished(
+            msg.sender,
+            sequence,
+            nonce,
+            topic,
+            payload,
+            consistencyLevel
+        );
+        return sequence;
+    }
+
+    /// @inheritdoc IMessageBus
+    function retrieveAllFunds(address receiver) external virtual override onlyOwner {
+        (bool ok, ) = receiver.call{value: address(this).balance}("");
+        require(ok, "failed sending value");
+    }
+}

contracts/src/cross_chain_messaging/common/CrossChainMessenger.sol

@@ -5,6 +5,7 @@ pragma solidity >=0.7.0 <0.9.0;
 import "./ICrossChainMessenger.sol";
 import "@openzeppelin/contracts/proxy/utils/Initializable.sol";
 import "../L1/IMerkleTreeMessageBus.sol";
+import "./IL2MessageBus.sol";
 
 
 /**
@@ -53,7 +54,7 @@ contract CrossChainMessenger is ICrossChainMessenger, Initializable {
         Structs.CrossChainMessage calldata message
     ) private {
         require(
-            IMessageBus(address(messageBusContract)).verifyMessageFinalized(message),
+            IL2MessageBus(address(messageBusContract)).verifyMessageFinalized(message),
             "Message not found or finalized."
         );
         bytes32 msgHash = keccak256(abi.encode(message));

contracts/src/cross_chain_messaging/common/IL2MessageBus.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache 2
+pragma solidity >=0.7.0 <0.9.0;
+
+import "../../common/Structs.sol";
+import "./IMessageBus.sol";
+
+/// @title IL2MessageBus
+/// @dev Extension of IMessageBus with Layer 2 specific functionality
+interface IL2MessageBus is IMessageBus {
+    /// @dev Verifies that a cross chain message has been submitted from the other network
+    ///      and that its challenge period has passed.
+    function verifyMessageFinalized(
+        Structs.CrossChainMessage calldata crossChainMessage
+    ) external view returns (bool);
+
+    /// @dev Returns the timestamp when a message becomes final. Reverts if the message
+    ///      was never submitted.
+    function getMessageTimeOfFinality(
+        Structs.CrossChainMessage calldata crossChainMessage
+    ) external view returns (uint256);
+
+    /// @dev Stores a message coming from the other linked layer.
+    /// @param crossChainMessage The message to store
+    /// @param finalAfterTimestamp Seconds after which the message is considered final
+    function storeCrossChainMessage(
+        Structs.CrossChainMessage calldata crossChainMessage,
+        uint256 finalAfterTimestamp
+    ) external;
+}

contracts/src/cross_chain_messaging/common/IMessageBus.sol

@@ -1,55 +1,31 @@
 // SPDX-License-Identifier: Apache 2
-
 pragma solidity >=0.7.0 <0.9.0;
 
-import "../../common/Structs.sol";
-
-// This represents the smart contract API that allows dApps and accounts to send and verify received messages 
-// between layer 1 and layer 2.
+// This interface defines the common publish API shared between
+// Layer 1 and Layer 2 message buses.
 interface IMessageBus {
-
-    // The event emitted whenever a message is published. 
-    // The enclave listens for it on the deployed message bus addresses.
-    event LogMessagePublished
-    (
-        address sender, 
-        uint64 sequence, 
-        uint32 nonce, 
-        uint32 topic, 
-        bytes payload, 
+    /// @dev Emitted whenever a message is published.
+    event LogMessagePublished(
+        address sender,
+        uint64 sequence,
+        uint32 nonce,
+        uint32 topic,
+        bytes payload,
         uint8 consistencyLevel
     );
 
-    // This method is called from contracts to publish messages to the other linked message bus.
-    // nonce - This is provided and serves as deduplication nonce. It can also be used to group a batch of messages together.
-    // topic - This is the topic for which the payload is published. 
-    // payload - This is the actual message.
-    // consistencyLevel - this is how many block confirmations to wait before publishing the message. 
-    // Notice that consistencyLevel == 0 is still secure, but might make your protocol result more prone to reorganizations.
-    // returns sequence - this is the unique id of the published message for the address calling the function. It can be used
-    // to determine the order of incoming messages on the other side and if something is missing.
+    /// @dev Publishes a message to the linked message bus.
+    /// @return sequence Unique id of the published message for the sender.
     function publishMessage(
         uint32 nonce,
         uint32 topic,
-        bytes calldata payload, 
+        bytes calldata payload,
         uint8 consistencyLevel
     ) external payable returns (uint64 sequence);
 
-    // This function verifies that a cross chain message provided by the caller has indeed been submitted from the other network
-    // and returns true only if the challenge period for the message has passed.
-    function verifyMessageFinalized(Structs.CrossChainMessage calldata crossChainMessage) external view returns (bool);
-
-    // Returns the time when a message is final (when the rollup challenge period has passed). If the message was never submitted the call will revert.
-    function getMessageTimeOfFinality(Structs.CrossChainMessage calldata crossChainMessage) external view returns (uint256);
-
-    // This is the smart contract function which is used to store messages sent from the other linked layer. 
-    // The function will be called by the ManagementContract on L1 and the enclave on L2. 
-    // It should be access controlled and called according to the consistencyLevel and Obscuro platform rules.
-    function storeCrossChainMessage(Structs.CrossChainMessage calldata crossChainMessage, uint256 finalAfterTimestamp) external;
-
-    // This is a testnet function which allows the bridge owner to retrieve all funds from the message bus.
+    /// @dev Testnet utility function to retrieve all funds from the message bus.
     function retrieveAllFunds(address receiver) external;
 
-    // the fee needed to be paid in msg.value to publish the value transfer
+    /// @dev Returns the fee required in msg.value to publish a message.
     function getPublishFee() external view returns (uint256);
 }

contracts/src/cross_chain_messaging/common/MessageBus.sol

@@ -5,33 +5,20 @@ import "../../common/Structs.sol";
 import "../../system/contracts/Fees.sol";
 
 import "../../system/interfaces/IFees.sol";
-import "./IMessageBus.sol";
-import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
+import "./BaseMessageBus.sol";
+import "./IL2MessageBus.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
-import "../../common/UnrenouncableOwnable2Step.sol";
 
 /**
  * @title MessageBus
  * @dev Implementation of the IMessageBus interface for cross-layer message handling.
  * Manages message publishing, verification, and value transfers between L1 and L2.
  */
-contract MessageBus is IMessageBus, Initializable, UnrenouncableOwnable2Step {
-    /// @custom:oz-upgrades-unsafe-allow constructor
+contract MessageBus is BaseMessageBus, IL2MessageBus {
     constructor() {
         _disableInitializers();
     }
 
-    /**
-     * @dev Initializes the contract with an owner and fees contract
-     * @param caller The address to set as the owner
-     * @param feesAddress The address of the fees contract
-     */
-    function initialize(address caller, address feesAddress) public virtual initializer {
-        __UnrenouncableOwnable2Step_init(caller);  // Initialize UnrenouncableOwnable2Step
-        fees = IFees(feesAddress);
-    }
-
-
     /**
      * @dev Modifier to restrict access to owner or self
      * Since this contract exists on L2, when messages are added from L1,
@@ -52,62 +39,6 @@ contract MessageBus is IMessageBus, Initializable, UnrenouncableOwnable2Step {
     // The stored messages, currently unconsumed.
     mapping(address sender => mapping(uint32 topic => Structs.CrossChainMessage[] messages)) messages;
 
-    // This stores the current sequence number that each address has reached.
-    // Whenever a message is published, this sequence number increments.
-    // This gives ordering to messages, guaranteed by us.
-    mapping(address sender => uint64 sequence) addressSequences;
-
-    IFees fees;
-
-    /**
-     * @dev Increments and returns the sequence number for a sender
-     * @param sender The address to increment the sequence for
-     * @return sequence The previous sequence number
-     */
-    function incrementSequence(
-        address sender
-    ) internal returns (uint64 sequence) {
-        sequence = addressSequences[sender];
-        addressSequences[sender] += 1;
-    }
-
-    function getPublishFee() public view returns (uint256) {
-        return fees.messageFee();
-    }
-
-    /**
-     * @dev Publishes a message to the other linked message bus
-     * @param nonce Deduplication nonce, can group messages together
-     * @param topic The topic for which the payload is published
-     * @param payload The actual message content
-     * @param consistencyLevel Block confirmations to wait. Level 0 is secure but more prone to reorganizations
-     * @return sequence Unique ID of the published message for the calling address
-     */
-    function publishMessage(
-        uint32 nonce,
-        uint32 topic,
-        bytes calldata payload,
-        uint8 consistencyLevel
-    ) external payable override returns (uint64 sequence) {
-        if (address(fees) != address(0)) { // No fee required for L1 to L2 messages.
-            uint256 fee = getPublishFee();
-            require(msg.value >= fee, "Insufficient funds to publish message");
-            (bool ok, ) = address(fees).call{value: fee}("");
-            require(ok, "Failed to send fees to fees contract");
-        }
-
-        sequence = incrementSequence(msg.sender);
-        emit LogMessagePublished(
-            msg.sender,
-            sequence,
-            nonce,
-            topic,
-            payload,
-            consistencyLevel
-        );
-        return sequence;
-    }
-
     /**
      * @dev Verifies  that a cross chain message provided by the caller has indeed been submitted from the other network
      *  and returns true only if the challenge period for the message has passed.
@@ -165,17 +96,6 @@ contract MessageBus is IMessageBus, Initializable, UnrenouncableOwnable2Step {
         );
     }
 
-    /**
-     * @dev Retrieves all funds from the contract (Testnet only - to be removed before mainnet deployment)
-     * @param receiver The address to receive the funds
-     */
-    function retrieveAllFunds(
-        address receiver
-    ) external onlyOwner {
-        (bool ok, ) = receiver.call{value: address(this).balance}("");
-        require(ok, "failed sending value");
-    }
-
     fallback() external {
         revert("unsupported");
     }