TREX-162 : bind and set module

CHANGELOG.md

@@ -37,6 +37,10 @@ All notable changes to this project will be documented in this file.
     - `IAFactory`
   - Each contract now implements the `supportsInterface` function to identify the supported interfaces, ensuring compliance with ERC-165 standards.
 
+- **Add and Set Module in a Single Transaction**:
+  - Introduced the `addAndSetModule` function in the `ModularCompliance` contract, allowing the contract owner to add a new compliance module and interact with it in a single transaction.
+  - This function supports adding a module and performing up to 5 interactions with the module in one call, streamlining the setup process for compliance modules.
+
 ### Updated
 
 - **Token Recovery Function**:

contracts/compliance/modular/IModularCompliance.sol

@@ -138,6 +138,40 @@ interface IModularCompliance {
      */
     function callModuleFunction(bytes calldata callData, address _module) external;
 
+    /**
+     * @dev Adds a module to the modular compliance contract and performs multiple interactions with it in a single transaction.
+     *
+     * This function allows the contract owner to add a new compliance module and immediately configure it by calling
+     * specified functions on the module. This can be useful for setting up the module with initial parameters or configurations
+     * right after it is added.
+     *
+     * @param _module The address of the module to add. The module must either be "plug and play"
+     * or be able to bind with the compliance contract.
+     * @param _interactions An array of bytecode representing function calls to be made on the module.
+     * These interactions are performed after the module is added.
+     *
+     * Requirements:
+     * - The caller must be the owner of the `ModularCompliance` contract.
+     * - The `_module` address must not be a zero address.
+     * - The `_module` must not already be bound to the contract.
+     * - The total number of modules must not exceed 25 after adding the new module.
+     * - The `_interactions` array must contain no more than 5 elements to prevent out-of-gas errors.
+     *
+     * Operations:
+     * - The function first adds the module using the `addModule` function.
+     * - Then, it iterates over the `_interactions` array, performing each
+     *   interaction on the module using the `callModuleFunction`.
+     *
+     * Emits:
+     * - A `ModuleAdded` event upon successful addition of the module.
+     * - A `ModuleInteraction` event for each function call made to the module during the interaction phase.
+     *
+     * Reverts if:
+     * - Any of the above requirements are not met.
+     * - Any of the module interactions fail during execution.
+     */
+    function addAndSetModule(address _module, bytes[] calldata _interactions) external;
+
     /**
      *  @dev function called whenever tokens are transferred
      *  from one wallet to another

contracts/compliance/modular/ModularCompliance.sol

@@ -68,6 +68,7 @@ import "./IModularCompliance.sol";
 import "./MCStorage.sol";
 import "./modules/IModule.sol";
 import "../../errors/ComplianceErrors.sol";
+import "../../errors/CommonErrors.sol";
 import "../../errors/InvalidArgumentErrors.sol";
 import "@openzeppelin/contracts/utils/introspection/IERC165.sol";
 import "../../roles/IERC173.sol";
@@ -128,23 +129,6 @@ contract ModularCompliance is IModularCompliance, OwnableUpgradeable, MCStorage,
         emit TokenUnbound(_token);
     }
 
-    /**
-     *  @dev See {IModularCompliance-addModule}.
-     */
-    function addModule(address _module) external override onlyOwner {
-        require(_module != address(0), ZeroAddress());
-        require(!_moduleBound[_module], ModuleAlreadyBound());
-        require(_modules.length <= 24, MaxModulesReached(25));
-        IModule module = IModule(_module);
-        require(module.isPlugAndPlay() || module.canComplianceBind(address(this)), 
-            ComplianceNotSuitableForBindingToModule(_module));
-
-        module.bindCompliance(address(this));
-        _modules.push(_module);
-        _moduleBound[_module] = true;
-        emit ModuleAdded(_module);
-    }
-
     /**
      *  @dev See {IModularCompliance-removeModule}.
      */
@@ -204,40 +188,14 @@ contract ModularCompliance is IModularCompliance, OwnableUpgradeable, MCStorage,
     }
 
     /**
-     *  @dev see {IModularCompliance-callModuleFunction}.
+     *  @dev See {IModularCompliance-addAndSetModule}.
      */
