Whitelist plugin (5afe#10)

* [5afe#7] Whitelist plugin

* [5afe#7] Format file, remove unused vars, update dataTypes.ts

* [5afe#7] Change logic

* [5afe#7] Update imports

* [5afe#7] Define OwnerManager interface

* Update contracts/contracts/WhitelistPlugin.sol

Co-authored-by: Mikhail <[email protected]>

* [5afe#7] Update docstring

---------

Co-authored-by: Mikhail <[email protected]>

contracts/contracts/WhitelistPlugin.sol

@@ -0,0 +1,84 @@
+// SPDX-License-Identifier: LGPL-3.0-only
+pragma solidity ^0.8.18;
+import {ISafe} from "@safe-global/safe-core-protocol/contracts/interfaces/Accounts.sol";
+import {ISafeProtocolPlugin} from "@safe-global/safe-core-protocol/contracts/interfaces/Integrations.sol";
+import {ISafeProtocolManager} from "@safe-global/safe-core-protocol/contracts/interfaces/Manager.sol";
+import {SafeTransaction, SafeRootAccess, SafeProtocolAction} from "@safe-global/safe-core-protocol/contracts/DataTypes.sol";
+import {BasePluginWithEventMetadata, PluginMetadata} from "./Base.sol";
+
+/**
+ * @title OwnerManager
+ * @dev This interface is defined for use in WhitelistPlugin contract.
+ */
+interface OwnerManager {
+    function isOwner(address owner) external view returns (bool);
+}
+
+/**
+ * @title WhitelistPlugin maintains a mapping that stores information about accounts that are
+ *        permitted to execute non-root transactions through a Safe account.
+ * @notice This plugin does not need Safe owner(s) confirmation(s) to execute Safe txs once enabled
+ *         through a Safe{Core} Protocol Manager.
+ */
+contract WhitelistPlugin is BasePluginWithEventMetadata {
+    // safe account => account => whitelist status
+    mapping(address => mapping(address => bool)) public whitelistedAddresses;
+
+    event AddressWhitelisted(address indexed account);
+    event AddressRemovedFromWhitelist(address indexed account);
+
+    error AddressNotWhiteListed(address account);
+    error CallerIsNotOwner(address safe, address caller);
+
+    constructor()
+        BasePluginWithEventMetadata(
+            PluginMetadata({name: "Whitelist Plugin", version: "1.0.0", requiresRootAccess: false, iconUrl: "", appUrl: ""})
+        )
+    {}
+
+    /**
+     * @notice Executes a Safe transaction if the caller is whitelisted for the given Safe account.
+     * @param manager Address of the Safe{Core} Protocol Manager.
+     * @param safe Safe account
+     * @param safetx SafeTransaction to be executed
+     */
+    function executeFromPlugin(
+        ISafeProtocolManager manager,
+        ISafe safe,
+        SafeTransaction calldata safetx
+    ) external returns (bytes[] memory data) {
+        address safeAddress = address(safe);
+        // Only Safe owners are allowed to execute transactions to whitelisted accounts.
+        if (!(OwnerManager(safeAddress).isOwner(msg.sender))) {
+            revert CallerIsNotOwner(safeAddress, msg.sender);
+        }
+
+        SafeProtocolAction[] memory actions = safetx.actions;
+        uint256 length = actions.length;
+        for (uint256 i = 0; i < length; i++) {
+            if (!whitelistedAddresses[safeAddress][actions[i].to]) revert AddressNotWhiteListed(actions[i].to);
+        }
+        // Test: Any tx that updates whitelist of this contract should be blocked
+        (data) = manager.executeTransaction(safe, safetx);
+    }
+
+    /**
+     * @notice Adds an account to whitelist mapping.
+     *         The caller should be a Safe account.
+     * @param account address of the account to be whitelisted
+     */
+    function addToWhitelist(address account) external {
+        whitelistedAddresses[msg.sender][account] = true;
+        emit AddressWhitelisted(account);
+    }
+
+    /**
+     * @notice Removes an account from whitelist mapping.
+     *         The caller should be a Safe account.
+     * @param account address of the account to be removed from the whitelist
+     */
+    function removeFromWhitelist(address account) external {
+        whitelistedAddresses[msg.sender][account] = false;
+        emit AddressRemovedFromWhitelist(account);
+    }
+}

contracts/src/deploy/deploy_plugin.ts

@@ -20,6 +20,13 @@ const deploy: DeployFunction = async function (hre: HardhatRuntimeEnvironment) {
         log: true,
         deterministicDeployment: true,
     });
+
+    await deploy("WhitelistPlugin", {
+        from: deployer,
+        args: [],
+        log: true,
+        deterministicDeployment: true,
+    });
 };
 
 deploy.tags = ["plugins"];

contracts/src/utils/builder.ts

@@ -1,6 +1,7 @@
 import { AddressLike } from "ethers";
 import { SafeRootAccess, SafeTransaction } from "./dataTypes";
-export const buildSingleTx = (address: AddressLike, value: bigint, data: string, nonce: bigint, metaHash: Uint8Array): SafeTransaction => {
+
+export const buildSingleTx = (address: AddressLike, value: bigint, data: string, nonce: bigint, metadataHash: Uint8Array | string): SafeTransaction => {
     return {
         actions: [
             {
@@ -10,18 +11,18 @@ export const buildSingleTx = (address: AddressLike, value: bigint, data: string,
             },
         ],
         nonce: nonce,
-        metaHash: metaHash,
+        metadataHash: metadataHash,
     };
 };
 
-export const buildRootTx = (address: AddressLike, value: bigint, data: string, nonce: bigint, metaHash: Uint8Array): SafeRootAccess => {
+export const buildRootTx = (address: AddressLike, value: bigint, data: string, nonce: bigint, metadataHash: Uint8Array | string): SafeRootAccess => {
     return {
         action: {
             to: address,
             value: value,
             data: data,
         },
         nonce: nonce,
-        metaHash: metaHash,
+        metadataHash: metadataHash,
     };
 };

contracts/src/utils/contracts.ts

@@ -1,9 +1,9 @@
-import { BaseContract } from "ethers";
-import { BasePlugin, RelayPlugin, TestSafeProtocolRegistryUnrestricted } from "../../typechain-types";
+import { Addressable, BaseContract } from "ethers";
+import { BasePlugin, RelayPlugin, TestSafeProtocolRegistryUnrestricted, WhitelistPlugin } from "../../typechain-types";
 import { HardhatRuntimeEnvironment } from "hardhat/types";
-import { getProtocolRegistryAddress } from "./protocol";
+import { getProtocolManagerAddress, getProtocolRegistryAddress } from "./protocol";
 
-export const getInstance = async <T extends BaseContract>(hre: HardhatRuntimeEnvironment, name: string, address: string): Promise<T> => {
+export const getInstance = async <T extends BaseContract>(hre: HardhatRuntimeEnvironment, name: string, address: string | Addressable): Promise<T> => {
     // TODO: this typecasting should be refactored
     return (await hre.ethers.getContractAt(name, address)) as unknown as T;
 };
@@ -16,3 +16,4 @@ export const getSingleton = async <T extends BaseContract>(hre: HardhatRuntimeEn
 export const getPlugin = (hre: HardhatRuntimeEnvironment, address: string) => getInstance<BasePlugin>(hre, "BasePlugin", address);
 export const getRelayPlugin = (hre: HardhatRuntimeEnvironment) => getSingleton<RelayPlugin>(hre, "RelayPlugin");
 export const getRegistry = async (hre: HardhatRuntimeEnvironment) => getInstance<TestSafeProtocolRegistryUnrestricted>(hre, "TestSafeProtocolRegistryUnrestricted", await getProtocolRegistryAddress(hre));
+export const getWhiteListPlugin = async (hre: HardhatRuntimeEnvironment) => getSingleton<WhitelistPlugin>(hre, "WhitelistPlugin");

contracts/src/utils/dataTypes.ts

@@ -9,11 +9,11 @@ export interface SafeProtocolAction {
 export interface SafeTransaction {
     actions: SafeProtocolAction[];
     nonce: bigint;
-    metaHash: Uint8Array;
+    metadataHash: Uint8Array | string;
 }
 
 export interface SafeRootAccess {
     action: SafeProtocolAction;
     nonce: bigint;
-    metaHash: Uint8Array;
+    metadataHash: Uint8Array | string;
 }

contracts/test/WhitelistPlugin.spec.ts

@@ -0,0 +1,133 @@
+import hre, { deployments, ethers } from "hardhat";
+import { expect } from "chai";
+import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";
+import { getWhiteListPlugin, getInstance } from "../src/utils/contracts";
+import { loadPluginMetadata } from "../src/utils/metadata";
+import { buildSingleTx } from "../src/utils/builder";
+import { ISafeProtocolManager__factory, MockContract } from "../typechain-types";
+import { MaxUint256, ZeroHash } from "ethers";
+import { getProtocolManagerAddress } from "../src/utils/protocol";
+
+describe("WhitelistPlugin", async () => {
+    let user1: SignerWithAddress, user2: SignerWithAddress;
+
+    before(async () => {
+        [user1, user2] = await hre.ethers.getSigners();
+    });
+
+    const setup = deployments.createFixture(async ({ deployments }) => {
+        await deployments.fixture();
+        const manager = await ethers.getContractAt("MockContract", await getProtocolManagerAddress(hre));
+
+        const account = await (await ethers.getContractFactory("ExecutableMockContract")).deploy();
+        const plugin = await getWhiteListPlugin(hre);
+        return {
+            account,
+            plugin,
+            manager,
+        };
+    });
+
+    it("should be initialized correctly", async () => {
+        const { plugin } = await setup();
+        expect(await plugin.name()).to.be.eq("Whitelist Plugin");
+        expect(await plugin.version()).to.be.eq("1.0.0");
+        expect(await plugin.requiresRootAccess()).to.be.false;
+    });
+
+    it("can retrieve meta data for module", async () => {
+        const { plugin } = await setup();
+        expect(await loadPluginMetadata(hre, plugin)).to.be.deep.eq({
+            name: "Whitelist Plugin",
+            version: "1.0.0",
+            requiresRootAccess: false,
+            iconUrl: "",
+            appUrl: "",
+        });
+    });
+
+    it("should emit AddressWhitelisted when account is whitelisted", async () => {
+        const { plugin, account } = await setup();
+        const data = plugin.interface.encodeFunctionData("addToWhitelist", [user1.address]);
+        expect(await account.executeCallViaMock(await plugin.getAddress(), 0, data, MaxUint256))
+            .to.emit(plugin, "AddressWhitelisted")
+            .withArgs(user1.address);
+    });
+
+    it("Should not allow calls to non-whitelist address", async () => {
+        const { plugin, account, manager } = await setup();
+
+        // Required for isOwner(address) to return true
+        account.givenMethodReturnBool("0x2f54bf6e", true);
+
+        const safeTx = buildSingleTx(user2.address, 0n, "0x", 0n, hre.ethers.randomBytes(32));
+
+        await expect(
+            plugin.executeFromPlugin(await manager.getAddress(), await account.getAddress(), safeTx),
+        ).to.be.revertedWithCustomError(plugin, "AddressNotWhiteListed");
+    });
+
+    it("Should not allow non-owner to execute transaction to whitelisted address", async () => {
+        const { plugin, account, manager } = await setup();
+        const safeAddress = await account.getAddress();
+        const data = plugin.interface.encodeFunctionData("addToWhitelist", [user2.address]);
+        await account.executeCallViaMock(await plugin.getAddress(), 0, data, MaxUint256);
+
+        // Required for isOwner(address) to return false
+        account.givenMethodReturnBool("0x2f54bf6e", false);
+
+        const safeTx = buildSingleTx(user2.address, 0n, "0x", 0n, ZeroHash);
+        await expect(plugin.connect(user1).executeFromPlugin(manager.target, safeAddress, safeTx))
+            .to.be.revertedWithCustomError(plugin, "CallerIsNotOwner")
+            .withArgs(safeAddress, user1.address);
+
+        const managerInterface = ISafeProtocolManager__factory.createInterface();
+        const expectedData = managerInterface.encodeFunctionData("executeTransaction", [account.target, safeTx]);
+
+        expect(await manager.invocationCount()).to.be.eq(0);
+        expect(await manager.invocationCountForMethod(expectedData)).to.be.eq(0);
+    });
+
+    it("Should allow to execute transaction to whitelisted address", async () => {
+        const { plugin, account, manager } = await setup();
+        const safeAddress = await account.getAddress();
+        const data = plugin.interface.encodeFunctionData("addToWhitelist", [user2.address]);
+        await account.executeCallViaMock(await plugin.getAddress(), 0, data, MaxUint256);
+        // Required for isOwner(address) to return true
+        account.givenMethodReturnBool("0x2f54bf6e", true);
+
+        const safeTx = buildSingleTx(user2.address, 0n, "0x", 0n, ZeroHash);
+        expect(await plugin.connect(user1).executeFromPlugin(manager.target, safeAddress, safeTx));
+
+        const managerInterface = ISafeProtocolManager__factory.createInterface();
+        const expectedData = managerInterface.encodeFunctionData("executeTransaction", [account.target, safeTx]);
+
+        expect(await manager.invocationCount()).to.be.eq(1);
+        expect(await manager.invocationCountForMethod(expectedData)).to.be.eq(1);
+    });
+
+    it("Should not allow to execute transaction after removing address from whitelist ", async () => {
+        const { plugin, account, manager } = await setup();
+        const safeAddress = await account.getAddress();
+
+        // Required for isOwner(address) to return true
+        account.givenMethodReturnBool("0x2f54bf6e", true);
+
+        const data = plugin.interface.encodeFunctionData("addToWhitelist", [user2.address]);
+        await account.executeCallViaMock(await plugin.getAddress(), 0, data, MaxUint256);
+
+        const data2 = plugin.interface.encodeFunctionData("removeFromWhitelist", [user2.address]);
+        expect(await account.executeCallViaMock(await plugin.getAddress(), 0, data2, MaxUint256))
+            .to.emit(plugin, "AddressRemovedFromWhitelist")
+            .withArgs(user1.address);
+
+        const safeTx = buildSingleTx(user2.address, 0n, "0x", 0n, ZeroHash);
+
+        await expect(plugin.connect(user1).executeFromPlugin(manager.target, safeAddress, safeTx))
+            .to.be.revertedWithCustomError(plugin, "AddressNotWhiteListed")
+            .withArgs(user2.address);
+
+        const mockInstance = await getInstance<MockContract>(hre, "MockContract", manager.target);
+        expect(await mockInstance.invocationCount()).to.be.eq(0);
+    });
+});

contracts/yarn.lock

@@ -791,9 +791,9 @@
   integrity sha512-AGuwhRRL+NaKx73WKRNzeCxOCOCxpaqF+kp8TJ89QzAipSwZy/NoflkWaL9bywXFRhIzXt8j38sfF7KBKCPWLw==
 
 "@openzeppelin/contracts@^4.9.1":
-  version "4.9.1"
-  resolved "https://registry.yarnpkg.com/@openzeppelin/contracts/-/contracts-4.9.1.tgz#afa804d2c68398704b0175acc94d91a54f203645"
-  integrity sha512-aLDTLu/If1qYIFW5g4ZibuQaUsFGWQPBq1mZKp/txaebUnGHDmmiBhRLY1tDNedN0m+fJtKZ1zAODS9Yk+V6uA==
+  version "4.9.2"
+  resolved "https://registry.yarnpkg.com/@openzeppelin/contracts/-/contracts-4.9.2.tgz#1cb2d5e4d3360141a17dbc45094a8cad6aac16c1"
+  integrity sha512-mO+y6JaqXjWeMh9glYVzVu8HYPGknAAnWyxTRhGeckOruyXQMNnlcW6w/Dx9ftLeIQk6N+ZJFuVmTwF7lEIFrg==
 
 "@safe-global/mock-contract@^4.0.0":
   version "4.0.0"