use @Zodomo/OAppUpgradeable

README.md

@@ -1,4 +1,4 @@
-## Upgradeable LayerZero V2 Foundry Starter Pack
+## Upgradeable LayerZero V2 Foundry Starter Pack 🛠️🚀
 
 This repository can be cloned to quickly start building upgradeable applications on top of LayerZero V2. It includes libraries required for development, contains test setup and working multichain deployment script written in Solidity.
 
@@ -33,7 +33,7 @@ $ forge build
 $ forge test
 ```
 
-Note: if you want to execute single test file you can add a flag, eg.: `--match-path ./test/CounterUpgradeability.t.sol`.
+*Note: if you want to execute single test file you can add a flag, eg.: `--match-path ./test/CounterUpgradeability.t.sol`.*
 
 ### Deploy
 
@@ -43,14 +43,14 @@ Dry run:
 forge script DeployCounter -s "deployCounterTestnet(uint256, uint256)" 1 1 --force --multi
 ```
 
-Note: `1 1` parameters are respectively: `uint256 _counterSalt, uint256 _counterProxySalt`. It affects generated addresses. If you have problem with deployment script failing try changing `1 1` to some random numbers instead. You can't deploy with the same salt twice - it fails with message: `script failed: <no data>`.
+*Note: `1 1` parameters are respectively: `uint256 _counterSalt, uint256 _counterProxySalt`. It affects generated addresses. If you have problem with deployment script failing try changing `1 1` to some random numbers instead. You can't deploy with the same salt twice - it fails with message: `script failed: <no data>`.*
 
 Deploy:
 ```shell
 forge script DeployCounter -s "deployCounterTestnet(uint256, uint256)" 1 1 --force --multi --broadcast
 ```
 
-Note: Don't use automatic `--verify` flag because it doesn't seem to work. Looks like Foundry error.
+*Note: Don't use automatic `--verify` flag because it doesn't seem to work. Looks like Foundry error.*
 
 ### Upgrade
 
@@ -129,7 +129,11 @@ This repository is, to a significant extent, a compilation of other people's wor
 
 LayerZero libraries and examples are based on: https://github.com/LayerZero-Labs/LayerZero-v2.
 
-Multichain script deployment setup is heavily based on: https://github.com/timurguvenkaya/foundry-multichain.
+Multichain script deployment setup is heavily based on: https://github.com/timurguvenkaya/foundry-multichain by @timurguvenkaya.
+
+LayerZero OApp Upgradeability is taken from: https://github.com/Zodomo/LayerZero-v2/tree/main by @Zodomo.
+
+*Note: Initially I have used my own implementation but I think @Zodomo version is slightly better structured. I've noticed that implementation after I created this repository.*
 
 ## License
 

src/OApp/OAppCoreInitializable.sol → lib/oapp-upgradeable/OAppCoreUpgradeable.sol

@@ -2,34 +2,27 @@
 
 pragma solidity ^0.8.22;
 
-import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
-import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import { OwnableUpgradeable } from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import { IOAppCore, ILayerZeroEndpointV2 } from "@layerzerolabs/lz-evm-oapp-v2/contracts/oapp/interfaces/IOAppCore.sol";
 
-import {IOAppCore, ILayerZeroEndpointV2} from "@layerzerolabs/lz-evm-oapp-v2/contracts/oapp/interfaces/IOAppCore.sol";
-import "forge-std/console2.sol";
 /**
- * @title OAppCoreInitializable
- * @dev Abstract contract implementing the IOAppCore interface with basic OApp configurations.
+ * @title OAppCoreUpgradeable
+ * @dev Abstract contract implementing the IOAppCore interface with basic OApp configurations and upgradeability.
+ * @author Zodomo, https://github.com/Zodomo/LayerZero-v2
  */
