[#47] Update RegistryManager

contracts/base/FunctionHandlerManager.sol

@@ -21,7 +21,6 @@ abstract contract FunctionHandlerManager is RegistryManager {
     event FunctionHandlerChanged(address indexed safe, bytes4 indexed selector, address indexed functionHandler);
 
     // Errors
-    error AddressDoesNotImplementFunctionHandlerInterface(address functionHandler);
     error FunctionHandlerNotSet(address safe, bytes4 functionSelector);
 
     /**
@@ -36,10 +35,13 @@ abstract contract FunctionHandlerManager is RegistryManager {
 
     /**
      * @notice Sets the function handler for a Safe account and function selector. The msg.sender must be the account.
+     *         This function checks if the functionHandler address is listed and not flagged in the registry.
      * @param selector bytes4 function selector
      * @param functionHandler Address of the contract to be set as a function handler
      */
-    function setFunctionHandler(bytes4 selector, address functionHandler) external onlyPermittedIntegration(functionHandler) {
+    function setFunctionHandler(bytes4 selector, address functionHandler) external {
+        if (functionHandler != address(0)) checkPermittedIntegration(functionHandler);
+
         // No need to check if functionHandler implements expected interfaceId as check will be done when adding to registry.
         functionHandlers[msg.sender][selector] = functionHandler;
         emit FunctionHandlerChanged(msg.sender, selector, functionHandler);

contracts/base/RegistryManager.sol

@@ -3,6 +3,7 @@ pragma solidity ^0.8.18;
 import {Ownable2Step} from "@openzeppelin/contracts/access/Ownable2Step.sol";
 import {ISafeProtocolRegistry} from "../interfaces/Registry.sol";
 import {IERC165} from "@openzeppelin/contracts/utils/introspection/IERC165.sol";
+import {RegistryReadable} from "./RegistryReadable.sol";
 
 contract RegistryManager is Ownable2Step {
     address public registry;
@@ -12,12 +13,8 @@ contract RegistryManager is Ownable2Step {
     error IntegrationNotPermitted(address plugin, uint64 listedAt, uint64 flaggedAt);
     error AccountDoesNotImplementValidInterfaceId(address account);
 
-    modifier onlyPermittedIntegration(address plugin) {
-        // Only allow registered and non-flagged plugins
-        (uint64 listedAt, uint64 flaggedAt) = ISafeProtocolRegistry(registry).check(plugin);
-        if (listedAt == 0 || flaggedAt != 0) {
-            revert IntegrationNotPermitted(plugin, listedAt, flaggedAt);
-        }
+    modifier onlyPermittedIntegration(address integration) {
+        checkPermittedIntegration(integration);
         _;
     }
 
@@ -29,6 +26,19 @@ contract RegistryManager is Ownable2Step {
         registry = _registry;
     }
 
+    /**
+     * @notice Checks if given integration address is listed and not flagged in the registry.
+     *         Reverts if given address is not-listed or flagged.
+     * @param integration Address of the integration
+     */
+    function checkPermittedIntegration(address integration) internal view {
+        // Only allow registered and non-flagged integrations
+        (uint64 listedAt, uint64 flaggedAt) = ISafeProtocolRegistry(registry).check(integration);
+        if (listedAt == 0 || flaggedAt != 0) {
+            revert IntegrationNotPermitted(integration, listedAt, flaggedAt);
+        }
+    }
+
     /**
      * @notice Allows only owner to update the address of a registry. Emits event RegistryChanged(egistry, newRegistry)
      * @param newRegistry Address of new registry

test/FunctionHandlerManager.spec.ts

@@ -4,7 +4,7 @@ import { getMockFunctionHandler } from "./utils/mockFunctionHandlerBuilder";
 import { IntegrationType } from "./utils/constants";
 import { expect } from "chai";
 import { getMockTestExecutorInstance, getInstance } from "./utils/contracts";
-import { MaxUint256 } from "ethers";
+import { MaxUint256, ZeroAddress } from "ethers";
 import { ISafeProtocolFunctionHandler__factory, MockContract } from "../typechain-types";
 
 describe("Test Function Handler", async () => {
@@ -57,11 +57,42 @@ describe("Test Function Handler", async () => {
         expect(await functionHandlerManager.getFunctionHandler.staticCall(safe.target, functionId)).to.be.equal(mockFunctionHandler.target);
     });
 
-    it("Should not allow non permitted function handler", async () => {
+    it("Should allow removing function handler", async () => {
+        const { safe, functionHandlerManager, mockFunctionHandler } = await setupTests();
+
+        // 0xf8a8fd6d -> function test() external {}
+        const functionId = "0xf8a8fd6d";
+        const dataSetFunctionHandler = functionHandlerManager.interface.encodeFunctionData("setFunctionHandler", [
+            functionId,
+            mockFunctionHandler.target,
+        ]);
+
+        await safe.executeCallViaMock(functionHandlerManager, 0n, dataSetFunctionHandler, MaxUint256);
+
+        const dataSetFunctionHandler2 = functionHandlerManager.interface.encodeFunctionData("setFunctionHandler", [
+            functionId,
+            ZeroAddress,
+        ]);
+        const tx = await safe.executeCallViaMock(functionHandlerManager, 0n, dataSetFunctionHandler2, MaxUint256);
+
+        const receipt = await tx.wait();
+        const events = (
+            await functionHandlerManager.queryFilter(
+                functionHandlerManager.filters.FunctionHandlerChanged,
+                receipt?.blockNumber,
+                receipt?.blockNumber,
+            )
+        )[0];
+        expect(events.args).to.deep.equal([safe.target, functionId, ZeroAddress]);
+
+        expect(await functionHandlerManager.getFunctionHandler.staticCall(safe.target, functionId)).to.be.equal(ZeroAddress);
+    });
+
+    it("Should not allow non-permitted function handler", async () => {
         const { functionHandlerManager } = await setupTests();
-        await expect(functionHandlerManager.setFunctionHandler("0x00000000", hre.ethers.ZeroAddress))
+        await expect(functionHandlerManager.setFunctionHandler("0x00000000", user1.address))
             .to.be.revertedWithCustomError(functionHandlerManager, "IntegrationNotPermitted")
-            .withArgs(hre.ethers.ZeroAddress, 0, 0);
+            .withArgs(user1.address, 0, 0);
     });
 
     it("Should revert with FunctionHandlerNotSet when function handler is not enabled", async () => {