-    function callModuleFunction(bytes calldata callData, address _module) external override onlyOwner {
-        require(_moduleBound[_module], ModuleNotBound());
-        // NOTE: Use assembly to call the interaction instead of a low level
-        // call for two reasons:
-        // - We don't want to copy the return data, since we discard it for
-        // interactions.
-        // - Solidity will under certain conditions generate code to copy input
-        // calldata twice to memory (the second being a "memcopy loop").
-        // solhint-disable-next-line no-inline-assembly
-        assembly {
-            let freeMemoryPointer := mload(0x40) // Load the free memory pointer from memory location 0x40
-
-            // Copy callData from calldata to the free memory location
-            calldatacopy(freeMemoryPointer, callData.offset, callData.length) 
-
-            if iszero( // Check if the call returns zero (indicating failure)
-                call( // Perform the external call
-                    gas(), // Provide all available gas
-                    _module, // Address of the target module
-                    0, // No ether is sent with the call
-                    freeMemoryPointer, // Input data starts at the free memory pointer
-                    callData.length, // Input data length
-                    0, // Output data location (not used)
-                    0 // Output data size (not used)
-                )
-            ) {
-                returndatacopy(0, 0, returndatasize()) // Copy return data to memory starting at position 0
-                revert(0, returndatasize()) // Revert the transaction with the return data
-            }
+    function addAndSetModule(address _module, bytes[] calldata _interactions) external onlyOwner override {
+        require(_interactions.length <= 5, ArraySizeLimited(5));
+        addModule(_module);
+        for (uint256 i = 0; i < _interactions.length; i++) {
+            callModuleFunction(_interactions[i], _module);
         }
-
-        emit ModuleInteraction(_module, _selector(callData));
     }
 
     /**
@@ -274,6 +232,60 @@ contract ModularCompliance is IModularCompliance, OwnableUpgradeable, MCStorage,
         return true;
     }
 
+    /**
+     *  @dev See {IModularCompliance-addModule}.
+     */
+    function addModule(address _module) public override onlyOwner {
+        require(_module != address(0), ZeroAddress());
+        require(!_moduleBound[_module], ModuleAlreadyBound());
+        require(_modules.length <= 24, MaxModulesReached(25));
+        IModule module = IModule(_module);
+        require(module.isPlugAndPlay() || module.canComplianceBind(address(this)),
+            ComplianceNotSuitableForBindingToModule(_module));
+
+        module.bindCompliance(address(this));
+        _modules.push(_module);
+        _moduleBound[_module] = true;
+        emit ModuleAdded(_module);
+    }
+
+    /**
+     *  @dev see {IModularCompliance-callModuleFunction}.
+     */
+    function callModuleFunction(bytes calldata callData, address _module) public override onlyOwner {
+        require(_moduleBound[_module], ModuleNotBound());
+        // NOTE: Use assembly to call the interaction instead of a low level
+        // call for two reasons:
+        // - We don't want to copy the return data, since we discard it for
+        // interactions.
+        // - Solidity will under certain conditions generate code to copy input
+        // calldata twice to memory (the second being a "memcopy loop").
+        // solhint-disable-next-line no-inline-assembly
+        assembly {
+            let freeMemoryPointer := mload(0x40) // Load the free memory pointer from memory location 0x40
+
+        // Copy callData from calldata to the free memory location
+            calldatacopy(freeMemoryPointer, callData.offset, callData.length)
+
+            if iszero( // Check if the call returns zero (indicating failure)
+                call( // Perform the external call
+                    gas(), // Provide all available gas
+                    _module, // Address of the target module
+                    0, // No ether is sent with the call
+                    freeMemoryPointer, // Input data starts at the free memory pointer
+                    callData.length, // Input data length
+                    0, // Output data location (not used)
+                    0 // Output data size (not used)
+                )
+            ) {
+                returndatacopy(0, 0, returndatasize()) // Copy return data to memory starting at position 0
+                revert(0, returndatasize()) // Revert the transaction with the return data
+            }
+        }
+
+        emit ModuleInteraction(_module, _selector(callData));
+    }
+
     /**
      *  @dev See {IERC165-supportsInterface}.
      */

test/compliances/module-country-allow.test.ts

@@ -523,4 +523,103 @@ describe('CountryAllowModule', () => {
       expect(await countryAllowModule.supportsInterface(ierc165InterfaceId)).to.equal(true);
     });
   });
