Add contracts, tests

Author: akshay-ap

contracts/NoDegegateCallGuard.sol

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.28;
+
+import { BaseGuard } from "@safe-global/safe-contracts/contracts/base/GuardManager.sol";
+import { Enum } from "@safe-global/safe-contracts/contracts/common/Enum.sol";
+
+contract NoDegegateCallGuard is BaseGuard {
+
+    error DelegateCallNotAllowed();
+
+    function checkTransaction(
+        address /*to*/,
+        uint256 /*value*/,
+        bytes memory /*data*/,
+        Enum.Operation operation,
+        uint256 /*safeTxGas*/,
+        uint256 /*baseGas*/,
+        uint256 /*gasPrice*/,
+        address /*gasToken*/,
+        address payable /*refundReceiver*/,
+        bytes memory /*signatures*/,
+        address /*msgSender*/
+    ) external {
+        if(operation == Enum.Operation.DelegateCall) {
+            revert DelegateCallNotAllowed();
+        }
+    }
+
+    function checkAfterExecution(bytes32 txHash, bool success) external {
+
+    }
+}

contracts/SafeGuard.sol

@@ -0,0 +1,104 @@
+import { ethers } from "hardhat";
+import { expect } from "chai";
+import { Contract, Signer, ZeroAddress } from "ethers";
+import { Safe, Safe__factory, SafeProxyFactory } from "../typechain-types";
+import { execTransaction } from "./utils/utils";
+
+describe("Example module tests", async function () {
+  let deployer: Signer;
+  let alice: Signer;
+  let masterCopy: Safe;
+  let proxyFactory: SafeProxyFactory;
+  let safeFactory: Safe__factory;
+  let safe: Safe;
+  let exampleGuard: Contract;
+
+  // Setup signers and deploy contracts before running tests
+  before(async () => {
+    [deployer, alice] = await ethers.getSigners();
+
+    safeFactory = await ethers.getContractFactory("Safe", deployer);
+    masterCopy = await safeFactory.deploy();
+
+    proxyFactory = await (
+      await ethers.getContractFactory("SafeProxyFactory", deployer)
+    ).deploy();
+  });
+
+  // Setup contracts: Deploy a new token contract, create a new Safe, deploy the TokenWithdrawModule contract, and enable the module in the Safe.
+  const setupContracts = async (
+    walletOwners: Signer[],
+    threshold: number
+  ) => {
+    const ownerAddresses = await Promise.all(
+      walletOwners.map(async (walletOwner) => await walletOwner.getAddress())
+    );
+
+    const safeData = masterCopy.interface.encodeFunctionData("setup", [
+      ownerAddresses,
+      threshold,
+      ZeroAddress,
+      "0x",
+      ZeroAddress,
+      ZeroAddress,
+      0,
+      ZeroAddress,
+    ]);
+
+    // Read the safe address by executing the static call to createProxyWithNonce function
+    const safeAddress = await proxyFactory.createProxyWithNonce.staticCall(
+      await masterCopy.getAddress(),
+      safeData,
+      0n
+    );
+
+    // Create the proxy with nonce
+    await proxyFactory.createProxyWithNonce(
+      await masterCopy.getAddress(),
+      safeData,
+      0n
+    );
+
+    if (safeAddress === ZeroAddress) {
+      throw new Error("Safe address not found");
+    }
+
+    // Deploy the TokenWithdrawModule contract
+    exampleGuard = await (
+      await ethers.getContractFactory("NoDegegateCallGuard", deployer)
+    ).deploy();
+
+    const safe = await ethers.getContractAt("Safe", safeAddress);
+
+    // Enable the module in the safe
+    const enableModuleData = masterCopy.interface.encodeFunctionData(
+      "setGuard",
+      [exampleGuard.target]
+    );
+
+    // Execute the transaction to enable the module
+    await execTransaction(
+      walletOwners.slice(0, threshold),
+      safe,
+      safe.target,
+      0,
+      enableModuleData,
+      0
+    );
+  };
+
+  // Test case to verify token transfer to bob
+  it("Should successfully transfer tokens to bob", async function () {
+    const wallets = [alice];
+    await setupContracts(wallets, 1);
+    // Execute the transaction to enable the module
+    await expect ( execTransaction(
+        wallets,
+        safe,
+        ZeroAddress,
+        0,
+        "0x",
+        1
+    )).to.be.revertedWithCustomError(exampleGuard, "DelegateCallNotAllowed");
+  });
+});

test/NoDelegateCall.test.ts

@@ -0,0 +1,77 @@
+
+import { ethers } from "hardhat";
+import { Signer, AddressLike, BigNumberish, ZeroAddress } from "ethers";
+import { Safe } from "../../typechain-types";
+
+/**
+ * Executes a transaction on the Safe contract.
+ * @param wallets - The signers of the transaction.
+ * @param safe - The Safe contract instance.
+ * @param to - The address to send the transaction to.
+ * @param value - The value to send with the transaction.
+ * @param data - The data to send with the transaction.
+ * @param operation - The operation type (0 for call, 1 for delegate call).
+ */
+const execTransaction = async function (
+  wallets: Signer[],
+  safe: Safe,
+  to: AddressLike,
+  value: BigNumberish,
+  data: string,
+  operation: number,
+): Promise<void> {
+  // Get the current nonce of the Safe contract
+  const nonce = await safe.nonce();
+
+  // Get the transaction hash for the Safe transaction
+  const transactionHash = await safe.getTransactionHash(
+    to,
+    value,
+    data,
+    operation,
+    0,
+    0,
+    0,
+    ZeroAddress,
+    ZeroAddress,
+    nonce
+  );
+
+  let signatureBytes = "0x";
+  const bytesDataHash = ethers.getBytes(transactionHash);
+
+  // Get the addresses of the signers
+  const addresses = await Promise.all(wallets.map(wallet => wallet.getAddress()));
+  // Sort the signers by their addresses
+  const sorted = wallets.sort((a, b) => {
+    const addressA = addresses[wallets.indexOf(a)];
+    const addressB = addresses[wallets.indexOf(b)];
+    return addressA.localeCompare(addressB, "en", { sensitivity: "base" });
+  });
+
+  // Sign the transaction hash with each signer
+  for (let i = 0; i < sorted.length; i++) {
+    const flatSig = (await sorted[i].signMessage(bytesDataHash))
+      .replace(/1b$/, "1f")
+      .replace(/1c$/, "20");
+    signatureBytes += flatSig.slice(2);
+  }
+
+  // Execute the transaction on the Safe contract
+  await safe.execTransaction(
+    to,
+    value,
+    data,
+    operation,
+    0,
+    0,
+    0,
+    ZeroAddress,
+    ZeroAddress,
+    signatureBytes
+  );
+};
+
+export {
+  execTransaction,
+};