Merge pull request #44 from datachainlab/initializable-erc20transfermodule

YukiTsuchida · web-flow · commit 8f9934f64956 · 2025-12-10T15:42:40.000+09:00
fix: Initializable ERC20TransferModule
diff --git a/src/example/ERC20TransferModule.sol b/src/example/ERC20TransferModule.sol
@@ -3,15 +3,21 @@ pragma solidity ^0.8.20;
 
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 import {ContractModuleBase} from "../core/ContractModuleBase.sol";
 import {CrossContext} from "../core/IContractModule.sol";
 
-abstract contract ERC20TransferModule is ContractModuleBase {
+abstract contract ERC20TransferModule is Initializable, ContractModuleBase, Ownable {
     using SafeERC20 for IERC20;
 
     error ERC20TransferModuleInvalidCallInfo();
     error ERC20TransferModuleTxAlreadyPending();
     error ERC20TransferModuleUnauthorized();
+    error ERC20TransferModuleNotInitialized();
+    error ERC20TransferModuleInvalidAddress();
+
+    event ERC20TransferModuleInitialized(address indexed crossModule, address indexed token);
 
     struct PendingTx {
         address from;
@@ -22,17 +28,25 @@ abstract contract ERC20TransferModule is ContractModuleBase {
     // txID => PendingTx
     mapping(bytes32 => PendingTx) public pendingTxs;
 
-    address public immutable CROSS_MODULE;
-    IERC20 public immutable TOKEN;
+    address public crossModule;
+    IERC20 public token;
+
+    constructor() Ownable(msg.sender) {}
 
     modifier onlyCrossModule() {
-        if (msg.sender != CROSS_MODULE) revert ERC20TransferModuleUnauthorized();
+        if (crossModule == address(0)) revert ERC20TransferModuleNotInitialized();
+        if (msg.sender != crossModule) revert ERC20TransferModuleUnauthorized();
         _;
     }
 
-    constructor(address _crossModule, address _token) {
-        CROSS_MODULE = _crossModule;
-        TOKEN = IERC20(_token);
+    function initialize(address _crossModule, address _token) external initializer onlyOwner {
+        if (_crossModule == address(0) || _token == address(0)) {
+            revert ERC20TransferModuleInvalidAddress();
+        }
+        crossModule = _crossModule;
+        token = IERC20(_token);
+
+        emit ERC20TransferModuleInitialized(_crossModule, _token);
     }
 
     function decodeCallInfo(bytes calldata callInfo)
@@ -60,7 +74,7 @@ abstract contract ERC20TransferModule is ContractModuleBase {
 
         // IMPORTANT: The implementing contract MUST ensure in `_authorize` that the `from` address corresponds to the authenticated signer.
         // slither-disable-next-line arbitrary-send-erc20
-        TOKEN.safeTransferFrom(from, to, amount);
+        token.safeTransferFrom(from, to, amount);
 
         return "";
     }
@@ -82,7 +96,7 @@ abstract contract ERC20TransferModule is ContractModuleBase {
 
         // IMPORTANT: The implementing contract MUST ensure in `_authorize` that the `from` address corresponds to the authenticated signer.
         // slither-disable-next-line arbitrary-send-erc20
-        TOKEN.safeTransferFrom(from, address(this), amount);
+        token.safeTransferFrom(from, address(this), amount);
 
         return "";
     }
@@ -94,7 +108,7 @@ abstract contract ERC20TransferModule is ContractModuleBase {
         if (pending.from != address(0)) {
             delete pendingTxs[txID];
             // Transfer locked tokens to the destination
-            TOKEN.safeTransfer(pending.to, pending.amount);
+            token.safeTransfer(pending.to, pending.amount);
         }
     }
 
@@ -105,7 +119,7 @@ abstract contract ERC20TransferModule is ContractModuleBase {
         if (pending.from != address(0)) {
             delete pendingTxs[txID];
             // Refund tokens to the sender
-            TOKEN.safeTransfer(pending.from, pending.amount);
+            token.safeTransfer(pending.from, pending.amount);
         }
     }
 }
diff --git a/test/ERC20TransferModule.t.sol b/test/ERC20TransferModule.t.sol
@@ -9,6 +9,8 @@ import {Account as AuthAccount, AuthType} from "../src/proto/cross/core/auth/Aut
 import {GoogleProtobufAny} from "@hyperledger-labs/yui-ibc-solidity/contracts/proto/GoogleProtobufAny.sol";
 import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 import {IERC20Errors} from "@openzeppelin/contracts/interfaces/draft-IERC6093.sol";
+import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 
 contract MockERC20 is ERC20 {
     constructor() ERC20("Mock Token", "MCK") {}
@@ -19,8 +21,6 @@ contract MockERC20 is ERC20 {
 }
 
 contract ERC20TransferModuleHarness is ERC20TransferModule {
-    constructor(address _crossModule, address _token) ERC20TransferModule(_crossModule, _token) {}
-
     function _authorize(CrossContext calldata, bytes calldata) internal pure override {
         // allow all for testing
     }
@@ -45,7 +45,8 @@ contract ERC20ModuleTest is Test {
 
         token = new MockERC20();
 
-        harness = new ERC20TransferModuleHarness(address(this), address(token));
+        harness = new ERC20TransferModuleHarness();
+        harness.initialize(address(this), address(token));
 
         token.mint(sender, INITIAL_BALANCE);
     }
@@ -74,6 +75,54 @@ contract ERC20ModuleTest is Test {
         return abi.encode(_from, _to, _amount);
     }
 
+    // --- Initialization Tests ---
+
+    function test_initialize_Success() public {
+        ERC20TransferModuleHarness newHarness = new ERC20TransferModuleHarness();
+        address crossModule = makeAddr("newCrossModule");
+        address newToken = makeAddr("newToken");
+
+        vm.expectEmit(address(newHarness));
+        emit ERC20TransferModule.ERC20TransferModuleInitialized(crossModule, newToken);
+
+        newHarness.initialize(crossModule, newToken);
+
+        assertEq(newHarness.crossModule(), crossModule);
+        assertEq(address(newHarness.token()), newToken);
+    }
+
+    function test_initialize_RevertWhen_InvalidAddress() public {
+        ERC20TransferModuleHarness newHarness = new ERC20TransferModuleHarness();
+        address validAddress = makeAddr("valid");
+
+        // 1. CrossModule is zero
+        vm.expectRevert(ERC20TransferModule.ERC20TransferModuleInvalidAddress.selector);
+        newHarness.initialize(address(0), validAddress);
+
+        // 2. Token is zero
+        vm.expectRevert(ERC20TransferModule.ERC20TransferModuleInvalidAddress.selector);
+        newHarness.initialize(validAddress, address(0));
+    }
+
+    function test_initialize_RevertWhen_AlreadyInitialized() public {
+        vm.expectRevert(Initializable.InvalidInitialization.selector);
+        harness.initialize(address(this), address(token));
+    }
+
+    function test_initialize_RevertWhen_CallerIsNotOwner() public {
+        ERC20TransferModuleHarness newHarness = new ERC20TransferModuleHarness();
+
+        address attacker = makeAddr("attacker");
+        address crossModule = makeAddr("crossModule");
+        address newToken = makeAddr("newToken");
+
+        vm.prank(attacker);
+
+        vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, attacker));
+
+        newHarness.initialize(crossModule, newToken);
+    }
+
     // --- Decode Tests ---
 
     function test_decodeCallInfo_Success() public view {
@@ -136,6 +185,15 @@ contract ERC20ModuleTest is Test {
         harness.onContractCommitImmediately(context, callInfo);
     }
 
+    function test_onContractCommitImmediately_RevertWhen_NotInitialized() public {
+        ERC20TransferModuleHarness uninitHarness = new ERC20TransferModuleHarness();
+        CrossContext memory context = _createContext(sender);
+        bytes memory callInfo = _createCallInfo(sender, receiver, AMOUNT);
+
+        vm.expectRevert(ERC20TransferModule.ERC20TransferModuleNotInitialized.selector);
+        uninitHarness.onContractCommitImmediately(context, callInfo);
+    }
+
     // --- Prepare Tests ---
 
     function test_onContractPrepare_Success() public {
@@ -216,6 +274,15 @@ contract ERC20ModuleTest is Test {
         harness.onContractPrepare(context, callInfo);
     }
 
+    function test_onContractPrepare_RevertWhen_NotInitialized() public {
+        ERC20TransferModuleHarness uninitHarness = new ERC20TransferModuleHarness();
+        CrossContext memory context = _createContext(sender);
+        bytes memory callInfo = _createCallInfo(sender, receiver, AMOUNT);
+
+        vm.expectRevert(ERC20TransferModule.ERC20TransferModuleNotInitialized.selector);
+        uninitHarness.onContractPrepare(context, callInfo);
+    }
+
     // --- Commit Tests ---
 
     function test_onCommit_Success() public {
@@ -260,6 +327,14 @@ contract ERC20ModuleTest is Test {
         harness.onCommit(context);
     }
 
+    function test_onCommit_RevertWhen_NotInitialized() public {
+        ERC20TransferModuleHarness uninitHarness = new ERC20TransferModuleHarness();
+        CrossContext memory context = _createContext(sender);
+
+        vm.expectRevert(ERC20TransferModule.ERC20TransferModuleNotInitialized.selector);
+        uninitHarness.onCommit(context);
+    }
+
     // --- Abort Tests ---
 
     function test_onAbort_Success() public {
@@ -305,4 +380,12 @@ contract ERC20ModuleTest is Test {
         vm.expectRevert(ERC20TransferModule.ERC20TransferModuleUnauthorized.selector);
         harness.onAbort(context);
     }
+
+    function test_onAbort_RevertWhen_NotInitialized() public {
+        ERC20TransferModuleHarness uninitHarness = new ERC20TransferModuleHarness();
+        CrossContext memory context = _createContext(sender);
+
+        vm.expectRevert(ERC20TransferModule.ERC20TransferModuleNotInitialized.selector);
+        uninitHarness.onAbort(context);
+    }
 }