+  describe('.addAndSetModule()', () => {
+    describe('when module is already bound', () => {
+      it('should revert', async () => {
+        const {
+          suite: { compliance, countryAllowModule },
+          accounts: { deployer },
+        } = await loadFixture(deployComplianceWithCountryAllowModule);
+
+        await expect(compliance.connect(deployer).addAndSetModule(countryAllowModule.target, [])).to.be.revertedWithCustomError(
+          compliance,
+          'ModuleAlreadyBound',
+        );
+      });
+    });
+    describe('when calling batchAllowCountries not as owner', () => {
+      it('should revert', async () => {
+        const {
+          suite: { compliance, countryAllowModule },
+          accounts: { deployer, anotherWallet },
+        } = await loadFixture(deployComplianceWithCountryAllowModule);
+
+        // Remove the module first
+        await compliance.connect(deployer).removeModule(countryAllowModule.target);
+
+        await expect(
+          compliance
+            .connect(anotherWallet)
+            .addAndSetModule(countryAllowModule.target, [
+              new ethers.Interface(['function batchAllowCountries(uint16[] calldata countries)']).encodeFunctionData('batchAllowCountries', [[42]]),
+            ]),
+        ).to.be.revertedWith('Ownable: caller is not the owner');
+      });
+    });
+    describe('when providing more than 5 interactions', () => {
+      it('should revert', async () => {
+        const {
+          suite: { compliance, countryAllowModule },
+          accounts: { deployer },
+        } = await loadFixture(deployComplianceWithCountryAllowModule);
+
+        // Remove the module first
+        await compliance.connect(deployer).removeModule(countryAllowModule.target);
+
+        const interactions = Array.from({ length: 6 }, () =>
+          new ethers.Interface(['function batchAllowCountries(uint16[] calldata countries)']).encodeFunctionData('batchAllowCountries', [[42]]),
+        );
+
+        await expect(compliance.connect(deployer).addAndSetModule(countryAllowModule.target, interactions)).to.be.revertedWithCustomError(
+          compliance,
+          'ArraySizeLimited',
+        );
+      });
+    });
+    describe('when calling addAllowedCountry twice with the same country code', () => {
+      it('should revert with CountryAlreadyAllowed', async () => {
+        const {
+          suite: { compliance, countryAllowModule },
+          accounts: { deployer },
+        } = await loadFixture(deployComplianceWithCountryAllowModule);
+
+        // Remove the module first to simulate a fresh start
+        await compliance.connect(deployer).removeModule(countryAllowModule.target);
+
+        await expect(
+          compliance
+            .connect(deployer)
+            .addAndSetModule(countryAllowModule.target, [
+              new ethers.Interface(['function addAllowedCountry(uint16 country)']).encodeFunctionData('addAllowedCountry', [42]),
+              new ethers.Interface(['function addAllowedCountry(uint16 country)']).encodeFunctionData('addAllowedCountry', [42]),
+            ]),
+        ).to.be.revertedWithCustomError(countryAllowModule, 'CountryAlreadyAllowed');
+      });
+    });
+    describe('when calling batchAllowCountries twice successfully', () => {
+      it('should add the module and interact with it', async () => {
+        const {
+          suite: { compliance, countryAllowModule },
+          accounts: { deployer },
+        } = await loadFixture(deployComplianceWithCountryAllowModule);
+
+        // Remove the module first
+        await compliance.connect(deployer).removeModule(countryAllowModule.target);
+
+        const tx = await compliance
+          .connect(deployer)
+          .addAndSetModule(countryAllowModule.target, [
+            new ethers.Interface(['function batchAllowCountries(uint16[] calldata countries)']).encodeFunctionData('batchAllowCountries', [[42]]),
+            new ethers.Interface(['function batchAllowCountries(uint16[] calldata countries)']).encodeFunctionData('batchAllowCountries', [[66]]),
+          ]);
+
+        await expect(tx).to.emit(compliance, 'ModuleAdded').withArgs(countryAllowModule.target);
+        await expect(tx).to.emit(countryAllowModule, 'CountryAllowed').withArgs(compliance.target, 42);
+        await expect(tx).to.emit(countryAllowModule, 'CountryAllowed').withArgs(compliance.target, 66);
+
+        expect(await countryAllowModule.isCountryAllowed(compliance.target, 42)).to.be.true;
+        expect(await countryAllowModule.isCountryAllowed(compliance.target, 66)).to.be.true;
+      });
+    });
+  });
 });