Registry checks for hooks

contracts/SafeProtocolManager.sol

@@ -327,7 +327,8 @@ contract SafeProtocolManager is ISafeProtocolManager, RegistryManager, HooksMana
         address msgSender
     ) external {
         // Store hooks address in tempHooksAddress so that checkAfterExecution(...) and checkModuleTransaction(...) can access it.
-        address tempHooksAddressForSafe = tempHooksAddress[msg.sender] = enabledHooks[msg.sender];
+        tempHooksAddress[msg.sender] = enabledHooks[msg.sender];
+        address tempHooksAddressForSafe = enabledHooks[msg.sender];
 
         if (tempHooksAddressForSafe == address(0)) return;
         bytes memory executionMetadata = abi.encode(

contracts/base/HooksManager.sol

@@ -1,9 +1,11 @@
 // SPDX-License-Identifier: LGPL-3.0-only
 pragma solidity ^0.8.18;
 import {ISafeProtocolHooks} from "../interfaces/Integrations.sol";
+
+import {RegistryManager} from "./RegistryManager.sol";
 import {OnlyAccountCallable} from "./OnlyAccountCallable.sol";
 
-contract HooksManager is OnlyAccountCallable {
+abstract contract HooksManager is RegistryManager, OnlyAccountCallable {
     mapping(address => address) public enabledHooks;
 
     /// @notice This variable should store the address of the hooks contract whenever
@@ -14,9 +16,6 @@ contract HooksManager is OnlyAccountCallable {
     // Events
     event HooksChanged(address indexed safe, address indexed hooksAddress);
 
-    // Errors
-    error AddressDoesNotImplementHooksInterface(address hooksAddress);
-
     /**
      * @notice Returns the address of hooks for a Safe account provided as a fucntion parameter.
      *         Returns address(0) is no hooks are enabled.
@@ -32,8 +31,10 @@ contract HooksManager is OnlyAccountCallable {
      * @param hooks Address of the hooks to be enabled for msg.sender.
      */
     function setHooks(address hooks) external onlyAccount {
-        if (hooks != address(0) && !ISafeProtocolHooks(hooks).supportsInterface(type(ISafeProtocolHooks).interfaceId)) {
-            revert AddressDoesNotImplementHooksInterface(hooks);
+        if (hooks != address(0)) {
+            checkPermittedIntegration(hooks);
+            if (!ISafeProtocolHooks(hooks).supportsInterface(type(ISafeProtocolHooks).interfaceId))
+                revert AccountDoesNotImplementValidInterfaceId(hooks);
         }
         enabledHooks[msg.sender] = hooks;
         emit HooksChanged(msg.sender, hooks);

contracts/test/TestExecutor.sol

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: LGPL-3.0-only
 pragma solidity ^0.8.18;
 import {ISafe} from "../interfaces/Accounts.sol";
-import {MockContract} from "@safe-global/mock-contract/contracts/MockContract.sol";
 
 contract TestExecutor is ISafe {
     address public module;

test/SafeProtocolManager.spec.ts

@@ -475,6 +475,7 @@ describe("SafeProtocolManager", async () => {
                 const { safeProtocolManager, safe, safeProtocolRegistry } = await loadFixture(deployContractsWithEnabledManagerFixture);
                 // Enable hooks on a safe
                 const hooks = await getHooksWithPassingChecks();
+                await safeProtocolRegistry.connect(owner).addIntegration(hooks.target, IntegrationType.Hooks);
                 const dataSetHooks = safeProtocolManager.interface.encodeFunctionData("setHooks", [await hooks.getAddress()]);
                 await safe.exec(safe.target, 0, dataSetHooks);
 
@@ -509,6 +510,7 @@ describe("SafeProtocolManager", async () => {
                 const { safeProtocolManager, safe, safeProtocolRegistry } = await loadFixture(deployContractsWithEnabledManagerFixture);
                 // Enable hooks on a safe
                 const hooks = await getHooksWithFailingPrechecks();
+                await safeProtocolRegistry.connect(owner).addIntegration(hooks.target, IntegrationType.Hooks);
 
                 const dataSetHooks = safeProtocolManager.interface.encodeFunctionData("setHooks", [await hooks.getAddress()]);
                 await safe.exec(safe.target, 0, dataSetHooks);
@@ -529,6 +531,7 @@ describe("SafeProtocolManager", async () => {
                 const { safeProtocolManager, safe, safeProtocolRegistry } = await loadFixture(deployContractsWithEnabledManagerFixture);
                 // Enable hooks on a safe
                 const hooks = await getHooksWithFailingPostCheck();
+                await safeProtocolRegistry.connect(owner).addIntegration(hooks.target, IntegrationType.Hooks);
 
                 const dataSetHooks = safeProtocolManager.interface.encodeFunctionData("setHooks", [await hooks.getAddress()]);
                 await safe.exec(safe.target, 0, dataSetHooks);
@@ -671,8 +674,11 @@ describe("SafeProtocolManager", async () => {
             it("Should execute a transaction from root access enabled plugin with hooks enabled", async () => {
                 const { safeProtocolManager, safe, safeProtocolRegistry } = await loadFixture(deployContractsWithEnabledManagerFixture);
                 const safeAddress = await safe.getAddress();
+
                 // Enable hooks on a safe
                 const hooks = await getHooksWithPassingChecks();
+                await safeProtocolRegistry.connect(owner).addIntegration(hooks.target, IntegrationType.Hooks);
+
                 const dataSetHooks = safeProtocolManager.interface.encodeFunctionData("setHooks", [await hooks.getAddress()]);
                 await safe.exec(safe.target, 0, dataSetHooks);
 
@@ -716,6 +722,7 @@ describe("SafeProtocolManager", async () => {
                 const { safeProtocolManager, safe, safeProtocolRegistry } = await loadFixture(deployContractsWithEnabledManagerFixture);
                 // Enable hooks on a safe
                 const hooks = await getHooksWithFailingPrechecks();
+                await safeProtocolRegistry.connect(owner).addIntegration(hooks.target, IntegrationType.Hooks);
 
                 const dataSetHooks = safeProtocolManager.interface.encodeFunctionData("setHooks", [await hooks.getAddress()]);
                 await safe.exec(safe.target, 0, dataSetHooks);
@@ -748,6 +755,8 @@ describe("SafeProtocolManager", async () => {
 
                 // Enable hooks on a safe
                 const hooks = await getHooksWithFailingPostCheck();
+                await safeProtocolRegistry.connect(owner).addIntegration(hooks.target, IntegrationType.Hooks);
+
                 const dataSetHooks = safeProtocolManager.interface.encodeFunctionData("setHooks", [await hooks.getAddress()]);
                 await safe.exec(safe.target, 0, dataSetHooks);
 
@@ -928,6 +937,7 @@ describe("SafeProtocolManager", async () => {
 
             await safeProtocolRegistry.connect(owner).addIntegration(hooks.target, IntegrationType.Hooks);
             await safeProtocolRegistry.connect(owner).addIntegration(hooksWithFailingPreChecks.target, IntegrationType.Hooks);
+            await safeProtocolRegistry.connect(owner).addIntegration(hooksWithFailingPostCheck.target, IntegrationType.Hooks);
 
             return { safe, safeProtocolManager, hooks, hooksWithFailingPreChecks, hooksWithFailingPostCheck };
         });

test/base/HooksManager.spec.ts

@@ -3,6 +3,7 @@ import { expect } from "chai";
 import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";
 import { getHooksWithPassingChecks } from "../utils/mockHooksBuilder";
 import { ZeroAddress } from "ethers";
+import { IntegrationType } from "../utils/constants";
 
 describe("HooksManager", async () => {
     let deployer: SignerWithAddress, user1: SignerWithAddress, owner: SignerWithAddress;
@@ -21,8 +22,9 @@ describe("HooksManager", async () => {
 
         const safe = await hre.ethers.deployContract("TestExecutor", [hooksManager.target], { signer: deployer });
         const hooks = await getHooksWithPassingChecks();
+        await safeProtocolRegistry.connect(owner).addIntegration(hooks.target, IntegrationType.Hooks);
 
-        return { hooksManager, hooks, safe };
+        return { hooksManager, hooks, safe, safeProtocolRegistry };
     });
 
     it("Should emit HooksChanged event when hooks are enabled", async () => {
@@ -65,15 +67,17 @@ describe("HooksManager", async () => {
         await expect(hooksManager.setHooks(hooksAddress)).to.be.reverted;
     });
 
-    it("Should revert AddressDoesNotImplementHooksInterface if user attempts address does not implement Hooks interface", async () => {
-        const { hooksManager, safe } = await setupTests();
+    it("Should revert AccountDoesNotImplementValidInterfaceId if user attempts address does not implement Hooks interface", async () => {
+        const { hooksManager, safe, safeProtocolRegistry } = await setupTests();
         const contractNotImplementingHooksInterface = await (await hre.ethers.getContractFactory("MockContract")).deploy();
-        await contractNotImplementingHooksInterface.givenMethodReturnBool("0x01ffc9a7", false);
+        await contractNotImplementingHooksInterface.givenMethodReturnBool("0x01ffc9a7", true);
+        await safeProtocolRegistry.connect(owner).addIntegration(contractNotImplementingHooksInterface.target, IntegrationType.Hooks);
 
+        await contractNotImplementingHooksInterface.givenMethodReturnBool("0x01ffc9a7", false);
         const calldata = hooksManager.interface.encodeFunctionData("setHooks", [contractNotImplementingHooksInterface.target]);
         await expect(safe.exec(safe.target, 0n, calldata)).to.be.revertedWithCustomError(
             hooksManager,
-            "AddressDoesNotImplementHooksInterface",
+            "AccountDoesNotImplementValidInterfaceId",
         );
     });
 });