-
-abstract contract OAppCoreInitializable is Initializable, IOAppCore, OwnableUpgradeable {
+abstract contract OAppCoreUpgradeable is IOAppCore, OwnableUpgradeable {
     // The LayerZero endpoint associated with the given OApp
     ILayerZeroEndpointV2 public endpoint;
 
     // Mapping to store peers associated with corresponding endpoints
     mapping(uint32 eid => bytes32 peer) public peers;
 
-    constructor() {
-        _disableInitializers();
-    }
-
     /**
-     * @dev Initialize the OAppCore with the provided endpoint and owner.
-     * @param _endpoint The address of the LOCAL LayerZero endpoint.
-     * @param _owner The address of the owner of the OApp.
+     * @dev Constructor to initialize the OAppCore with the provided endpoint and owner.
+     * @param _endpoint The address of the LOCAL Layer Zero endpoint.
+     * @param _owner The address of the owner of the OAppCore.
      */
-    function initialize(address _endpoint, address _owner) public initializer {
-        __Ownable_init();
+    function _initializeOAppCore(address _endpoint, address _owner) internal onlyInitializing {
         _transferOwnership(_owner);
         endpoint = ILayerZeroEndpointV2(_endpoint);
         endpoint.setDelegate(_owner); // @dev By default, the owner is the delegate
@@ -73,4 +66,4 @@ abstract contract OAppCoreInitializable is Initializable, IOAppCore, OwnableUpgr
     function setDelegate(address _delegate) public onlyOwner {
         endpoint.setDelegate(_delegate);
     }
-}
+}

src/OApp/OAppReceiver.sol → lib/oapp-upgradeable/OAppReceiverUpgradeable.sol

@@ -2,16 +2,16 @@
 
 pragma solidity ^0.8.22;
 
-import {
-    ILayerZeroReceiver, Origin
-} from "@layerzerolabs/lz-evm-protocol-v2/contracts/interfaces/ILayerZeroReceiver.sol";
-import {OAppCoreInitializable} from "./OAppCoreInitializable.sol";
+import { ILayerZeroReceiver, Origin } from "@layerzerolabs/lz-evm-protocol-v2/contracts/interfaces/ILayerZeroReceiver.sol";
+import { OAppCoreUpgradeable } from "./OAppCoreUpgradeable.sol";
 
 /**
- * @title OAppReceiver
- * @dev Abstract contract implementing the ILayerZeroReceiver interface and extending OAppCore for OApp receivers.
+ * @title OAppReceiverUpgradeable
+ * @dev Abstract upgradeable contract implementing the ILayerZeroReceiver interface and extending OAppCore for OApp
+ * receivers.
+ * @author Zodomo, https://github.com/Zodomo/LayerZero-v2
  */
-abstract contract OAppReceiver is ILayerZeroReceiver, OAppCoreInitializable {
+abstract contract OAppReceiverUpgradeable is ILayerZeroReceiver, OAppCoreUpgradeable {
     // Custom error message for when the caller is not the registered endpoint/
     error OnlyEndpoint(address addr);
 
@@ -56,7 +56,7 @@ abstract contract OAppReceiver is ILayerZeroReceiver, OAppCoreInitializable {
      * @dev This is also enforced by the OApp.
      * @dev By default this is NOT enabled. ie. nextNonce is hardcoded to return 0.
      */
-    function nextNonce(uint32, /*_srcEid*/ bytes32 /*_sender*/ ) public view virtual returns (uint64 nonce) {
+    function nextNonce(uint32 /*_srcEid*/, bytes32 /*_sender*/) public view virtual returns (uint64 nonce) {
         return 0;
     }
 

src/OApp/OAppSender.sol → lib/oapp-upgradeable/OAppSenderUpgradeable.sol

@@ -2,19 +2,17 @@
 
 pragma solidity ^0.8.22;
 
-import {SafeERC20, IERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-import {
-    MessagingParams,
-    MessagingFee,
-    MessagingReceipt
-} from "@layerzerolabs/lz-evm-protocol-v2/contracts/interfaces/ILayerZeroEndpointV2.sol";
-import {OAppCoreInitializable} from "./OAppCoreInitializable.sol";
+import { SafeERC20, IERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import { MessagingParams, MessagingFee, MessagingReceipt } from "@layerzerolabs/lz-evm-protocol-v2/contracts/interfaces/ILayerZeroEndpointV2.sol";
+import { OAppCoreUpgradeable } from "./OAppCoreUpgradeable.sol";
 
 /**
- * @title OAppSender
- * @dev Abstract contract implementing the OAppSender functionality for sending messages to a LayerZero endpoint.
+ * @title OAppSenderUpgradeable
+ * @dev Abstract upgradeable contract implementing the OAppSender functionality for sending messages to a LayerZero
+ * endpoint.
+ * @author Zodomo, https://github.com/Zodomo/LayerZero-v2
  */
-abstract contract OAppSender is OAppCoreInitializable {
+abstract contract OAppSenderUpgradeable is OAppCoreUpgradeable {
     using SafeERC20 for IERC20;
 
     // Custom error messages
@@ -48,15 +46,17 @@ abstract contract OAppSender is OAppCoreInitializable {
      *      - nativeFee: The native fee for the message.
      *      - lzTokenFee: The LZ token fee for the message.
      */
-    function _quote(uint32 _dstEid, bytes memory _message, bytes memory _options, bool _payInLzToken)
-        internal
-        view
-        virtual
-        returns (MessagingFee memory fee)
-    {
-        return endpoint.quote(
-            MessagingParams(_dstEid, _getPeerOrRevert(_dstEid), _message, _options, _payInLzToken), address(this)
-        );
+    function _quote(
+        uint32 _dstEid,
+        bytes memory _message,
+        bytes memory _options,
+        bool _payInLzToken
+    ) internal view virtual returns (MessagingFee memory fee) {
+        return
+            endpoint.quote(
+                MessagingParams(_dstEid, _getPeerOrRevert(_dstEid), _message, _options, _payInLzToken),
+                address(this)
+            );
     }
 
     /**
@@ -84,11 +84,12 @@ abstract contract OAppSender is OAppCoreInitializable {
         uint256 messageValue = _payNative(_fee.nativeFee);
         if (_fee.lzTokenFee > 0) _payLzToken(_fee.lzTokenFee);
 
-        return endpoint
+        return
             // solhint-disable-next-line check-send-result
-            .send{value: messageValue}(
-            MessagingParams(_dstEid, _getPeerOrRevert(_dstEid), _message, _options, _fee.lzTokenFee > 0), _refundAddress
-        );
+            endpoint.send{ value: messageValue }(
+                MessagingParams(_dstEid, _getPeerOrRevert(_dstEid), _message, _options, _fee.lzTokenFee > 0),
+                _refundAddress
+            );
     }
 
     /**

lib/oapp-upgradeable/OAppUpgradeable.sol

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.22;
+
+// @dev Import the 'MessagingFee' so it's exposed to OApp implementers
+// solhint-disable-next-line no-unused-import
+import { OAppSenderUpgradeable, MessagingFee } from "./OAppSenderUpgradeable.sol";
+// @dev Import the 'Origin' so it's exposed to OApp implementers
+// solhint-disable-next-line no-unused-import
+import { OAppReceiverUpgradeable, Origin } from "./OAppReceiverUpgradeable.sol";
+import { OAppCoreUpgradeable } from "./OAppCoreUpgradeable.sol";
+
+/**
+ * @title OAppUpgradeable
+ * @dev Abstract contract serving as the base for OAppUpgradeable implementation, combining OAppSenderUpgradeable and
+ * OAppReceiverUpgradeable functionality.
+ * @author Zodomo, https://github.com/Zodomo/LayerZero-v2
+ */
+abstract contract OAppUpgradeable is OAppSenderUpgradeable, OAppReceiverUpgradeable {
+    /**
+     * @dev Initializer for the upgradeable OApp with the provided endpoint and owner.
+     * @param _endpoint The address of the LOCAL LayerZero endpoint.
+     * @param _owner The address of the owner of the OApp.
+     */
+    function _initializeOApp(address _endpoint, address _owner) internal virtual onlyInitializing {
+        _initializeOAppCore(_endpoint, _owner);
+    }
+
+    /**
+     * @notice Retrieves the OApp version information.
+     * @return senderVersion The version of the OAppSender.sol implementation.
+     * @return receiverVersion The version of the OAppReceiver.sol implementation.
+     */
+    function oAppVersion()
+        public
+        pure
+        virtual
+        override(OAppSenderUpgradeable, OAppReceiverUpgradeable)
+        returns (uint64 senderVersion, uint64 receiverVersion)
+    {
+        return (SENDER_VERSION, RECEIVER_VERSION);
+    }
+}

remappings.txt

@@ -1,8 +1,8 @@
 @openzeppelin/contracts/=lib/openzeppelin-contracts/contracts/
 @openzeppelin/contracts-upgradeable/=lib/openzeppelin-contracts-upgradeable/contracts/
-@openzeppelin/contracts-upgradeable/=lib/openzeppelin-contracts-upgradeable/contracts/
 @layerzerolabs/lz-evm-oapp-v2/=lib/LayerZero-v2/oapp/
 @layerzerolabs/lz-evm-protocol-v2/=lib/LayerZero-v2/protocol/
 @layerzerolabs/lz-evm-messagelib-v2/=lib/LayerZero-v2/messagelib/
 @layerzerolabs/lz-evm-v1-0.7/=lib/LayerZero/
-solidity-bytes-utils/=lib/solidity-bytes-utils/
+solidity-bytes-utils/=lib/solidity-bytes-utils/
+@zodomo/oapp-upgradeable/=lib/oapp-upgradeable/

script/Counter.s.sol

@@ -6,7 +6,6 @@ import {Script, console2} from "forge-std/Script.sol";
 import {BaseDeployer} from "./BaseDeployer.s.sol";
 import {Counter} from "../src/Counter.sol";
 import {UUPSProxy} from "../src/UUPSProxy.sol";
-import {OAppCoreInitializable} from "../src/OApp/OAppCoreInitializable.sol";
 
 contract DeployCounter is Script, BaseDeployer {
     address private create2addrCounter;
@@ -61,7 +60,7 @@ contract DeployCounter is Script, BaseDeployer {
                     type(UUPSProxy).creationCode,
                     abi.encode(
                         create2addrCounter,
-                        abi.encodeWithSelector(OAppCoreInitializable.initialize.selector, lzEndpoints[i], ownerAddress)
+                        abi.encodeWithSelector(Counter.initialize.selector, lzEndpoints[i], ownerAddress)
                     )
                 )
             );
@@ -96,8 +95,7 @@ contract DeployCounter is Script, BaseDeployer {
         console2.log("Counter address:", address(counter), "\n");
 
         proxyCounter = new UUPSProxy{salt: counterProxySalt}(
-            address(counter),
-            abi.encodeWithSelector(OAppCoreInitializable.initialize.selector, lzEndpoint, ownerAddress)
+            address(counter), abi.encodeWithSelector(Counter.initialize.selector, lzEndpoint, ownerAddress)
         );
 
         proxyCounterAddress = address(proxyCounter);

src/Counter.sol

@@ -2,14 +2,26 @@
 pragma solidity ^0.8.19;
 
 import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import {OAppUpgradeable, MessagingFee, Origin} from "@zodomo/oapp-upgradeable/OAppUpgradeable.sol";
 
-import {OAppInitializable, MessagingFee, Origin} from "./OApp/OAppInitializable.sol";
-
-contract Counter is OAppInitializable, UUPSUpgradeable {
+contract Counter is OAppUpgradeable, UUPSUpgradeable {
     bytes public constant MESSAGE = "";
 
     uint256 public count;
 
+    constructor() {
+        _disableInitializers();
+    }
+
+    /**
+     * @dev Initialize the OApp with the provided endpoint and owner.
+     * @param _endpoint The address of the LOCAL LayerZero endpoint.
+     * @param _owner The address of the owner of the OApp.
+     */
+    function initialize(address _endpoint, address _owner) public initializer {
+        _initializeOApp(_endpoint, _owner);
+    }
+
     function increment(uint32 _dstEid, bytes calldata _options) public payable {
         _lzSend(
             _dstEid, // Destination chain's endpoint ID.

test/Counter.t.sol

@@ -7,7 +7,6 @@ import {OptionsBuilder} from "@layerzerolabs/lz-evm-oapp-v2/contracts/oapp/libs/
 
 import {Counter} from "../src/Counter.sol";
 import {UUPSProxy} from "../src/UUPSProxy.sol";
-import {OAppCoreInitializable} from "../src/OApp/OAppCoreInitializable.sol";
 import {ProxyTestHelper} from "./utils/ProxyTestHelper.sol";
 
 contract CounterTest is ProxyTestHelper {

test/CounterUpgradeability.t.sol

@@ -8,7 +8,6 @@ import {OptionsBuilder} from "@layerzerolabs/lz-evm-oapp-v2/contracts/oapp/libs/
 import {Counter} from "../src/Counter.sol";
 import {Counter2} from "./mocks/Counter2.sol";
 import {UUPSProxy} from "../src/UUPSProxy.sol";
-import {OAppCoreInitializable} from "../src/OApp/OAppCoreInitializable.sol";
 import {ProxyTestHelper} from "./utils/ProxyTestHelper.sol";
 
 contract CounterUpgradeabilityTest is ProxyTestHelper {